Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields
cs.LG, cond-mat.mtrl-sci, physics.chem-ph, q-bio.BM
Submitted: 2025-03-11
Updated: 2025-05-29
DOI: 10.1039/d5dd00260e
Project page: https://tkreiman.github.io/projects/mlff_distribution_shifts
License: http://creativecommons.org/licenses/by/4.0/
The gist: Machine Learning Force Fields (MLFFs) are a promising alternative to expensive ab initio quantum mechanical molecular simulations.
Terminology
Abstract
Machine Learning Force Fields (MLFFs) are a promising alternative to expensive ab initio quantum mechanical molecular simulations. Given the diversity of chemical spaces that are of interest and the cost of generating new data, it is important to understand how MLFFs generalize beyond their training distributions. In order to characterize and better understand distribution shifts in MLFFs, we conduct diagnostic experiments on chemical datasets, revealing common shifts that pose significant challenges, even for large foundation models trained on extensive data. Based on these observations, we hypothesize that current supervised training methods inadequately regularize MLFFs, resulting in overfitting and learning poor representations of out-of-distribution systems. We then propose two new methods as initial steps for mitigating distribution shifts for MLFFs. Our methods focus on test-time refinement strategies that incur minimal computational cost and do not use expensive ab initio reference labels. The first strategy, based on spectral graph theory, modifies the edges of test graphs to align with graph structures seen during training. Our second strategy improves representations for out-of-distribution systems at test-time by taking gradient steps using an auxiliary objective, such as a cheap physical prior. Our test-time refinement strategies significantly reduce errors on out-of-distribution systems, suggesting that MLFFs are capable of and can move towards modeling diverse chemical spaces, but are not being effectively trained to do so. Our experiments establish clear benchmarks for evaluating the generalization capabilities of the next generation of MLFFs. Our code is available at https://tkreiman.github.io/projects/mlff distribution shifts/.
Sources
- A foundation model for atomistic materials chemistry
- Graph Neural Networks Use Graphs When They Shouldn't
- Overcoming systematic softening in universal machine learning interatomic potentials by fine-tuning
- Nutmeg and SPICE: Models and Data for Biomolecular Machine Learning
- Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction
- GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets
- On Over-Squashing in Message Passing Neural Networks: The Impact of Width, Depth, and Topology
- Implicit Bias of Gradient Descent on Linear Convolutional Networks
- Test-Time Adaptation via Self-Training with Nearest Neighbor Information
- EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
- Transition1x -- a Dataset for Building Generalizable Reactive Machine Learning Potentials
- From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks