Physics-Informed Weakly Supervised Learning for Interatomic Potentials

Authors: Makoto Takamoto, Viktor Zaverkin, Mathias Niepert
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: machine learning interatomic potential, weakly supervised learning, machine learning for science

Abstract:
Machine learning plays an increasingly important role in computational chemistry and materials science, complementing computationally intensive ab initio and first-principles methods. Despite their utility, machine-learning models often lack generalization capability and robustness during atomistic simulations, yielding unphysical energy and force predictions that hinder their real-world applications. We address this challenge by introducing a physics-informed, weakly supervised approach for training machine-learned interatomic potentials (MLIPs). We introduce two novel loss functions, extrapolating the potential energy via a Taylor expansion and using the concept of conservative forces. Our approach improves the accuracy of MLIPs applied to training tasks with sparse training data sets and reduces the need for pre-training computationally demanding models with large data sets. Particularly, we perform extensive experiments demonstrating reduced energy and force errors---often lower by a factor of two---for various baseline models and benchmark data sets. Moreover, we demonstrate improved robustness during MD simulations of the MLIP models trained with the proposed weakly supervised loss.
    Finally, we show that our approach facilitates MLIPs' training in a setting where the computation of forces is infeasible at the reference level, such as those employing complete-basis-set extrapolation. An implementation of our method and scripts for executing experiments are available at \url{https://anonymous.4open.science/r/PICPS-ML4Sci-1E8F}.

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

Reviews and discussion:
<thread>
creation date:
1729585670257

thread_type:
official_review

<thread_object>
creator: reviewer_xngz

<content>
summary:
This work presents a physics-informed, weakly supervised approach for training machine-learned interatomic potentials. The method is innovative, and the ideas are inspiring. However, it may not be practical for real-world applications, highlighting the need for further validation and case studies. Additionally, the improvements in experimental results are minimal, particularly with the rMD17 dataset, a common benchmark.

strengths:
1. The work introduces a physics-informed, weakly supervised approach for training machine-learned interatomic potentials.
2. The method is innovative, and the ideas are truly inspiring.

weaknesses:
The proposed method may not be practical for real-world applications, highlighting the need for further validation and case studies. Additionally, the improvement in experimental results is minimal, particularly with the rMD17 dataset, which is a common benchmark.

questions:
In training MLIPs for small molecules, there is usually a substantial amount of data available (over 1,000 samples). How effective is the proposed method with a larger dataset of small molecules?

How effective is training on datasets for large molecules, such as the MD22 dataset?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730035191800

thread_type:
official_review

<thread_object>
creator: reviewer_ji8n

<content>
summary:
The paper proposes an extended loss for the training of machine learning interatomic potentials. Two additional terms are added in the loss, the physics-informed Taylor-expansion-based consistency (PITC) loss and the physics-informed spatial-consistency (PISC) loss. The idea behind these terms is to generate additional approximate training data points by a first-order Taylor expansion of the energy at a given training data point. The goal is to increase the robustness of the machine learning interatomic potential for applications in this training approach. The approach is validated in several benchmarks.

strengths:
- The approach appears to provide consistently better results.
- The approach is independent on the machine learning potential model.
- The approach can be easily integrated in other projects.

weaknesses:
- The potential energy curve of the C-H bond is badly represented in both cases, with and without PIWSL. For me, the argument is not strong if it still fails but not as badly. Maybe this has just to be trained to better/more (400, 800) training data and then both will give good results. To show the advantage of PIWSL, I think you have to find the sweet spot where PIWSL can already deal with a low amount of data and the conventional approach fails due to the lack of data. You could then also add a graph, where both approaches are trained on more data and both succeed to demonstrate that it is really the more efficient data handling of PIWSL. I think, the reference and predicted curves should at least agree in the range where the relative energy is below about 2 kcal/mol, because this region will be sampled in molecular dynamics simulations. However, especially the error for bond increases in plot d appears very large (>1 kcal/mol), currently.
- The approach increases the computational time required for training significantly.

questions:
- Are the reported E and F RMSE values for training or test data?
- If I understand it correctly, the training was stopped after a per-defined number of training epochs. However, the PIWSL method requires 2-3 times more training time due to the additional calculation. I think a fair comparison would be to reduce the number of training epochs for PIWSL by a factor of 2-3 or otherwise increase the number of training epochs for the baseline approach. You could also just use the wall-clock time as stopping criterion instead of a given number of epochs. Did you check the convergence of the baseline approach? You could add learning curves to validate the convergence.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730054243313

thread_type:
official_review

