F ORCES ARE NOT E NOUGH : B ENCHMARK AND C RITICAL E VALUATION FOR M ACHINE L EARNING F ORCE F IELDS WITH M OLECULAR S IMULATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Molecular dynamics (MD) simulation techniques are widely used for various natural science applications.
Increasingly, machine learning (ML) force field (FF) models begin to replace ab-initio simulations by predicting forces directly from atomic structures.
Despite significant progress in this area, such techniques are primarily benchmarked by their force/energy prediction errors, even though the practical use case would be to produce realistic MD trajectories.
We aim to fill this gap by introducing a novel benchmark suite for ML MD simulation.
We curate representative MD systems, including water, organic molecules, peptide, and materials, and design evaluation metrics corresponding to the scientific objectives of respective systems.
We benchmark a collection of state-of-the-art (SOTA) ML
FF models and illustrate, in particular, how the commonly benchmarked force accuracy is not well aligned with relevant simulation metrics.
We demonstrate when and how selected SOTA methods fail, along with offering directions for further improvement.
Specifically, we identify stability as a key metric for ML models to improve.
Our benchmark suite comes with a comprehensive open source codebase for training and simulation with ML FFs to facilitate further work.
I NTRODUCTION
F o r c e N e t S c h N e t
D eep P o t - SE N equ I P D
i m e N e t
G e m N e t - d T G e m N e t - T
F o r c e M AE
RD F M AE
S t ab ili t y
High stability and low
MAE mean
Molecular Dynamics (MD) simulations provide atomistic insights into physical phenomena in materials and biological systems.
Such simulations are typically based on force fields (FFs) that characterize the underlying potential energy surface (PES) of the system and then use Newtonian forces to simulate long trajectories (Frenkel &
Smit, 2001).
The PES itself is challenging to compute and would ideally be done through quantum chemistry which is computationally expensive.
Traditionally, the alternative has been parameterized force fields that are surrogate models built from empirically chosen functional forms (Halgren, 1996).
Recently, machine learning (ML) force fields (Unke et al., 2021b) have shown promise to accelerate MD simulations by orders of magnitude while being quantum chemically accurate.
The evidence supporting the utility of ML FFs is often based on their accuracy in reconstituting forces across test cases (Faber et al., 2017).
The evaluations invariably do not involve simulations.
However, we show that force accuracy alone does not suffice for effective simulation (Figure 1).
MD simulation not only describes microscopic details on how the system evolves, but also entails macroscopic observables that characterize system properties.
Calculating meaningful observables often requires long simulations to sample the underlying equilibrium distribution.
These observables are designed to be predictive of material properties such as diffusivity in electrolyte materials (Webb et al., 2015), and reveal detailed physical mechanisms, such as the folding kinetics of protein dynamics (Lane et al., 2011;
Lindorff-Larsen et al., 2011)
Although these observables are critical products of MD simulations, sys- tematic evaluations have not been sufficiently studied in existing ML FF literature.
To gain insight into the performance of existing models in a simulation setting, we propose a series of simulation-based benchmark protocols.
Compared to the popular multistep prediction task in the learned simulator community (Sanchez-Gonzalez et al., 2020), MD observables focus on distributional quantities.
The exact recovery of the trajectories given the initial conditions is not the ultimate goal.
Evaluating learned models through MD simulations requires careful design over the selection of systems, the simulation protocol, and the evaluation metrics : (1) A benchmark suite should cover diverse and representative systems to reflect the various challenges in different MD applications.
(2) Simulations can be computationally expensive.
An ideal benchmark needs to balance the cost of evaluation and the complexity of the system so that meaningful metrics can be obtained with reasonable time and computation.
(3) Selected systems should be well studied in the simulation domain, and chosen metrics should characterize the system’s important degrees of freedom or geometric features.
Are current state-of-the-art (SOTA)
ML FFs capable of simulating a variety of MD systems?
What might cause a model to fail in simulations?
In this paper, our aim is to answer these questions with a novel benchmark study.
The contributions of this paper include:
• We introduce a novel benchmark suite for ML MD simulation with simulation protocols and quantitative metrics.
We perform extensive experiments to benchmark a collection of SOTA ML models.
We provide a complete codebase for training and simulating MD with ML FFs to lower the barrier to entry and facilitate future work.
• We show that many existing models are inadequate when they are evaluated on simulationbased benchmarks, even when they show accurate force prediction (as shown in Figure 1).
• By performing and analyzing MD simulations, we summarize common failure modes and discuss the causes and potential solutions to motivate future research.
R ELATED W ORK
ML force fields learn the potential energy surface (PES) from the data by applying expressive regressors such as kernel methods (Chmiela et al., 2017) and neural networks on symmetry-preserving representations of atomic environments (Behler & Parrinello, 2007; Khorshidi & Peterson, 2016;
Smith et al., 2017; Artrith et al., 2017; Unke & Meuwly, 2018; Zhang et al., 2018b;a; Kovács et al., 2021; Thölke & De Fabritiis, 2021).
Recently, graph neural network architectures (Gilmer et al., 2017; Schütt et al., 2017; Gasteiger et al., 2020) have gained popularity as they provide a systematic strategy for building many-body correlation functions to capture highly complex PES.
In particular, equivariant representations have been shown powerful in representing atomic environments (Thomas et al., 2018; Qiao et al., 2021; Schütt et al., 2021; Batzner et al., 2022; Gasteiger et al., 2021; Liao & Smidt, 2022), leading to significant improvements in benchmarks such as MD17 and OC22/20.
Some works presented simulation-based results (Unke et al., 2021a;
Park et al., 2021; Batzner et al., 2022; Musaelian et al., 2022) but do not compare different models with simulation-based metrics.
Existing benchmarks for ML force fields (Ramakrishnan et al., 2014; Chmiela et al., 2017) mostly focus on force/energy prediction, with small molecules being the most typical systems.
The catalystfocused OC20 (Chanussot et al., 2021) and OC22 (Tran et al., 2022) benchmarks proposed the
Initial Structure to Relaxed Structure/Energy tasks, which aim to predict the relaxed structure/energy through structural optimizations.
These tasks do not characterize system properties under a structural ensemble which requires simulations.
Several recent works (Rosenberger et al., 2021) have also studied the utility of certain ML FFs in MD simulations.
In particular, Stocker et al.
2022 uses
GemNet (Gasteiger et al., 2021) to simulate small molecules in the QM7-x (Hoja et al., 2021) dataset, with a focus on simulation stability.
Zhai et al.
2022 applies the DeepMD (Zhang et al., 2018a) architecture to simulate water and demonstrates its shortcoming in generalization across different phases.
However, systematic benchmarks for simulation-based metrics are lacking in the existing literature, which obscures the challenges in applying ML FF for MD applications.
P RELIMINARIES
Training.
An ML FF aims to learn the potential energy surface ˆ E ( x ) ∈ R as a function of atomic coordinates x ∈ R N × 3 ( N is the number of atoms), by fitting atom-wise forces ˆ F ( x ) and energies
[Figure]
[Table]
from a training dataset: { x i , F i , E i } N data i =1 , where x i ∈ R N × 3 , F ∈ R N × 3 , E ∈ R .
For evaluation, the test force prediction accuracy is used as a proxy to quantify the quality of the learned PES.
The force field learning protocol has been well established (Unke et al., 2021b).
MD simulation.
Simulating molecular behaviors requires integrating a Newtonian equation of motion with forces obtained by differentiating ˆ E ( x ) :
F ( x ) = − ∂ ˆ E ( x ) /∂
x .
To mimic desired thermodynamic conditions, an appropriate thermostat and barostat are chosen to augment the equation of motion with extended variables.
The simulation produces a time series of positions: { x t ∈ R N × 3 } Tt =0 , where t is the temporal order index, and T is the total simulation steps.
Observables.
From the time series observations of positions and velocities, observables O ( x t ) can be computed to characterize the state of the system at different granularities.
Under the ergodic hypothesis, the time averages of the simulation observables converge to distributional averages under the Gibbs measure: ⟨ O ⟩ = 1 T lim T →∞ (cid:80)
Tt O ( x t )
=
(cid:82) dx p ( x ) O ( x ) ,
where p ( x ) ∝ exp( − ˆ E ( x ) k B T ) with T as the bath temperature and k B as the Boltzmann constant.
Calculations of such observables require the system to reach equilibrium.
Simulation observables connect simulations to experimental measurements and are predictive of macroscopic properties of matter.
Common observables include radial distribution functions (RDFs), virial stress tensor, mean-squared displacement (MSD), etc.
The metrics used in this work are all well-established observables in the respective types of systems.
Definitions of benchmarked observables are in Appendix A.
D ATASETS AND M ETRICS
Popular benchmark datasets, such as MD17, focus on the force prediction task for gas-phase small molecules.
However, successes in these tasks are not sufficient evidence for (1) capturing complex interatomic interactions that are critical for condensed phase systems; and (2) recovery of critical simulation observables that cannot be directly indicated by force prediction accuracy.
We wish to extend existing benchmarks to expand coverage of the types of systems that manifest complex intermolecular interactions at multiple scales.
Furthermore, beyond force predictions, we conduct simulations and benchmark observables that reflect the actual simulation quality, along with stability and computational efficiency.
We choose four representative MD systems (summarized in Table 1): small molecules, liquid water, peptide, and solid-state materials.
To facilitate relatively easy and fast benchmarking, we limit the size of the systems to less than 200 atoms.
Quantifying simulation stability.
ML FFs can produce unstable dynamics, as the learned force field may not extrapolate robustly to the undersampled configuration space.
As a result, the trajectories can enter nonphysical states that are not meaningful for observable calculations.
Therefore, we closely monitor how much the simulated structure deviates from the physical configurations, with the RDF for condensed phase systems and bond lengths for flexible molecules.
We say that a simulation becomes “unstable” when the deviation exceeds a threshold, which implies sampling of
[Table]
highly nonphysical structures.
We then use the time duration for which a model remains stable in simulations to measure its stability.
The ensemble statistics are only computed over the stable part of the simulated trajectories.
Details on the stability criterion are included in Appendix A.
MD17 (Chmiela et al., 2017) dataset contains AIMD calculations for eight small organic molecules and is widely used as a force prediction benchmark for ML FFs.
We adopt four molecules from
MD17 and benchmark the simulation performance.
In addition to force error, we evaluate the stability and the distribution of interatomic distances h ( r ) .
For each molecule, we randomly sample 9, configurations for training and 500 for validation from the MD17 database.
We randomly sample 10,000 configurations from the rest of the data for force error evaluation.
We perform five simulations of 300 ps for each model/molecule by initializing from 5 randomly sampled testing configurations, with a time step of 0.5 fs, at 500 K temperature, under a Nosé–Hoover thermostat.
Water is arguably the most important molecular fluid in biological and chemical processes.
Due to its complex thermodynamic and phase behavior, it poses great challenges for molecular simulations.
In addition to force error, we evaluate simulation stability and recovery of both equilibrium statistics and dynamical statistics, namely the element-conditioned RDFs and liquid diffusion coefficient.
Our dataset consists of 100,000 structures collected every 10 fs from a 1 ns trajectory sampled at equilibrium and a temperature of 300 K. We benchmarked all models with various training+validation dataset sizes (1k/10k/90k randomly sampled structures) and used the remaining 10,000 structures for testing.
We performed 5 simulations of 500 ps by initializing from 5 randomly sampled testing configurations, with a time step of 1 fs, at 300 K temperature, under a Nosé–Hoover thermostat.
Alanine dipeptide features multiple metastable states, making it a classic benchmark for the development of MD sampling methods and force field development (Head-Gordon et al., 1989; Kaminski et al., 2001).
Its geometric flexibility is well represented by the central dihedral (torsional) angles ϕ and ψ .
Our reference data are obtained from simulations with explicit water molecules, with detailed protocols described in Appendix B. For faster simulation, we learn an implicitly solvated FF following a protocol similar to Chen et al.
Our task is more challenging in that it aims to learn the implicitly solvated atomistic FF rather than the implicit solvation correction in Chen et al.
To facilitate accelerated sampling, we apply metadynamics with ϕ and ψ as the collective variables.
We evaluate force prediction, simulation stability, and free energy surface (FES) reconstruction F ( ϕ, ψ ) .
Our dataset consists of 50,000 structures dumped every 2 ps from a 100 ns trajectory at a temperature of 300 K. We used 38,000 randomly sampled structures for training, 2,000 for validation, and the rest as a test set.
We performed 6 simulations of 5 ns by initializing from six local minima on the FES (Figure 5) with a time step of 2 fs at 300 K, and under a Langevin thermostat to mimic random noise from solvation effects.
More information on our simulation protocols can be found in Appendix B.
LiPS is a crystalline superionic lithium conductor relevant to battery development and a representative system for MD simulation usage in studying kinetic properties in materials.
We adopt this dataset from Batzner et al. 2022, and benchmark all models on their force error, stability, RDF recovery, and
Li-ion diffusivity coefficient.
The dataset has 25,000 structures in total, from which we use 19, randomly sampled structures for training, 1,000 structures for validation, and the rest for computing force error.
We conduct 5 simulations of 50 ps by initializing from 5 randomly sampled testing configurations, with a time step of 0.25 fs, at 520 K temperature, under a Nosé–Hoover thermostat.
E XPERIMENTS
Benchmarked models.
We adopt the Open Catalyst Project implementation of SchNet (Schütt et al., 2017), DimeNet (Gasteiger et al., 2020), ForceNet (Hu et al., 2021), GemNet-T/dT (Gasteiger et al.,
[Table]
1 Stability 246 (109) 18 (2) 85 (68) 300 (0) 8 (2) 25 (10) 300 (0)
h
( r ) 0 .
FPS 109 .
2 Salicylic Acid Force 14 .26 .
6 Stability 300 (0) 300 (0) 73 (82) 1 (0) 26 (24) 94 (109) 300 (0) h ( r ) 0 .
FPS 94 .111 .
, and the official implementation of DeepPot-SE (Zhang et al., 2018b) and NequIP (Batzner et al., 2022).
A summary of all benchmarked models is in Table 2.
These models have been popular in previous benchmark studies for force/energy prediction.
They use different representations for atomistic interactions and respect different levels of euclidean symmetry.
We follow all original hyperparameters introduced in the respective papers and only make minimal adjustments when the training is unstable.
More details on the hyperparameters can be found in
Appendix C.
Key observations.
We make two key observations as evidenced in the experimental results:
Despite being widely used, force prediction is not sufficient for evaluating ML FFs.
It generally does not align with simulation stability and performance in estimating ensemble properties.
While often neglected, stability can be a major bottleneck for ML FFs.
Lower force error and more training data does not necessarily give rise to more stable simulations, suggesting stability as a fundamental consideration for comparison and model design.
We next go through experimental results of all four datasets in detail to demonstrate the key observations while making other observations.
MD17.
As shown in Table 3, more recent models that lie on the right side of the table generally achieve a lower force error, but may lack stability.
Figure 3 selects results from Table 3 to demonstrate the non-aligned trends of force prediction performance vs. simulation performance, which supports key observation 1 .
GemNet-T/dT can attain a very low force error for all four molecules, but often collapse before the simulation finishes.
This observation constitutes key observation 2 .
We note that
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
although the stable portion of simulated trajectories produced by GemNet-T/dT can recover the h ( r ) relatively accurately, stability will become a bigger issue when the statistics of interest require long simulations, as demonstrated in other experiments.
On the other hand, despite having a relatively high force error
[Figure]
DeepPot-SE performs very well on simulation-based metrics on all molecules except for Aspirin (Figure 3).
With the highest molecular weight
[Figure]
Aspirin is indeed the hardest task in
MD17 in the sense that all models attain high force prediction errors on it.
We further observe that good stability does not imply accurate recovery of trajectory statistics.
Although ForceNet remains stable for Ethanol and Naphthalene , the extracted h ( r ) deviates a lot from the reference (Table 3), indicating that ForceNet does not learn the underlying PES correctly, possibly due to its lack of energy conservation and rotational equivariance.
Overall, NequIP is the best-performing model on MD17.
It achieves the best performance in both force prediction and simulation-based metrics for all molecules while requiring the highest computational cost.
More detailed results on MD17 including individual h ( r ) are included in Appendix C.
Water.
Under different challenges posed by a condensed phase system, key observation 1 and 2 are still evident according to Table 4: GemNet-T/dT and DimeNet are the top-3 models in terms of force prediction, but all lack stability.
With insufficient simulation time due to early collapse, the RDFs computed from their simulated trajectories are of high variance.
The water diffusivity coefficient requires long (100 ps in our experiments) trajectories to estimate and thus cannot be extracted for unstable models.
Like MD17, DeepPot-SE does not achieve the best force prediction performance but demonstrates decent stability and highly accurate recovery of simulation statistics.
Figure 4 further compares model performance with different training dataset sizes.
Key observation 1 and 2 are clearly shown: Models located on the left of each scatter plot have very low force error but may have poor stability or high error in simulation statistics.
More specifically, although more training data almost always improve force prediction performance, its effect on simulation performance is not entirely clear.
On the one hand, GemNet-T/dT, DimeNet, and ForceNet are not
[Table]
[Table] [Table]
stable even when under the highest training data budget.
On the other hand, we observe a clear improvement of DeepPot-SE when more training data is used.
NequIP is again the best performing model, achieving very low force error, excellent stability, and accurate recovery of ensemble statistics, even under the lowest data budget of 1,000 training+validation structures.
However, when the training dataset is sufficiently large (90k), DeepPot-SE has equally good results as NequIP while being more than 20 times faster – dataset size also influences the model of choice for a certain dataset.
More results, including tables for water-1k/90k and a study on model size, are included in Appendix C.
Alanine dipeptide poses unique challenges in sampling different metastable states separated by high free energy barriers.
Table 5 shows all models have high force errors due to the random forces introduced by the lack of explicit account of water molecules.
Although the force errors are in the same order of magnitude, all models except NequIP fail to simulate stably.
The FES reconstruction task requires stable simulation for the entire 5 ns.
NequIP is the only model that manages to finish five simulations out of six but produces inaccurate statistics.
All other models are not stable enough to produce meaningful results.
We further analyze the results on this task in Section 6.
LiPS.
Compared to flexible molecules and liquid water, this solid material system features slower kinetics.
From Table 6 we observe that most models are capable of finishing 50-ns simulations stably.
In this dataset, the performance on diffusivity estimation and force prediction align well.
We observe that both GemNet-T and GemNet-dT show excellent force prediction, stability, and recovery of observables, while GemNet-dT is 2.6 times faster.
The better efficiency comes from the direct prediction of atomic forces F instead of taking the derivative F = ∂E/∂ x , which also makes
GemNet-dT not energy-conserving – a potential issue we further discuss in Section 6.
Implications on model architecture.
More recent models utilizing SE(3)/E(3)-equivariant representations and operations such as GemNet and NequIP are more expressive and can capture interatomic interactions more accurately.
This is reflected by their very low force error and accurate recovery of ensemble statistics when not bottlenecked by stability.
Moreover, NequIP shows that excellent accuracy and stability can be simultaneously achieved.
The stability may come from parity-equivariance and the explicit architecture in manipulating higher-order geometric tensors.
We believe further investigations into the extrapolation behavior induced by different equivariant geometric representations and operations (Batatia et al., 2022) is a fruitful direction in designing more powerful ML FFs.
F AILURE M ODES : C AUSES AND F UTURE D IRECTIONS
A case study on alanine dipeptide simulation.
NequIP achieves decent performance on all our tasks but fails on alanine dipeptide.
It is also the only model that can simulate stably for 5 ns.
Figure 5 (a) demonstrates how NequIP fails to reconstruct the FES: it does not manage to sample much of the
[Figure]
transition regions and the configuration space with ϕ ∈ [0 , 180 ◦ ] .
Figure 5 (b) demonstrates the reconstructed FES, which significantly deviates from the reference.
This failure can be partially explained by Figure 5 (c), the training data distribution produced by the reference potential.
The relatively high-energy (low-density) regions are exactly those that are not reachable by NequIP.
Even though our MD trajectory is well-equilibrated, the relative difference in populations of different meta-stable states creates data imbalance, making it more challenging for the model to learn PES for higher-energy configurations where density is relatively low.
In our experiments, we observe that simulations starting from the low-density meta-stable state (e.g., black star marked in Figure 5 (c)) tend to fail.
This implies that generalization across different regions in the conformational space is an important challenge for ML FFs.
To prevent ML FFs from sampling nonphysical regions, which is a common precursor to failed simulation (Figure 6), one can deliberately include distorted and off-equilibrium geometries in the training data to improve model robustness(Stocker et al., 2022).
Alternatively, one can resort to active learning (Wang et al., 2020b; Vandermause et al., 2020;
Schwalbe-Koda et al., 2021) to acquire new data points based on model uncertainty.
Energy conservation.
Models that directly predict forces may not conserve energy.
Figure 5 (d) demonstrates the evolution of model-predicted total energy for selected models on LiPS, in a microcanonical (NVE) ensemble.
The energy of an isolated system in the NVE ensemble is in principle conserved.
We observe that GemNet-T conserves energy, whereas GemNet-dT fails to conserve the predicted total energy.
The existence of non-conservative forces breaks the time reversal symmetry and, therefore, may not be able to reach equilibrium for observable calculation.
However, in our experiment, GemNet-dT performs well on the LiPS dataset when coupled with a thermostat.
Previous works (Kolluru et al., 2022) also found that energy conservation is not required for SOTA performance on OC20.
The usability of non-conservative FFs in simulations requires further careful investigations.
Simulation instability is a major bottleneck for highly accurate models such as GemNet-T to fail on several simulation tasks.
Moreover, in our water experiments, we find a larger amount of training data does not resolve this issue for GemNet-T/dT and DimeNet (Figure 4).
We further experiment with smaller simulation time steps for GemNet-T on water (Figure 5 (e)), but stability still does not improve.
On the other hand, Stocker et al.
(2022) demonstrates that the stability of GemNet improves with larger training sets on QM7-x, which includes high-energy off-equilibrium geometries obtained from normal mode sampling.
We hypothesize that including these distorted geometries may improve the model’s robustness against going into nonphysical configurations.
We also observe that t = 3702 fs t = 3704 fs t = 3706 fs t = 3622 fs t = 3716 fs t = 28330 fs t = 28370 fs t = 28372 fs t = 28373 fs t
=
fs
[Figure]
Figure cases, observe that the nonphysical configurations first emerge at local regions (circled) which cascade to the entire system very quickly as extremely large forces are being produced and subsequently integrated into the dynamics.
At the end of the visualization, the bonds in the alanine dipeptide system are broken.
Therefore, the local-descriptor-based NequIP model predicts very small forces.
For the water system, the particles are packed in a finite periodic box.
The nonphysical configurations exhibit incorrect coordination structures and extremely large forces.
Regarding stability, past works found adding noise paired with a denoising objective during training helpful in improving out-of-distribution test performance on OC20 (Godwin et al., 2021), and in stabilizing learned simulations (Sanchez-Gonzalez et al., 2020).
Another relevant line of work in coarse-grained MD simulation has studied regularization with an empirical “prior energy”
(Wang et al., 2019b) and post-prediction refinement (Fu et al., 2022) to battle simulation instability.
C ONCLUSION AND O UTLOOK
We have introduced a diverse suite of MD simulation tasks and conducted a thorough comparison of
SOTA ML FFs to reveal novel insights into ML for MD simulation.
As shown in our experiments, benchmarking only force error is not sufficient, and simulation-based metrics should be used to reflect the practical utility of a model.
We demonstrate case studies on the failure of existing training schemes/models to better understand their limitations and emphasize the importance of simulation stability.
Our experiments also show that the performance of a model can be highly case-dependent.
For more challenging MD systems, more expressive atomistic representations may be required.
For example, recent work has explored non-local descriptors (Kabylda et al., 2022) aiming at capturing long-range interactions in large molecules.
Strictly local equivariant representations (Musaelian et al., 2022) are studied for very large systems where computational scalability is critical.
New datasets (Eastman et al., 2022) and benchmarks have been playing an important role in inspiring future work in the thriving field of ML MD simulation.
The possibility of ML in advancing MD simulation is not limited to ML force fields.
Enhanced sampling methods enable fast sampling of rare events and have been augmented with ML techniques (Schneider et al., 2017;
Sultan et al., 2018; Holdijk et al., 2022).
Differentiable simulations (Schoenholz & Cubuk, 2020; Wang et al., 2020a; Doerr et al., 2021; Ingraham et al., 2018;
Greener & Jones, 2021) offer a principled way of learning the force field by directly training the simulation process to reproduce experimental observables (Wang et al., 2020a; 2022; Thaler &
Zavadlav, 2021).
We hope our datasets and benchmarks will encourage future developments in all related aspects to push the frontier in ML for MD simulations.
R
A E VALUATION M ETRICS
Distribution of interatomic distances is a low-dimensional description of the 3D structure and has been studied in previous work (Zhang et al., 2018a).
For a given configuration x , the distribution of interatomic distances h ( r ) is computed with:
[Equation]
where r is the distance from a reference particle; N is the total number of particles; i, j indicates the pairs of atoms that contribute to the distance statistics; δ is the Dirac Delta function to extract value distributions.
To calculate the ensemble average, h ( r ) is calculated and averaged over frames from equilibrated trajectories.
For learned simulations, we compute h ( r ) using only the stable part of the simulation.
RDF.
As one of the most informative simulation observables, the radial distribution function (RDF) describes the structural/thermodynamic properties of the system and is also experimentally measurable.
By definition, the RDF describes how density varies as a function of distance from a particle (illustrated in Figure 7 (a)).
For a given configuration x , the RDF can be computed with the following formula:
[Equation]
where r is the distance from a reference particle; N is the total number of particles; i, j indicates the pairs of atoms that contribute to the distance statistics; ρ is the density of the system; δ is the Dirac
Delta function to extract value distributions.
To calculate the ensemble average, RDF( r ) is calculated and averaged over frames from equilibrated trajectories.
The final RDF MAE is then calculated by integrating r:
[Equation]
where ⟨·⟩ is the averaging operator, ⟨ RDF( r ) ⟩ is the reference equilibrium RDF, and ⟨ RDF( r ) ⟩ is the model-predicted RDF.
Diffusivity coefficient.
The diffusivity coeffcient D quantifies the time-correlation of the translational displacement (illustrated in Figure 7 (b)), and can be computed from the mean square displacement:
[Equation]
where x i ( t ) is the coordinate of particle i at time t , and N ′ is the number of particles being tracked in the system.
For the water system we monitor the liquid diffusivity coefficient and track alloxygen atoms.
For the LiPS system, we monitor Li-ion Diffusivity and track all 27 Li-ions.
As the definition implies, D is a quantity that converges with longer simulation time.
Accurate recovery of
D requires sufficient long trajectories sampled from the Hamiltonian with ML FFs.
In this paper, we only compute diffusivity for stable trajectories of at least 100 ps for water and 40 ps for LiPS.
As we simulate multiple random runs for each system/model, we average the diffusivity coefficient extracted from each valid trajectory to obtain the final prediction.
For learned simulations, we compute h ( r ) using only the stable part of the simulation.
Free energy surface.
Given the probability distributions over configurations p ( x ) and a chosen geometric coordinate ξ transformed from x .
Based on the marginalized density p ( ξ ) , the free energy can be calculated from the following:
[Equation]
Diffusion t t+
Radial Distribution r ϕ ψ
Dihedral Angles
[Figure] [Figure]
) (c)
In the specific case of alanine dipeptide, there are two main conformational DOF: dihedral angle ϕ of
C − N − C α − C and dihedral angle ψ of N − C α − C − N (illustrated in Figure 7 (c)).
Therefore, the FES w.r.t ϕ and ψ is the most physically informative.
We propose our quantitative metric
MAE F ( ϕ ) ,F ( ψ ) based on the absolute error in reconstructing the FES along the ϕ and ψ coordinates.
We integrate the absolute difference between the reference free energy F and the model predicted ˆ F from [ − π, π ) :
[Equation]
Stability criterion.
Abstractly speaking, stability is a notion of staying within physical (low-energy) configuration spaces.
Since all MD systems studied in this paper are in equilibrium, practically we keep track of stability by closely monitoring equilibrium statistics.
For systems with periodic boundary conditions, we monitor the RDF and say a simulation becomes “unstable” at time T when
[Equation]
8) where ⟨·⟩ is the averaging operator, τ is a short time window, and ∆ is the stability threshold.
In this paper we use τ = 1 ps , ∆ = 3 . 0 for water, and τ = 1 ps , ∆ = 1 . 0 for LiPS.
For water, we assert unstable if any of the three element-conditioned RDFs: RDF ( O,O ) , RDF ( H,H ) , RDF ( H,O ) exceeds the threshold.
For flexible molecules, we keep track of stability through the bond lengths and say a simulation becomes “unstable” at time T when:
[Equation]
where 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.
We use ∆ = 0 .
5 Å for both MD17 molecules and alanine dipeptide.
We measure the stability of a learned model through the simulation time it remains stable.
For each dataset, the threshold ∆ we adopt is rather relaxed that an “unstable” equilibrium statistic usually indicates the system is already in a highly nonphysical configuration.
FPS.
All frames per second (FPS) metrics are measured with an NVIDIA Tesla V100-PCIe GPU.
We present FPS as a reference for models’ computational efficiency but also note that code speed can be affected by many factors, and likely has room for improvement.
[Figure]
B D ATASET DETAILS
The MD17 dataset 1 (Chmiela et al., 2017) and the LiPS dataset 2 (Batzner et al., 2022) are adapted from previous works and are publicly available.
We refer interested readers to the respective papers for more details on the data generation process.
The water dataset and alanine dipeptide dataset are generated by ourselves, and will be made publicly available after the reviewing process.
Water.
Our water dataset is generated from molecular dynamics simulations of a simple classical water model, namely, the flexible version of the Extended Simple Point Charge water model (SPC/Efw) (Wu et al., 2006) at temperature T = 300 K and pressure P = 1 atm.
This classical model has been well-studied in previous work (Wu et al., 2006; Yue et al., 2021) and has shown reasonable predictions of the physical properties of liquid water.
It provides a computationally inexpensive way to generate a large amount of training data.
The experience and knowledge gained from the benchmark based on the simple model can be readily extended to systems with higher accuracy, such as the ab-initio models.
Alanine dipeptide.
Our dataset is generated from the MD simulation of an alanine dipeptide molecule solvated in explicit water (1164 water molecules) performed in GROMACS (Abraham et al., 2015) using the AMBER-03 (Ponder & Case, 2003) force-field.
The NPT ensemble is applied in simulations, with hydrogen bond length constraints using LINear Constraint Solver (LINCS) and a time step offs.
The temperature and pressure of the system are controlled at T = 300 K and P = 1 bar using a stochastic velocity rescaling thermostat with damping frequency t v = 0 .
1 ps and Parrinello-Rahman barostat with coupling frequency t p = 2 .
ps, respectively.
The Particle Mesh Ewald approach is used to compute long-range electrostatics with periodic boundary conditions applied to the x, y, and z directions.
Implicit solvation.
The explicit solvent of 1164 water molecules is not the subject of study but adds a significant computational burden.
In this task, we attempt to learn an implicit solvent model (ISM) of the alanine dipeptide, in which the explicit solvent environment is incorporated in the learned FF.
The
ISM is commonly used in drug design (Wang et al., 2019a) because it can speed up the computation by dramatically decreasing the number of particles required for simulation.
In general, the mean-field estimation in ISM ignores the effect of solvent, thermal fluctuations, and solvent friction (Feig, 2007).
Thus, molecular kinetics is not directly comparable to the explicit solvation simulation.
However, the equilibrium configurations can be explicitly compared, as conducted in Chen et al.
Metadynamics simulation.
Simulating energy barrier jump usually requires a long sampling of the trajectory in MD simulations.
The conformational change of alanine dipeptide in water involves such a process, making it difficult to extract the complete free energy surface, i.e., the
Ramachandran plot, in normal MD.
In order to examine the learned ML FFs within a reasonable time limit, metadynamics (Laio & Parrinello, 2002) is employed to explore the learned FES of the solvated alanine dipeptide.
Metadynamics is a widely used technique in atomistic simulations to accelerate the sampling of rare events and estimate the FES of a certain set of degrees of freedom.
It is based on iteratively “filling” the potential energy using a sum of Gaussians deposited on a set of suitable collective variables (CVs) along the trajectory.
At evaluation time, we perform metadynamics
[Table] [Table]
with dihedral angles ϕ and ψ as CVs 3 , starting from the configurations located at one of the six energy minimums in the free energy surface indicated in Figure 5 (c).
The Gaussians with height h = 1 . 2 and sigma σ = 0 .
35 are deposited every 1 ps centered on ψ and ϕ .
As shown in Figure 8, the estimated FES of both ϕ and ψ do not significantly change after 5 ns.
In addition, the height of the bias gaussian potential smoothly converges to ∼ 0 in the time limit of 5 ns.
Therefore, a simulation time of 5 ns is sufficient for the convergence of the metadynamics.
This metadynamics simulation of alanine dipeptide with AMBER force-fields is carried out using GROMACS (Abraham et al., 2015) integrated with the PLUMED library (Tribello et al., 2014; Colón-Ramos et al., 2019) of version 2.8.
C E XPERIMENTAL DETAILS
The Open Catalyst Project 4 codebase and the official codebases of DeepPot-SE 5 and NequIP 6 are all publicly available.
We build our MD simulation framework based on the Atomic Simulation
Environment (ASE) library 7 (Larsen et al., 2017).
Our code will be made publicly available after the reviewing process.
Hyperparameters.
We adopt the original model hyperparameters in the respective papers and find they can produce good force prediction results that match the trend and numbers for MD17 reported in previous work.
As we introduce new datasets, we set training hyperparameters such as the batch size and summarize them in Table 7.
For water and LiPS, we use a batch size of 1 as each structure already contains a reasonable number of atoms and interactions.
Following previous work, we use an initial learning rate of 0.001 for all experiments except for NequIP, which uses 0.005 as the initial learning rate in the original paper.
For models that minimize a mixture of force loss and energy loss, we set the force loss coefficient λ F to be 1000 and the energy loss coefficient λ E to be 1 , if not specified in the original paper for a dataset.
A higher force loss coefficient is common in previous work (Zhang et al., 2018a; Batzner et al., 2022) as simulations do not directly rely on the energy.
[Table] [Table]
Width=32, r=
[Table]
=64,
[Table]
=
[Table]
Notably, NequIP proposed several sets of hyperparameters for different datasets, including MD17, a water+ice dataset from Zhang et al.
2018a, LiPS, etc.
We follow the MD17 hyperparameters of
NequIP for our MD17 and alanine dipeptide datasets; the water+ice hyperparameters of NequIP for our water dataset; and the LiPS hyperparameters of NequIP for the same LiPS dataset.
The only architectural adjustment we made is because we observed training instability for ForceNet on water using the original hyperparameters.
We resolve this issue by reducing the network width from 512 to 128 for ForceNet in our water experiments.
To facilitate fair benchmarking, we stop the training of a model if either of the following conditions is met: (1) a maximum training time of 7 days is reached on an NVIDIA Tesla V100-PCIe GPU; (2) a maximum number of epochs specified in Table 7 is reached; (3) The learning rate drops below 10 −with a ReduceLROnPlateau scheduler with factor 0.8 and learning rate (LR) patience specified in Table 7.
We also report the longest time for an ML model to finish our benchmark simulation in Table 7.
All numbers are results of NequIP.
The high computational cost for evaluating MD simulations has been a major consideration in designing our benchmark datasets and metrics.
Complete water results.
We present results on water-1k in Table 8 and results on water-90k in
Table 9.
Results on water-10k is presented in Table 4 in the main text.
All models generally achieve lower force error when trained with more data, but stability and estimation of ensemble statistics don’t necessarily improve.
Notably, DeepPot-SE shows clear improvement with more training data and becomes as good as NequIP on water-90k.
SchNet demonstrates significant improvement in stability, but the estimation of ensemble statistics does not improve.
This may be due to the limited accuracy of SchNet coming from the limited expressiveness of the invariant atomic representation.
Influence of model size.
Table 10 shows an ablation study over the model size and radius cutoff of NequIP over water-1k.
We observe that all models are highly stable and attain equally good performance in simulation-based metrics.
Although a small radius cutoff of 4 leads to worse performance in force prediction, it is more computationally efficient and preserves the trajectory statistics.
These results show that there exists a trade-off between accuracy and efficiency when choosing the hyperparameters of an ML force field, and force error may not be the preferred criterion for model selection.
Distribution of interatomic distances for MD17.
Figure 9 shows the h ( r ) curves for all models and molecules benchmarked.
We randomly selected one simulation out of the five simulations we conducted for each model and molecule.
We observe that due to lack of stability, DeepPotSE produces noisy h ( r ) on Aspirin.
ForceNet does not manage to learn the correct interatomic interactions and produces incorrect h ( r ) curves.
Most models are able to produce h ( r ) that match well with the reference, with SchNet being less accurate on Aspirin and Ethanol .
RDFs for water.
RDF curves for water-1k/10k/90k are in Figure 10, Figure 11, and Figure 12.
Most noisy curves are due to insufficient sampling time, which results in a small number of frames to be averaged in obtaining the RDF curves.
We observe that SchNet and ForceNet produce inaccurate curves that are not very noisy, showing that their failure is not entirely due to lack of stability but because of inaccurate modeling of interactions caused by limited expressiveness and lower sample efficiency.
Further, we note that the reference curves have zero values below a certain threshold, as any pair of atoms cannot get too close to each other.
However, DimeNet and GemNet-T exhibit abnormal high values for very small distances, indicating the simulations have gone into nonphysical configurations and collapsed.
RDFs for LiPS.
As shown in Figure 13, DeepPot-SE does not manage to stay stable on LiPS.
ForceNet learns inaccurate interactions and produces inaccurate RDFs.
All other models can produce highly accurate RDF and can reproduce Li-ion diffusivity relatively accurately, as demonstrated in
Table 6.
[Figure]
RD F (r)
[ Å 1 ]
DeepPot-SE, O-O
Reference Prediction
DeepPot-SE, H-H
Reference Prediction
DeepPot-SE, H-O
Reference Prediction
RD F (r)
[ Å 1 ]
SchNet,
[Figure]
-O
Reference Prediction
[Figure] [Figure]
[Table]
Prediction
[Figure]
[Table]
Prediction
RD F (r)
[ Å 1 ]
DimeNet, O-O
Reference Prediction
[Figure] [Figure]
[Table]
Prediction
[Figure] [Figure]
Reference Prediction
RD F (r)
[ Å 1 ]
ForceNet, O-O
Reference Prediction
ForceNet, H-H
Reference Prediction
ForceNet, H-O
Reference Prediction
RD F (r)
[ Å 1 ]
GemNet-T, O-O
Reference Prediction
[Figure]
[Figure] [Figure]
RD
[Figure]
(r)
[ Å
[Figure] [Figure]
RD
[Figure]
(r)
[ Å
[Figure]
RD F (r)
[ Å 1 ]
[Figure]
RD F (r)
[ Å 1 ]
[Figure]
RD F (r) [ Å
[Figure]
[Figure] [Figure]
[Figure] [Figure]
RD F (r)
[ Å 1 ]
DeepPot
[Figure]
SE, O-O
Reference Prediction
[Figure]
[Table]
[Figure] [Figure]
[Figure]
RD F (r)
[ Å 1 ]
[Figure]
Prediction
[Figure] [Figure]
[Figure] [Figure]
RD F (r)
[ Å 1 ]
[Figure]
Reference Prediction
[Figure] [Figure]
[Figure]
RD F (r)
[ Å 1 ]
ForceNet
[Figure]
O-O
Reference Prediction
[Figure]
Prediction
[Figure]
Prediction
RD F (r)
[ Å 1 ]
GemNet
[Figure] [Figure]