# Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 26 of its 80 references that are in the map.
- First posted on arXiv 21 Feb 2024.

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

- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [OptNet: Differentiable Optimization as a Layer in Neural Networks](https://mlsanity.com/paper/1703.00443/) (2017)
- [Directional Message Passing for Molecular Graphs](https://mlsanity.com/paper/2003.03123/) (2020)
- [MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields](https://mlsanity.com/paper/2206.07697/) (2022)
- [Equivariant message passing for the prediction of tensorial properties and molecular spectra](https://mlsanity.com/paper/2102.03150/) (2021)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Spherical Message Passing for 3D Graph Networks](https://mlsanity.com/paper/2102.05013/) (2021)
- [On Differentiating Parameterized Argmin and Argmax Problems with Application to Bi-level Optimization](https://mlsanity.com/paper/1607.05447/) (2016)
- [Less is more: sampling chemical space with active learning](https://mlsanity.com/paper/1801.09319/) (2018)
- [ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations](https://mlsanity.com/paper/2103.01436/) (2021)
- [Machine Learning Coarse-Grained Potentials of Protein Thermodynamics](https://mlsanity.com/paper/2212.07492/) (2022)
- [Learning differentiable solvers for systems with hard constraints](https://mlsanity.com/paper/2207.08675/) (2022)
- [Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size](https://mlsanity.com/paper/2304.10061/) (2023)
- [Learning continuous models for continuous physics](https://mlsanity.com/paper/2202.08494/) (2022)
- [Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting](https://mlsanity.com/paper/2106.01138/) (2021)
- [Differentiable Molecular Simulations for Control and Learning](https://mlsanity.com/paper/2003.00868/) (2020)
- [Committee neural network potentials control generalization errors and enable active learning](https://mlsanity.com/paper/2006.01541/) (2020)
- [EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations](https://mlsanity.com/paper/2310.02428/) (2023)
- [Accurate machine learning force fields via experimental and simulation data fusion](https://mlsanity.com/paper/2308.09142/) (2023)
- [Data efficiency and extrapolation trends in neural network interatomic potentials](https://mlsanity.com/paper/2302.05823/) (2023)
- [Learned Force Fields Are Ready For Ground State Catalyst Discovery](https://mlsanity.com/paper/2209.12466/) (2022)
- [Differentiable Simulations for Enhanced Sampling of Rare Events](https://mlsanity.com/paper/2301.03480/) (2023)
- [Learning Pair Potentials using Differentiable Simulations](https://mlsanity.com/paper/2209.07679/) (2022)
- [Top-down machine learning of coarse-grained protein force-fields](https://mlsanity.com/paper/2306.11375/) (2023)

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

- [Orb: A Fast, Scalable Neural Network Potential](https://mlsanity.com/paper/2410.22570/) (2024)
- [Transformers Discover Molecular Structure Without Graph Priors](https://mlsanity.com/paper/2510.02259/) (2025)
- [chemtrain-deploy: A parallel and scalable framework for machine learning potentials in million-atom MD simulations](https://mlsanity.com/paper/2506.04055/) (2025)
- [Understanding and Mitigating Distribution Shifts For Machine Learning Force Fields](https://mlsanity.com/paper/2503.08674/) (2025)
- [Beyond Force Metrics: Pre-Training MLFFs for Stable MD Simulations](https://mlsanity.com/paper/2506.14850/) (2025)
- [Bridging the Simulation-to-Experiment Gap with Generative Models using Adversarial Distribution Alignment](https://mlsanity.com/paper/2604.01169/) (2026)
- [From Evaluation to Design: Using Potential Energy Surface Smoothness Metrics to Guide Machine Learning Interatomic Potential Architectures](https://mlsanity.com/paper/2602.04861/) (2026)
- [Generative Quasi-Continuum Modeling of Confined Fluids at the Nanoscale](https://mlsanity.com/paper/2509.08223/) (2025)

## Abstract (arXiv)

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.