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

Sanjeev Raja  Department of Electrical Engineering and Computer Science,
University of California, Berkeley  Ishan Amin  Department of Electrical
Engineering and Computer Science, University of California, Berkeley
Department of Physics, University of California, Berkeley  Fabian Pedregosa
Google DeepMind  Aditi S. Krishnapriyan Corresponding author:
<EMAIL_ADDRESS>Department of Chemical and Biomolecular Engineering,
University of California, Berkeley  Department of Electrical Engineering and
Computer Science, University of California, Berkeley

###### Abstract

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 multi-modal 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].
Neural network interatomic potentials (NNIP) have recently emerged as a highly
successful class of ML methods, demonstrating the ability to capture complex
many-body interactions, and in some cases transfer flexibly across chemical
space [58, 29, 42, 24, 25, 5, 47, 4, 59]. 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].

MD simulations aim to accurately estimate system observables like the radial
distribution function, virial stress tensor, and diffusivity coefficient. This
often requires long simulation to fully explore the underlying PES and
minimize sampling error. Unfortunately, NNIPs are known to produce unstable
simulations, meaning that they can irreversibly enter unphysical regions of
phase space (e.g., a bond breaking event in a non-reactive system at low
temperature) [23, 63, 69, 8, 46, 72]. Sampling of such regions often quickly
leads to unrecoverable simulation collapse as the NNIP drifts far from the
distribution of its training data. As a result, computed observables may have
high statistical error due to insufficient sampling, or be altogether
inaccurate. This can limit the ability of NNIP-based MD simulations to
investigate long-timescale phenomena like ion diffusion and protein folding,
as well as rare events that may require extensive sampling to encounter.
Figure 1 illustrates typical NNIP instability behaviors and provides examples
of unstable structures sampled during MD simulation.

Insufficient coverage of phase space and inaccuracies within the reference
datasets utilized for training NNIPs have been suggested as factors
contributing to instability [23, 63, 69, 8, 46, 72]. Expanding the dataset
requires additional QM calculations, and improving its accuracy requires more
exact QM approaches [39]. Both options are very computationally expensive,
particularly for larger systems like macromolecules. Recent work has even
suggested that expanding the dataset may not always improve stability [23],
consistent with the notion that phase space is sampled with exponentially
vanishing probability based on the local potential energy. This suggests the
need for additional sources of information beyond energies and forces to train
stable NNIPs.

Figure 1: Neural network interatomic potential (NNIP) failure modes.  (a)
Illustrative examples of true and learned potential energy surfaces (PES) and
resulting dynamics for unstable NNIPs. NNIPs can be highly accurate in
approximating most of the PES, but contain localized ”holes” where energy and
forces estimates deviate significantly from the true PES, leading to sampling
of highly unphysical regimes. Once such a regime is reached, further
simulation typically does not re-enter physically plausible regions. As a
result, observables computed from MD simulation may be inaccurate or have high
statistical error. (b) Examples of stable versus unstable configurations
sampled by NNIPs during molecular dynamics simulation of systems considered in
this work. The unstable configurations contain unphysical bond lengths,
angles, and coordination structures which are not meaningful for observable
calculations. (c) Selected states from an unstable MD trajectory of aspirin. A
unidirectional progression is demonstrated, in which a stable, physical
configuration evolves into a highly unstable, unphysical configuration.

We posit that system observables derived from experiment or high-fidelity
simulation can serve this purpose. Parameterizing NNIPs to match observables
is an alternative approach to conventional energy-and-forces-training [65,
70]. This approach has been successful in learning potentials of condensed-
phase [71] and titanium [54] systems, enhanced sampling of rare events [61],
and protein folding simulation [30, 48, 34]. However, using observables to
improve the stability of NNIPs in downstream MD simulations remains
unexplored.

In this work, we introduce Stability-Aware Boltzmann Estimator (StABlE)
Training, a procedure designed to produce NNIPs that are both stable and
accurate. The core idea behind StABlE is to use MD simulation to rapidly
explore regions of molecular phase space where the NNIP becomes unstable,
followed by a targeted refinement of these regions using reference system
observables. The key to efficient and numerically stable training lies in the
Boltzmann Estimator, which avoids computing gradients directly through MD
simulations. StABlE Training is a self-contained, efficient process that
leverages learning signals from both reference observables and existing QM
data, with no reliance on performing additional QM reference calculations. To
our knowledge, StABlE Training is the first method that combines these data
modalities to improve the stability of neural network potentials in MD
simulations. Its flexibility allows for integration with any NNIP architecture
or molecular system, making it a valuable addition to existing NNIP workflows.

We demonstrate StABlE Training on three systems and NNIP architectures: 1)
simulation of aspirin with SchNet [58], 2) simulation of the Ac-Ala3-NHMe
tetrapeptide with NequIP [5], and 3) simulation of an all-atom water system
with GemNet-T [24]. Relative to NNIPs trained solely on reference energy and
forces data, our StABlE-trained models produce significantly more stable
downstream MD simulations and recover both structural and dynamical
observables more accurately. For an aspirin simulation with a SchNet model
trained on 1,000 reference structures, StABlE Training increases the median
stability from 35 to 140 picoseconds, exhibits better generalization to unseen
simulation temperatures, and outperforms models trained on 50 times larger,
labeled datasets. For an all-atom water simulation with a GemNet-T model
trained on 1,000 reference structures, StABlE Training improves the median
stability from 23 to 52 picoseconds and yields an 87% accuracy improvement in
estimation of the diffusivity coefficient. Our results suggest that utilizing
both quantum-mechanical and observable-based modalities is required to fully
exploit the available learning signal in reference datasets and train stable
and accurate NNIPs.

##  1 Preliminaries

#### Molecular Dynamics.

Molecular dynamics simulation is used to evolve the positions and momenta of
an atomistic system. Consider a system of $N$ atoms. The state of the system
at time $t$ is defined by $\Gamma(t)=\\{r(t),p(t)\\}$, where
$r(t),p(t)\in\mathbb{R}^{N\times 3}$ are the position and momenta of the
atoms. We assume the system has a scalar-valued Hamiltonian function
$\mathcal{H}:\mathbb{R}^{N\times 3}\times\mathbb{R}^{N\times
3}\rightarrow\mathbb{R}$ of the form,

$\mathcal{H}(\Gamma)=\sum_{i=1}^{N}\frac{p_{i}^{2}}{2m_{i}}+U(r)\,,$ (1)

where $U:\mathbb{R}^{N\times 3}\rightarrow\mathbb{R}$ is a potential energy
function and $m_{i}$ and $p_{i}$ are the mass and momentum of atom $i$. The
system obeys Hamilton’s laws of motion,

$\dot{r_{i}}=\frac{\partial H}{\partial
p_{i}}=\frac{p_{i}}{m_{i}}\,,\quad\dot{p_{i}}=-\frac{\partial H}{\partial
r_{i}}=-\frac{\partial U(r_{i})}{\partial r_{i}}\,.$ (2)

Evolving these equations in time for $T_{f}$ steps produces a sequence of
simulation states, $\bar{\Gamma}=\\{\Gamma(t)\\}_{t=0}^{T_{f}}$, which
conserve the Hamiltonian and constitute samples from the microcanonical
ensemble. It is also common to couple a thermostat to the simulation to
conserve temperature instead of the Hamiltonian, thereby sampling from the
canonical ensemble. In this case, the state distribution has a Boltzmann form:
$P(\Gamma)\propto e^{-\frac{H(\Gamma)}{k_{B}T}}$ where $T$ is the temperature
and $k_{B}$ is Boltzmann’s constant.

#### Training Neural Network Potentials with Reference Energies and Forces.

In the context of neural network interatomic potentials (NNIPs), the potential
energy is approximated by a function $U_{\theta}$ with learnable parameters
$\theta$ given by a neural network. When training a model to learn from
quantum mechanical reference data, $\theta$ is obtained by using stochastic
gradient descent (SGD) to minimize the following energy and forces loss
function:

$\displaystyle\mathcal{L}_{\text{QM}}(\theta)\stackrel{{\scriptstyle\text{def}}}{{=}}$
$\displaystyle\mathbb{E}_{\Gamma\sim
P_{\text{QM}}(\Gamma)}\Big{[}\lambda_{U}\|U_{\text{ref}}(\Gamma)-U_{\theta}(\Gamma)\|_{2}^{2}$
(3)
$\displaystyle\quad+\lambda_{F}||F_{\text{ref}}(\Gamma)+\nabla_{r}{U_{\theta}(\Gamma)}||_{2}^{2}\Big{]}\,,$

where $P_{\text{QM}}(\Gamma)$ is the equilibrium distribution induced by ab-
initio (or high-fidelity) simulation, $U_{\text{ref}}$ and $F_{\text{ref}}$
are the high-fidelity reference energy and force functions, and $\lambda_{U}$
and $\lambda_{F}$ are hyperparameters controlling the relative weights of the
energy and force loss terms. Note that while we have adopted the notation
$U(\Gamma)$ for convenience, the potential energy is a function only of the
position component, $r$, of the state $\Gamma$. In this work,
$P_{\text{QM}}(\Gamma)$ is approximated by a reference dataset
$\mathcal{D}_{\text{train}}$ of states sampled from a high-fidelity
simulation, such as ab-initio molecular dynamics. We note that in some cases,
$\mathcal{D}_{\text{train}}$ may not be fully representative of
$P_{\text{QM}}(\Gamma)$, due to biased sampling procedures favoring rare or
interesting events.

#### Estimation of System Observables from MD Simulation.

Observables, $g(\Gamma(t))$, characterize the state of a MD simulation at time
$t$, and relate to macroscopic properties or experimental measurements of the
system. Examples include the radial and angular distribution functions,
velocity autocorrelation function, and diffusivity coefficient. Observables
can be computed conveniently as an empirical average over states from a MD
simulation. This is justified by the ergodicity hypothesis, under which a
time-average over a sufficiently long simulation converges to a distributional
average over the Boltzmann distribution. For more detail on the observables
used in this work, please refer to Supplementary Section A.4.

##  2 Methods

We present StABlE Training, our proposed procedure to train stable and
accurate NNIPs by leveraging both reference quantum mechanical data and system
observables. Our approach is enabled by Boltzmann Estimators, which allows
gradient-based optimization of NNIPs based on system observables.

###  2.1 Boltzmann Estimator

Reference system observables can be estimated from high-fidelity MD
simulations or experimental measurements. Since observables are linked to the
NNIP parameters through a molecular simulation, training with this source of
information requires a reliable way to optimize through MD trajectories.

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 (in the loss,
$\hat{U}$ refers to the NNIP-predicted energy). This dataset remains fixed
throughout the procedure, and is never expanded with new calculations. Upon
convergence of pre-training, StABlE then alternates between two main phases,
simulation and learning. In the simulation phase, we perform many molecular
dynamics simulations in parallel with the NNIP and find regions of
instability. When a sufficient fraction of simulations become unstable, we
enter the learning phase, where the NNIP is further trained to match known
system observables from ab-initio simulation or experiment (in the loss,
$\hat{g}$ refers to the NNIP-predicted observable). 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. Sample stable and unstable molecules are included next to
the corresponding graphs.

Formally, we define $g(\Gamma)$ to be a vector-valued observable of a state
$\Gamma$, and $g_{\text{ref}}$ to be the reference value of the observable. To
train our NNIP with information about the system observable, we seek to
minimize the following loss function:

$\mathcal{L}_{\text{obs}}(\theta)\stackrel{{\scriptstyle\text{def}}}{{=}}\|\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]-g_{\text{ref}}\|_{2}^{2},$ (4)