<thread_object>
creator: reviewer_xuz7

<content>
summary:
The authors propose a new pre-training method for MLIPs by augmenting the typical supervised loss with a physics-inspired loss (PIWSL). The proposed objective has a part that encourages the model to learn a conservative force field, and another part that ensures consistency of the potential energy surface based on a Taylor expansion around structure in the dataset. The authors conduct experiments on small datasets showing that their physics-informed loss improves performance compared to training just with supervision on energies and forces.

strengths:
Strengths:
- The paper is written clearly and the method is motivated and presented well.
- Pre-training / self-supervised methods are important for the field of ML and could be particularly impactful for the field of MLIPs, where data is computationally expensive to generate.
- The proposed method is novel and physically motivated. It seems to be easy to implement programmatically and could be applied generally to any MLIP.

weaknesses:
The main weaknesses are in the evaluation of the proposed method. In most other areas of ML, weakly-supervised learning is done with large datasets to learn general representations that can then be transferred / fine tuned downstream. However, in this paper, the experiments are done on extremely small datasets ranging in size from 100 - 1000 samples. This makes it hard to argue that the PIWSL serves as a useful objective, especially as the field migrates to larger datasets (OC20, MPTrj, SPICE, etc.). Concretely (roughly in order of importance):
- The experiments are done on extremely small datasets, making it hard to argue that the PIWSL helps the model learn useful representations (I suspect that any form of regularization would improve performance on datasets this small). On these small datasets, a non-deep learning based model like sGDML [3] outperforms all of the models tested in this paper in Table A9 by a large margin. It would be much more compelling if PIWSL still improved performance on the larger (1M+ samples) more common datasets (OC20, SPICE, MPTrj, full ANI-1x dataset, etc.).  
- Since the models are trained on such small datasets, the errors are currently huge (more than 10 kcal / mol across in almost all cases in table 1 for instance). Therefore, the improvement of ~1 kcal / mol with PIWSL still leaves the models with a huge error (reflected in the poor simulation performance in Figure 3.) As the dataset size increases, the improvement from PIWSL seems to decrease as well (mentioned by the authors in Fig. 2).
- The baseline NoisyNodes isn’t really a valid comparison since it is meant to only work for equilibrium structures. A much more valid comparison would be something like DeNS [1] that does work for the non-equilibrium structures that the authors test on.
- The models tested do not necessarily cover the state-of-the-art models for these types of molecular datasets. Equiformer and eSCN for instance are tuned for OC20 and are often outperformed by models like MACE [2] on molecular datasets. Something like MACE is also often able to get errors well below 1 kcal / mol with just supervision from energies and forces. If PIWSL still helped with these better models, that would be more compelling.

Weakly/self-supervised strategies have shown to be extremely effective in other fields of ML, and I am optimistic that methods like the ones presented by the authors will be important for MLIPs. However, the current evaluations in the paper do not provide convincing evidence that the current instantiation of PIWSL would actually help the best MLIPs when trained on large datasets. I would be happy to raise my score with more rigorous evaluations.

questions:
- Why were the simulations done with the NVE ensemble (it makes it harder to compare to the work in Fu et. al. (2023) since they do an NVT with NoseHoover)? This is briefly discussed around line 400 but I feel like some results with NVT simulations could be helpful, since at the moment the NVE simulations yield very low stability.
- What happens if the models are trained explicitly to minimize curl as opposed to the PISC loss?

References:
[1] Yi-Lun Liao, Tess Smidt, Muhammed Shuaibi, & Abhishek Das. (2024). Generalizing Denoising to Non-Equilibrium Structures Improves Equivariant Force Fields. 
[2] Ilyes Batatia, Dávid Péter Kovács, Gregor N. C. Simm, Christoph Ortner, & Gábor Csányi. (2023). MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields
[3] Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T. Unke, Adil Kabylda, Huziel E. Sauceda, Alexandre Tkatchenko, & Klaus-Robert Müller (2023). Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9(2), eadf0873.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731117684490

thread_type:
official_review

<thread_object>
creator: reviewer_kyhe

<content>
summary:
The authors propose two self-supervised losses to train machine learning interatomic potentials. One is based on a Taylor-expansion approach, and the other is based on a consistency loss. They use these losses to train MLIPs on ANI, a TiO2 dataset, and some molecules in rMD17. They look at training models with mostly n=100 to n=1000 samples, and performance primarily on energy and forces prediction, as well as an example of running the molecular dynamics simulation.

strengths:
- Self-supervised learning strategies for MLIPs are not very common, so it is nice to see some attempts towards developing these.
- Using these losses could be useful in cases where only energy labels are available, and not forces

