Understanding and Mitigating Distribution Shifts for Machine Learning Force Fields

Authors: Tobias Kreiman, Aditi S. Krishnapriyan
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: machine learning force fields, test-time training, distribution shifts

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 ab initio labels. The first strategy, based on spectral graph theory, modifies the edges of test graphs to align with graph structures seen during training. It can be applied to any existing pre-trained model to mitigate connectivity distribution shifts. Our second strategy improves representations for out-of-distribution systems at test-time by taking gradient steps using an auxiliary objective. Inspired by previous test-time training works in computer vision, we replace self-supervised objectives at test time with an objective that uses an efficient prior to address distribution shifts. Our test-time refinement strategies can reduce force errors by an order of magnitude 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.

Decision: reject
Reviewer scores: 6, 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1729871662928

thread_type:
official_review

<thread_object>
creator: reviewer_ube3

<content>
summary:
This paper addresses the test-time domain-shift problem in machine learning force fields (MLFFs), a common issue when applying pretrained MLFFs to new materials in real-world applications. The authors identify three typical types of domain shifts and propose two test-time refinement strategies to address them. The authors provide partial experimental validation of their claims.

strengths:
S1: The authors identify three primary origins of domain shift in MLFFs, most of which would not be explicitly appeared in the previous literature. 

S2: They propose two strategies to mitigate these shifts: test-time radius refinement (RR) and test-time training (TTT) using inexpensive priors. As far as I know, this work would be the first one applying test time refinement strategies to MLFFs study. 

S3: Experiments were conducted across various datasets, including MD17, MD22, SPICE, SPICEv2, and OC20.

weaknesses:
Major Weaknesses:

W1. $\textbf{Assumption of Prior Knowledge on Test Data}$. The method assumes prior knowledge of test data (e.g., target materials, force labels). It conflicts with machine learning principles, where test data should ideally be inaccessible prior to testing, while this assumption may partially align with certain use cases in ab initio molecular dynamics simulations (MDs) studies (e.g., at least target material name and typical structure would be known beforehand). The scenario targeted in the paper should be more thoroughly explained in the abstract and introduction, including limitations and specific applications of the method.

W2. $\textbf{Lack of Clear Motivation}$. Related to the first point, the utility of the proposed method in practical settings is unclear. In computational chemistry, force MAE/RMSE values typically need to be below 1 kcal/mol/Å (~43 meV/Å) for reliability, but most results provided do not reach this threshold. Thus, practitioners may prefer simple data generation or active learning via DFT/CCSD for rigorously reliable results, leaving the proposed method of more academic than practical interest. This should be approached together with W1 in introduction, abstract, and conclusion (and maybe in limitation). 

W3. $\textbf{Unclear Test Set Assumptions}$. Also connected to the first weakness, the assumptions regarding test set data (e.g., accessibility of force/potential labels, material structure details, availability of classical force fields) lack clarity. These setup should be explicitly defined in Section 2.

Minor Weaknesses:

W4. $\textbf{Potential Physical Limitations of Radius Refinement (RR)}$. Adjusting the cutoff radius in RR may introduce unexpected potential discontinuities, causing sudden and potentially destabilizing force changes (which is well-known issue for DFT/MLFFs). This modification can affect MD simulation results, and thus, a careful theoretical and empirical examination of the cutoff radius effect on MD simulations is advised.

W5. $\textbf{Insufficient MD Simulation Analysis.}$
        
(5.1) The authors conducted NVT simulations, which can artificially stabilize simulations due to the thermostat. A fair comparison between MLFFs with and without RR/TTT would benefit rather from NVE simulations.
       
(5.2) The MD simulation analysis includes only TTT, with no validation of RR effects. Based on point W4, it is strongly recommended to test RR's impact on MD simulations.

questions:
Q1. How can the authors evaluate force distribution shifts without access to test dataset labels?
    
Q2. Why are there no experiments combining MACE+TTT and GemNet-T+RR?
    
Q3. Is TTT more accurate than models trained directly using prior labels, such as those from classical force fields?
    
Q4. Why were MACE, GemNet-T, and NequIP chosen as the MLFF models? Do they represent the broader MLFF landscape?
    