where $P_{\theta}(\Gamma)$ is the equilibrium distribution induced by the NNIP
$U_{\theta}$. In this work, we consider systems with a fixed volume,
temperature, and number of particles, corresponding to the canonical ensemble.
This means that microstates obey a Boltzmann distribution,
$P_{\theta}(\Gamma)={\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma)\right)}/{C(\theta)}$,
where the Hamiltonian
$\mathcal{H}_{\theta}(\Gamma)=\sum_{i=1}^{N}\frac{p_{i}^{2}}{2m_{i}}+U_{\theta}(\Gamma)$,
$T$ is the sampling temperature, $k_{B}$ is Boltzmann’s constant, and
$C(\theta)=\int\exp(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma^{\prime}))\mathop{}\\!\mathrm{d}\Gamma^{\prime}$
is the normalizing partition function. MD simulation is required at each
training iteration to sample from this distribution and estimate the loss.

We optimize the parameters of the NNIP, $\theta$, using stochastic gradient
descent. This requires computing gradients of Equation 4 with respect to
$\theta$. The required gradient, $\nabla_{\theta}\mathcal{L}_{\text{obs}}$,
can be decomposed via the chain rule as follows:

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{obs}}^{\top}$
$\displaystyle=\frac{\partial\mathcal{L}_{\text{obs}}}{\partial{\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)[g(\Gamma)]}}}\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$ (5)
$\displaystyle=2(\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]-g_{\text{ref}})^{\top}\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}\,.$

The non-trivial quantity to compute is the Jacobian,
${\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}/{\partial\theta}$. One way to estimate this
quantity is by using a chain rule expansion that corresponds to each step of
the unrolled MD simulation [70]. However, this approach is memory-intensive
and can lead to numerical instability for long trajectories [71, 61]. We can
instead leverage the known Boltzmann form of the state distribution to
construct an unbiased estimator of the Jacobian.

######  Definition 1 ($N$-sample Boltzmann estimator).

Given $N$ independent samples $\Gamma^{(1)},\ldots,\Gamma^{(N)}$ from a
Boltzmann distribution $P_{\theta}(\Gamma)$, we define the _$N$ -sample
Boltzmann estimator_ $\mathcal{E}(\Gamma^{(1)},\ldots,\Gamma^{(N)})$ as,

$\displaystyle\mathcal{E}(\Gamma^{(1)},\ldots,\Gamma^{(N)})=$ (6)
$\displaystyle\frac{N}{k_{B}T(N-1)}\left[\hat{\mathbb{E}}\left[g(\Gamma)\right]\hat{\mathbb{E}}\big{[}{\nabla_{\theta}U_{\theta}(\Gamma)}\right]^{\top}\\!-\hat{\mathbb{E}}\left[{g(\Gamma)\cdot\nabla_{\theta}U_{\theta}(\Gamma)}^{\top}\right]\big{]}\,,$

where $\hat{\mathbb{E}}[f(\Gamma)]=\frac{1}{N}\sum_{i=1}^{N}f(\Gamma^{(i)})$
denotes an empirical mean over the samples.

This estimator provides an unbiased estimate of the Jacobian
$\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$. Our proof is provided in
Supplementary Section A.1. This result can also be derived using thermodynamic
perturbation theory [74, 65].

#### Localized Boltzmann Estimator for Spatial Specificity.

In some scenarios, unphysical configurations can occur within localized
regions of the simulation domain, such as collisions between two molecules in
a large condensed-phase system. Due to spatial averaging, global observables
$g(\Gamma)$ may be insensitive to these localized events, limiting the ability
to identify unphysical states. To address this, we propose the $N$-sample
Localized Boltzmann Estimator. Here, the global energy $U_{\theta}(\Gamma)$
and observable $g(\Gamma)$ are replaced with local versions
$U_{\theta}(\gamma)$ and $g(\gamma)$, where $\gamma$ denotes a local
neighborhood of $n<N$ atoms. NNIPs compute global energies by summing
individual atomic energies, and we can leverage this to compute
$U_{\theta}(\gamma)$ by summing only over atoms in the local neighborhood. See
Supplementary Section A.2 for a formal definition of the localized estimator.
A concrete example of localized instability, along with the use of our
localized Boltzmann estimator to correct it, will be presented in the context
of an all-atom water system in Results, Section 3.4.

#### Key Advantages.

The Boltzmann estimators defined earlier have several advantages over explicit
differentiation through the MD simulation to train the NNIP. Due to the use of
independent samples, the estimator avoids numerical instability associated
with differentiating through long, continuous trajectories. Further, the most
computationally expensive component of the estimator, the gradient of the
potential energy $\nabla_{\theta}U_{\theta}(\Gamma)$, is independent of the
observable $g(\Gamma)$. This means it can be reused, allowing efficient
training to match multiple observables simultaneously. Finally, our localized
formulation of the estimator enables finer control over isolated
instabilities, providing a more informative gradient with which to optimize
the NNIP.

###  2.2 Stability-Aware Boltzmann Estimator (StABlE) Training

Stability-Aware Boltzmann Estimator (StABlE) Training begins with conventional
supervised pre-training of a NNIP on a dataset of structures from a high-
fidelity, constant temperature MD simulation, paired with reference energy and
forces values. The method then proceeds by alternating between two major
phases, simulation and learning, as illustrated in Figure 2.

#### Simulation Phase.

The simulation phase aims to explore the molecular phase space and pinpoint
regions where the NNIP becomes unstable. We sample $R$ equilibrium states from
the training dataset as initial conditions for separate MD trajectories
(replicas). Using the pre-trained NNIP $U_{\theta}$, we run MD simulations on
these replicas in parallel for $t$ timesteps. We apply a predetermined
stability criterion (see Supplementary Section A.5 for details) to each
replica, freezing those marked unstable at their current states. Simulations
continue for the remaining replicas, with stability checks every $t$
timesteps. Once a specified fraction of replicas become unstable, we rewind
all unstable replicas by $t$ timesteps. This ensures that further simulation
for $t$ timesteps will trigger instability, which is then corrected in the
next phase.

#### Learning Phase.

The learning phase aims to refine the NNIP to correct the previously
encountered instabilities. Starting from the near-unstable configurations
obtained in the simulation phase, we perform MD simulation for $t$ timesteps,
sampling every $S^{th}$ state to obtain uncorrelated samples. Using the
sampled states $\Gamma^{(1)},\ldots,\Gamma^{(N_{d})}$, where
$N_{d}=\frac{t}{S}$, we compute the observable loss function (Equation 4). To
update the NNIP parameters $\theta$, we compute an unbiased estimator of the
loss gradient via the Boltzmann Estimator (Section 2.1) and use it to perform
a single step of gradient descent. We then reset all replicas to their
original near-unstable states, simulate with the updated NNIP, recompute the
loss and gradient estimator with the newly sampled states, and again update
the NNIP weights. This process is repeated until the fraction of unstable
replicas drops below a predetermined threshold. When this occurs, the learning
phase has concluded, and we begin a new simulation phase starting from the
endpoints of the last learning phase. We continue alternating between
simulation and learning phases until a predetermined computational budget is
reached, at which point StABlE Training has concluded. See Supplementary
Section A.3 for a formal algorithmic description of the StABlE Training
procedure.

#### Regularizing StABlE Training with Energy and Forces Reference Data.

In practice, the mapping between a sparse set of system observables and a
potential energy function is non-unique [50]. Consequently, learning with
observables alone is underconstrained. To combat this, we regularize the
observable loss function (Equation 4) with the energy and forces loss function
(Equation 3). The final StABlE loss function thus becomes,

$\mathcal{L}_{\text{StABlE}}(\theta)\stackrel{{\scriptstyle\text{def}}}{{=}}\mathcal{L}_{\text{obs}}+\lambda\mathcal{L}_{\text{QM}},$
(7)

where $\mathcal{L}_{\text{obs}}$ and $\mathcal{L}_{\text{QM}}$ were defined in
Equations 4 and 3 respectively, and the hyperparameter $\lambda$ controls the
strength of the regularization. Crucially, $\mathcal{L}_{\text{QM}}$ is still
only computed over the static training dataset $\mathcal{D}_{\text{train}}$,
and not over new structures explored during MD simulation. Therefore, the
regularization introduces no additional quantum-mechanical supervision.

##  3 Results

Figure 3: Aspirin simulation with StABlE Training. Results of StABlE Training
a SchNet [58] NNIP on the 21-atom aspirin molecule from the MD17 [18] dataset
at 500 Kelvin. (a) Alternation of simulation and learning phases during StABlE
Training with 128 parallel replicas. The simulation phases correspond to
regions where no learning occurs and 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 a SchNet NNIP trained on 1,000
(blue), 10,000 (magenta), 50,000 (grey) quantum mechanical (QM) energy and
forces reference data points. We apply StABlE Training on the SchNet NNIP
trained with 1,000 QM data points. This results in significantly more stable
simulations (orange), surpassing models trained on 50$\times$ more QM data.
(c) Distribution of interatomic distances ($h(r)$) from an ab-initio molecular
dynamics (AIMD) simulation (black), a SchNet model trained on 1,000 QM data
points (blue), and a SchNet model trained on 1,000 QM data points, followed by
our StABlE Training procedure (orange). 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) relative to an ab-
initio simulation from a SchNet model trained on 1,000 QM data points (blue)
and a SchNet model trained on 1,000 QM datapoints, followed by our StABlE
Training procedure (orange). The VACF is an observable computed from
simulation dynamics, which the model is not trained on. While the mean
accuracy of the StABlE-trained SchNet model is similar to the QM-only
baseline, there is 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.

We present the results of our StABlE training procedure on three different
systems and NNIP architectures: 1) an aspirin molecule with a SchNet [58]
architecture (Sections 3.1 \- 3.2), 2) an Ac-Ala3-NHMe tetrapeptide with a
NequIP [5] architecture (Section 3.3), and 3) an all-atom water system with a
GemNet-T [24] architecture (Section 3.4). For systems 1) and 2), we use the
Boltzmann estimator presented in Equation LABEL:eq:estimator. For system 3),
we employ our localized Boltzmann estimator (Equation 8) to correct isolated
instabilities. We focus our evaluation of NNIPs on their ability to produce
stable MD simulations which accurately recover key system observables. For all
systems, stability thresholds are chosen to be very large relative to the
scale of typical fluctuations in reference simulations at the chosen
temperature. Simulations which cross the threshold thus cannot reasonably re-
enter the manifold of stable configurations (see Supplementary Section A.5 for
more details). We compute simulation observables by averaging over sampled
states from the trajectory corresponding to the median stability improvement
(over all replicas) between the baseline and StABlE-trained NNIP. More
detailed information about NNIP architectures, training, simulation, and
evaluation can be found in Supplementary Sections A.6, A.7, and A.8.

###  3.1 Aspirin Molecule

Aspirin (chemical formula $C_{9}H_{8}O_{4}$) is the largest molecule from MD17
[18], a widely used benchmark dataset for atomistic simulations which contains
energy and forces calculations computed at the PBE+vdW-TS [66] level of
theory. Consisting of 21 atoms, aspirin has been shown to be the most
challenging molecule in MD17 for state-of-the-art NNIPs to simulate stably
[23].

#### Problem setup.

We pre-train a SchNet [58] model on the energy and forces matching objective
(Equation 3) using a subset of 1,000 aspirin structures from the reference
dataset. After convergence, we begin StABlE Training with the global Boltzmann
estimator, using 128 parallel replicas at a temperature of 500K. By simulating
hundreds of replicas in parallel, we expose the NNIP to a comprehensive range
of failure modes much more quickly than possible with a single replica.
Following [23], we use a maximum bond length deviation criterion to detect
instability in the simulations (see Supplementary Section A.5 for details). We
use the distribution of interatomic distances, $h(r)$, as our training
observable, with the $h(r)$ computed over structures in the training dataset
serving as the reference (see Supplementary Section A.4 for details). This
choice is motivated by the observation that unphysical bond stretches
constitute the majority of failure cases in aspirin simulations. The $h(r)$ is
sensitive to these abnormalities, and is thus an informative optimization
criterion. We perform four complete StABlE cycles of simulation and learning
with a SGD optimizer.