weaknesses:
There are many concerns I have with this method and actual utility in the broader scope of the MLIP literature.

- The results are not very convincing that this is a method that would actually get used. The authors are looking at a low number of data points, when larger datasets now exist. It is true that they are looking at more data points in the ANI dataset, though the level of accuracy and data generation of this dataset is suspect.

- In some cases, models can be trained by taking gradients of the energy to get a conservative force field, including cases the authors looked at. If this is the case, why do we need these losses? How do these losses compare in accuracy to taking gradients of the energy? This would make sense to do in datasets like rMD17 and other molecular datasets that have the forces.

- Stability analysis is only done for 3-6 ps which is very, very short. The models are also barely improving and still not usable, and it is not a good sign that it does not do well for such a simple molecule like aspirin. Why are the authors doing NVE simulations instead of NVT? Why not try NVT?

- The poor molecular dynamics performance means that these potentials are not very good, so they cannot be used for actual simulations.

- The authors say that they run NVE simulations and that this means the total energy is conserved. This is not necessarily true, because the authors may not have a potential that actually conserves energy. The authors can test actual energy conservation with their model.

- There is a line on 361 about “this observation indicates strong evidence of the effectiveness of PIWSL applied to bulk materials.” This is a very specific dataset (TiO2) and such claims are much greater than what is actually being presented. If the authors want to go towards such claims, they should be looking at datasets like the Materials Project, which is much more representative of bulk materials.

- Besides the above example, there are too many strong conclusions are being drawn from not great results to back these up, as these experiments are not being done on more representative datasets and sizes.

questions:
- Do these loss terms help for more representative benchmarked datasets like Materials Project? Or MD22, which is a more representative dataset than MD17

- How does using these losses compare to taking energy gradients when the forces are known?

- Do the losses actually improve energy conservation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732266139289

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Remark 2

comment:
$\textbf{Further analysis on ML ab-initio MD simulations (NVT and NVE with gradient-based force model): }$

We would like to clarify the purpose of Figure 3. Its intention is not to claim that the trained model is entirely suitable for real-world applications but to demonstrate that PIWSL improves prediction stability during simulations. Indeed, the stability of MLIP models depends strongly on the model architecture. Figure 3 shows models employing a direct force branch with relatively poor stability.

To further validate the effectiveness of PIWSL, we added a new stability analysis using the PaiNN model with gradient-based force predictions. The results are shown in  Figure 3, on the right panel in the revised manuscript. This analysis shows that simulations can be performed stably for up to 100 ps, even when trained on only 200 samples, and that PIWSL further enhances stability in this scenario.

Additionally, in response to the reviewer's feedback, we incorporated a new stability measure using MD simulations in the canonical statistical ensemble (NVT), presented in Figure A2. These results highlight that the thermostat, which is required to maintain a constant temperature, significantly stabilizes simulations and can obscure the inherent stability characteristics of the MLIP model.

Finally, we would like to emphasize that the effectiveness of PIWSL has been thoroughly validated across five diverse datasets (ANI-1x, rMD17, MD22, TiO2, LMNTO) using six different models (SchNet, SpinConv, PaiNN, eSCN, Equiformer v2, MACE), including several state-of-the-art (SOTA) approaches. These validations span a wide range of data sizes, from as few as 50 samples to as many as 5,000,000, and include molecular dynamics (MD) simulations under both NVE and NVT conditions. We hope the reviewers will recognize and appropriately value the scope and rigor of these experiments.

Detailed responses to each review are provided below. We hope that our responses address your question satisfactorily. Please let us know if further clarification is needed. 

Best regards,
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1732266331635

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Remark

comment:
Dear Reviewers, 

Thank you very much for your very kind and encouraging suggestions. Based on your valuable suggestions, we uploaded the refined new manuscript, where the newly added or modified text was written in red. Before addressing each review individually, we want to summarize the key points.  

$\textbf{The objective of Our Paper:}$

As we emphasized in the abstract, introduction, and conclusion, the primary objective of our work is to introduce a simple yet widely applicable method that reduces errors in the sparse data regime. This is particularly important for real-world computational chemistry applications, where training data for MLIP models should be generated from scratch using DFT solvers. Due to the high computational cost of DFT calculations, minimizing the overall dataset size is crucial. Reducing errors in the sparse data regimes is essential to accelerate this process and enhance sample efficiency. We emphasize that a reduction of the error in sparse data regimes (even though it could be considered as making a "bad" ML model less "bad") is crucial in such scenarios because it minimizes the number of required DFT calculations, increases the length of atomistic simulations with the surrogate (i.e., increases the number of sampled MD frames), and decreases the overall training time for a reliable MLIP model. Thus, this capability to efficiently reduce error represents a significant contribution from a practical point of view. We showed that the PIWSL method can reduce errors across various models, including recent SOTA models, within these sparse data regimes across multiple material datasets. 