Q5. Why are "Necessity of Proper Pre-training for Test-Time Training" (Sec. C) and "Notes on the Prior" (Sec. F.4) in the Appendix? These seem critical and perhaps better to belong in the main text.
    
Q6. What does "the same element with a different charge" mean at line 150? If it refers to ionization, electrostatic forces should be considered, as this represents a significant problem setting change, not merely a distribution shift.
    
Q7. At lines 261-262, what does “for many molecular datasets” mean specifically? Additional information on the datasets would clarify this statement.
    
Q8. In Figure 1’s caption, “(middle)” appears to be missing after “high force norms (…).”

For additional feedback, refer to the "Weaknesses" section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730359459406

thread_type:
official_review

<thread_object>
creator: reviewer_bnuv

<content>
summary:
This paper aims to analyze and mitigate the out-of-distribution problems for machine learning force fields. This paper proposes pertaining and test-time-training (TTT) using data generated from classical physical priors instead of QM calculations. The proposed method is evaluated on the SPICE dataset and is shown to improve OOD performance, including experiments on MD simulations. the experiments also show an analysis of different types of "OOD" in the context of molecular simulations.

strengths:
- The ML force field is a direction of significant interest to the AI for Science community and ICLR audience.
- The OOD challenge is an important one for ML FFs, as we expect to use ML FFs in extrapolation tasks such as relaxation and MD simulations.
- The proposed method improves model OOD performance on a variety of tasks.
- The analysis of different flavors of OOD-ness and how the proposed method helps is valuable to the community.

weaknesses:
The authors include an MD benchmark on the effectiveness of TTT which is great. However, it would be more interesting to test out the SPICE-trained models on MD simulation as a generalization to unseen molecules will be a more suited task for a SPICE-trained model. 

Further, it will be interesting to see how the OOD-ness of a test molecule impacts MD performance. The authors show a reduction in Force MAE with TTT and its correlation with the OOD-ness metrics such as force norm and graph Laplacian eigenvalues. How would TTT impact MD stability under these different types of ODD-ness?

questions:
Why does the author train the GemNet-T model with only 10% of the data? Clearly, the impact of TTT is perhaps higher with less training data. It would be ideal if the authors could include the results of a GemNet-T trained on the same data as the MACE model.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730628348503

thread_type:
official_review

<thread_object>
creator: reviewer_2dvj

<content>
summary:
The paper studies the problem of training machine learning force fields. Specifically, the paper aims to investigate the generalization ability of current MLFFs. To address the generalization problem, the paper proposes two methods: 1) use radius refinement to identify the test graph structure most consistent with the training graph distribution; 2) use cheap priors to perform test-time-training on the MLFF. The paper shows that the proposed method can effectively improve the generalization ability of current MLFFs on unseen molecules.

strengths:
1. The paper studies existing foundation MLFFs and investigates their generalization ability, which could be helpful for the community.

2. The paper proposes an interesting approach to search for the most similar graph structure by tuning the radius for the test unseen molecules.

weaknesses:
1. Even though the paper shows some improvements on unseen molecules, the performance is still magnitudes away from the desirable chemical accuracy. In other words, the proposed method could hardly be useful in practice. As the downstream task could be much more expensive (e.g., wet lab experiments or subject studies) than the simulations, accuracy is still the top priority for MLFF.

2. The actual distance for radius refinement is quite simple, and may not capture more detailed information in the graph connectivity. Also, for the train eigenvalues, since we could have multiple molecules during training, do we aggregate their overall eigenvalue distribution? It seems to me a more reasonable assumption would be that as long as there is one training molecule graph that is similar to the test, the generalization performance could be better. Also, I think more case studies should be provided for this part. I.e., how does the radius change for certain test molecules change its connectivity and make the graph similar to training molecules?

3. The idea of test-time-training is not novel, and many previous works have utilized low-cost priors for MLFF training/fine-tuning.

4. Some assumptions about the generalization issue may not be practical. For example, if the test dataset has unseen elements or the test molecule has a specific sub-structure that has never been encountered, it is probably not ideal to use the pre-trained MLFF for such a task. 

Minor:
1. The main paper should include at least a brief summary of the related work.