#### Results.

We evaluate our final models by performing constant-temperature MD
simulations, starting from 256 initial structures not seen during training.
Figure 3a shows the progression of StABlE Training. During the simulation
phase, the learning rate is zero and the fraction of unstable replicas
steadily increases. When the upper instability threshold is reached, learning
commences. The fraction of unstable replicas steadily decreases, confirming
that optimization of the interatomic distance distribution produces the
desired stability improvement. A decaying learning rate schedule is employed
to avoid overfitting to the current batch of replicas.

A SchNet model trained via our stability-aware approach is significantly more
stable during MD simulation than models trained only on the conventional QM
energy and forces objective function (Figure 3b). As a result of our training
procedure, the mean stable simulation time increases from 35 to 140
picoseconds. Even when StABlE training is employed after training on only on
1,000 reference structures, it is more stable than models trained
conventionally on 10,000 and 50,000 reference structures, highlighting the
effectiveness of our method in improving stability without reliance on any
additional QM reference data. We also observe that simply reducing the
simulation timestep does not reliably lead to an improvement in the stability
of simulations produced by conventionally trained models (see Supplementary
Section A.11 for details).

Simulations from our model closely recover the true distribution of
interatomic distances, while the distributions produced by models solely
trained on QM data are noisier, due to the limited stable simulation time
(Figure 3c). We also test the ability of our StABlE-trained model to recover
the velocity autocorrelation function (VACF), a fundamental dynamical
observable not seen during training (Figure 33). We see that the quality of
the recovered VACF remains similar after StABlE Training, with lower variance
over the simulated replicas (see Supplementary Section A.12 for a sample VACF
computed from the simulations). The preservation of a held-out dynamical
observable suggests that StABlE Training does not achieve stability
improvements by simply restricting the model around a narrow regime
corresponding to the reference $h(r)$. If this were the case, characteristic
dynamic fluctuations around the reference $h(r)$ would be suppressed, leading
to an inaccurate VACF. Figure 3e depicts aspirin structures sampled from
simulations during the course of a single learning phase of StABlE Training,
demonstrating that unphysical bond stretching in these structures is resolved
by the training procedure.

###  3.2 Temperature generalization