$\textbf{Training on Larger Data sets (ANI-1x: 1,000,000 and 5,000,000 Samples)}$

In response to reviewers' recommendations, we included a new assessment of the PIWSL method on a large dataset. Specifically, we trained the PaiNN model on 1,000,000 and 5,000,000 samples from the ANI-1x dataset. The results are as follows:

$$\\begin{array} {|r|r|}\\hline  & N_{train} & Farce MAE & Force RMSE & Energy MAE & Energy RMSE \\\ \\hline Baseline & 1,000,000 & 0.67 & 2.74 & 4.90 & 6.56 \\\ \\hline PIWSL &  & 0.68 & 2.77 & \textbf{4.48} & \textbf{6.06} \\\ \\hline Baseline & 5,000,000 & 0.53 & 2.18 & 3.94 & 5.24 \\\ \\hline PIWSL &  & 0.55 & 2.24 & \textbf{3.25} & \textbf{4.75} \\\ \\hline  \\end{array}$$

These results demonstrate that PIWSL remains effective in the large dataset regime (1,000,000 and 5,000,000 samples), reducing the energy error by approximately 10% while maintaining the force prediction error. This table is also included as Table A5 in the revised manuscript.

$\textbf{A New result with MACE model on MD22 data set (buckyball catcher)}$

To address the reviewers' suggestions, we conducted an additional experiment on a large molecule dataset using the state-of-the-art MACE model. For this purpose, we trained the MACE model on the MD22 dataset, focusing on the "buckyball catcher" molecule due to its large size (148 atoms) and complex structure. The results are as follows:

$$\\begin{array} {|r|r|}\\hline  &  & N_{train}=50 &  & N_{train}=600 & \\\ \\hline Model &  & Baseline & PIWSL & Baseline & PIWSL \\\ \\hline MACE & E & 0.163 & \textbf{0.155} & 0.091 & \textbf{0.085} \\\ \\hline  & F & 6.780 & 6.758 & 1.886 & 1.899 \\\ \\hline MACE w.t. SWA & E & 0.150 & \textbf{0.106} & 0.080 & \textbf{0.069} \\\ \\hline  & F & 6.844 & 6.807 & 2.056 & 2.052 \\\ \\hline  \\end{array}$$

These results confirm that PIWSL is effective in the large molecule regime when applied with MACE, particularly in improving energy predictions in the sparse data regime, as expected. This table is also included as Table A10 in the revised manuscript.
(Here we used the set up introduced in the official document [2] with slight modification on max_L and mini-batch size, to accelerate training).

[1] Ilyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. Adv. Neural Inf. Process. Syst., 35:11423–11436, 2022.

[2] https://mace-docs.readthedocs.io/en/latest/examples/training_examples.html
</content>
</thread_object>

</thread>
<thread>
creation date:
1732867747127

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experiments

comment:
Dear Reviewers, 

Thank you very much for your insightful feedback. To address your comments, we conducted additional experiments to further demonstrate the effectiveness of PIWSL in the following scenarios: (1) fine-tuning foundation models and (2) cases where the MLIP model has already achieved an accuracy of 1 kcal/mol. The detailed results are as follows:

$$\textbf{(1) fine-tuning foundation models}$$

To illustrate that PIWSL remains valuable in the foundation model + fine-tuning framework, we conducted an experiment on the MD22 dataset using MACE as the foundation model, trained with 50 samples. The results are: 

$$\\begin{array} {|r|r|}\\hline Model &  & F-MAE & E-MAE \\\ \\hline Baseline & MP & 16.503 \\pm 0.206 & 0.329 \\pm 0.071 \\\ \\hline PIWSL & MP & \\mathbf{16.150 \\pm 0.244} & 0.314 \\pm 0.023 \\\ \\hline Baselie & OFF & 7.981 \\pm 0.090 & 0.181 \\pm 0.024 \\\ \\hline PIWSL & OFF & \\mathbf{7.781 \\pm 0.015} & \\mathbf{0.155 \\pm 0.008} \\\ \\hline  \\end{array}$$

