Understanding and Mitigating Distribution Shifts For Machine Learning
  Force Fields
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/.
67d323231d5d1fd92c35ad1c
Tobias Kreiman
extracted_confirmed
2025-03-13T18:32:48.171Z
64bece3e01f1983a86a3e043
/avatars/7f6372d409e8eb0474f7ff6282cfd4e9.svg
Toby Kreiman
user
tkreiman
67d323231d5d1fd92c35ad1d
Aditi S. Krishnapriyan
claimed_verified
2025-05-14T07:37:00.927Z
6823adb00368164fae85147c
/avatars/f9b46f9415be40d5b42cb88dc38b9d75.svg
Aditi Krishnapriyan
user
a1k12
Understanding and Mitigating Distribution Shifts For Machine Learning
  Force Fields
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/.
Test-time refinement strategies using spectral graph theory and auxiliary objectives enhance the generalization of MLFFs to out-of-distribution chemical systems.
MLFFs ab initio quantum mechanical molecular simulations distribution shifts spectral graph theory auxiliary objective out-of-distribution systems