Figure 4: Testing temperature generalization for aspirin. We test the
generalization of StABlE Training by performing 256 parallel MD simulations at
350 Kelvin and 700 Kelvin with a SchNet model which was StABlE-trained at 500
Kelvin on aspirin from MD17. (a, b) Fraction of unstable replicas as a
function of simulation time for a SchNet NNIP trained on 1,000 mechanical (QM)
energy and forces reference data points (blue). Applying StABlE Training to
this model yields significantly more stable simulations (orange). At 350K,
before StABlE Training, $70\%$ of replicas are unstable after 300 picoseconds
of simulation. After StABlE Training, $<$5% of replicas are unstable. (c, d)
Distribution of interatomic distances ($h(r)$) from an ab-initio molecular
dynamics (AIMD) simulation (black), a SchNet model trained on 1,000 QM data
points (blue), and a SchNet model trained on 1,000 QM datapoints, followed by
our StABlE Training procedure (orange). The StABlE-trained SchNet model more
accurately recovers the true $h(r)$ at both 350K and 700K. At 350K, the
baseline model trained only on QM data incorrectly estimates several peaks in
the $h(r)$, while at 700K the baseline model produces a noisy $h(r)$ due to
its limited stable simulation time (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.

#### Problem setup.

We explore our method’s generalization to different thermodynamic conditions.
To do so, we perform 256 parallel MD simulations at 350K and 700K using the
SchNet model which was trained with our StABlE Training procedure at 500K. We
use the same criterion based on bond length deviation to detect instability as
in Section 3.1. We estimate the reference $h(r)$ at 350K and 700K by
Boltzmann-reweighting samples from the original 500K training dataset (see
Supplementary Section A.9 for details).

#### Results.

Aspirin simulations produced by a StABlE-Trained SchNet model are
significantly more stable at both temperatures than baseline models trained
only on QM reference data (Figure 4a, b). At 350K, the difference is
particularly large. After 300 picoseconds of simulation, less than 5% of
replicas are unstable in the StABlE-trained NNIP simulation, while over 75%
replicas are unstable in the baseline NNIP simulation. The StABlE-trained
model also recovers the true distribution of interatomic distances more
closely than the baseline NNIP at both temperatures (Figure 4c, d). At 350K,
the baseline NNIP is stable enough produce a smooth $h(r)$, but the structure
is inaccurate, indicating incorrect sampling of the phase space. At 700K, the
baseline NNIP produces a $h(r)$ which is close to the AIMD-produced
distribution, but is noisy due to insufficient stable simulation time. This
underscores the necessity for a NNIP to be both stable and accurate in order
to be practically useful in MD simulation. StABlE Training improves both of
these metrics for this system. As generalization to new thermodynamic
conditions is of crucial importance for NNIPs to be useful in practical
applications [35], we see the temperature-transferability of StABlE Training
as an important strength of our method.

###  3.3 Ac-Ala3-NHMe Tetrapeptide

Ac-Ala3-NHMe (chemical formula $C_{12}H_{22}N_{4}O_{4}$) is a tetrapeptide
from the MD22 dataset [17], a challenging benchmark consisting of reference
energy and forces computed at the PBE+MBD [51, 67] level of theory for
considerably larger molecules than those in the MD17 dataset. Relative to MD17
molecules, Ac-Ala3-NHMe poses unique challenges for atomistic simulations due
to its larger size and flexibility [32]. As a result, more expressive
NNIPs—often incorporating E(3) equivariance [27]—are required to accurately
model the underlying potential energy surface.

#### Problem setup.

In our experiments, we use a NequIP [5] model due to its promising accuracy
and data efficiency on challenging atomistic simulation tasks [37, 45]. We
pre-train a NequIP model on the QM energy and forces matching objective using
14,890 (or 25%) of available structures from the training dataset. After
convergence, we begin StABlE Training with the global Boltzmann estimator,
using 48 parallel replicas at a temperature of 500K. As with aspirin (Sections
3.1 \- 3.2), we detect instability using the maximum bond length deviation
criterion, and choose the distribution of interatomic distances, $h(r)$, as
our training observable.

#### Results.

The StABlE-trained NequIP model gives an improvement in stability relative to
a baseline model trained only on QM reference data (Figure 5a). Additionally,
the StABlE-trained model closely recovers the true $h(r)$, while the baseline
NNIP produces a highly noisy $h(r)$. This is because of insufficient sampling,
due to a lower stable simulation time (Figure 5b). The baseline NNIP also
underestimates the latter portion of the $h(r)$ distribution, indicating poor
modelling of long-range interactions. Figure 5c shows selected structures
throughout the course of a single learning phase of StABlE Training,
demonstrating a progression from very unstable to stable structures.

###  3.4 Water

Figure 5: Ac-Ala3-NHMe tetrapeptide simulation with StABlE Training. Results
of StABlE Training a NequIP [5] NNIP on the 42-atom Ac-Ala3-NHMe tetrapeptide
from the MD22 [17] dataset at 500 Kelvin. (a) Fraction of unstable replicas as
a function of simulation time for a NequIP NNIP trained on 1,000 energy and
forces reference data points (blue). Applying StABlE Training yields a model
which can simulate more MD replicas stably over time (orange). (b)
Distribution of interatomic distances ($h(r)$) from an ab-initio molecular
dynamics (AIMD) simulation (black), a NequIP model trained on 14,890 QM data
points (blue), and a NequIP model trained on 14,890 QM datapoints, followed by
our StABlE Training procedure (orange). 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.

As a final evaluation of StABlE Training, we consider liquid-phase water.
Water has historically posed unique challenges for atomistic simulations due
to the presence of long-range interactions, proton disorder, and nuclear
quantum effects [16, 52, 44]. Additionally, estimating transport properties,
such as the diffusivity coefficient, requires long, continuous trajectories to
minimize statistical error. We use the dataset produced in [23], which
contains reference all-atom simulations of 64 water molecules using the
flexible Extended Simple Point Charge model [73] at 300K and 1 atm. For this
system, we consider a GemNet-T model [24]. GemNet-T displays a failure mode in
which unphysical configurations (e.g., bond stretching) first arise in highly
localized ”pockets” of 1-2 molecules (Figure 6a), which then gradually cascade
to the rest of the simulation domain [23]. Such configurations do not
noticeably affect global system observables such as the radial distribution
function. Therefore, in this setting it is more appropriate to use our
localized version of the Boltzmann gradient estimator (Equation 8) in order to
guide the optimization process towards specific spatial domains of
instability.

#### Problem setup.

We pre-train a GemNet-T model to the energy and forces matching objective
using a subset of 1,000 reference structures from the training dataset. After
convergence, we begin StABlE Training with the localized Boltzmann estimator,
simulating 10 parallel replicas at a temperature of 300K. We use a minimum
intermolecular distance criterion (see Supplementary Section A.5) to detect
instabilities. To estimate gradients for optimization, from each global state
we extract all possible local neighborhoods containing a single water molecule
(see Supplementary Section A.6 for more details on the local neighborhood
sampling scheme). The most prevalent failure mode in this system is unphysical
bond stretching, which eventually leads to collisions with nearby water
molecules. Motivated by this observation, we use the mean hydrogen-oxygen bond
length as our training observable. We perform a single cycle of simulation and
learning.

#### Results.

We evaluate the performance of the StABlE-Trained model by performing 10
parallel MD simulations at 300K, starting from held-out initial conditions. We
achieve significant improvements in stability and accuracy using our training
approach. We nearly double the stable simulation time relative to a baseline
NNIP model trained only on QM reference data (Figure 6b). Compared to a
reference value of $2.3\times 10^{-9}\frac{m^{2}}{s}$, the mean diffusivity
produced by StABlE-trained model simulations is $1.6\times
10^{-9}\frac{m^{2}}{s}$, while the mean diffusivity produced by baseline model
simulations is $5.0\times 10^{-9}\frac{m^{2}}{s}$ (Figure 6c). This
corresponds to an 87% improvement in the accuracy of diffusivity coefficient
estimation. We note that the StABlE-trained model systemically underestimates
the reference diffusivity, while the baseline model overestimates the value.
We hypothesize that in this particular case, StABlE Training has a retarding
effect on system dynamics. As the StABlE optimization process assigns high
probability to microstates closely matching the equilibrium bond length, small
fluctuations about this value may be excessively penalized, inhibiting the
motion of molecules in the simulation. Finally, the spikes in the short-range
region of the element-conditioned RDFs (Figures 6d-f), which indicate
unphysical bond stretching and collision events, are largely removed in
StABlE-trained model simulations, while long-range correlations are also
better captured. We reiterate that in this experiment, neither the element-
conditioned RDFs nor the diffusivity coefficient were explicitly seen during
StABlE Training.

Figure 6: All-Atom Water Simulation with StABlE Training. Results of StABlE
Training a GemNet-T [24] NNIP on an all-atom water system with 64 molecules at
300 Kelvin. (a) Example unstable configuration arising in a GemNet-T
simulation of water. Instabilities are highly localized and manifest as
unphysical bond stretches or intermolecular collisions, motivating the use of
our localized Boltzmann estimator (Section 22.1) during training (b) Fraction
of unstable replicas as a function of simulation time for a GemNet-T NNIP
trained on energy and forces reference data points (blue), followed by StABlE
Training (orange). StABlE Training yields a clear stability improvement
relative to training on only energy and forces data, more than doubling the
median stable simulation time from 23 to 52 picoseconds. (c) Left: Convergence
of the diffusivity coefficient over the simulations produced by the reference
Extended Simple Point Charge model [73] (black), a GemNet-T model trained on
1,000 energy and forces datapoints (blue), and a GemNet-T model trained on
1,000 energy and forces datapoints followed by StABlE Training (orange).
Right: Relative to the reference value of $2.3\times
10^{-9}\frac{m^{2}}{s}$The StABlE-trained model gives a more accurate final
estimate of the diffusivity coefficient, with a mean predicted value of
$1.6\times 10^{-9}\frac{m^{2}}{s}$ as compared with a value of $5.0\times
10^{-9}\frac{m^{2}}{s}$ produced by the baseline model. (d - f) Element-
conditioned (OH, OO, HH) radial distribution functions (RDF) produced by
simulations from the reference model (black), a GemNet-T model trained on
1,000 energy and forces data points (blue), and a GemNet-T model trained on
1,000 energy and forces data points followed by StABlE Training (orange). The
StABlE-trained model produces fewer unphysical collisions and more accurately
captures long-range interactions.

##  4 Discussion

We have introduced StABlE Training, a strategy for training stable and
accurate neural network interatomic potentials. We have demonstrated StABlE
Training on a flexible organic molecule, a tetrapeptide, and a liquid-phase
system, as well as with three modern NNIP architectures. Our training
procedure results in NNIPs which remain stable for considerably longer periods
in MD simulation compared to models trained only on QM data, and as a result,
more accurately reproduce key simulation observables. We also note that we
have improved the recovery of observables that the NNIPs were not explicitly
trained on. These include dynamical observables such as the diffusivity
coefficient of water and velocity autocorrelation function of aspirin. We
showed that StABlE-trained models exhibit better generalization to unseen
temperatures for the aspirin molecule. Finally, we demonstrated that StABlE-
trained models can outperform NNIP models trained on up to 50 times more QM
data. We thus expect StABlE Training to be a particularly compelling
alternative to traditional QM training in cases where producing comprehensive
ab-initio datasets is prohibitively expensive, such as for large system sizes.

#### Advantages of StABlE Training.

One advantage of StABlE Training is that it can utilize both quantum-
mechanical energies and forces and system observables to supervise NNIP
training. As reference observables can be easily computed by averaging over
existing datasets, the stability, accuracy, and data efficiency gains brought
by StABlE Training require no additional reference calculations or data.
StABlE Training is thus self-contained and efficient, requiring minimal
additional computational expense beyond conventional NNIP training (see
Supplementary Section A.6 for details on training times).

StABlE Training can also be flexibly applied across atomistic systems due to
its ability to handle a diverse set of failure modes arising in MD simulation,
including both global and local instabilities. Specifically, due to our
localized Boltzmann estimator, we can target and resolve isolated
instabilities found in simulations of large condensed-phase systems. The
adaptability of our method is further exemplified by the lack of reliance on a
manually specified prior potential, a common requirement of observable-based
learning schemes. The pre-training step which precedes StABlE Training
effectively learns a QM-informed prior, which is then further refined in the
alternating simulation and learning stages.

The flexibility of StABlE Training extends to the neural network
parameterization of the potential energy function. The three NNIP
architectures we considered in this work vary widely in their expressivity, as
measured by the number of total learnable parameters, as well as in the
symmetry principles governing their internal feature transformations (see
Supplementary Section A.6 for details). The effectiveness of StABlE Training
across these models suggests that our approach will remain applicable as NNIP
architectures continue to evolve.

Finally, we note that the underlying Boltzmann gradient estimator used to
train the NNIPs differentiates through the potential energy of uncorrelated
states. It is thus agnostic to the method of sampling from the Boltzmann
distribution $P_{\theta}(\Gamma)$, and is readily applicable to sampling
schemes other than MD simulation (e.g., Monte Carlo methods). This is related
to recent work in machine learning on implicit layers, which enables
differentiation of optimization problem solutions without knowledge of the
numerical solver used [1, 28, 53, 9, 49].

#### Limitations.

A limitation of StABlE Training is that training NNIPs to match reference
observables lacks the convergence guarantees provided by traditional QM
learning in the large data limit. Specifically, the observable-matching
objective is under-constrained, as the mapping between a potential energy
function and a sparse set of simulation observables is in general non-unique
[50]. Although we mitigate this problem by pre-training and regularizing
StABlE Training with the conventional energy and forces loss function, we
observe that stability improvements resulting from StABlE Training are
accompanied by a minor increase in the energy and forces error on a held-out
test set, indicating a misalignment between the observable-matching and QM
objectives. Adjusting the strength of the QM regularization and the learning
rate used for StABlE Training are two primary ways in which to navigate this
tradeoff. We note, however, that our observed error increases are typically
within the range of DFT error (see Supplementary Section A.10 for details). We
also note that StABlE Training is currently incompatible with dynamical
observables, due to the use of uncorrelated states to compute the Boltzmann
Estimator. Overcoming this limitation is non-trivial, as it requires
optimization over long paths, which is a significantly higher-dimensional
problem with large-variance gradients [10].

#### Future Work.

In this work, we took observables calculated from high-fidelity simulations as
reference values. Many of the observables we considered, such as the radial
distribution function, diffusivity coefficient, and equilibrium bond length,
are also experimentally measurable. Future work could explore training with
experimental observables from multiple thermodynamic conditions (e.g.,
temperatures and pressures) simultaneously to yield more generalizable and
robust potentials. As an independent source of information, experimental
observables could also be used to refine NNIPs trained on data obtained from
lower levels of theory. Future work could also explore incorporating
additional observables into StABlE Training, particularly those which are
dynamical. This could restrict the set of learnable functions and address the
under-constrained nature of learning with observables. Finally, we note that
established techniques, such as active learning [62, 68, 41], are not mutually
exclusive with StABlE Training, and could be used in tandem to further improve
NNIP stability and robustness.

##  5 Acknowledgements

This work was initially supported by Laboratory Directed Research and
Development (LDRD) funding under Contract Number DE-AC02-05CH11231. It was
then supported by the U.S. Department of Energy, Office of Science, Office of
Advanced Scientific Computing Research, Scientific Discovery through Advanced
Computing (SciDAC) program under contract No. DE-AC02-05CH11231. This research
used resources of the National Energy Research Scientific Computing Center
(NERSC), a U.S. Department of Energy Office of Science User Facility located
at Lawrence Berkeley National Laboratory, operated under Contract No. DE-
AC02-05CH11231. We thank Rasmus Lindrup, Toby Kreiman, Geoffrey Negiar,
Muhammad Hasyim, Ritwik Gupta, Aayush Singh, David Limmer, Kranthi Mandadapu,
David Prendergast, Bryan McCloskey, and Muratahan Aykol for fruitful
discussions and comments on the manuscript.

##  6 Author Contributions

S.R and A.S.K. jointly conceived and initiated the project. S.R.
conceptualized and implemented the Localized Boltzmann Estimator and StABlE
Training method, and performed experiments. I.A. trained the baseline models
and contributed to the Boltzmann Estimator theory. F.P contributed to a unique
derivation and analysis of the Boltzmann Estimator. A.S.K. supervised and
guided the project from conception to design of experiments, implementation,
theory, and analysis. All authors analyzed results and wrote the paper.

##  7 Code Availability

Our code will be made available at https://github.com/ASK-Berkeley/StABlE-
Training.

## References

  * [1] Brandon Amos and J. Kolter  “OptNet: Differentiable optimization as a layer in neural networks”  In _Proceedings of the International Conference on Machine Learning (ICML)_ , 2017, pp. 136–145
  * [2] A.. Bartók, M.. Payne, R. Kondor and G. Csányi  “Gaussian approximation potentials: The accuracy of quantum mechanics, without the electrons”  In _Phys. Rev. Lett._ 104, 2010, pp. 136403
  * [3] Ilyes Batatia et al.  “A foundation model for atomistic materials chemistry”, 2023  arXiv:2401.00096 [physics.chem-ph]
  * [4] Ilyes Batatia et al.  “MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields”  In _Neural Information Processing Systems (NeurIPS)_ , 2022
  * [5] Simon Batzner et al.  “E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials”  In _Nature Communications_ 13.1 Springer ScienceBusiness Media LLC, 2022  DOI: 10.1038/s41467-022-29939-5
  * [6] J. Behler  “Neural network potential-energy surfaces in chemistry: A tool for large-scale simulations”  In _Phys. Chem. Chem. Phys._ 13, 2011, pp. 17930–17955
  * [7] Jörg Behler and Michele Parrinello  “Generalized Neural-Network Representation of High-Dimensional Potential-Energy Surfaces”  In _Phys. Rev. Lett._ 98 American Physical Society, 2007, pp. 146401  DOI: 10.1103/PhysRevLett.98.146401
  * [8] Vaibhav Bihani et al.  “EGraFFBench: Evaluation of Equivariant Graph Neural Network Force Fields for Atomistic Simulations”, 2023  arXiv:2310.02428 [cs.LG]
  * [9] Mathieu Blondel et al.  “Efficient and Modular Implicit Differentiation”  In _Neural Information Processing Systems (NeurIPS)_ , 2022  Google Research
  * [10] Peter G. Bolhuis, Z. Brotzakis and Bettina G. Keller  “Optimizing molecular potential models by imposing kinetic constraints with path reweighting”  In _J. Chem. Phys._ 159, 2023, pp. 074102  DOI: 10.1063/5.0151166
  * [11] Max Born and J. Oppenheimer  “On the Quantum Theory of Molecules” In German  In _Annalen der Physik_ 389.20, 1927, pp. 457–484  DOI: 10.1002/andp.19273892002
  * [12] R. Car and M. Parrinello  “Unified Approach for Molecular Dynamics and Density-Functional Theory” Published 25 November 1985  In _Phys. Rev. Lett._ 55 American Physical Society, 1985, pp. 2471  DOI: 10.1103/PhysRevLett.55.2471
  * [13] Scott P. Carmichael and M. Shell  “A New Multiscale Algorithm and Its Application to Coarse-Grained Peptide Models for Self-Assembly”  In _J. Phys. Chem. B_ 116.29 American Chemical Society, 2012, pp. 8383–8393  DOI: 10.1021/jp2114994
  * [14] Nicholas E. Charron et al.  “Navigating protein landscapes with a machine-learned transferable coarse-grained model”, 2023  arXiv:2310.18278 [q-bio.BM]
  * [15] C. Chen and S.P. Ong  “A universal graph deep learning interatomic potential for the periodic table”  In _Nat Comput Sci_ 2, 2022, pp. 718–728  DOI: 10.1038/s43588-022-00349-3
  * [16] Bingqing Cheng, Edgar A. Engel, Jörg Behler and Michele Ceriotti  “Ab initio thermodynamics of liquid and solid water” Edited by Pablo G. Debenedetti, Princeton University, Princeton, NJ, approved December 3, 2018 (received for review September 4, 2018)  In _Proc. Natl. Acad. Sci. U.S.A._ 116.4, 2019, pp. 1110–1115  DOI: 10.1073/pnas.1815117116
  * [17] Stefan Chmiela et al.  “Accurate global machine learning force fields for molecules with hundreds of atoms”, 2022  arXiv:2209.14865 [physics.chem-ph]
  * [18] Stefan Chmiela et al.  “Machine learning of accurate energy-conserving molecular force fields”  In _Science Advances_ 3.5, 2017  DOI: 10.1126/sciadv.1603015
  * [19] Bowen Deng et al.  “CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling”  In _Nature Machine Intelligence_ 5.9 Nature Publishing Group UK London, 2023, pp. 1031–1041
  * [20] FA Faber et al.  “Prediction Errors of Molecular Machine Learning Models Lower than Hybrid DFT Error” Epub 2017 Oct 10  In _J Chem Theory Comput_ 13.11, 2017, pp. 5255–5264  DOI: 10.1021/acs.jctc.7b00577
  * [21] Daan Frenkel and Berend Smit  “Understanding Molecular Simulation: From Algorithms to Applications”  Academic Press, 2001
  * [22] Richard A. Friesner  “Ab initio quantum chemistry: Methodology and applications”  In _Proceedings of the National Academy of Sciences_ 102.19, 2005, pp. 6648–6653  DOI: 10.1073/pnas.0408036102
  * [23] Xiang Fu et al.  “Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations”, 2022  arXiv:2210.07237 [physics.comp-ph]
  * [24] Johannes Gasteiger, Florian Becker and Stephan Günnemann  In _Conference on Neural Information Processing Systems (NeurIPS)_ , 2021
  * [25] Johannes Gasteiger, Janek Groß and Stephan Günnemann  “Directional message passing for molecular graphs”  In _8th International Conference on Learning Representations, ICLR_ , 2020
  * [26] Johannes Gasteiger et al.  “GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation Datasets”, 2022  arXiv:2204.02782 [cs.LG]
  * [27] Mario Geiger and Tess Smidt  “e3nn: Euclidean Neural Networks”, 2022  arXiv:2207.09453 [cs.LG]
  * [28] S. Gould et al.  “On differentiating parameterized argmin and argmax problems with application to bi-level optimization”  In _arXiv preprint arXiv:1607.05447_ , 2016
  * [29] Weihua Hu et al.  “ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations”, 2021  arXiv:2103.01436 [cs.LG]
  * [30] John Ingraham, Adam Riesselman, Chris Sander and Debora Marks  “Learning Protein Structure with a Differentiable Simulator”  In _International Conference on Learning Representations, ICLR_ , 2019
  * [31] R.. Jones  “Density functional theory: Its origins, rise to prominence, and future”  In _Rev. Mod. Phys._ 87 American Physical Society, 2015, pp. 897–923  DOI: 10.1103/RevModPhys.87.897
  * [32] A. Kabylda et al.  “Efficient interatomic descriptors for accurate machine learning force fields of extended molecules” Erratum in: Nat Commun. 2023 Jul 11;14(1):4116. PMID: 37322039; PMCID: PMC10272221.  In _Nat Commun_ 14.1, 2023, pp. 3562  DOI: 10.1038/s41467-023-39214-w
  * [33] Thomas N. Kipf and Max Welling  “Semi-Supervised Classification with Graph Convolutional Networks”, 2017  arXiv:1609.02907 [cs.LG]
  * [34] Christopher Kolloff and Simon Olsson  “Rescuing off-equilibrium simulation data through dynamic experimental data with dynAMMo” © 2023 The Author(s)  In _Mach. Learn.: Sci. Technol._ 4.4 IOP Publishing Ltd, 2023, pp. 045050  DOI: 10.1088/2632-2153/ac7b4f
  * [35] Dávid Péter Kovács et al.  “Linear Atomic Cluster Expansion Force Fields for Organic Molecules: Beyond RMSE” PMID: 34735161  In _Journal of Chemical Theory and Computation_ 17.12, 2021, pp. 7696–7711  DOI: 10.1021/acs.jctc.1c00647
  * [36] Dávid Péter Kovács et al.  “MACE-OFF23: Transferable Machine Learning Force Fields for Organic Molecules”, 2023  arXiv:2312.15211 [physics.chem-ph]
  * [37] Boris Kozinsky, Albert Musaelian, Anders Johansson and Simon Batzner  “Scaling the Leading Accuracy of Deep Equivariant Models to Biomolecular Simulations of Realistic Size”  In _Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis_ , SC ’23  Denver, CO, USA: Association for Computing Machinery, 2023  DOI: 10.1145/3581784.3627041
  * [38] Aditi S Krishnapriyan, Alejandro F Queiruga, N Benjamin Erichson and Michael W Mahoney  “Learning continuous models for continuous physics”  In _Communications Physics_ 6.1 Nature Publishing Group UK London, 2023, pp. 319
  * [39] H.. Kümmel  “A biography of the coupled cluster method”  In _Recent Progress in Many-Body Theories_ Singapore: World Scientific Publishing, 2002, pp. 334–348
  * [40] Ask Hjorth Larsen et al.  “The atomic simulation environment — a python library for working with atoms”  In _Journal of Physics: Condensed Matter_ , 2017  DOI: 10.1088/1361-648X/aa680e
  * [41] Qidong Lin, Liang Zhang, Yaolong Zhang and Bin Jiang  “Searching Configurations in Uncertainty Space: Active Learning of High-Dimensional Neural Network Reactive Potentials”  In _J. Chem. Theory Comput._ 17.5, 2021, pp. 2691–2701  DOI: 10.1021/acs.jctc.0c01193
  * [42] Yi Liu et al.  “Spherical Message Passing for 3D Molecular Graphs”  In _International Conference on Learning Representations (ICLR)_ , 2022
  * [43] M. Majewski, A. Pérez and P. Thölke  “Machine learning coarse-grained potentials of protein thermodynamics”  In _Nat Commun_ 14 Nature Publishing Group, 2023, pp. 5739  DOI: 10.1038/s41467-023-41343-1
  * [44] TE Markland and M Ceriotti  “Nuclear quantum effects enter the mainstream”  In _Nat Rev Chem_ 2, 2018, pp. 0109  DOI: 10.1038/s41570-017-0109
  * [45] A. Merchant, S. Batzner and S.S. Schoenholz  “Scaling deep learning for materials discovery”  In _Nature_ 624 Nature Publishing Group, 2023, pp. 80–85  DOI: 10.1038/s41586-023-06735-9
  * [46] Joe D. Morrow, John L.. Gardner and Volker L. Deringer  “How to validate machine-learned interatomic potentials”  In _J. Chem. Phys._ 158, 2023, pp. 121501
  * [47] Albert Musaelian et al.  “Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics”, 2022  arXiv:2204.05249 [physics.comp-ph]
  * [48] Carles Navarro, Maciej Majewski and Gianni De Fabritiis  “Top-Down Machine Learning of Coarse-Grained Protein Force Fields”  In _J. Chem. Theory Comput._ , 2023
  * [49] Geoffrey Négiar, Michael W. Mahoney and Aditi Krishnapriyan  “Learning differentiable solvers for systems with hard constraints”  In _The Eleventh International Conference on Learning Representations_ , 2023
  * [50] W.. Noid  “Perspective: coarse-grained models for biomolecular systems”  In _J. Chem. Phys._ 139 American Institute of Physics, 2013, pp. 090901
  * [51] John P. Perdew, Kieron Burke and Matthias Ernzerhof  “Generalized Gradient Approximation Made Simple”  In _Phys. Rev. Lett._ 77 American Physical Society, 1996, pp. 3865–3868  DOI: 10.1103/PhysRevLett.77.3865
  * [52] Z Raza  “Proton ordering in cubic ice and hexagonal ice; a potential new ice phase–XIc”  In _Phys. Chem. Chem. Phys._ 13, 2011, pp. 19788–19795  DOI: 10.1039/c1cp22167a
  * [53] Jie Ren et al.  “TorchOpt: An Efficient Library for Differentiable Optimization”, 2022  arXiv:2211.06934 [cs.MS]
  * [54] Sebastien Röcken and Julija Zavadlav  “Accurate machine learning force fields via experimental and simulation data fusion”, 2023  arXiv:2308.09142 [physics.chem-ph]
  * [55] M. Rupp, A. Tkatchenko, K.-R. Müller and O.. Lilienfeld  “Fast and accurate modeling of molecular atomization energies with machine learning”  In _Phys. Rev. Lett._ 108, 2012, pp. 058301
  * [56] Franco Scarselli et al.  “The Graph Neural Network Model”  In _IEEE Transactions on Neural Networks_ 20.1, 2009, pp. 61–80  DOI: 10.1109/TNN.2008.2005605
  * [57] Michael Schaarschmidt et al.  “Learned Force Fields Are Ready For Ground State Catalyst Discovery”, 2022  arXiv:2209.12466 [cond-mat.mtrl-sci]
  * [58] K.. Schütt et al.  “SchNet – A deep learning architecture for molecules and materials”  In _J. Chem. Phys._ 148.24, 2018, pp. 241722  DOI: 10.1063/1.5019779
  * [59] Kristof T. Schütt, Oliver T. Unke and Michael Gastegger  “Equivariant message passing for the prediction of tensorial properties and molecular spectra”  In _International Conference on Machine Learning (ICML)_ , 2021
  * [60] Nima Shoghi et al.  “From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction”, 2023  arXiv:2310.16802 [cs.LG]
  * [61] Martin Šípka, Johannes C.. Dietschreit, Lukáš Grajciar and Rafael Gómez-Bombarelli  “Differentiable Simulations for Enhanced Sampling of Rare Events”, 2023  arXiv:2301.03480 [physics.chem-ph]
  * [62] Justin S. Smith et al.  “Less is more: Sampling chemical space with active learning”  In _J. Chem. Phys._ 148, 2018, pp. 241733  DOI: 10.1063/1.5023802
  * [63] Sina Stocker et al.  “How robust are modern graph neural network potentials in long and hot molecular dynamics simulations?”  In _Machine Learning: Science and Technology_ 3.4 IOP Publishing, 2022, pp. 045010  DOI: 10.1088/2632-2153/ac9955
  * [64] S. Takamoto, C. Shinagawa and D. Motoki  “Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements”  In _Nat Commun_ 13, 2022, pp. 2991  DOI: 10.1038/s41467-022-30687-9
  * [65] Sebastian Thaler and Julija Zavadlav  “Learning neural network potentials from experimental data via Differentiable Trajectory Reweighting”  In _Nature Communications_ 12, 2021, pp. 6884  DOI: 10.1038/s41467-021-27241-4
  * [66] Alexandre Tkatchenko and Matthias Scheffler  “Accurate Molecular Van Der Waals Interactions from Ground-State Electron Density and Free-Atom Reference Data”  In _Phys. Rev. Lett._ 102 American Physical Society, 2009, pp. 073005  DOI: 10.1103/PhysRevLett.102.073005
  * [67] Alexandre Tkatchenko, Robert A. DiStasio, Roberto Car and Matthias Scheffler  “Accurate and Efficient Method for Many-Body van der Waals Interactions”  In _Phys. Rev. Lett._ 108 American Physical Society, 2012, pp. 236402  DOI: 10.1103/PhysRevLett.108.236402
  * [68] J. Vandermause, S.B. Torrisi and S. Batzner  “On-the-fly active learning of interpretable Bayesian force fields for atomistic rare events”  In _npj Comput Mater_ 6, 2020, pp. 20  DOI: 10.1038/s41524-020-0283-z
  * [69] Joshua A. Vita and Daniel Schwalbe-Koda  “Data efficiency and extrapolation trends in neural network interatomic potentials”  In _Mach. Learn.: Sci. Technol._ 4 IOP Publishing, 2023, pp. 035031
  * [70] Wujie Wang, Simon Axelrod and Rafael Gómez-Bombarelli  “Differentiable Molecular Simulations for Control and Learning”, 2020  arXiv:2003.00868 [physics.comp-ph]
  * [71] Wujie Wang, Zhenghao Wu, Johannes C.. Dietschreit and Rafael Gó mez-Bombarelli  “Learning pair potentials using differentiable simulations”  In _The Journal of Chemical Physics_ 158.4 AIP Publishing, 2023, pp. 044113  DOI: 10.1063/5.0126475
  * [72] Zun Wang et al.  “Improving machine learning force fields for molecular dynamics simulations with fine-grained force metrics”  In _J. Chem. Phys._ 159, 2023, pp. 035101  DOI: 10.1063/5.0147023
  * [73] Yujie Wu, Harald L. Tepper and Gregory A. Voth  “Flexible simple point-charge water model with improved liquid-state properties”  In _The Journal of Chemical Physics_ 124.2, 2006, pp. 024503  DOI: 10.1063/1.2136877
  * [74] Robert W. Zwanzig  “High-Temperature Equation of State by a Perturbation Method. I. Nonpolar Gases”  In _J. Chem. Phys._ 22, 1954, pp. 1420–1426

##  Appendix A Supplementary Information

###  A.1 Derivation of Boltzmann Estimator

We provide a full derivation of our Boltzmann estimator, which we use to train
our NNIP as part of StABlE Training. Consider a vector-valued observable
$g(\Gamma)$ of a state $\Gamma$, and a reference value of the observable
$g_{\text{ref}}$. Training a NNIP $U_{\theta}$ to match $g_{\text{ref}}$
requires minimizing the loss function,

$\mathcal{L}_{\text{obs}}(\theta)\stackrel{{\scriptstyle\text{def}}}{{=}}\|\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]-g_{\text{ref}}\|_{2}^{2},$

