Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators
Authors: Sanjeev Raja
Ishan Amin
Fabian Pedregosa
Aditi S. Krishnapriyan
Venue: Accepted by TMLR ()

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

Published in Transactions on Machine Learning Research (02/2025)
Stability-Aware Training of Machine Learning Force Fields with Differentiable Boltzmann Estimators
Sanjeev Raja
sanjeeor@berrelle..dd edu
Deportment of Computer Science, UC Berkeley
Ishan Amin
Deportment of Computer Science, Deportment of Physics, UC Berkeley
Fabian Pedregosa
pedregosa.@gnngle.com
Google Deepmind
Aditi Krishnapriyan
aditik:1 @berkeleyy edu
Deportment of Computer Science and Department of Chemioal Engineering. UC Berkeley; LBNL
Reviewed on OpenReview: https: // optpr/veeereeeeeeee net/ forum? id= ZcKLMG0000
net/forum?
Abstract
Machine learning force fields (MLFFs) are an attractive alternative to ab-initio methods for molecular dynamics (MD) simulations. However, they can produce unstable simulations, limiting their ability model phenomena occurring over longer timescales and compromising the quality of estimated observables. To address these challenges, we present Stability-Aaare Boltzmann Estimator (StABIE) Training, multi-modal training procedure which leverages joint supervision from reference quantun-mechanical calculations and system observables. StABIE Training iteratively runs many MD simulations in parallel to seek out unstable regions, and corrects the instabilities via supervision with a reference observable. We achieve efficient end-to-end automatic differentiation through MD simulations using our Boltzmann Estimator, a generalization of implicit differentiation techniques to a broader class of stochastic algorithms. Unlike existing techniques based on active learning, our approach requires additional ab-inilio energy and forces calculations to correct instabilities. We demonstrate molecules, tetrapeptides, and condensed phase systems, using three modern MLFF architccturos. StABIE-traines models achieve significant improvements in simulation stability, data efficiency, and agreement, with reference observables. Crucially, the stability improvements cannot be matched by simply reducing the simulation timestep, meaning that StABIE Training effectively allows for larger timesteps. By incorporating observables into the training process alongside fist--priccppes calculations, SLABIE Training can be viewed general semi-empirical framework applicable across MLFF architectures and systems. This makes it powerful tool for training stable and accurate MLFFs, particularly in the absence of large reference datasets. Our code is publicly available at https://githun com/ cww/ASK--e///// ASAAASK-BBBt SttLe-eeeeee///ABBEEE----- -ttp://tttB--T-riiinn
methodology across organic
1 Introduction
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 (Frenkel be obtained on-the-fly via. ab-inétio Srantum-mechaniee (QM) calculations (Car & Parrinello 1985), this is prohibitively expensive for realistic length and time scales (Friesner 2005). Machine learning force fields (MLFFs) have recently emerged as a promising option to serve as surrogate models for QM calculations, demonstrating the ability to capture complex many- many-body interactions, and in some cases transfer flexibly across chemical space (Schütt ct al. 2018 Hu et 2021 Liu al.
& Smit 2001 While the atomic forces needed for MD simulation
2022
Gasteiger al. 2021 2020 Batzner et al. 2022 Musaclian
2022 Batatia al. 2022 Schitt et al.
2021
Graph neural network (GNN)-basee MLFFF FFs trained on large ab-imilio datasets increasingly being used

---

Published in Transactions on Machine Learning Research (02/2025)
(a)
(b)
Stable
Unstable
U(x)
U(x)
Learned PES
True PES
Learned Trajectory
True Trajectory
(c)
Stable
Unstable
Figure 1: Machine learning force field (MLFF) failure modes. (a) Illustrative examples of true and learned potential energy surfaces (PES) and resulting dynamics for unstable MLFFs. MLFFs can be accurate in approximat- ing much of the PES, but contain regions where energy and forces estimates deviate significantly from the true PES, leading to sampling of highly unphysical regimes. As a result, observables computed from MD simulation may be biased by the oversampling of unphysical states, or have high statistical error in the extreme case of unrecoverable simulation collapse. (b) Examples of stable versus unstable configurations sampled by MLFFs during molecular dy- namics simulattio systems considered in this work. (c) Selected states from an unstable MD trajectory of aspirin.
to model challenging and important chemical systems with favorable results (Merchant et al.
2023
Kozinsky et al,
2023
Chen & Ong 2022, Schaarschmidt et al. 2022 Takamoto et al. [0022] Majewski ct al. 2023
Charron et al.
2023
Batatia
2023
Kovács et al. 2023, Shoghi al. 2023 Gasteiger al. 2022 Deng al. 2023
MD simulations aim to accuratoly estimate system obscrvables like the radial distribution function, virial stress tensor, and diffusivity coefficient. This often requires long simulation to fully explore the underlying PES Unfortunately, MLFFss are known to produce unstable simulations, meaning that they can irreversibly enter umphysical regions
of phase space (e.g., a bond breaking event. in non-reactive system at low temperature) (Fu et al. 2022 Stocker
ct al. 2022 Vita & Schwalbe- such regions can lead to inaccuracies in computed observables as the MLFF gradually drifts from the distribution
2023 Bihani et al. 2023 Morrow et al. 2023 Wang et al. 2023b Sampling of
of its training data. In extreme cases, instabilities can lead to unrecoverable simulation collapse, in which case computed obscrvables may have high statistical error due to insufficient sampling. In cither case, instability can limit the ability of MLFF-based MD simulations investigale long-timescale phenomena like ion diffusion and protein folding, as well as rare events that may require extensive sampling encounter. Figure 1 illustrates typical MLFF instability behaviors and provides examples of unstable structures sampled during MD simulation.
Simulation instability has been shown to have an unreliable correlation with the energy and force error metrics
typically used to train and evaluate ML.FFs (uu et al. 2022 Stocker et al. 2022 Bihani et al. 2023 Recent works have introduced alternative (Ge et al.] 2024) energy and force errors, to achieve better downstream stability. However, these approaches do not allow the MLFF to visit new configurations, limiting the distribution over which it is trained and thus the potential improvement. Alternatively, expanding the phase space coverage of the datasot can be an effective way to combat MLFF instability (Fu et al. 2022 Stocker al. 2022 Vita & Schwalbe-
innulation-based objectives, such as localized (Wang et al. 2023b) or reweighted
2023 Bihani et al. 2023 This typically accomplished via active learning (Smith et al. 2018
Morrow et al. 2023 Wang et al. 2023b).
Vandermause et al. 2020 Schran et al. 2020 Lin et al. 2021 Kulichenko et al. 2023 approaches, where new
atomistic configurations are selected, ab-initio QM calculations are performed, and the MLFF is retrained on the
expanded dataset. However, these techniques rely on performing additional ab-initio calculations to expand the
dataset. With MLFFs being trained on increasingly diverse datasets (Shoghi et al. 2023 Batatia et al. 2023 Kovács et al. (022)) and larger atomistic systems, the expense of these calculations may hamper the practicality of active learning workflows. This suggests the need for additional sources of information beyond energies and forces to train stable ML.FFs.
2

---

Published in Transactions on Machine Learning Research (02/2025)
In this work, we bridge this gap by using both system observables and ab-initio QM data to improve MLFF stability, We introduce Stability-Aware Boltzmann Estimator BIE) Training, woa procedure designed to produce MLFFs that are both stable and accurate. The core idea behind StABlE is to which uses efficient, parallelized MD sim- ulations to rapidly explore regions of molccular phase space where the MLFF 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 enables end-to-end gradient- based learning without backpropagating directly through MD simulations. We also introduce a localized version of the Boltzmann Estimator, which enables targeted refinement of local instabilities. This is particularly important for stabilizing simmlations of large, condensed- phase systems. StABIE 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.
We demonstrate StABIE Training on three systems and MLFF architectures: 1) simmlation of aspirin with SchNet (Schütt ct al. 2018), 2) simulation of the Ac-Ala3NHMe tetrapeptide with NoqulP (Batzner et al. 2022), and 3) sinulation of an all-atom water system with GemNet-T (Gasteiger et al. 2021). Relative to MLF FI's trained solely on energies and forces, our StABiE-trained models produce significantly more stable MD simulations, recover observables more accurately, exhibit better generalization to unseen simulation temperatures, and outperform models trained on 50 times larger, labeled datascts. Our results suggest that utilizing both quantum-mechaniec and observable-based modalities is required to fully exploit. the available learning signal in reference datasets and train stable and accurate MLFFs. To our knowledge, StABIE Training is the first method that combines these data modalities to improve the stability of neural network potentials in MD simulations.
2 Preliminaries
Molecular Dynamics. Molecular dynamics simulation is used to evolve the positions and momenta of atomistic system. Given system of N atoms, its state at time t is defined by r(t) = (r(t),p(t)), where r(t),p(t) E Nx3 are the position and momenta of the atoms. We assume the system has a scalar-valued
Hamiltonian function H: Nx3 x Nx3 +R of the form H(T)= = E 2m(1) Ur) where U: Nx3 - R is potential energy function and m)) and p(t) are the mass and momentum of atom i. By updating in the direction of the per-atom forces 0r using numerical integration scheme such as Langevin dynamics (Bussi & Parrinello 2007), a sequence of K simulation states {T()){< is produced.
av
Machine Learning Force Fields. A MLFF is a function approximator which learns a potential energy Uo (())) from QM reference data. MLFFs are trained to minimize the following
and forces Fg = - V (_.,.("))))))) regression loss, with supervision from a dataset of reference energies and forces Derain = (()
N LQM= - DI
(1)
N
System Observables. Observables, g(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 distri- bution functions, velocity autocorrelation function, and diffusivity coefficient. Observables can be computed conve- niently 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. More details on the observables used in this work can be found in Supplementary Section A.4
Training MLFFss with Observables. Observables have been used extensively in the historical development of classical potentials (Cornell et al. 1995 Marrink et al) 2007 Li et al. 2018 and more recently are gaining traction
complementary data source for training MLFFs (Wang et al. 2020 Thaler & Zavadlav 2021 Fuchs et al. 2025).
Observables have boen used to train MLFFs in the context of condensed-phase (Wang et al. 2023a and titanium
Röcken & Zavadlav 2023) systems, enhanced sampling of rare events (Sipka et al. 2023). and protein folding
simulations (Ingraham et al. 2019 Navaro et al. 2023 Kolloff & Olsson| 2023). Training with observables requires an efficient way to compute gradients through MD simulations while avoiding numerical instability and memory limitations. The approach is appealing due to the lack of reliance on expensive ab-initio quantum mechanical calcula- tions, and the possibility of improved empirical consistency in settings where the underlying ab-inaitoo method may be
3

---

Published in Transactions on Machine Learning Research (02/2025)
Step Quantum Mechanical Pre-Training
QM Dataset
Pre-Trained NNIP
n
LQM Σ Au 10, - U//2
Ui. F
00%
+ Arll
111
Dr
Step 2: StABIE Training
Paraleel Replicas
Upper instabilit Threshold
Simulate
Fraction Unstable
Unstable
Lower instability Threshald
Find Unstable
Correct Instability
Simuiaiio Time
Regions
Simuliation Time
Upper Instability Threshold
Leam
Fraction Unstable
Unstable
Lobs = l1
9ref
Lower Instabilly Threshold
Simulaatin Time
Simulition Time
Observable 9ref
Figure 2: Schematic of Stability-Aware Boltzmann Estimator (StABIE) Training procedure. Our proposed StABIE Training procedure begins with conventional pre-training on a small reference dataset of QM calculations. This datasct remains fixed throughout the procodure, and is never expanded with new calculations. Upon convergence of pre-training, SLABIE alternates between two main phases, simulaaion and learning. In the simulation phase, we perform many molecular dynamics simulations in parallel with the MLFF and find regions of instability. When sufficient fraction of simulations become unstable, enter the learning phase, where the MLFF is further trained to match known system observables from ab-initio simulation experiment. Gradients are computed efficiently through the MD simulation via our Boltzmann Estimator. After sufficient reduction in the portion of unstable trajectories, we re-enter the simulation phase, and repeat the training cycle until a predetermined computational budget reached.
unreliable (Cheetham & Seshadri 2024 However, past works do not consider MLFFs with stability problems, and dispense entirely with using QM data, thus losing the first-principles guarantocs of matching an ab-initio PES.
3 Methods
We present StABIE Training, our proposod procedure to train stable and accurate MLFFs by leveraging both reference quantuu mechanical data and system observables. Our approach is enabled by Boltzmann Estimators, which allows gradient-based optimization of MLFFs based on system observables.
3.1 Boltzmann Estimator
Reference system observables can be estimated from high-fidelity MD simulations or experimental measurements. Since observables are linked to the MLFF parameters through a molecular simulation, training with this source of information requires a reliable way to optimize through MD trajectories.
Formally, we define g(T)) be a vector-valued observable of state r, and gref to be the reference value of the observable. To train our MLFF to match reference observable, we minimize the following loss function:
Lobe (0)
E (TTTorr))(TT)) gretl
(2)
4

---

Published in Transactions on Machine Learning Research (02/2025)
where (T) is the equilibrium distribution induced by the MLFF Ug. In this work, we primarily consider systems with a fixed volume, temperature, and number of particles, corresponding to the canonical (NVT) ensemble. We note, however, that the estimator is readily applicable to other ensembles, including the isothermal/isolee /isobaiic (NPT) and grand canonical (nVT) ensembles (sce Supplementary Sections A.2 and A.10 for details). In the canonical ensemble, microstates obey a Boltzmann distribution, Px(T)==epp( k Ho(()) ((())/CCC)))))))))))))))) where Ho(L)) is the Hamiltonian, is the sampling temperature, kB is Boltzmann's 's constant, and C(+) -(0)=foxp(-- 1 (TT is the normalizing partition function. MD simulation required at each training ileration to sample from this distribution and estimate the Joss.
Kg
Optimizing the MLFF requires computing gradients of Equation with respect to 8. The required gradient, VoLobe: can be decomposod via the chain rule as follows:
OLobs
OE Barys(
P (TT)
Vo
-2(Er~ra(r) (g())))(((])))))r))) -ree))
(3)
00
The non-trivial quantity to compute is the Jacobian, @Er-Pa(T) [F) One way to estimate it is by using a chain rule expansion that corresponds to each step of the unrolled MD simmlation (Wang et al. 2020). The adjoint mothod (Chen et al. 2019 can be used to limit the memory footprint at the expense of increasod computation,
but can still lead lo numerical instability for long trajectories (Wang et al. 2023a Sípka. al. 2023 (see
Supplementary Section A.13 for an empirical comparison with direct and adjoint- based backpropagation).
We can avoid direct backpropagation through the simulation by noting that the equilibrium state distribution (T) is independent of the algorithm (i.e., MD integrator) used to sample the distribution. This is analogous to implicit differentiation techniques for differentiable optimization (Amos & Kolter 2017 Gould et al. 2016 Ren
el al. 2022 Blondel el al. 2022 Négiar et al.) 2023). in which the solution to oplimization problem is decoupled from the numerical solver used to obtain it. We leverage the known Boltzmann form of Po((') to construct an
unbiasod estimator of the Jacobian.
Definition 1 (N-sample Boltzmann estimator). Given N independent samples I1,...,I'N from a Boltzmann distribution (T',, we define the N-sample Boltzmann estimator E(II....['NN of the Jacobian (Er- as,
KnT(N- [Bo(()) ( -Elo]]-)-[[[[[]-00000 E(N)
N
(4)
where B(((T))= E (T+)) denotes an empérical meen over the samples.
Po A proof is provided in Supplementary 210
This estimator provides an unbiased estimate the Jacobian Section A.1 The estimator related to the REINFORCE trick Williams
1992] and policy gradient estimators from reinforcement learning (Silver et al., 2014) when Boltzmann state distribution is assumed. The result can
also be derived using thermodynamic perturbation theory ZZwanzig
1954 Thaler & Zavadlav
2021
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 large condensed- phase system. Due to spatial averaging, global observables g(T)) may be insensitive these localized events, limiting the ability to identify unphysical states. To address this, propose the N N-sample Localized Boltzmann Estimator. Here, the global energy UT) and obscrvable g(T) are replaced with local versions Ug(y) and g(7), where denotes a local neighborhood of n < N atoms. Formally, define a neighborhood of n aloms W N= VV==(1,00000 (21,,222... a E EZ,1 <a; S N for all i= I.). The local state y is defined as y=
The local energy is then defined as Un( - - ( ())), where y()) = [NN))),,pNN))] contains the position and momenta of the 7th atom in the local neighborhood. The local energy is easily obtained by noting that MLFFs parameterize their global energy prediction Up as sum individual atomic energies. The localized estimator is thus given follows:
5

---

Published in Transactions on Machine Learning Research (02/2025)
Definition 2 (N-sample localized Boltzmann estimator). Given N ii.d. samples of local states "n n,.,,, where each 7/s is extracted from a global state T ~ Po(L'), we define the N N-sample localized estimator of the Jacobian Putssstt)) 80
as
def
KB [lo( -B -B(0 ]]
N
E(A
The localized estimator follows from the original Boltzmann estimator due to the fact that as a subset of the larger Boltzmann-
system, any local neighborhood also obeys a battaman--dsstibuttaan we extract multiple local neighborhoods from each global state to increase the sample size and state space coverage (Section A. 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 44
In practice,
Key Advantages. Unlike active learning approaches, additional quantum mochanical energy and forces calculations are required to compule our Boltanann Estimators (all energy terms are model predictions). As a result, learning with the Boltzmann estimator is computationally efficient. Due to the use of independent samples, the estimator also avoids numerical instability and memory demands associated with differentiating through long, contimous trajectorics. Further, the most compatationally expensive component of the estimator, the gradient of the potential energy VTUg(T), is independent of the observable g(T). This means it can be reused, allowing efficient training to match multiple observables simultaneously.
3.2 Stability-Aware Boltzmann Estimator (StABIE) Training
Stability-Aware Boltzmann Estimator (StABIE) Training begins with conventional supervised pre-training of a MLFF on a reference dataset of energy and forces, using the loss defined in Equation [ 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 MLFF bocomes unstable. We sample R equilibrium states from the training dataset as initial conditions for separate MD trajectories (replicas). Using the pre-trained MLFF Us, we run MD simulations on these replicas in parallel for t timesteps. Due to our use of efficient vectorized GPU operations, simulating many replicas in parallel has similar computational cost to simulating a single replica. We apply predetermined stability criterion (Supplementary Section A.5) to each replica, freezing those marked unstable at their current states. Simulations contimue 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 timesteps will trigger instability, which is then corrected in the next phase. We note that many other strategies exist to explore phase space and find unstable regions, including Diffusion Monte Carlo with fictitious masses (Li et al. 2 or uncertainty-based sampling (Kulichenko et al. 2023 StABIE Training can be flexibly used with any of these stralegies without impacting the learning stage (described next) of the procedure.
Learning Phase. The learning phase aims to refine the MLFF to correct the previously encountered instabilities. Starting from the near-unstable configurations obtained in the simulation phase, we perform MD simulation for timesteps, sampling every 5th state to obtain uncorrelated samples. Using the sampled states Ti,..,,UN,, where ; we compute the observable loss function (Equation 2. To update the MLFF parameters 0, we compute an unbiased estimator of the loss gradient via the Boltzmann Estimator (Soction 3.1 and use to perform a single step of gradient descent. We then reset all replicas to their original near-unstabe states, simulate with the updated MLFF, recompute the loss and gradient estimator with the newly sampled states, and again update the MLFF weights. This process is repeated until the fraction of unstable replicas drops below 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 contimue alternating between simulation and learning phases until a predetermined computational budget is reached, at which point StABIE Training has concluded. Soc Supplementary Section A.3 for a formal algorithmic description of the SLABIE Training procedure.
Na= 18
Regularizing StABIE 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 (Noid 2013). Consequently, learning with observables alone is raieed. To combat this, we regularize the observable loss function (Equation
6

---

Published in Transactions on Machine Learning Research (02/2025)
(b)
(c)
Number Simu atiote
80
0.61
-075
Lewning Bate
2020
60-
MP)
0.50
40
0.2
02-
0.0
100
200
300
400
500
1.000
0-0
200
400
1000
0.01
10
Training Epoch
(oo)
(Angstromsl
(d)
(e)
Number Simutssions
60-
40--
20-
2
0-
0.05
0.10
0.15
0.26
VACF MAE
Unstable
Stable
Reference
QM Training
(1K) Training
(10K) Training
(50k) Training
Figure 3: Aspirin simulation with StABIE Training. (a) Alternation of simulation and learning phascs during StABIE Training with 128 parallel replicas. The simulation phases correspond to regions where the fraction of unstable replicas increases, while the learning phases correspond to regions where learning occurs and the fraction of unstable replicas decreascs. (b) Stable simulation time of 256 parallel aspirin trajectorics from SchNet MLFFs. Applying StABIE Training yields significantly more stable simulations than models trained only on energies and forces, surpassing models trained on 50x more QM data. (c) Distribution of interatomic distances (h(r)) from MLFF simulations. A StAiBlE-trained SchNet model closely recovers the true distribution of interatomic distances, while the model trained only on QM reference data produces 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). StABIE Training yields a reduction in variance across replicas. (e) Aspirin structures sampled over epochs of a single learning phase of StABIE Training. There is a clear progression as unstable configurations become stable.
2 with the energy and forces loss function (Equation The final StABIE loss function thus becomes,
LSAABIE (0) def Lobs +ACQM>
(6)
where Lobe and LQm were defined in Equations 2 and respectively, and À controls the strength of the regularization. Crucially, LQM is only computed over the original training dataset Dtrain, and not over new structures explored during MD simulation. Therefore, the regularization requires no additional ab-initio calculations.
4 Results
We present the results of StABlE Training on the aspirin molecule with SchNet (Section 4.1, Al-Ala3-NHMe tetrapeptide with NequIP (Section 4.3) and an all-atom water system with GemNet-T (Section 4.4
4.1
Aspirin Molecule
Aspirin (chemical formula CgH5O4) is the largest molecule from MD17 (Chmiela et al. 2017), a widely used benchmark dataset for atomistic simulations which contains energy and forces calculations computed at the PBE+vdW-TS (Tkatchenko & Scheffler. 2209)) level of theory. Consisting of 21 atoms, aspirin has been shown to be the most challenging molecule in MD17 for state-of-the-art MLFF's simulate stably (Fu et al. 2022). We pre-train SchNet (Schütt et al.) 2018) model on the energy and forces matching objective (Equation 1 using subset of 1,000 aspirin structures from the reference dataset. After convergence, we begin StABIE 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 MLFF to a comprehensive range of failure modes much more quickly than possible with a single replica. Following (Fu et al. 2022), we use a maximum bond length deviation criterion to detoct instability in the simulations (Soction A.5). We use the distribution of interatomic distanoes, h(r), as our training observable, with the h()) computed over structures in the training dataset serving the reference (Section A.4h. This choice is motivated by the observation that unphysical bond stretches constitute the majority of failure
7

---

Published in Transactions on Machine Learning Research (02/2025)
350 K
(b)
700
Fraction of Unstable Beplicas
0.8
Unicable Replicas
1.00
0.6
0.75
0.50
0.25
0.0
0.00
200
400
600
800
1000
0
50
100
150
200
(c)
Simulation Time (ps)
(d)
Simulation Time (ps)
0.4-
0.4
MF)
Me)
0.2
0.2
0.0
0.0
10
10
(Angstroms)
(Angstroms) QM (1k) StABIE Training
AIMD Reference
QM (1k) Training
Figure 4: Testing temperature generalization for aspirin. (a, b) Fraction of unstable replicas function of simulation time for SchNet MLFFs at 350K and 700K. Applying StABIE Training yields significantly more stable simulations than a bascline SchNet model. (c, d) Distribution of interatomic distances (h(r)) from SchNet MLFFs. The StABIE-Lrained SchNet model more accurately recovers the true h(r) at both 350K and 700K (the inset shows difference between the predicted and reference h(r)). The AIMD reference h(r) is computed by Boltzinant-ewweighting of independent samples from the 500K training dataset.
cases in aspirin simulations. The h(r) is sensitive to these abmormalities, and is thus an informative optimization criterion. We perform four complete StABIE cycles simulation and learming with a SGD optimizer. We evaluate our final models by performing nanosecond -entaant-eemperaturre MD simulations, starting from 256 initial structures not seen during training. Figure 3h shows the progression of StABIE 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 Sch model trained via our stability-aware approach 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 median stable simulation time increases from 42 602 picoseconds, Even when StABIE training is employed after training 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 in Supplementary Section A.12 that simply reducing the simulation timestep by a factor of ten does not eliminate instability in sinulations produced by conventionally trained models, and reducing instability further may require impractically small timesteps. In Supplementary Section A.12 we demonstrate the potential StABIE Training to accelerate simulations by enabling larger timesteps. Simulations from our model closely recover the true distribution of interatomic distances, while the distributions produced by models solely trained QM data noisier, due to the limited stable simulation time (Figure 3+). We also test the ability of our StAiBlEtrainnd model to recover the velocity autocorrelation function (VACF), fundamental dynamical obscrvable not seen during training (Figure 31). We see that the quality of the recovered VACF remains similar after SLABIE Training, with lower variance over the sinulated replicas (see Supplementary Section A.14 for a sample VACF computed from the simulations). The preservation a held-out dynamical obscrvable suggests that StABIE Training does not achieve stability improvements by simply restricting the model around a narrow regime corresponding the reference h((r). If this were the case, characteristic dynamic fluct around the reference h(r) would be suppressed, leading to an inaccurate VACF. Figure 3 depicts
8

---

Published in Transactions on Machine Learning Research (02/2025)
(a)
(b)
0.8
0.6
0.2
Facton Unstabii
0.4
0.1
02
0.0
0.0
50
100
150
200
250
0.0
2.5
10.0
12.5
15.0
Simulation Time (ps)
(Angstroms))
(c)
AIMD Reference
QM Training
QM StABIE Training
Unstable
Stable
Figure 5: Ac-Ala3-NHMe tetrapeptide simulation with StABIE Training. Fraction of unstable replicas as function of simulation time for NequlP MLFFs. Applying StABIE Training yields model which can simulate more MD replicas stably over time than baseline trained only energies and forces. (b) Distribution of interatoric distances (h(r)) from NequIP MLFF simulations. The StABiE-trained NequlP 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 StABIE Training. There is a clear progression as very unstable configurations becomn stable.
aspirin structures sampled from sinulations during the course of a single learning phase of StABIE Training, demonstrating that unphysical bond strchinngg in these structures is resolved by the training procedure.
4.2 Temperature generalization
We explore our method's generalization to different thermodynamic conditions. To do so, we perform 256 parallel MD sinulations at 350K and 700K using the SchNet model which was trained with our StABIE Training procedure at 500K. We use the same criterion based on bond length deviation to detect instability as in Section 4.1 We estimate the reference h(r) at 350K and 700K by Roltamann-veweighinn samples from the original 500K training dataset (Section A.9). Aspirin simulations produced by a StABIE-Trained Sch model are significantly more stable at both temperatures than baseline models trained only on QM reference data (Figure 4, b). At. 350K, the difference is particularly large. After nanosecond of simulation, virtually no replicas unstable in the StABIE- trained MLFF simulation, while approximately 80% of replicas are unstable in the bascline MLFF simulation. The SLABIE-trained model also recovers the true distribution of interatomic distances more closely than the baseline MLFF at both temperatures (Figure + d). At 350K, the baseline MLFF stable enough produce a smooth A(r), but the structure inaccurate, indicating incorrect sampling of the phase space. At 700K, the baseline MLFF produces a h(r) which is close the AIMD-produech distribution, but is noisy due insufficient stable simulation time. This underscores the necessity for a MLFF 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 thermodynamie conditions is of crucial importance for MLFFs to be useful in practical applications (Kovács et al. of StABIE Training as an important strength of our method.
2021,, we see the
4.3
Tetrapeptide
Ala3 NHMe (chemical formula O1))Hg2NN,Oi is tetrapeptide from the MD22 dataset CChmiela et al. 2022 challenging benchmark consisting of reference energy and forces computed at the PBE+MBD (Perdew et al. Tkaichenko et al.) 2012) level of theory for considerably larger molecules Ihan those in the MD17 dataset. Relative to MD17 molecules, Ac-Ala3-NHMM poses unique challenges for atomistic simulations due to its larger size and ficxibility (Kabylda al. 2023). As a result, more expressive STLFFF-sommtin incorporating E(3) cquivariance (Geiger & Smidt] 2022 are required to accurately model the underlying potential energy surface.
1996
our experiments, use NequlP Batzzer et. al. 2022) model due to its promising accuracy and data efficiency

---

Published in Transactions on Machine Learning Research (02/2025)
(a)
(b)
(c)
Fection of limtabl Replect stltet
0.8
Diffusiiity (10 2m²
0.6
6x
0.4
0.
0.0
2x12
Referance
800
1000
5óo
1000
(d)
(e)
Simu ation Time (pst
Simulation Time (p5)
HO vl2
00 y62
gi9)
0
(Angstroms)
(Angaroma)
(Angalromm
Reference Simulation
QM (1k) Training
(1K) SLABIE Training
Figure 6: All-Atom Water Simulation with StABIE Training. (a) Example unstable configuration arising in GemNet-T simulation of water. Instabilities are highly localized and manifest unphysical bond stretches or intermolecular collisions, motivating the use of our localized Boltzmann estimator (Section 3.1 during training (b) Fraction of unstable replicas function of simulation time for GemNet-T MLFFs. SLABIE Training yields a significant stability improvement relative to training on only energy and forces data. (c) Convergence of the diffusivity coefficient over MLFF simulations. The StABlE-trained model gives a considerably more accurate final estimate of the diffusivity coefficient. (d- f) Filement-condditionee (OH, O0, HH) radial distribution functions (RDF) produced by MLFF simulations. The StABIE-trainod model more accurately captures long-range interactions.
Merchant et al. 2023 We pre-train NequlP model the QM energy and forces matching objective using 14,890 (or 25%) of available structures from the training datasct. After convergence, we begin StABIE Training with the global Boltzmann estimator, using 128 parallel replicas al a lemperature of 500K. As with aspirin (Sections 4.1 4.2 we delect instability using the maximum bond length deviation criterion, and choose the distribution of interatomic distances, h(r), as our training observable. The StABBE-trained NequlP model gives an improvement in stability relative to a baseline model trained only on QM reference data (Figure 5p). Additionally, the -traiid model closely recovers the true h(r), while the baseline MLFF produces a highly noisy h(()) This is because of insufficient sampling, due lower stable simulation time (Figure pp). The baseline MLFF also underestinates the latter portion of the h(r) distribution, indicating poor modelling of long-range interactions. Figure 5p shows selected structures throughout the course of single learning phase of StABIE Training, demonstrating a progression from very unstable to stable structures.
challenging atomistic simulation tasks (Kozinsky et al. 2023
4.4 Water
As final evaluation of StABIE Training, consider liquid-phase water. Water has historically posed unique challenges for atomistic innulations due to the presence of long-range interactions, proton disorder, and nuclear quantum effects (Cheng et al. 2019 Raza et al. transport properties, such as the diffusivity coefficient, requires long, continuous trajectories to minimize statistical error. We use the dataset produced in Fu et al. (Fu et al.
2011 Markland & Ceriotti 2018). Additionally, estimating
2022 which contains reference all-atom simulations of 64
water molecules using the flexible Extended Simple Point Charge model (Wu al. 2006) 300K and atm. For this system, we consider a GemNet-T model (Gasteiger et al. 2021). GemnNetTT displays failure mode in which umphysical configurations (e.g., bond stretching) first arise in highly localized 'pockets" of 1-2 molecules (Figure G which then gradually cascade to the rest of the simulation domain (Fu et al. 2022 Such configurations do not noticcably affect global system observables such as the radial distribution function. Therefore, in this sctting it is more appropriate to use our localized version of the Boltzmann gradient estimator (Equation 5 in order to guide the optimization process towards specific spatial domains of instability, After pre- pre-training a GemNet-TT model using subsct of 1,000 reference structures, we begin StABIE Training with the localized Boltzmann estimator, simulating parallel replicas at a temperature of 300K. We use minimum intermolecular distance criterion (Section A.5 to detect instabilities during training. To estimate gradients for optimization, from each global state we extract all
10

---

Published in Transactions on Machine Learning Research (02/2025)
possible local neighborhoods containing a single water molecule (Supplementary Section A.6 The most prevalent failure mode produced by GemNet-T in this system is umphysical bond stretching, which eventually leads to umphysical coordination structures. 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. We then evaluate the performance of the ALA Trained model by performing parallel nanosecond (ns) MD simulations at 300K, starting from held-out, initial conditions and using a RDF MAE critcrion (defined in Supplementary Section A.5 to measure stability. We achieve significant improvements in stability and accuracy using our training approach. Our StABIE-trained model can simulate stably for 1 ns for all 5 initial conditions, while the baseline MLFF model trained only on QM reference
data can only do so for 1 initial condition, and has a median stability of just 312 ps (Figure G p)) Compared to a
reference value of 2 10 m the mean diffusivity produced by StABiE-trained model simmlations is 2.4 10 10-90m 9m"
while the mean diffusivity produced by baseline model simulations is 5.0 10 10-9 (Figure (C). The long-range correlations in the ceement-conditions RDFs (Figures Bi-f) are also captured more accurately by the StABIE- trained model. We reiterate that neither the element-condiitonee RDF's nor the diffuusivtty coefficient were explicitly seen during StABIE Training. We also highlight that the RDF MAE stability critcrion used for evaluation is different from the minimum intermolecular distance criterion used during training (see Supplementary Section A.5 for more details), suggesting that StABIE Training does not overfit to the training stability criterion. In Supplementary Section A0 we demonstrate that StABIE Training leads to similar stability improvements when water simulations are conducted in the nothermal/isobaric (NPT) ensemble, instead of the canonical (NVT) ensemble.
5 Conclusion and Future Work
We have introduced StABIE Training, strategy for training stable and accurate machine learning force fields. Our training procedure results in MLFFs which are significantly more stable in MD simulation, more accurately repro- duce key simulation observables, (including those that were not explicitly trained on), exhibit better generalization to unseen temperatures, and have superior data, efficiency relative to ML.FFs trained only on QM data.
Key Takeaways. StABIE Training can utilize both maattum-eehhanee
energies and forces and system observables to supervise MLFF training. We highlight that the reference observables need not be acquired from experimental measurements, and can be computed by averaging over existing QM datasets. This makes our approach applicable in realistic computational discovery scenarios in which experimental characterization is not always available for hypothetical systems. The stability, accuracy, and data efficiency gains brought. by SLABIE Training require no additional reference calculations or data. The procodure is thus self-contained and efficient, requiring minimal additional computational expense beyond a single iteration of conventional MLFF training (see Supplementary Section A.6 for details on training times). This suggests that StABIE Training may scale more gracefully to larger MLFFs and atomistic systems than existing active learning approaches, which require repeated QM reference calculations and retraining. StABIE 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 due to our localized Boltzmann estimator. This flexibility extends to the choice of MLFF: the offoctiveness of StABIE Training across the three diverse architectures considered in this work suggests that our approach will remain applicable as MLFF architectures continue to evolve.
Limitations. Unlike traditional QM learning, training MLFF's to match reference observables lacks convergence guarantees in the large data limit. Specifically, the abservable-matchingg objective is under-constraineed as the mapping between a potential energy function and a sparse set of simulation observables is in general non-unique (Noid 2013). Although we mitigate this problem by pre-training and regularizing StABIE Training with the conventional energy and forces loss function, we observe that stability improvements resulting from StABIE Training are accompanied by minor increase in the energy and forces error a held-out test set, indicating a misalignment between the observable- matchinn and QM objectives. Adjusting the strength of the QM regularization and the learning rate used for StABIE Training are two primary ways in which to navigate this tradcoff. We note, however, that our observed error increases are lypically within the range of DFT error (Supplementary Section A. We also note that StABIE Training is currently incompatible with dynamical observables, due to the use of uncorrelated states to compute the Boltzmann Estimator. Overcoming this limitation nontrivial as it requires optimization over long paths, but may become tractable with recent advances in path gradient computation
(Bolhuis et. al. 2023 Greener 2024
Han & Yu 2025
11

---

Published in Transactions on Machine Learning Research (02/2025)
Future Work. In this work, we took observables calculated from high-fidelity simmlations 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 simultancously to yield more generalizable and robust polentials. Future work could also explore incorporating additional observables into SLABIE Training, particularly those which are dynamical. This could restrict the set of learnable functions and address the under-constrained nature of learning with observables. We also note that the reference datasets used to pre-train MLFFs in this work were sampled uniformly from high-fidelity simulations. More sophisticated sampling strategies (Deringer et al.
2018 Karabin & Perez 2020 Yoo et al. 2021 Fonseca et al. 2021 Kulichenko et al. 2023 Qi et al. 2024) for diverse dataset generation, as well as during phase space exploration to find unstable regions, could be employed along with SLABIE Training lo achieve further stability and accuracy gains. Closely related techniques such as active learning are compatible with StABIE Training, and could be used in tandem to further improve MLFF stability and robustness. In this sctting, system observables can serve as a cheap source of information with which to augmentt the more expensive supervision provided by ab-ônilãoo calculations.
Acknowledgments
This work supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Rescarch, Scientific Discovery through Advanced Computing (SciDAC) program under contract No. DE-AC02-05CH1123 and the U.S. Department of Energy, Office of Science, Energy Earthshot initiatives as part of the Center for Tonomer-based Water Flectrolysis at Lawrence Berkeley National Laboratory under Award Number NE-AA02-05CH11231 This research used resources the National Energy Research Scientific Computing Center (NERSC), a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley We thank Rasmus Lindrup, Toby Kreiman, Gooffrey Negiar, Muhammad Hosyim, Ritwik Gupta, Martin Sipka, Johannes Dictschreit, Aayush Singh, David Limuner, Kranthi Mandadapu, David Prendergast, Bryan McCloskey, and Muratahan Aykol for fruitful discussions and comments on the manuscript.
National Laboratory, operated under Contract No. DF-ACC02-05CH1122
References
Brandon Amos and J. Zico Kolter. Optnet: Differentiable optimization as a layer in neural networks. In Proceedings of the International Conference on Machine Learning (ICML), pp. 136-145, 2017.
llyes Batatia, David P. Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. In Neural Information Processing Systems (NewrIPS), 2022.
Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kovács, Janosh Riebesell, Xavier R. Advincula, Mark Asta, William J. Baldwin, Noam Bornstein, Arghya Bhowmik, Samuel M. Blau, Vlad Cärare, James P. Darby, Sandip De, Flaviano Della Pia, Volker L. Deringer, Rokas Elijoöius, Zakariya El-Machachi, Fdvin Fako, Andrea C. Ferrari, Annalena Genreith--Shrrieer, Janine George, Rhys F. A. Goodall, Clare P. Grey, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Christian Holm, Jad Jaafar, Stephan Hofmann, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron D. Kaplan, Nima Karimitari, Namu Kroupa, Jolla Kullgren, Matthew C. Kuner, Domantas Kuryla, Guoda Liepuoniute, Johannes T. Margraf, Toan-Bogdan Magdän, Angelos Michaclides, J. Harry Moore, Aakash A. Naik, Samuel P. Niblett, Sam Walton Norwood, Niamh O'Neill, Christoph Ortner, Kristin A. Persson, Karsten Reuter, Andrew S. Rosen, Lars L. Schaaf, Christoph Schran, Eric Sivonxay, Tamás K. Stenczel, Viktor Svahn, Christopher Sutton, Cas van der Oord, Esster VargacUm Tejs Vegge, Martin Vondrák, Yangshuai Wang, William C. Witt, Fabian Zills, and Gábor Csányi. A foundation model for atomistic materials chemistry, 2023.
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess F. Smidt, and Boris Kozinsky. F(3)-equivariaat graph neural networks for data-efficient, and accurate interatomic potentials. Nature Communications, 13(1), 5 2022. doi:
Vaibhav Bihani, Utkarsh Pratiush, Sajid Mannan, Tao Du, Zhimin Chen, Santiago Miret, Matthicu Micoulant, Morten M Smodskjaer, Sayan Ramu, and N M Anoop Krishnan. Egraffbench: Evaluation of equivariant graph neural network force fields for atomistic simulations, 2023.
12

---

Published in Transactions on Machine Learning Research (02/2025)
Mathieu Blondel, Quentin Berthet, Marco Cuturi, Roy Frostig, Stephan Hoyer, Felipe Llinares-Lpezz, Fabian Pedregosa, and Jean-Philippe Vert. Efficiient and modular implicit differentiation. Tn Neural Information Processing Systems (NewrIPS). Google Research, 2022.
Peter G. Bolhuis, Z. Faidon Brotzakis, and Bettina G. Keller. Optimizing molccular potential models by imposing kinetic constraints with path reweighting. J. Chem. Phys., 159:074102, 2023. doi: 10.1063/5.0151166. URL
https://doi. http:/.................. www.:///................... 0151166
Giovanni Bussi and Michele Parrinello. Accurate sampling using langevin dynamics. Physical Review E Statistioal, Nondinear, and Soft Maller Physics, 75(5)0056707, 2007.
R. Car and M. Parrinello. Unified approach for molecular dynamics and density-fuuctiinall theory. Phys. Rew. Lett. 55:2471, 1985. doi: 10.1103 /ttss//hyRee.....5........ Published 25 November 1985.
Scott P. Carmichaol and M. Scott Shell. new multiscale algorithm and application to coarso-grained peptide models for self-aseebly. J. Phys. Chem. B, 116((9):8383 8393, 2012. doi: 10.1021/1p2114994.
Nicholas E. Charron, Felix Musil, Andrea Guljas, Yaoyi Chen, Klara Bonneau, Aldo S. Pasos- Pasos-Trejo, Jacopo Venturin, Daria Gusew, Tryna Zaporozhets, Andreas Krämer, Clark Templeton, Atharva Kelkar, Aleksander E. P. Durumerie, Simon Olsson, Adrià Pérez, Maciej Majewski, Brooke E. Husic, Ankit Patel, Gianni De Fabritiis, Frank Noé, and Cecilia Clementi. Navigating protein landscapes with a machine-learned transferable coarse-grained model, 2023.
Anthony K Cheetham and Ram Seshadri. Artificial intelligence driving materials discovery? perspective on the article: Scaling deep learning for materials discovery. Chemistry of Materials, 36(8):3490 3495, 2024.
C. Chen and S.P. Ong. A universal graph deep learning interatomic potential for the periodic table. Nat Comput Sci, 2:718-722 728, 2022. doi: 10. (0.1038/s13588-002--00033 URL
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2019.
Bingging Cheng, Edgar A. Engel, Jörg Behler, and Michele Ceriolti. Ab initio thermodynamics of liquid and solid water. Proc. Natl. Acad. Sci. U.S.A., (16(():1100-11)) 1115, 2019. doi: 111773/pnas. 1815117116. URL hattps://dol. org/10. 10///....0 ww.11077ppaas http://paa..11111111.1 1815117116 Edited by Pablo G. Debenedetti, Princeton University, Princeton, NJ, approved December 3, 2018 (received for review September 4, 2018).
Stefan Chmicla, Alexandre Tkatchenko, Huziel Saucoda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Miller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 3(5), 5 2017. doi: 10.1126/sciadv.1100 16000115.
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T. Unke, Adil Kabylda, Huziel E. Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms, 2022.
Wendy D. Cornell, Piotr Cieplak, Christopher 1. Bayly, Ian R. Gould, Kenneth M. Merz, David M. Ferguson, David C. Spellmeyer, Thomas Fox, James W. Caldwell, and Peter A. Kollman. A sccond generation force field for the simulation of proteins, nucleic acids, and organic molecules. J. Am. Chem. Soc., 117:5179 5197, 1995.
Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J Bartel, and Gerbrand Coder. Chgnet as a pretrained universal noural network potential for charge-informed atomistic modelling. Nature Machine Intelligence, 5(9):1031 1041, 2023.
Volker L Deringer, Chris Pickard, and Gábor Csányi. Data-driven learning of total and local energies in elemental boron. Physieal review letters, 120(15):156001 2018.
FA Faber, L Hutchison, B Huang, Gilmer, SS Schoenholz, GE Dahl, O Vinyals, S Kearnes, PF Riley, and OA von Lilienfeld. Prediction errors of molccular machine learning models lower than hybrid dit error. J Chem Theory Comput, 13(1 (3(11):5255 5264, 11 2017. doi: 10.1021 /.111///cs.ccccc./000777. Epub 2017 Oct 10.
13

---

Published in Transactions on Machine Learning Research (02/2025)
Gregory Fonseca, Igor Poltavsky, Valentin Dassilev-Galinde and Alexandre Tkatchenko. Improving molecular force fields across configurational space by combining supervised and unsupervised machine learning. The Journal of Chemical Physics, 154(12), 2021.
Daan Frenkel and Berend Smit. Understanding Molecular Simulation: From Algorithoms to Applications. Academic Press, 2 edition, 11 2001.
Richard A. Friesner. Ab initio quantum chemistry: Methodology and applications. Academy of Sciences, 102(19):6648 6664-6633, 2005. doi: 10.1073/ppaa.04000000000
Proceedings of the National
Xiang Fu, Zhenghao Wu, Wujie Wang, Tian Xie, Sinan Keten, Rafacl Komez-Bombarclie and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations, 2022.
Paul Fuchs, Stephan Thaler, Sebastien Röcken, and Julija Zavadlav. chemtrain: Learning deep potential models via automatic differentiation and statistical physics. Computer Physics Commanications, pp. 109512, 2025.
Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. In 8th International Conference Learning Representations, ICLR, 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. In Conference on Neural Information Processing Systems (NewrIPS), 2021.
Johannes Gasteiger, Muhammed Shuaibi, Amuroop Sriram, Stephan Günnemann, Zachary Ulsssi, C. Lawrence Zitnick, and Abhishek Das. Gemnet-oc: Developing graph neural networks for large and diverse molecular simulation datasets, 2022.
Fuchun Ge, Ran Wang, Chen Qu, Peikun Zheng, Apurba Nandi, Riccardo Conte, Paul L. Houston, Joel M. Bowman, and Pavlo O. Dral. Tell machine learning potentials what they are needed for: Simulation-orientel training excmplified for glycine. J. Phys. Chem. Lett., 15(16):4451-4400, 2024. doi: 10.1021 10.1021 /cctt.pppc.t...4000...... URL huttps://doi.com
Publication Date: April 16, 2024.
Mario Geiger and Tess Smidt. c3nn: Euclidean neural networks, 2022.
S. Gould, B. Fernando, A. Cherian, P. Anderson, R. S. Cruz, and E. Guo. On differentiating parameterized argmin and argmax problems with application to bi-level optimization. arXiv preprinl. Xix::607.00447 2016.
Joe G Greener. Reversible molecular simulation for training classical and machine learning force fields. arXiv preprint arXiv:2412.0437, 2024.
Bin Han and Kuang Yu. Refining potential energy surface through dynamical properties via differentiable molecular simulation. Nature Communications, 16(1):816, 2025.
Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Amuropp Sriram, Jure Leskovec, Devi Parikh, and Lawrence Zitnick. Forcenet: A graph neural network for largo-scale quantum calculations, 2021.
John Ingraham, Adam Riesselman, Chris Sander, and Debora Marks. Learning protein structure with a differentiable simulator. In International Conference on Learning Representations, ICLR, 2019.
A. Kabylda, V. Vassilov-Galindo. S. Clmiicla I. Poltavsky, and A. Tkatchenko. Efficient interatomic descriptors for accurate machine learning force fields of exlended molecules. Nal. Commun,, 14(1):3562, 6 2023. doi: 10. 10.100//$4466--10 PMC10272221.
Erratum in: Nat Commun. 2023 Jul 11:14(1):4116. PMID: 37322039; PMCID:
Mariia. Karabin and Danny Perez. An entropy-maxiiizztion approach to automated training set generation for interatomic potentials. The Journal of Chemical Physics, 153(9), 2020.
Christopher Kolloff and Simon Olsson. Rescuing off-cquilibrium simulation data through dynamic experimental data with dynammo. Mach. Learn.: Sci. Technol., 4(4):045050, 2023. doi: 10.1088/2668 200////662211///ae/2 2153/ac7b4f. c 2023 The Author(s).
14

---

Published in Transactions on Machine Learning Research (02/2025)
Dávid Péter Kovács, Cas van der Oord, Jiri Kucera, Alice E. A. Allen, Daniel J. Cole, Christoph Ortner, and Gábor Csányi. Linear atomic cluster expansion force fields for organic molecules: Beyond rmse. Journal of Chemical Theory and Computation, 17(12):7696 7711, 2021. doi: 10.1021/acc.jetter jct. URL PMID: 34735161.
https://doi. org/ 10. 1021/acs. jet.jjtttc.io 1c00647
Dávid Péter Kovács, J. Harry Moore, Nicholas Browning, llyes Batatia, Joshua T. Horton, Venkat Kapil, William C. Witt, foan-Bogdan Magdän, Daniel J. Cole, and Gábor Csányi. Mace-off23: Transferable machine learning force fields for organic molecules, 2023.
Boris Kozinsky, Albert Musaclian, Anders Johansson, and Simon Batzner. Scaling the leading accuracy of deep equivariant models to biomolccular simulations of realistic size. In Proceedings of the International Conference for High Performance Compuling, Nelworking, Storage and Analysis, SC '23, New York, NY, USA, 2023. Association for Computing Machinery, doi: /5115//3551184.327011.. URL https://doft org/ 10.
Aditi S Krishnapriyan, Alejandro F Queiruga, N Benjarnin Erichson, and Michael W Mahoney. Learning continuous models for continuous physics. Communicatiins Physics, 6(1):319, 2023.
M. Kulichenko, K. Barros, N. Lubbers, el al.
teratomic potentials. Nat Comput Sci, Uncertainty-driven dynamics for active learning of in- https://doi.oww orgo 10. 1038/s43588 00444588-023-0000
2023. doi: 10. (0.1088//4338880223-0000 023-0006-..... URL
00000003--0000 00406-5
Ask Hjorth Larsen, Jens Jørgen Mortensen, Jakob Blomqvist, Ivano E. Castelli, Rune Christensen, Marcin Dulak, Jes- per Friis, Michael N. Groves, Bjork Hammerr Cory Hargus, et The atomic simulation environment - a python library for working with atoms. Journal of Physics: Condensed Matter, 2017. doi: 10.108/13661-660 -aa48x/a66800. /aa680e.
J. Li, C. Qu, and J. M. Bowman. Diffusion monte carlo with fictitious masses finds holes in poten- tial energy surfaces. Molecular Physics, 119(17-18), 2021. doi: 10.1080 001100/00220977.22221199997222 URL https://doi.. org/o0//// 10. 1080/00262976.... 2021. 2974.197426
Xiang-Guo Li, Chongze Hu, Chi Chen, Zhi Deng, Jian Luo, and Shyue Ping Ong- Quanturn-cccurate spectral neighbor analysis potential models for ni-mo binary alloys and foe metals. Phys. Rea. B, 98:094104, 2018. doi:
h..1003/PlysReevB... .4944440... URL httpt::// wwwiink......... aps otpp .8.. 08.094104
Qidong Lin, Liang Zhang, Yaolong Zhang, and Bin Jiang, Searching configurations in uncertainty space: Active learning of high-dimensional neural network reactive potentials. J. Chem. Theory Comput., 17(5):2691-2701, 2021. doi: 10.1021 acs. acc.jetcc.000119.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molccular graphs. In International Conference on Learning Representations (ICLR), 2022.
M. Majewski, A. Pérez, P. Thölke, et al. Machine learning coarse-grained potentials of pro- tein thermodynamies. Not Commun, 14:5739, 2023. doi: 10.1088/s41446---33-413333 https://doi.org orgg 10.
URL
TE Markland and M Ceriotti. Nuclear quantum effects enter the mainstream. Nat Rev Chem, 2:0109, 2018. doi: 10. 00.1088/411500-07-0000
S. Marrink, H. J. Rissclada, S. Yefimov, D. P. Ticleman, and A. H. De Vries. The martini force field: coarse grained model for biomolecular simulations. J. Phys. Chem. B, 111:7812 7824, 2007. doi: /1971177.
A. Merchant, S. Batzner, S.S. Schoenholz, et al. Scaling deep learning for materials discovery. Nature, 624:80-88 85, 2023. doi: 00.1038/441580-02---10
Joe D. Morrow, John L. A. Gardner, and Volker L. Deringer. How to validate machine-leamod interatomic potentials. J. Chem. Phys., 158:121501, 2023.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics, 2022.
15

---

Published in Transactions on Machine Learning Research (02/2025)
Carles Navarro, Maciej Majewski, and Gianni De Fabritiis. Top-down machine learning of coarse-grained protein force fields. J. Chem. Theory Compul. 2023.
Geoftrey Négiar, Michael W. Mahoney, and Aditi Krishnapriyan. Learning differentiable solvers for systems with hard constraints. In The Eleventh International Conference Learning Representations, 2023.
W. G. Noid. Perspective: coarse-grained models for biomolecular systems. J. Chem Phys., 139:090901, 2013.
John P. Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made doi:
simple. Phys. Rec. Lett., 77:3865-3866 10 1996.
URL
https://link aps doi/10. ottp./////....11////s/eeele....
77. 37.365
Ji Qi, Tsz Wai Ko, Brandon C Wood, Tuan Anh Pham, and Shyue Ping Ong. Robust training of machine learning interatomic potentials with dimensionality reduction and stratified sampling. npj Computational Materials, 10(1):43, 2024.
Z Raza et al. Proton ordering in cubic ioe and hexagonal ice; a potential new ice phase xic. Plys. Chem. Chem. Phys., 13:19788 19795, 2011. doi: 10. /clp//26222.
Jie Ren, Xidong Feng, Bo Liu, Xuehai Pan, Yao Fu, Luo Mai, and Yaodong Yang. Torchopt: An efficient library for differentiable optimization, 2022.
Sebastien Röcken and Julija Zavadlav. Accurate machine learning force fields via experimental and simulation data fusion, 2023.
Michael Schaarschmidt, Morgane Riviere, Alex M. Ganose, James S. Spencer, Alexander L. Gaunt, James Kirkpatrick, Simon Axelrod, Peter W. Battaglia, and Jonathan Godwin. Learned force fields are ready for ground state catalyst discovery, 2022.
Christoph Schran, Krystof Brezina, and Ondrej Marsalek. Committee neural network potentials control gener- alization errors and enable active learning. J. Chem. Phys., 153(10))104105, Sep 2020. doi: 10.1063/5.0016004. URL https:///oi..org//
1063/5.001600
K. T. Schütt, H. E. Sanceda, P.-. Kindermans, A. Tkatchenko, and K.-R. Müller. Schnet - a deep learning architecture for molocules and materials. J. Chem. Phys., 148(24))241722, 2018. doi: 1.1106/1.5019777
URL htttps://bii.rg/// 100.:/.......1.1111.....5050
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.
Nima Shoghi, Adeesh Kolluru, John R. Kitchin, Zachary W. Ulissi, C. Lawrence Zitnick, and Brandon M. Wood. From molecules to materials: Pre-training large generalizable models for atomic property prediction, 2023.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International conference on machine learning, pp. 387-395. Pmlr, 2014.
Justin S. Smith, Ben Nebgen, Nicholas Lubbers, Olexandr Isayev, and Adrian E. Roitberg. Less is more: Sampling chomical space with active learning. Chem. Phys., 148:241733, 2018. doi: 10.1063/1.502388
Sina Stocker, Johannes Gasteiger, Florian Bocker, Stephan Günnemann, and Johannes T Margraf. How
robust are modern graph neural nelwork polentials in long and hot molecular dynamics sinulations? Machine Learning: Science and Technology, 3(4)00450 11 2022. doi: 10.1088/2632 220000026322113///// 1088/08883333 222222213//aaeeee2 2153/ac9955
URL
https://dx. http/xxxdoo.... wro.o.ogg/1.... 100///00.101//111
S. Takamoto, C. Shinagawa, D. Motoki, et al. Towards universal neural network potential for material discovery ap- plicable to arbitrary combination of 45 elements. Nut Commun, 13:2991, 2022. doi: 10.1038 0011008//11467-002-30677777
URL https::/aoi.orgg// 0001://0000...00000///1114700020 922--6
Sebastian Thaler and Julija Zavadlav. Learning neural network potentials from experimental data via differentiable trajectory reweighting. Nature Communicutions, 12:6884, 2021. doi: 10. 40100227/444
16

---

Published in Transactions on Machine Learning Research (02/2025)
Alexandre Tkatchenko and Matthias Scheffler. Accurate molecular der waals interactions from ground-state electron density and free-atom reference data. Phys. Rew. Lett., 102:073005, 2 2009. doi: 1tt.io//Pppseeeeeee.00..00.00000
URL https://linkk..ee aps org/ hwrg/doi/10.oo http://...11///haaeeeeeo wtw.eo/eeeeeeeeeett....0.... 073005
Alexandre Tkatchenko, Robert A. DiStasio, Roberto Car, and Matthias Scheffler. Accurate and efficient method for many-body van der waals interactions. Phys. Rev. Lett., 108:236402, 6 2012. doi:
108. 108.2364022
J. Vandermause, S.B. Torrisi, Batzner, et al. On-the-fly active learning of interpretable bayesian force fields for atomistic rare events. upj Compul. Maler, 6:20, 2020. doi: 10.1088//111244--220022224 URL http::///abb..rrr/....o00////2/2
Joshua A. Vita and Daniel Koda.. Data efficiency and extrapolation trends in neural network interatomic potentials. Mach. Learn: Sci. Technol, 4:035031, 2023.
Wujie Wang, Simon Axelrod, and Rafael Homnez-oombarelli Differentiable molecular simulations for control and learning, 2020.
Wujic Wang, Zhenghao Wi,, Johannes C. B. Dictschreit, and Rafacl Gó mez-Bombarelli. Learning pair potentials using iiflerentiable simulations. The Journal of Chemical Physics, 158(4):044113, 1 2023a. doi: (0003/501126775.
Zun Wang, Honglei Wu, Lixin Sun, Xinheng He, Zhirong Liu, Bin Shao, Tong Wang, and Tie Yan Liu. Improving machine learning force fields for molecular dynamics simulations with fine-grained force metrics. J. Chem.
Phys., 159:035101, 2023b. doi: 10.1063/5.0144023..
Ronald Williams. Simple statistical Machine learning, 1992.
algorithms for connectionist reinforcement learning.
Yujie Wu, Harald L. Tepper, and Gregory A. Voth. Flexible simple point-charge water model with improved liquid-state properties. The Journal of Chemical Physics, 124(2):024503, 1 2006. ISSN 0021-9606, 1089-7690. doi: 10.1063/1 /...//....1677
Dongsun Yoo, Jisu Jung, Wonscok Jeong, and Seungwu Han. Metadynamics sampling in atomic environment space for collecting training data for machine learning potentials. upj Compulational Materials, 7(1):131, 2021.
Robert W. Zwanzig. High-temperature equation of state by a perturbation method. nonpolar gases. Chem. Phys., 22:1420-1422 1954.
Martin Sipka, Johannes C. B. Dietschreit, Lukáš Grajciar, and Rafael (6mez-Bombarelli. Differentiable simulations for enhanced sampling of rare events, 2023.
17


Reviews:
[{"review_id": "HJIyhmbDQm", "reviewer": "Reviewer_wTtv", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The proposed method addresses a critical limitation of MLFFs—simulation instability—through a well-founded theoretical approach and innovative training procedures.\n2. The theoretical foundation of the method is well described, providing a clear description of Boltzmann estimators and their capabilities and limitations. Furthermore, detailed derivations and proofs are provided in the supplement.\n3. The evaluation is well done, covering three distinct MLFF models and datasets that include varied systems, showcasing the method's applicability in differing scenarios. Throughout all these experiments the StABlE training method shows significant improvements in the stability measure chosen by the authors, while at the same time resulting in only a minor reduction in energy and force prediction accuracy.\n4. The besides the proofs in the supplementary materials, the evaluations there are detailed as well, offering in-depth analyses of different aspects of the method and it's performance on the different models and datasets.\n\nWeaknesses:\n1. The instability criterion chosen is usually closely related to the observable being trained on. While the authors demonstrate that the observable based training does not negatively impact or even improves other observables such as VACF for aspirine and diffusivity for water, the authors have failed to show whether StABlE training also improves or at least does not deteriorate stability with respect to other stability criteria as defined in Fu et. al (2022).\n2. On page 8, the authors state that because the stability gains by StABlE training are better than those gained by taking smaller simulation steps with the model trained just on energies and forces, this means that StABlE training could lead to speeding up MD through taking larger time steps during the simulation. This claim is not substantiated or supported at just based on the stability comparison for smaller time steps, since the main challenge when taking larger time steps it the error due to the larger approximation in the integration procedure, regardless of the quality of the force field itself. The authors should either remove this claim, or show experiments demonstrating that this is actually the case and that models trained with StABlE can produce stable and accurate simulations with MD steps larger than 1fs for example.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This manuscript introduces Stability-Aware Boltzmann Estimator (StABlE) Training, a novel training approach with the goal of specifically enhancing the stability of machine learning force fields (MLFFs) in molecular dynamics (MD) simulations, based on system observables. \n\nStABlE training integrates ML-based MD simulations in unstable regions of configuration space and optimized the model to obtain the correct observables in this unstable region to improve both simulation stability and accuracy. It leverages an innovative Boltzmann Estimator that enabling efficient gradient-based optimization of the observables without requiring additional computationally expensive quantum-mechanical calculations in these unstable regions, a great contrast and advantage when compared to active learning approaches. This approach is combined with using the 'standard' force and energy loss terms as additional or regularizing loss terms, enabling the model to avoid the one-to-many mapping challenge that would arise when training purely based on observables. \n\nOverall, this approach paved the way toward the efficient training of more reliable MLFFs that are applicable to long-timescale simulations and diverse systems.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. On Page 2, \"instability Simulation instability\" should include a period to separate the sentences properly.\n2. On Page 8, the claim that \"training on observables can enable larger timesteps for simulation\" should be supported with empirical tests or removed if unsubstantiated. As mentioned in the weaknesses section, the current evidence comparing timestep reductions and StABlE training does not adequately support this claim. (critical)\n3. Additionally, in the corresponding supplementary section A.12, specifically in supplementary Figure 11, the x-axis only contains one labeling for the time value, making it difficult to judge what are the other timestep values that the model was evaluated on. \n4. On Page 8, the statement \"less than 5% of virtually no replicas are unstable\" does not quite make sense as it is. Maybe the authors wanted to say \"less than 5% or virtually no replicas are unstable\". Either way this needs to be corrected for clarity.\n5. Page 11, it seems that a plus-minus symbol (e.g., ±) is missing when stating the diffusivity coefficients, since it seems that the authors are trying to include the error margins of the estimated values.\n6. As mentioned in the weaknesses, the authors should include some others of the usual stability measures for each of their simulations to demonstrate that StABlE training targeted at a specific observable which is closely related to the chosen stability criterion does not actively lead to a worse performance with respect to other stability criteria. (would strengthen the work)", "broader_impact_concerns": "There are no broader impact concerns for this work as far as I am aware.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "RKC5pJX2w0", "reviewer": "Reviewer_otu9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Main text and supplementary are well-written and easy to follow. The problem statement, motivation, and solution are well addressed.\n* Experiments are comprehensive. It is convincing to perform the experiments across multiple systems, architectures, and temperatures.\n* Boltzmann Estimators are a nice idea and good technical contribution.\n* Generalizing to new temperature is a promising result of this method.\n\nWeaknesses:\n* It would have been nice to compare the Boltzmann Estimator compared to the adjoint method or chain rule expansion of the unrolled MD simulation. This would show empirical insights of the estimator on a toy problem.\n* I find some hyper parameters such as the number of replicas and timesteps magically chosen. It would be nice to include a sweep over them and details of how one would choose these hyperparameters for new systems.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The present work develops a training procedure, Stability-Aware Boltzmann Estimator (StABlE) Training, that uses joint supervision from reference QM calculations and system observables to train MLFFs to improve the stability and accuracy of MD simulations. The method is intuitive by having two stages: (1) simulating the learned MLFF until it reaches a state with unstable observables and then (2) learning to correct the unstable observables. Matching the observables is challenging due to the need of backpropping through MD simulations or requiring costly QM calculations. The technical novelty of the work is developing a Boltzmann Estimator using the REINFORCE trick and inspiration from thermodynamic perturbation theory. The Boltzmann Estimator is an unbiased estimator of the Jacobin needed for matching observables. An extension called the localized Boltzmann Estimator helps scale up to larger systems. Experiments are comprehensive as they are performed across multiple systems (organic molecules, tetra peptides, and condensed phase systems) as well as multiple MLFF architectures. The results show clear improvement in stability across extended simulations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The only change I request is to discuss how one would choose the hyper parameters for new systems. Could all the hyper parameters be listed somewhere? I may have missed it.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "asAcHagUNq", "reviewer": "Reviewer_XXPs", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\n- The paper presents a practical way to incorporate system observables into MLFF training without requiring extensive additional QM calculations, addressing a critical need in the field.\n\n- The emphasis on simulation stability, often overlooked in favor of energy and force accuracy, fills an important gap in MLFF development.\n\n- The experiments span diverse systems (organic molecules, peptides, and water) and architectures, showcasing the method’s broad applicability.\n\n- By avoiding direct backpropagation through MD simulations and leveraging the Boltzmann Estimator, the method is computationally efficient and avoids instability issues inherent in other approaches.\n\n**Weaknesses:**\n\n- The dataset scope is somewhat limited. Experiments are restricted to specific examples rather than a broader dataset where training and testing are performed across distinct systems, limiting conclusions about generalization. The choice of observables needs to be specific to different systems.\n\n- While stability and observables are targeted, the methodology does not yet address dynamical observables, a crucial aspect for many MD applications.\n\n- Tradeoff with Energy/Force Accuracy: Stability improvements come with a slight degradation in energy and force accuracy, which might concern users prioritizing these metrics. It is also an extra hyperparameter to consider.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces Stability-Aware Boltzmann Estimator (StABlE) Training, a novel methodology for training machine learning force fields (MLFFs) that combines quantum mechanical (QM) data and system observables as training signals. The paper proposes using a Boltzmann Estimator interpretation of the MLFF to enable efficient optimization without backpropagating through MD simulations. The authors demonstrate the method's efficacy on diverse systems, including Aspirin, a tetrapeptide, and a water system, across three MLFF architectures. Experimental results indicate improvements in simulation stability and accuracy in recovering observables.\n\nThis paper should be interesting to the AI for science community that is part of the TMLR audience. The claims are largely supported by the evidence presented in the paper.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- The authors mentioned the performance drop in force/energy accuracy with their approach. Is this sensitive to the loss coefficient? It would be good to quantify and contextualize the energy and force accuracy tradeoff observed during training with a figure/table, to better inform readers of practical implications.\n\n- The observables that are feasible to obtain are often limited for larger-scale systems, sometimes even infeasible for DFT calculations. Can the authors discuss the applicability of their approach in larger-scale problems?\n\n- It would strengthen the paper if the authors could expand experiments to include training/testing across different systems to evaluate generalization capabilities, but I think this is not required for paper acceptance.", "broader_impact_concerns": "There are no concerns about ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall, this paper presents novel algorithmic contributions with respect to stability of molecular field generations using neural networks. The results are reasonably good, the writing is clear, and visaulization clear. The paper is clearly a vaiuable contribution to the literature.

The authors did a good job in answering all questions by reviewers including larger simulation steps, clarifications about the stability criteria used, comparisons to adjoint methods. According to new experiments, StABlE Training is stable also for longer rollouts, as and does not overfit to the criteria defined and used during training. Also, as clarified by authors and confirmed experimentally, the proposed algorithm is more stable than adjoint methods in longer rollouts, while having lower memory requirements.
Decision: Accept (poster)