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

TL;DR

<3-5 sentence high-level summary>StABlE Training addresses the instability challenges of neural network-based machine learning force fields (MLFFs) in molecular dynamics by coupling traditional quantum-mechanical (QM) supervision with system observables via a Boltzmann-gradient mechanism. The method introduces the Boltzmann Estimator and its localized variant to enable end-to-end differentiable learning without backpropagating through long MD trajectories, and it alternates simulation and learning phases to identify and rectify unstable regions. Demonstrations on aspirin, Ac-Ala3-NHMe, and all-atom water across multiple architectures show substantial gains in stability, data efficiency, and fidelity of observables not explicitly trained, including improved RDFs and diffusivity. The approach is self-contained, broadly applicable, and promises practical impact by enabling stable, long-timescale MD with MLFFs even when large QM datasets are unavailable.

Abstract

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.

Figures (14)

• Figure 1: Machine learning force field (MLFF) failure modes.(a) Illustrative examples of true and learned potential energy surfaces (PES) and resulting dynamics for unstable MLFFs. MLFFs can be accurate in approximating much of the PES, but contain regions where energy and forces estimates deviate significantly from the true PES, leading to sampling of highly unphysical regimes. As a result, observables computed from MD simulation may be biased by the oversampling of unphysical states, or have high statistical error in the extreme case of unrecoverable simulation collapse. (b) Examples of stable versus unstable configurations sampled by MLFFs during molecular dynamics simulation of systems considered in this work. (c) Selected states from an unstable MD trajectory of aspirin.
• Figure 2: Schematic of Stability-Aware Boltzmann Estimator (StABlE) Training procedure. Our proposed StABlE Training procedure begins with conventional pre-training on a small reference dataset of QM calculations. This dataset remains fixed throughout the procedure, and is never expanded with new calculations. Upon convergence of pre-training, StABlE alternates between two main phases, simulation and learning. In the simulation phase, we perform many molecular dynamics simulations in parallel with the MLFF and find regions of instability. When a sufficient fraction of simulations become unstable, we enter the learning phase, where the MLFF is further trained to match known system observables from ab-initio simulation or experiment. Gradients are computed efficiently through the MD simulation via our Boltzmann Estimator. After a sufficient reduction in the portion of unstable trajectories, we re-enter the simulation phase, and repeat the training cycle until a predetermined computational budget is reached.
• Figure 3: Aspirin simulation with StABlE Training.(a) Alternation of simulation and learning phases during StABlE Training with 128 parallel replicas. The simulation phases correspond to regions where the fraction of unstable replicas increases, while the learning phases correspond to regions where learning occurs and the fraction of unstable replicas decreases. (b) Stable simulation time of 256 parallel aspirin trajectories from SchNet MLFFs. Applying StABlE Training yields significantly more stable simulations than models trained only on energies and forces, surpassing models trained on 50$\times$ more QM data. (c) Distribution of interatomic distances ($h(r)$) from MLFF simulations. A StABlE-trained SchNet model closely recovers the true distribution of interatomic distances, while the model trained only on QM reference data produces a noisier $h(r)$ because it cannot stably simulate the system for longer time periods. Inset shows difference between predicted and reference $h(r)$. (d) Distribution of velocity autocorrelation function (VACF) mean absolute error (MAE). StABlE Training yields a reduction in variance across replicas. (e) Aspirin structures sampled over epochs of a single learning phase of StABlE Training. There is a clear progression as unstable configurations become stable.
• Figure 4: Testing temperature generalization for aspirin.(a, b) Fraction of unstable replicas as a function of simulation time for SchNet MLFFs at 350K and 700K. Applying StABlE Training yields significantly more stable simulations than a baseline SchNet model. (c, d) Distribution of interatomic distances ($h(r)$) from SchNet MLFFs. The StABlE-trained SchNet model more accurately recovers the true $h(r)$ at both 350K and 700K (the inset shows difference between the predicted and reference $h(r)$). The AIMD reference $h(r)$ is computed by Boltzmann-reweighting of independent samples from the 500K training dataset.
• Figure 5: Ac-Ala3-NHMe tetrapeptide simulation with StABlE Training.(a) Fraction of unstable replicas as a function of simulation time for NequIP MLFFs. Applying StABlE Training yields a model which can simulate more MD replicas stably over time than a baseline trained only on energies and forces. (b) Distribution of interatomic distances ($h(r)$) from NequIP MLFF simulations. The StABlE-trained NequIP model much more closely recovers the true $h(r)$, while the $h(r)$ produced by the model trained only on QM reference data is noisy and inaccurate due to insufficient sampling time. (c) Ac-Ala3-NHMe structures sampled over epochs of a single learning phase of StABlE Training. There is a clear progression as very unstable configurations become stable.

Theorems & Definitions (3)

• Definition 1: $N$-sample Boltzmann estimator
• Definition 2: $N$-sample localized Boltzmann estimator
• proof