Stability-Aware Training of Neural Network Interatomic Potentials with Differentiable Boltzmann Estimators

Neural network interatomic potentials (NNIPs) are an attractive alternative to ab-initio methods for molecular dynamics (MD) simulations. However, they can produce unstable simulations which sample unphysical states, limiting their usefulness for modeling phenomena occurring over longer timescales. To address these challenges, we present Stability-Aware Boltzmann Estimator (StABlE) Training, a multimodal training procedure which combines conventional supervised training from quantum-mechanical energies and forces with reference system observables, to produce stable and accurate NNIPs. StABlE Training iteratively runs MD simulations to seek out unstable regions, and corrects the instabilities via supervision with a reference observable. The training procedure is enabled by the Boltzmann Estimator, which allows efficient computation of gradients required to train neural networks to system observables, and can detect both global and local instabilities. We demonstrate our methodology across organic molecules, tetrapeptides, and condensed phase systems, along with using three modern NNIP architectures. In all three cases, StABlE-trained models achieve significant improvements in simulation stability and recovery of structural and dynamic observables. In some cases, StABlE-trained models outperform conventional models trained on datasets 50 times larger. As a general framework applicable across NNIP architectures and systems, StABlE Training is a powerful tool for training stable and accurate NNIPs, particularly in the absence of large reference datasets. Molecular dynamics (MD) simulation is a staple method of computational science, enabling high-resolution spatiotemporal modeling of atomistic systems throughout biology, chemistry, and materials science [21]. Under the Born-Oppenheimer approximation, system evolution is governed by the underlying potential energy surface (PES), which is a function of the nuclear Cartesian coordinates [11]. While the atomic forces needed for MD simulation can be obtained on-the-fly via ab-initio quantum-mechanical (QM) calculations [12], the unfavorable scaling of this approach makes it prohibitively expensive for realistic system sizes and timescales [22]. There is a long history of using machine learning (ML) approaches in place of ab-initio methods to efficiently approximate the global PES [7, 6, 2, 55]. NNIPs, typically parameterized as graph neural networks [56, 33], are trained by matching energy and forces of a molecule or material from a reference dataset of QM calculations, such as Density Functional Theory (DFT) [31]. NNIPs trained on large ab-initio datasets are increasingly being used to model challenging and important chemical systems with favorable results [45, 37, 15, 57, 64, 43, 14, 3, 36, 60, 26, 19].