Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators
This addresses the critical issue of simulation instability in MLFFs for molecular dynamics, enabling more reliable modeling of longer timescales and better observable estimation, particularly in data-scarce scenarios.
The paper tackles the problem of unstable simulations in machine learning force fields (MLFFs) by introducing Stability-Aware Boltzmann Estimator (StABlE) Training, which uses joint supervision from quantum-mechanical calculations and system observables to correct instabilities without additional ab-initio calculations. The method achieves significant improvements in simulation stability, data efficiency, and agreement with reference observables across various systems and MLFF architectures, allowing for larger timesteps.
Machine learning force fields (MLFFs) are an attractive alternative to ab-initio methods for molecular dynamics (MD) simulations. However, they can produce unstable simulations, limiting their ability to model phenomena occurring over longer timescales and compromising the quality of estimated observables. To address these challenges, we present Stability-Aware Boltzmann Estimator (StABlE) Training, a multi-modal training procedure which leverages joint supervision from reference quantum-mechanical calculations and system observables. StABlE Training iteratively runs many MD simulations in parallel to seek out unstable regions, and corrects the instabilities via supervision with a reference observable. We achieve efficient end-to-end automatic differentiation through MD simulations using our Boltzmann Estimator, a generalization of implicit differentiation techniques to a broader class of stochastic algorithms. Unlike existing techniques based on active learning, our approach requires no additional ab-initio energy and forces calculations to correct instabilities. We demonstrate our methodology across organic molecules, tetrapeptides, and condensed phase systems, using three modern MLFF architectures. StABlE-trained models achieve significant improvements in simulation stability, data efficiency, and agreement with reference observables. The stability improvements cannot be matched by reducing the simulation timestep; thus, StABlE Training effectively allows for larger timesteps. By incorporating observables into the training process alongside first-principles calculations, StABlE Training can be viewed as a general semi-empirical framework applicable across MLFF architectures and systems. This makes it a powerful tool for training stable and accurate MLFFs, particularly in the absence of large reference datasets. Our code is available at https://github.com/ASK-Berkeley/StABlE-Training.