where $P_{\theta}(\Gamma)$ is the equilibrium distribution induced by the NNIP
$U_{\theta}$. This requires computing the gradient,
$\nabla_{\theta}\mathcal{L}_{\text{obs}}$, which can be decomposed via the
chain rule as follows:

$\displaystyle\nabla_{\theta}\mathcal{L}_{\text{obs}}^{\top}$
$\displaystyle=\frac{\partial\mathcal{L}_{\text{obs}}}{\partial{\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)[g(\Gamma)]}}}\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$
$\displaystyle=2(\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]-g_{\text{ref}})^{\top}\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}\,.$

We derive the $N$-sample estimator, presented in Equation LABEL:eq:estimator,
of the Jacobian, ${\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}/{\partial\theta}$. The estimator is repeated
below for convenience.

$\mathcal{E}(\Gamma^{(1)},\ldots,\Gamma^{(N)})\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{N}{k_{B}T(N-1)}\left[\hat{\mathbb{E}}\left[g(\Gamma)\right]\hat{\mathbb{E}}\left[{\nabla_{\theta}U_{\theta}(\Gamma)}\right]^{\top}-\hat{\mathbb{E}}\left[{g(\Gamma)\nabla_{\theta}U_{\theta}(\Gamma)}^{\top}\right]\right],$