questions:
Please address the weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730646250114

thread_type:
official_review

<thread_object>
creator: reviewer_xyhq

<content>
summary:
The paper proposes two low-cost test-time refinement strategies to address distribution shifts in Machine Learning Force Field (MLFF) models, a significant challenge even for large foundation models trained on extensive datasets. Specifically, the first strategy leverages spectral graph theory to adjust the edges of test-time graphs, aligning them with structures observed during training, while the second strategy adapts representations for out-of-distribution (OOD) systems at test-time through gradient steps based on an auxiliary objective. The authors provide empirical results on several OOD benchmarks, demonstrating the effectiveness of these approaches.

strengths:
* Given that expanding MLFF applicability to diverse chemical spaces is a primary goal within AI-for-Science, this approach is well-motivated and offers a promising step towards this objective.
* The paper establishes practical criteria for identifying distribution shifts, which may inspire future work, and the proposed refinement strategies demonstrate meaningful performance improvements over large-scale foundation models and pre-trained baselines on established benchmarks.
* The paper is well-organized and easy to follow.

weaknesses:
* The proposed strategies yield only modest performance gains over pre-trained models, falling short of significantly improving OOD sample prediction accuracy to levels comparable with in-distribution (ID) samples or chemical accuracy. This limitation may hinder the practical applicability of the methods and affect the perceived technical contribution.
* The prior model used in the test-time training strategy (sGDML) may also face generalization issues on OOD samples, potentially limiting the effectiveness of pseudo-labels derived for fine-tuning. Did the authors consider using semi-empirical electronic structure methods (e.g., DFTB) as a prior model? Semi-empirical methods might offer better generalization in broader chemical spaces.

questions:
* The idea of using the Laplacian spectrum to characterize and align graph structures for OOD and ID samples is interesting but would benefit from more theoretical or intuitive insights. Could the authors clarify the motivation behind aligning spectra as a strategy for managing distribution shifts?
* The accuracy differences reported in Table 1 (e.g., 26.75 vs. 26.0) appear slight. Could the authors provide statistical analyses to ensure these differences are not due to random fluctuations in test data?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732067704274

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We appreciate the detailed and thoughtful feedback from all the reviewers. We summarize important updates and comments about the paper here. We have uploaded a new manuscript and updates to the manuscript are in blue. We also reply individually to reviewers with further updates and clarifications. 

**At a high level, we want to emphasize that our paper is about characterizing and better understanding distribution shifts in machine learning force fields**. We also present some initial strategies to mitigate these shifts that **establish a lower bound on the performance gains that can only be improved with further architectural, data, and training innovations.** These topics are especially timely given that more large “foundation models” are being developed in this field, and we think that it is important to learn and characterize where such models are not doing well, in order to make them better. One of our goals is to also help us understand how MLFFs generalize, and these initial training strategies provide evidence for the hypothesis that MLFFs can be trained to generalize more effectively.

**We want to clarify that for out-of-distribution tasks (OOD), matching in-distribution performance (ID) is challenging for machine learning models [1][2].** While it is difficult in machine learning in general to exactly match ID performance on OOD tasks, our experiments on SPICE (section 4.1), MD17 (section 4.2), and MD22 (section C.4) show that TTT and RR are able to mitigate distribution shifts, and the majority of errors on the molecules are below 25 meV / Å (more than 8,000 / 10,000 for SPICE, and all the MD17 and section C.4 MD22 results). **We also have added Fig. 10 (to supplement Tables 2 and 3) to the paper to explicitly quantify that many molecular systems start with errors above 40 meV / Å and get errors brought down close to ID performance with TTT and RR.** 




**We have included results with a GemNet-T model trained on 100% of the SPICE data (as suggested by reviewer bnuv).** We show that the same patterns hold in terms of increased errors with distribution shifts, helping us understand how MLFFs fail to generalize. TTT still significantly helps with OOD performance, and all the absolute errors are lower, providing further evidence that our results can only be improved with further data, training, and architecture innovations. 

**We have added more molecular dynamics simulation experiments.** We show that TTT enables stable simulations with an NVE integrator. We also run MD simulations with models trained on the SPICE dataset, as well as with models that use our radius refinement technique, showing that our methods improve stability of the simulations. These can be found in Tab. 2, Tab. 3, and section C.2.