Here, MP refers to MACE-MP ("large"), a universal foundation model, and OFF refers to MACE-OFF ("large\_OFF"), a foundation model specialized for organic materials. We used the setup provided by official repository with increasing total epoch from 10 to 100. The results demonstrate that PIWSL accelerates the learning of force labels, achieving consistently lower errors even under the best-case scenario. Notably, MACE-OFF outperforms MACE-MP, highlighting that MACE-MP struggles to learn from organic material data. This indicates that while foundation models can be effective, their performance is highly dependent on the pretraining dataset, and they are not universally applicable. We believe this observation supports our assumption that scratch training remains valuable for studies involving new materials. 

Additionally, we observed that incorporating the second-order term of the Taylor expansion in the PITC loss is critical in this setup. Through simple mathematical derivations,  $\textbf{we extended the PITC loss formulation to approximately account for the second-order term}$, by simply adding $(F(S_{dr}) - F(S))/2 \\cdot dr$ to PITC potential energy estimation. This also demonstrates the flexibility of the PIWSL method based on mathematics and physics.

$$\textbf{(2) 1kcal/mol Accuracy:}$$

We appreciate the reviewer’s suggestion and acknowledge that we may have misunderstood the intention behind their proposal. Regarding accuracy, it is inherently related to the model size, as larger models are better equipped to capture the diversity of larger datasets. In our case, the PaiNN model used was not sufficiently large. This choice was constrained by the rebuttal timeframe, as training larger models, such as Equiformer or eSCN, would have required significantly more time.
To address this point, we performed an additional experiment using eSCN trained on the TiO₂ dataset with 2,000 samples. The results are as follows:

$$\\begin{array} {|r|r|}\\hline  & F-MAE & F-RMSE & E-MAE & E-RMSE \\\ \\hline Baseline & 0.061 \\pm 0.010 & 0.415 \\pm 0.169 & 0.634 \\pm 0.017 & 1.775 \\pm 0.070 \\\ \\hline PIWSL & \\mathbf{0.028 \\pm 0.000} & \\mathbf{0.163 \\pm 0.015} & \\mathbf{0.329 \\pm 0.026} & \\mathbf{0.899 \\pm 0.085} \\\ \\hline  \\end{array}$$
Notably, the energy errors are less than 1 kcal/mol. Moreover, PIWSL achieves an error reduction of approximately 20–50% for both forces and energies.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The submission presents a physics-informed weakly supervised method to improve the generalization of machine-learning force field models (MLFFs), especially in situations where labeled data are scarce. Two losses are proposed based on the physics fact that the force is the negative gradient of potential energy. Comprehensive experiments are conducted over datasets and MLFF architectures, and evaluation on realistic use cases of MD simulation is also appreciated. I agree with the reviewers on the appreciation of the contribution to improve MLFF by leveraging physical laws that bridge different quantities and tasks beyond standard supervised training. Nevertheless, there still remain a few insufficiencies that multiple reviewers have raised:
* Since the proposed method does not provide more information than asking the model to be conservative, what are the advantages over learning a gradient-based/conservative force model, and what could be the rationale that the proposed method outperforms conservative force models. I also asked this question myself to the authors. The authors provided multiple pieces of replies, but are mainly empirical demonstrations, and further clarifications on the roles of the losses do not seem to be helpful for answering this question, since a conservative force model is also expected to achieve the effects. I expect more rationales and evidence on why a conservative force model cannot achieve the same effects, hence ruling out the doubt that the improvement comes from other factors.
* Another common challenge is the use case. The major demonstration is the case of scarce labeled data, while multiple reviewers challenged that there are large data in the domain. The authors further provided additional results on large datasets in the rebuttal. Some reviewers (ji8n and xnGz) appreciated the effort and results, hence increased their scores, while Reviewer xuZ7 challenged the significance of the improvements in the large data cases. I agree with the reviewer on this point. There also remains the doubt that the improvement is insignificant when using better model architectures or conservative force models, or finetuning a large foundation model. Nevertheless, while showing the benefit in the large-data case is one way to show the merit, I don't think it is not common that molecular tasks have scarce data, but the authors may need to find a proper realistic setting, e.g., very accurate energy/force labels, hard-to-compute quantities, or costly experimental data.
* Reviewer kYHe also challenged the stability analysis in MD is short and that the MD performance and settings are not satisfying. These two points seem to have been addressed by additional experiments in the rebuttal.

Considering both sides, while the authors have done a hard work that systematically studied the method in various cases, the remaining insufficiencies are still non-negligible. I recommend a reject to the submission for the current form, and encourage the authors to further strengthen the insufficiencies accordingly for a stronger paper.