where $\hat{\mathbb{E}}$ denotes sample averages. This is an unbiased
estimator, that is, $\mathbb{E}_{(\Gamma^{(1)},\ldots,\Gamma^{(N)}\sim
P_{\theta}(\Gamma)}[\mathcal{E}(\Gamma^{(1)},\ldots,\Gamma^{(N)})]=\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$.

###### Proof.

Given
$P_{\theta}(\Gamma)\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma)\right)}{C(\theta)}$,
where
$\mathcal{H}(\Gamma)=\sum_{i=1}^{N}\frac{p_{i}^{2}}{2m_{i}}+U_{\theta}(r)$,
and
$C(\theta)\stackrel{{\scriptstyle\text{def}}}{{=}}\int\exp(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma^{\prime}))\mathop{}\\!\mathrm{d}\Gamma^{\prime}$
is the partition function, we have,

$\displaystyle\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$
$\displaystyle=\frac{\partial}{\partial\theta}\int
g(\Gamma^{\prime})P_{\theta}(\Gamma^{\prime})\mathop{}\\!\mathrm{d}\Gamma^{\prime}$
$\displaystyle=\int
g(\Gamma^{\prime})\nabla_{\theta}P_{\theta}(\Gamma^{\prime})^{\top}\mathop{}\\!\mathrm{d}\Gamma^{\prime}$
$\displaystyle=\int
g(\Gamma^{\prime})\nabla_{\theta}(\frac{\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma)\right)}{C(\theta)})^{\top}\mathop{}\\!\mathrm{d}\Gamma^{\prime}.$

Expanding via the chain rule for Jacobians, and noting that
$\nabla_{\theta}\mathcal{H}_{\theta}(\Gamma)=\nabla_{\theta}U_{\theta}(r)$ (we
will write $U_{\theta}(\Gamma)$ for convenience), we get,

$\displaystyle=\int
g(\Gamma^{\prime})\frac{-\frac{1}{k_{B}T}\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma^{\prime})\right)\nabla_{\theta}U_{\theta}(\Gamma^{\prime})^{\top}C(\theta)-\nabla_{\theta}C(\theta)^{\top}\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma^{\prime})\right)}{C(\theta)^{2}}\mathrm{d}\Gamma^{\prime}$
$\displaystyle=\int
g(\Gamma^{\prime})\left[-\frac{1}{k_{B}T}\nabla_{\theta}U_{\theta}(\Gamma^{\prime})^{\top}P_{\theta}(\Gamma^{\prime})-\frac{\nabla_{\theta}C(\theta)^{\top}}{C(\theta)}P_{\theta}(\Gamma^{\prime})\right]\mathrm{d}\Gamma^{\prime}\,.$

By definition of $C(\theta)$, the quotient
$\frac{\nabla_{\theta}C(\theta)}{C(\theta)}$ can be simplified as,

$\displaystyle\frac{\nabla_{\theta}C(\theta)}{C(\theta)}$
$\displaystyle=\frac{-\int\nabla_{\theta}U_{\theta}(\Gamma^{\prime})\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma^{\prime})\right)\mathrm{d}\Gamma^{\prime}}{k_{B}T\cdot
C(\theta)}$
$\displaystyle=-\frac{1}{k_{B}T}\int\nabla_{\theta}U_{\theta}(\Gamma^{\prime})P_{\theta}(\Gamma^{\prime})\mathop{}\\!\mathrm{d}\Gamma^{\prime}$
$\displaystyle=-\frac{1}{k_{B}T}\mathbb{E}_{\Gamma}[\nabla_{\theta}U_{\theta}(\Gamma)]\,.$

Putting it all together, we have,

$\displaystyle\frac{\partial\mathbb{E}_{\Gamma\sim
P_{\theta}(\Gamma)}[g(\Gamma)]}{\partial\theta}$
$\displaystyle=-\frac{1}{k_{B}T}\left(\underbrace{\int
g(\Gamma^{\prime})\nabla_{\theta}U_{\theta}(\Gamma^{\prime})^{\top}P_{\theta}(\Gamma^{\prime})\mathop{}\\!\mathrm{d}\Gamma^{\prime}}_{\mathbb{E}\left[g(\Gamma)\cdot\nabla_{\theta}U_{\theta}(\Gamma)\right]}-\underbrace{\left(\int\cdot
g(\Gamma^{\prime})\cdot
P_{\theta}(\Gamma^{\prime})\cdot\mathrm{d}\Gamma^{\prime}\right)}_{\mathbb{E}[g(\Gamma)]}\mathbb{E}[\nabla_{\theta}U_{\theta}(\Gamma)]^{\top}\right)$
$\displaystyle=-\frac{1}{k_{B}T}\left(\mathbb{E}\left[g(\Gamma)\cdot\nabla_{\theta}U_{\theta}(\Gamma)^{\top}\right]-\mathbb{E}[g(\Gamma)]\mathbb{E}[\nabla_{\theta}U_{\theta}(\Gamma)]^{\top}\right)$
$\displaystyle=-\frac{1}{k_{B}T}\operatorname{Cov}\left(g(\Gamma),\nabla_{\theta}U_{\theta}(\Gamma)\right).$

In this work, we use the following unbiased estimator for covariance. Given 2
random vectors $X,Y$, and samples of these vectors
$\\{(X^{(1)},Y^{(1)}),...(X^{(n)},Y^{(n)})\\}$, our estimator is given by:

$\widehat{\operatorname{Cov}}(X,Y)=\frac{1}{N-1}\sum_{i=1}^{N}X^{(i)}{Y^{(i)}}^{\top}-\frac{1}{N(N-1)}\left(\sum_{j=1}^{N}X^{(j)}\right)\left(\sum_{k=1}^{N}{Y^{(k)}}^{\top}\right).$

We can prove it is unbiased by taking the expectation of the right hand side
and applying linearity of expectation:

$\displaystyle=\frac{1}{N-1}\sum_{i=1}^{N}\mathbb{E}[X^{(i)}{Y^{(i)}}^{\top}]-\frac{1}{N(N-1)}\left(\sum_{i=1}^{N}\mathbb{E}[X^{(i)}{Y^{(i)}}^{\top}]+\sum_{j\neq
k}\mathbb{E}[X^{(j)}{Y^{(k)}}^{\top}]\right)$
$\displaystyle=\frac{N}{N-1}\mathbb{E}[XY^{\top}]-\frac{1}{N(N-1)}\left(N\cdot\mathbb{E}[XY^{\top}]+N(N-1)\cdot\mathbb{E}[X]\mathbb{E}[Y]^{\top}\right)$
$\displaystyle=\mathbb{E}[XY^{\top}]-\mathbb{E}[X]\mathbb{E}[Y]^{\top}=\operatorname{Cov}(X,Y).$

Given the samples $\\{\Gamma^{(1)},,...\Gamma^{(N)}\\}$, we can then use the
estimator above to get,

$\displaystyle\widehat{\operatorname{Cov}}\left(g(\Gamma),\nabla_{\theta}U_{\theta}(\Gamma)\right)$
$\displaystyle=\frac{1}{N-1}\sum_{i=1}^{N}g(\Gamma^{(i)})\nabla_{\theta}U_{\theta}(\Gamma^{(i)})^{\top}-\frac{1}{N(N-1)}\left(\sum_{j=1}^{N}g(\Gamma^{(j)})\right)\left(\sum_{k=1}^{N}\nabla_{\theta}U_{\theta}(\Gamma^{(k)})\right)^{\top}$
$\displaystyle=\frac{N}{N-1}\left[\hat{\mathbb{E}}\left[{g(\Gamma)\nabla_{\theta}U_{\theta}(\Gamma)^{\top}}\right]-\hat{\mathbb{E}}\left[g(\Gamma)\right]\hat{\mathbb{E}}\left[{\nabla_{\theta}U_{\theta}(\Gamma)}\right]^{\top}\right].$

The final Boltzmann estimator is thus given as,

$\mathcal{E}(\Gamma^{(1)},\ldots,\Gamma^{(N)})\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{N}{k_{B}T(N-1)}\left[\hat{\mathbb{E}}\left[g(\Gamma)\right]\hat{\mathbb{E}}\left[{\nabla_{\theta}U_{\theta}(\Gamma)}\right]^{\top}-\hat{\mathbb{E}}\left[{g(\Gamma)\nabla_{\theta}U_{\theta}(\Gamma)}^{\top}\right]\right].$

∎

We note that a similar estimator is obtained in [65] by differentiating
through a reweighting scheme arising from thermodynamic perturbation theory.
We have presented an alternative derivation that does not require reweighting.

###  A.2 Localized Boltzmann Estimator