**We have added additional foundation models (EquiformerV2 and JMP) to section 2** to further help identify distribution shifts that foundation models face. We have also shown that test-time radius refinement (RR) works on the JMP model (which is built off of a GemNet backbone). See section 2 and Tab. 6 and 7 in the appendix.

**We have included further motivation and clarifying details about our RR approach.** To summarize, we were inspired by the theoretical analyses presented in [3] which show that GNNs easily overfit to regularly connected graphs. Our RR approach adjusts graph connectivity at test time to match the connectivity of graphs seen during training. We do so by aligning graph Laplacian eigenvalue distributions using the heuristic described in the paper. Although we explored more involved distance measures for aligning connectivities, we found that they had negligible performance improvements compared to our heuristic. We have added these clarifications to the manuscript in appendix F.     

We thank the reviewers again for the helpful feedback. We believe these additional experiments and clarifications address the major issues brought up by the reviewers. 

[1] Yu Sun, et. al. Test-time training with self-supervision for generalization under distribution shifts, 2020.

[2] Yossi Gandelsman,et. al. Test-time training with masked autoencoders. Advances in Neural Information Processing Systems, 2022.

[3] Maya Bechler-Speicher, et. al. Graph neural networks use graphs when they shouldn’t, 2024. URL https://arxiv.org/abs/2309.04332.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes an analysis on types of distribution shift for learning and using machine learning force field (MLFF) and two ways to bridge them. This is an important topic of practical consideration, and the methods are investigated over a variety of datasets. While most reviewers appreciate the contributions, there still remain some concerns and insufficiencies. While I also appreciate the effort on this topic, I still feel that the paper seems a bit heuristic and is expected to provide more serious justifications on some points.
* The characterizations of distributional shift seems a bit superficial. The first type is more like an instance/sample-level difference: they are just different inputs. The second one needs further justifications considering that typical GNNs may already be able to output forces with a larger norm when the input atoms are closer. The authors may need to consider the built-in inductive bias of a model/architecture when discussing generalization. The third one also seems superficial: if the difference in the input graph structure can be accounted as a distribution shift, then why won't any change in the input, e.g., a different 3D conformation, is considered?
* The paper seems to not sufficiently distinguish instance-level difference and distribution-level difference. I hope to see more analysis on the distribution shift across different datasets, while the paper seems to present instance-to-distribution difference at most (for the cutoff radius selection). Moreover, how do the authors ensure that it is not due to the in-distribution generalization gap that including more in-distribution data still cannot bridge the gap (Fig. 2)?
* For the first method, why is graph Laplacian sufficient to represent the data distribution? This is only a one-dimensional projection of the data distribution, and might lose distribution in other aspects. I read from the authors' rebuttal and would like to further ask, does the ineffectiveness of using KL-divergence between eigenvalue distributions indicate that the proposed method may not be working in the designed way?
* Also regarding the first method, why changing cutoff radius won't break the model's knowledge on inter-atomic interaction? More specifically, in training, the model may have aggregated effects of atoms outside the cutoff radius to those within the radius. Will the model double-count the effects when using a larger radius? Or if the model would miss necessary interactions when using a smaller radius?
* Pretraining by aligning with a classical force field is not essentially a novel one. It could be very similar to pretraining a model on a large number of cheaper labels. The authors may also be interested in discussing why this won't make the model inherit artifacts (e.g., overly flat energy landscape) of the classical force field, which would make a noticeable difference on off-equilibrium structures.
* Some reviewers also complained the final performance not being useful as there is still a gap to the in-distribution case. This is indeed an insufficiency, but I won't put too much criticism on this due the intrinsic distribution shift gap (as the authors argued). As the method does not require direct supervision on the test domain, an improvement could already be beneficial.

I have gone through reviewers' comments. Although more reviewers have proposed a score suggesting acceptance and I'm aware of the contributions, and one reviewer has raised the score, they still raised a few insufficiencies, with some still remain. Moreover, I did not find my concerns sufficiently addressed throughout the rebuttals. I hence recommend a reject and hope the authors could further improve the paper for a better version.