# Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields

arXiv: https://arxiv.org/abs/2503.08674 · HTML: https://mlsanity.com/paper/2503.08674/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 66 references that are in the map.
- First posted on arXiv 11 Mar 2025.

## Builds on (most cited first, 30 of 30)

- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [SchNet: A continuous-filter convolutional neural network for modeling quantum interactions](https://mlsanity.com/paper/1706.08566/) (2017)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Measuring Robustness to Natural Distribution Shifts in Image Classification](https://mlsanity.com/paper/2007.00644/) (2020)
- [GemNet: Universal Directional Graph Neural Networks for Molecules](https://mlsanity.com/paper/2106.08903/) (2021)
- [Implicit Bias of Gradient Descent on Linear Convolutional Networks](https://mlsanity.com/paper/1806.00468/) (2018)
- [The Open Catalyst 2020 (OC20) Dataset and Community Challenges](https://mlsanity.com/paper/2010.09990/) (2020)
- [Deep Potential Molecular Dynamics: a scalable model with the accuracy of quantum mechanics](https://mlsanity.com/paper/1707.09571/) (2017)
- [Test-Time Training with Masked Autoencoders](https://mlsanity.com/paper/2209.07522/) (2022)
- [EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations](https://mlsanity.com/paper/2306.12059/) (2023)
- [Domain Adaptive Ensemble Learning](https://mlsanity.com/paper/2003.07325/) (2020)
- [The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts](https://mlsanity.com/paper/2206.08917/) (2022)
- [Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations](https://mlsanity.com/paper/2210.07237/) (2022)
- [Test-Time Training on Nearest Neighbors for Large Language Models](https://mlsanity.com/paper/2305.18466/) (2023)
- [Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction](https://mlsanity.com/paper/2502.12147/) (2025)
- [Test-Time Adaptation via Self-Training with Nearest Neighbor Information](https://mlsanity.com/paper/2207.10792/) (2022)
- [GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets](https://mlsanity.com/paper/2204.02782/) (2022)
- [OOD-CV: A Benchmark for Robustness to Out-of-Distribution Shifts of Individual Nuisances in Natural Images](https://mlsanity.com/paper/2111.14341/) (2021)
- [The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains](https://mlsanity.com/paper/2410.24169/) (2024)
- [From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction](https://mlsanity.com/paper/2310.16802/) (2023)
- [Transition1x -- a Dataset for Building Generalizable Reactive Machine Learning Potentials](https://mlsanity.com/paper/2207.12858/) (2022)
- [Graph Neural Networks Use Graphs When They Shouldn't](https://mlsanity.com/paper/2309.04332/) (2023)
- [EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations](https://mlsanity.com/paper/2310.02428/) (2023)
- [Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians](https://mlsanity.com/paper/2501.09009/) (2025)
- [chemtrain: Learning Deep Potential Models via Automatic Differentiation and Statistical Physics](https://mlsanity.com/paper/2408.15852/) (2024)
- [Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators](https://mlsanity.com/paper/2402.13984/) (2024)
- [Synthetic pre-training for neural-network interatomic potentials](https://mlsanity.com/paper/2307.15714/) (2023)
- [Overcoming systematic softening in universal machine learning interatomic potentials by fine-tuning](https://mlsanity.com/paper/2405.07105/) (2024)
- [Injecting Domain Knowledge from Empirical Interatomic Potentials to Neural Networks for Predicting Material Properties](https://mlsanity.com/paper/2210.08047/) (2022)

## Built on by (most cited first, 3 of 3)

- [Space Group Equivariant Crystal Diffusion](https://mlsanity.com/paper/2505.10994/) (2025)
- [A Survey of Large Language Models for Text-Guided Molecular Discovery: from Molecule Generation to Optimization](https://mlsanity.com/paper/2505.16094/) (2025)
- [BLIPs: Bayesian Learned Interatomic Potentials](https://mlsanity.com/paper/2508.14022/) (2025)

## Abstract (arXiv)

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/.