We provide additional details on the Localized Boltzmann Estimator. Formally,
define a neighborhood of $n$ atoms
$\mathcal{N}=\\{x_{1},x_{2},\ldots,x_{n}\mid x_{i}\in\mathbb{Z},1\leq
x_{i}\leq N\ \text{for all}\ i=1,2,\ldots,n\\}$. The local state $\gamma$ is
defined as
$\gamma=\\{[r_{\mathcal{N}_{1}};\ldots;r_{\mathcal{N}_{n}}],[P_{\mathcal{N}_{1}};\ldots;P_{\mathcal{N}_{n}}]\\}$.
The local energy is then defined as
$U_{\theta}(\gamma)=\sum_{i=1}^{n}U_{\theta}(\gamma_{i})$, where
$\gamma_{i}=\\{r_{\mathcal{N}_{i}},P_{\mathcal{N}_{i}}\\}$ contains the
position and momenta of the $i^{th}$ atom in the local neighborhood. We can
obtain this by noting that NNIPs first compute individual atomic energies
before summing over all atoms to obtain the global energy. Finally, given $N$
i.i.d. samples of local states $\gamma^{(1)},\ldots,\gamma^{(N)}$, where each
$\gamma^{(i)}$ is extracted from a global state $\Gamma^{(i)}\sim
P_{\theta}(\Gamma)$, we define the $N$-sample localized estimator of the
Jacobian $\frac{\partial\mathbb{E}_{\gamma\sim
P_{\theta}(\gamma)}[g(\gamma)]}{\partial\theta}$ as

$\mathcal{E}(\gamma^{(1)},\ldots,\gamma^{(N)})\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{N}{k_{B}T(N-1)}\left[\hat{\mathbb{E}}\left[g(\gamma)\right]\hat{\mathbb{E}}\left[{\nabla_{\theta}U_{\theta}(\gamma)}\right]^{\top}-\hat{\mathbb{E}}\left[{g(\gamma)\cdot\nabla_{\theta}U_{\theta}(\gamma)}^{\top}\right]\right]$
(8)

The localized estimator follows from the original Boltzmann estimator due to
the fact that as a subset of the larger Boltzmann-distributed system, any
local neighborhood also obeys a Boltzmann-distribution. In practice, we
extract multiple local neighborhoods from each global state to increase the
sample size and state space coverage. See Supplementary Section A.6 for a
description of the neighborhood extraction procedure used for the Water
system.

###  A.3 StABlE Training Algorithm

We provide an algorithmic description of our StABlE-Training procedure. The
algorithm proceeds in two main phases: simulation, which seeks out unstable
regions of molecular phase space, and learning, which corrects the instability
via supervision from the reference observables.

Algorithm 1 StABlE Training Procedure

1:Initialize:

2:Pre-trained Neural Network Potential $U_{\theta}$

3:Reference energy and forces dataset $\mathcal{D}_{\text{train}}$ and
observables $\\{g_{\text{ref}}^{(i)}\\}_{i=1}^{N}$

4:Simulation length $t$, number of parallel replicas $R$, minimum unstable
threshold $f_{min}$,

5:maximum unstable threshold $f_{max}$, energy and forces loss weight
$\lambda$, learning rate $\alpha$

6:

7:$\bar{\Gamma}_{curr}\leftarrow\\{\Gamma^{(1)}(0),\Gamma^{(2)}(0),\ldots\Gamma^{(R)}(0)\\}\sim\mathcal{D}_{\text{train}}$

8:Current fraction of unstable replicas $f_{unst}\leftarrow 0$

9:Total simulated time $T_{f_{i}}\leftarrow 0,\forall i=1,\ldots,R$

10:Mark all replicas as active for simulation

11:repeat Simulation Phase

12: while $f_{unst}<f_{max}$ do

13: Simulate active replicas w/ $U_{\theta}$ for $t$ steps starting from
$\bar{\Gamma}_{curr}$

14: $T_{f_{i}}\leftarrow T_{f_{i}}+t$, $\forall i$ corresponding to stable
replicas

15:
$\bar{\Gamma}_{curr}\leftarrow\\{\Gamma^{(1)}({T_{f_{1}}}),\Gamma^{(2)}({T_{f_{2}}}),\ldots\Gamma^{(R)}({T_{f_{R}}})\\}$

16: Update $f_{unst}$ and mark unstable replicas inactive

17: end while

18:

19: // $\bar{\Gamma}_{curr}$ now contains
$\\{\Gamma^{(1)}({T_{f_{1}}}),\Gamma^{(2)}({T_{f_{2}}}),\ldots\Gamma^{(R)}({T_{R}})\\}$,
where $T_{f_{i}}$’s are per-replica total simulation times

20: // At least $f_{max}$ fraction of replicas are unstable at this point
Learning Phase

21: while $f_{unst}>f_{min}$ do

22: Rewind all trajectories by $t$ timesteps:
$\bar{\Gamma}_{curr}\leftarrow\\{\Gamma^{(1)}({T_{f_{1}}-t}),\ldots,\Gamma^{(R)}({T_{f_{R}}-t})\\}$

23: Simulate all replicas w/ $U_{\theta}$ for $t$ steps starting from
$\bar{\Gamma}_{curr}$

24: Update $f_{unst}$

25: Compute observables
$\\{\widehat{\mathbb{E}}_{\Gamma}[g^{(i)}(\Gamma)]\\}_{i=1}^{N}$ over all
length-$t$ trajectories

26:
$\mathcal{L}_{obs}\leftarrow\sum_{i=1}^{N}\|\widehat{\mathbb{E}}_{\Gamma}[g^{(i)}(\Gamma)]-g_{\text{ref}}^{(i)}\|^{2}$

27: $\mathcal{L}_{\text{QM}}\leftarrow$ energy and forces loss of $U_{\theta}$
on dataset $\mathcal{D}_{\text{train}}$

28:
$\theta\leftarrow\theta-\alpha\cdot\nabla_{\theta}(\mathcal{L}_{obs}+\lambda\mathcal{L}_{\text{QM}})$
// Compute Boltzmann Estimator

29: end while

30: Mark stable replicas active, unstable replicas inactive for next
simulation phase

31: // At most $f_{min}$ fraction of replicas are unstable at this point

32:until Convergence or maximum cycles reached

###  A.4 Observables

We provide definitions and details of the observables considered in this work.

#### Distribution of Interatomic Distances.

The distribution of interatomic distances serves as a low-dimensional
description of 3D structure. For a configuration
$r^{\prime}\in\mathbb{R}^{N\times 3}$, it is defined as,

$h(r)=\frac{1}{N(N-1)}\sum_{i=1}^{N}\sum_{\begin{subarray}{c}j\neq
i\end{subarray}}^{N}\delta\left(r-\left\|{r_{i}}{{}^{\prime}}-{r_{j}}{{}^{\prime}}\right\|\right),$

where $\delta$ is the Dirac-Delta function. Although the observable need not
be differentiable with respect to $r$ for our Boltzmann learning framework, we
compute a differentiable version of $h(r)$ via Gaussian smearing as in [71] to
facilitate comparison with differentiable simulation methods.

#### Radial Distribution Function.

The radial distribution function (RDF) captures how density (relative to the
bulk) varies as a function of distance from a reference particle, and thus
characterizes the structural/thermodynamic properties of the system. The RDF
is defined as,

$RDF(r)=\frac{V}{N^{2}4\pi r^{2}}h(r).$

where $V$ is the volume of the simulation domain, $N$ is the number of
particles, $r$ is the radial distance from a reference particle, and $h(r)$ is
a histogram of pairwise distances. As with $h(r)$, we use Gaussian smearing to
make the RDF differentiable.

#### Velocity Autocorrelation Function.

The velocity autocorrelation function (VACF) is an important dynamical
observable. Many fundamental properties, such as the diffusion coefficient and
vibrational spectra, are functions of this observable. Computing the VACF
requires a window of consecutive simulation states to compute. The VACF at a
given time lag $\Delta t$ is given by,

$VACF(\Delta
t)=\frac{1}{S}\sum_{t_{0}}\sum_{i}<v_{i}(t_{0}),v_{i}(t_{0}+\Delta t)>,$

where $t_{0}$ is an initial time, $v_{i}(t)$ is the velocity of the $i^{th}$
atom at timestep $t$, $<\cdot,\cdot>$ is an inner product, and $S$ is the
total number of samples considered given the summations over initial times and
atoms. In this work, we compute the VACF over a window of 100 consecutive
simulation timesteps, and normalize the values by the autocorrelation at
$\Delta t=0$ to restrict the range to $[0,1]$.

#### Diffusivity Coefficient

The diffusivity coefficient is a fundamental transport property with crucial
implications on the performance of energy storage systems, among other
applications. Related to the time-derivative of the mean squared displacement,
the diffusivity coefficient is defined as,

$D=\lim_{t\to\infty}\frac{1}{6t}\frac{1}{N}\sum_{i=1}^{N}\left|r_{i}(t)-r_{i}(0)\right|^{2}$

, where $r_{i}(t)$ is the coordinate of the $i^{th}$ particle at time $t$ and
$N$ is the number of atoms considered. For the water system considered in this
work, we measure the diffusivity of all 64 oxygen atoms.

###  A.5 Stability Criteria

We provide definitions and details on the stability criteria considered in
this work.

#### Maximum Bond Length Deviation.

Adapted from [23], the maximum bond length deviation metric captures
unphysical bond stretching or collapse in small flexible molecules. According
to this criterion, a simulation becomes unstable at time $T$ if,

$\max_{(i,j)\in\mathcal{B}}\left|\left(\left\|r_{i}(T)-r_{j}(T)\right\|-b_{i,j}\right)\right|>\Delta,$

where $\mathcal{B}$ is the set of all bonds, $i,j$ are the two endpoint atoms
of the bond, and $b_{i,j}$ is the equilibrium bond length computed from the
reference simulation. Following [23], we set $\Delta=0.5A$ for final stability
evaluation. However, we adopt a more conservative value of $\Delta=0.25A$
during training in order to detect and correct instability earlier. We use
this criterion for the MD17 and MD22 datasets.

#### Minimum Intermolecular Distance.

The minimum intermolecular distance metric is used to detect unphysical
coordination structures or collisions between molecules in the water system.
According to this criterion, a simulation becomes unstable at time $T$ if,

$\min_{(i,j)\notin\mathcal{B}}\left\|r_{i}(T)-r_{j}(T)\right\|<\Delta,$

where $\mathcal{B}$ is the set of all bonds, and $i,j$ are the endpoint
indices of two non-bonded atoms. We set $\Delta=1.2A$ to detect instability
during training.

#### Radial Distribution Function Mean Absolute Error.

The minimum intermolecular distance metric is appropriate at train-time to
detect instability early before it cascades to the rest of the system.
However, it is too sensitive to use for evaluation, as realistic simulation
can still be achieved for some time after the occurrence of a highly localized
instability. Therefore, following [23], we adopt an instability metric based
on the radial distribution function, defined as,

$\int_{r=0}^{\infty}\left\|RDF_{\text{ref}}(r)-\left\langle
RDF^{t}(r)\right\rangle_{t=T}^{T+\tau}\right\|\,dr>\Delta,$

where $\langle\cdot\rangle$ is the averaging operator, $\tau$ is a short time
window, and $\Delta$ is the stability threshold. We use $\tau=1$ ps and
$\Delta=3.0$ for water. The stability criterion is triggered if any of the
three element-conditioned water RDFs (H-H, O-O, or H-O) exceeds the threshold.

#### Justification of Chosen Stability Thresholds.

Following [23], we choose stability thresholds such that a realistic, high-
fidelity simulation at the chosen temperature would virtually never cross the
threshold. This means that if a simulation does cross the threshold, this is
indicative of catastrophic failure. Thresholds are set more conservatively
during training in order to facilitate early detection of potential collapse.
Figure 7 shows the distribution of values of the stability criterion over
high-fidelity reference simulations for the three systems considered in this
work, along with thresholds chosen to denote instability for training and
evaluation.

Figure 7: Distribution of stability criterion over reference simulations.
Instability thresholds are chosen to be very relaxed, such that crossing of
the threshold signifies catastrophic, unrecoverable instability.

###  A.6 Architecture and Training Details

We provide details on the model architectures and training procedures used in
this work. MD simulations and NNIPs are written in the PyTorch framework and
are built upon the MDsim [23] and Atomic Simulation Environment [40] packages.
All training is performed on a single NVIDIA A100 GPU.

Table 1 provides details on the NNIP architectures. $r_{max}$ is the cutoff
distance used to construct the radius graph. $l_{max}$ denotes the level of
E(3) equivariance used in the network.

| Symmetry | Parameter | $r_{max}$ | $l_{max}$
---|---|---|---|---
| Principle | Count | (A) |
SchNet [58] | E(3)-invariant | 0.12M | 5.0 | -
NequIP [5] | E(3)-equivariant | 0.12M | 5.0 | 1
GemNet-T [24] | SE(3)-equivariant | 1.89M | 5.0 | -
Table 1: NNIP Architecture Details.

For energy and forces pre-training, we follow the protocols in [23]. In order
to isolate the effect of observable-based learning, we begin StABlE Training
only after pre-training has fully converged (that is, when
$\mathcal{L}_{\text{QM}}$ has reached a plateau). This means that any
improvements in stability or accuracy as a result of StABlE Training can be
attributed to the learning signal from the reference observable, as opposed to
the regularization from the QM energy and forces data.

We include relevant settings used for StABlE Training in Table 2. $\alpha$ is
the learning rate, $\lambda$ is the strength of energy and forces
regularization, $t$ is the number of simulation timesteps per epoch, and $R$
is the number of parallel replicas. We note that in practice, we compute the
outer products and empirical means in the Boltzmann estimator in batched
fashion. Thus, to limit memory usage, we compute $\frac{N}{B}$ separate
Boltzmann estimators from minibatches of $B<N$ states and subsequently average
them to produce a final estimator.

| NNIP | Stability | Training | Estimator | $\alpha$ | $\lambda$ | t | $R$ | $B$
---|---|---|---|---|---|---|---|---|---
|  | Criterion, Threshold | Observable | Type |  |  |  |  |
Aspirin | SchNet | Bond Len. Dev., 0.25 A | $h(r)$ | Global | 0.001 | 10 | 2000 | 128 | 40
Ac-Ala3-NHMe | NequIP | Bond Len. Dev., 0.25 A | $h(r)$ | Global | 0.001 | 10 | 2000 | 128 | 40
Water | GemNet-T | IMD, 1.2 A | O-H Bond Length | Local | 0.003 | 0 | 1000 | 8 | 4
Table 2: StABlE Training Settings.

#### Details on Local Neighborhood Sampling Scheme for Water.

For water training, from each sampled global state we consider all possible
local neighborhoods containing a single water molecule. Since instability is
highly localized, the majority of local neighborhoods are stable, leading to
an imbalance in the sampled states. To provide more informative samples for
optimization, we assign the individual sampled molecules to one of 100 bins
based on their deviation from the equilibrium O-H bond length. Then, we choose
a fixed number (in this case 5) molecules from each bin from which to compute
the final localized estimator.

#### Wall Clock Time of StABlE Training.

In Table 3, we provide the total wall clock time spent on QM pre-training, as
well as the subsequent StABlE Training. All runtimes were measured on an
NVIDIA A100 GPU. We note that especially for Ac-Ala3-NHMe and Water, StABlE
Training incurs a relatively small marginal computational cost beyond that of
QM pre-training.

| QM pre-training | StABlE Training
---|---|---
Aspirin | 2 | 4
Ac-Ala3-NHMe | 16 | 4.7
Water | 31 | 3.5
Table 3: Wall clock time, in hours, of QM pre-training and StABlE Training.

###  A.7 Simulation Details

We provide MD simulation details in Table 4. All systems considered in this
work are simulated with a Nose-Hoover thermostat. All simulation parameters
are chosen to be consistent with prior literature [18, 17, 23].

| Temperature | Nose-Hoover | Timestep | Periodic Boundary
---|---|---|---|---
| (K) | Effective Mass (cm-1) | (fs) | Conditions
Aspirin | 500 | 2000 | 0.5 | No
Ac-Ala3-NHMe | 500 | 2000 | 0.5 | No
Water | 300 | 2000 | 1 | Yes
Table 4: Simulation Settings.

###  A.8 Evaluation Details

We provide further details on the protocol used to evaluate the NNIPs
considered in this work. Our evaluation protocol is centered around MD
simulations. To facilitate direct comparison between StaBlE-trained NNIPs and
those trained only on energy and forces reference data, for a given molecular
system we perform MD simulations starting from the same initial configurations
for all models. We choose the number of parallel replicas and total simulation
time on a per-system basis so as to saturate GPU memory usage while remaining
within a reasonable computational budget. Simulation conditions are the same
as described in Section A.7 except for the temperature generalization
experiment (Section 3.2), in which the temperature of simulation is varied.
Table 5 summarizes the relevant evaluation parameters for each system.

| Num. Parallel Replicas | Max. Simulation Time (ps) | Stability Criterion, Threshold
---|---|---|---
Aspirin | 256 | 300 | Bond Length Deviation, 0.5 A
Ac-Ala3-NHMe | 48 | 300 | Bond Length Deviation, 0.5 A
Water | 10 | 200 | RDF MAE, 3.0
Table 5: StABlE Evaluation Settings.

###  A.9 Temperature-Reweighting of Observables

We provide further details on the reweighting process used to estimate the
reference distribution of interatomic distances at 350K and 700K.

Under the canonical ensemble, microstates follow a Boltzmann distribution
$P_{\theta}(\Gamma)\stackrel{{\scriptstyle\text{def}}}{{=}}\frac{\exp\left(-\frac{1}{k_{B}T}\mathcal{H}_{\theta}(\Gamma)\right)}{C(\theta)}$.
Consider states $\Gamma^{(1)},\ldots,\Gamma^{(N)}$ sampled at temperature
$T_{1}$. Define a reweighting factor for each sample as follows,

$\displaystyle w_{i}$
$\displaystyle=\frac{\frac{P_{\theta}(\Gamma;T_{1})}{P_{\theta}(\Gamma;T_{2})}}{\sum_{i=1}^{N}\frac{P_{\theta}(\Gamma^{i};T_{1})}{P_{\theta}(\Gamma^{i};T_{2})}}$
$\displaystyle=\frac{\exp(-\frac{\mathcal{H}_{\theta}(\Gamma^{i})}{k_{B}}(\frac{1}{T_{2}}-\frac{1}{T_{1}}))}{\sum_{i=1}^{N}\exp(-\frac{\mathcal{H}_{\theta}(\Gamma^{i})}{k_{B}}(\frac{1}{T_{2}}-\frac{1}{T_{1}}))}.$

We can then compute a reweighted Monte Carlo estimate of the observable at
$T_{2}$ as follows [65],

$g_{true,T_{2}}=\sum_{i=1}^{N}w_{i}g(\Gamma^{(i)})$

.

The statistical error of the reweighted Monte Carlo estimate is captured by
the effective sample size, $N_{eff}\approx e^{-\sum_{i=1}^{N}w_{i}ln(w_{i})}$
[13]. A small effective sample size indicates that a few samples with high
weights dominate the average; this occurs for large differences between
$T_{1}$ and $T_{2}$. To select lower and upper temperatures at which to
perform the temperature generalization experiment, we set a minimum
$N_{eff}=1000$, leading us to choose 350K and 700K (Supplemental Figure 8).

Figure 8: Boltzmann-reweighting of aspirin samples. The effective sample size
($N_{eff}$) as a function of reweighting temperature for aspirin dataset.
$N_{eff}$ is maximized when the reweighting temperature is equal to the
original temperature (500K). Using a minimum sample size of 1000, we choose
upper and lower temperatures of 700K and 350K at which to perform temperature
generalization experiments.

###  A.10 Analysis of Energy and Forces Errors

We study the effect of two hyperparameters of StABlE Training, the learning
rate $\alpha$ and the strength of QM regularization $\lambda$, on the energy
and forces errors of a SchNet NNIP on a held-out test set of aspirin
structures. We perform StABlE Training as described in Section 2 for learning
rates ranging from $10^{-5}$ to $10^{-3}$ and QM regularization coefficients
ranging from $10^{0}$ to $10^{2}$. We perform evaluation of each trained model
via MD simulation of 256 parallel replicas at 500K.

Figure 9: Effect of training hyperparameters on stability and force error
improvements. (a) Models trained with higher learning rate and lower QM loss
coefficient achieve better stability gains relative to a baseline model
trained only on QM reference data. (b) Models trained with higher learning
rate and lower QM loss coefficient incur higher increases in force mean
absolute error (MAE) on a held-out test dataset relative to the baseline
model. (c) A Pareto frontier of Stability vs Force MAE arises. Some choices of
learning rate and loss coefficient are Pareto-suboptimal, while choosing
others moves along the Pareto frontier.

Training runs with high learning rate and low QM loss coefficient achieve
greater improvements in stable simulation time relative to a baseline model
trained only on QM reference data (Figure 9a). However, these training runs
also incur a greater increase in the Mean Absolute Error (MAE) of force
prediction on a held out test dataset (Figure 9b). Due to training on a single
structural observable, the observable-matching component of the StABlE
objective is ill-posed: a NNIP which collapses simulations onto a sparse set
of states exactly matching the reference observable would globally minimize
the observable-matching loss and yield indefinitely stable simulations, while
incurring a large QM/force error. As the learning rate of the StABlE Training
procedure is increased, the optimization is increasingly pushed towards this
degenerate mode. Increasing the weight of the QM objective counteracts this
tendency. Consequently, a Pareto frontier arises between stability and force
prediction accuracy (Figure 9c). Some settings of learning rate and QM loss
coefficient are Pareto suboptimal, while choosing among the remaining
combinations causes one to move along the frontier. Incorporating additional
training observables, particularly those which are dynamical in nature (e.g
velocity autocorrelation functions), could counteract the degeneracy and push
the Pareto frontier outwards.

Finally, we note that the observed energy MAE increase on aspirin for our
chosen combination of learning rate ($\alpha=0.001$) and loss coefficient
($\lambda=10$) is from 0.87 to 1.4 $\mathrm{kcal\,mol}^{-1}$, while DFT error
for energies on MD17 can be as high as 2.3 $\mathrm{kcal\,mol}^{-1}$ [20].
Thus, some of the error in the NNIP predictions could be attributable to
inaccuracies in the underlying DFT data.

###  A.11 Effect of Simulation Timestep on Stability

To demonstrate that the instability problem in NNIPs is not simply due to an
inappropriate choice of simulation timestep, we perform simulations at various
timesteps for each pair of system and model (Figure 10).

Figure 10: Effect of timestep on simulation stability. Reducing the timestep
does not reliably lead to an improvement in stability, and can sometimes
worsen stability.

We observe that stability does not necessarily improve as the timestep is
reduced. In fact, for some systems, stability seems to consistently decrease
as the timestep is reduced. Similar behavior has been observed in neural
network based solvers for ordinary differential equations [38].

###  A.12 Velocity Autocorrelation Function of Aspirin

We show the aspirin velocity autocorrelation function (VACF) corresponding to
the trajectory with the median stability improvement between conventional and
StABlE Training. A StABlE-trained model produces a similar VACF relative to a
conventionally trained model. This suggests that the StABlE procedure does not
significantly interfere with dynamic properties of the simulation, despite
only training with a structural observable ($h(r)$).

Figure 11: Velocity autocorrelation function of aspirin.