Learning Large-Time-Step Molecular Dynamics with Graph Neural Networks
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Molecular dynamics (MD) simulation predicts the trajectory of atoms by solving
Newton’s equation of motion with a numeric integrator.
Due to physical constraints,the time step of the integrator need to be small to maintain sufﬁcient precision.
This limits the efﬁciency of simulation.
To this end, we introduce a graph neuralnetwork (GNN) based model, MDNet, to predict the evolution of coordinates andmomentum with large time steps.
In addition, MDNet can easily scale to a largersystem, due to its linear complexity with respect to the system size.
We demonstratethe performance of MDNet on a 4000-atom system with large time steps, and showthat MDNet can predict good equilibrium and transport properties, well alignedwith standard MD simulations.
Introduction
Molecular dynamics (MD) is a simulation technique for analyzing the movement of many-bodysystems by solving Newton’s equation of motion.
The resulting trajectories provide a view of dynamicevolution of atoms and molecules, and can be used to analyze the equilibrium and transport propertiesof the system [1].
MD plays a key role in many scientiﬁc ﬁelds, including materials science [2, 3],biochemistry [4], drug discovery [5], etc.
The core of an MD algorithm is to solve Newton’s equation of motion at discrete time steps.
One ofthe most frequently used algorithm is the ordinary Verlet integrator [6], which predicts the coordinateof atom i at time t , q i ( t ) , after time step ∆ t by:
[Equation]
where the potential energy U of the system is a function of the coordinates q .
The time step ∆ tmust be set small enough to avoid the contribution of O (∆ t 4 ) .
In practice, ∆ t is usually chosenas 0.5 to 2 femtoseconds (1 fs = 10 − 15 s).
Recently, efforts have been made to solve the classicalmechanics of many-body systems via machine learning methods [7–12].
However, all these methodsare limited to small systems containing only dozens of atoms, whereas practical MD simulationsusually involve thousands to millions of atoms.
Therefore, a scalable machine learning approach withlinear complexity is highly desired.
In this paper, we propose MDNet, a scalable graph neural network (GNN) model to predict thedynamics of many-body systems.
The key contribution is to predict the movement of an atom basedon its local environment encoded by a series of descriptors similar as DeePMD [13, 14].
In addition,we introduce momenta in the local environment descriptor and make the model comply with the lawof conservation of momenta.
We show that MDNet can simulate a system containing thousands ofatoms with a time step 128 times larger than that of the Verlet integrator.
The resulting equilibriumand transport properties are shown to be consistent with the simulations using the Verlet integrator.
1.1 Related Work
The trajectories of MD simulations can be viewed as sequence data hence could be learned bysequence models, such as recurrent neural networks (RNNs).
Tsai et al. showed that one-dimensionalstochastic trajectories generated from higher-dimensional dynamics can be learned by long shortterm memory network (LSTM).
Kadupitiya et al.
demonstrated that LSTM was able to predict MDtrajectory for systems containing 16 atoms.
Another type of models including HNN [10], SRNN [11],
GFNN [12] and others [15–18], manged to incorporate physical knowledge and solve mechanicalproblems by learning Hamiltonian or other physical quantities.
However, all the aforementionedmodels used the information of all atoms to predict the dynamics, which limits their scalability.
In traditional MD simulations, a common assumption is that atoms only interact with their neighbors.
Motivated by this assumption, Satorras et al.
proposed E(n) equivariant graph neural networks(EGNNs), which model the system by a graph and achieve linear complexity.
Compared to EGNN,
MDNet use more sophisticated descriptors of local environments and achieve better performance.
Background
We brieﬂy introduce some basic notations in Hamiltonian mechanics that MDNet is built upon.
The
Hamiltonian for an N-particle system in d dimensions can be written as:
[Equation]
where p , q are the momenta and coordinates of the atoms.
The momentum is deﬁned as p i
= m i
v i ,where m i and v are the mass and velocity of atom i , respectively.
T ( p ) =
(cid:80) Ni =1 p
i 2 m i is the kineticenergy of the system and U ( q ) is the potential energy induced by interaction between atoms.
Themotion of atoms is determined by Hamilton’s equation:
[Equation]
which is a generalization of Newton’s equation of motion.
The equation is then solved at discretetime steps via numeric integrators, where the velocity Verlet (appendix A)
[20] is the most widelyused one.
Model Architecture
In Hamilton’s equation (3), the variation of q only depends on p and vice versa.
However, with alarger time step, atoms interact with each other, thus, ∆ q depends not only on p , but also q (and so is∆ p ).
In MDNet, we model ∆ q and ∆ p via neural networks:
[Equation]
and the integration algorithm becomes:
[Equation]
In MDNet, the system is represented by a graph where each atom is represented by a node.
Twonodes are connected by an edge if the distance between them is less than a cut-off distance, r c .
Foran edge between node i and j , its edge feature and message is calculated as:
[Equation]
where r ij = ||
q i − q j || 2 is the Euclidean distance between node i and j , e ij ∈ R m 1 and m
ij ∈
R m 2 .
In three-dimensional space, m 1 = 7.
The features 1 r i,j and q i − q j r 2 ij are motivated from DeePMD [13, 14]and p i − p
j brings the information of momenta.
Relative coordinates and momenta are used topreserve translation invariance.
The edge operator φ e is an MLP.
Then, edge features and messagesare gathered to nodes.
For node i with neighbors { j k } nk =0 , the local embedding L i is computed as
[Equation]
i
[Equation]
Here permutation invariance is preserved by the multiplication between E i and M i .
Finally, thevariation of coordinates and momenta are predicted by aggregating the contribution of each edge,
[Equation]
To ensure the conservation of momentum, for any edge ( i, j ) , we only compute one of ∆ p ij and let∆ p ji = − ∆ p ji .
This treatment is also applied to the calculation of ∆ q ij ’s.
Experiments
We test the performance of MDNet on a liquid argon system containing 4000 atoms.
The interatomicinteraction is modeled by Lennard-Jones potential [21]:
[Equation]
(10) where r is the distance between two atoms, (cid:15) and σ are set to 0.2374 kcal/mol and 3.432 Å respectively.
Although there is only one kind of interaction in the system, the dynamics could be rather complicated,with increasing space and time scale.
Dataset: To create the dataset, we run classical MD simulations with LAMMPS using Verletintegrator with a time step of 1 fs.
We independently sample 20 trajectories from different initializedconﬁgurations, where 15 are used for training, 3 are used for testing and 2 are used for validation.
For each trajectory, we randomly initialize 4000 atoms in a simulation box with periodic boundarycondition in all directions.
The system is ﬁrst relaxed under 80 K and 1 bar in the NPT ensemble with
Nosé-Hoover (NH) thermostat
[22] and barostat for 20000 steps.
Then, the system is further relaxedunder 80 K in the NVT ensemble with NH thermostat for 20000 steps.
Finally, a NVE simulation isperformed for 2560 steps and only the trajectory of this ﬁnal stage is used.
In each trajectory, wesample 10 frames { q ( t i ) , p ( t i ) } 10 i =1 of the system and use { q ( t i + ∆ t ) , p ( t i + ∆ t ) } 10 i =1 as targets.
According to the time reversibility of MD, the reversed data are also added to the dataset.
Besides, anextra 40960-steps MD simulation in the NVT ensemble is performed to validate roll-out performance.
Implementation details:
The cutoff radius r c for constructing the nearest neighbor graph is 7.0Å. The MLPs φ e and φ o both have 4 layers and use ReLU [23] as activation function.
φ e has 40neurons in each hidden layer and φ o has 200 neurons in each hidden layer.
Training is carried out for1000 epoch, batch size 1 with Adam optimizer [24].
Learning rate starts from 5 × 10 − 3 and decaysexponentially at a rate of 0.8.
The loss function is the mean squared error (MSE) of coordinatesand momenta, normalized by their standard deviations respectively.
We compare MDNet to EGNN[19] composed of 4 EGCL layers with 64 neurons per layer and Swish activation function [25].
Hyperparameters of EGNN are consist with the original paper.
One-step results: We present the one-step root mean squared errors (RMSE) obtained by the Verletintegrator, EGNN and MDNet on validation set in Table 1.
It is shown that MDNet reports comparable
RMSE with Verlet with a small time step (16 fs), and signiﬁcantly outperforms Verlet with a large
[Table]
time step.
As the time step increases, the RMSE of Verlet increases faster than MDNet, due to thecontribution of the fourth order error.
The performance of EGNN is poor with all the time steps.
Roll-out performance: MD simulations can be performed by adopting one-step prediction iteratively,which is known as roll-out.
We use two metrics to evaluate the roll-out performance of MDNet:radial distribution function (RDF) [26] and mean-squared displacement (MSD) [27, 28].
where RDFsprovide good descriptions of equilibrium structure and MSD is a measure of transport propertiesof simulation systems.
To compute RDFs, we conduct 40000 fs simulation and take 8 frames fromthe subsequent trajectory with a gap of 128 fs.
The ﬁnal RDFs are averaged over these 8 frames.
As shown in Figure 2 (left), RDFs obtained from MDNet with different time steps coincide withground truth, which show correct equilibrium structures.
To compute MSD, the trajectory is savedwith a resolution of 256 fs and calculate displacement against the ﬁrst frame.
As shown in Figure 2(right), MSDs obtained from MDNet are linear with time and are in good agreement with groundtruth, despite that the slopes with large time steps are slightly smaller.
In summary, MDNet canreproduce the equilibrium structures and transport properties of the simulation system with large timesteps.
We also compare the time cost of roll-out and show that MDNet enjoys faster roll-out speedcompared to Verlet, thanks to the large time step.
The comparison of time cost is in the appendix B.
Conclusion and Discussion
We introduce MDNet, a GNN model that can predict the dynamics of many body system with largetime steps.
We demonstrate that MDNet exhibites good precision on one-step prediction with a largetime step, and is able to generate MD trajectories.
The generated trajectories are characterized by
RDF and MSD, which showed good agreement with the ground truth.
It is worth mentioning that
MDNet only takes local environments as feature, so it could be generalized to larger systems with thesame species of atoms without further training and we plan to investigate this in our future work.
There are several future directions.
Firstly, the total energy is not guaranteed to be conserved duringroll-out of MDNet.
Incorporating thermostat into MDNet is necessary to control the temperatureand total energy.
Secondly, The Lennard-Jones system presented in this paper is the simplest casein MD applications.
With the introduction of chemical bonds, the system will exhibit multiplecharacteristic times and the dynamics will be much more difﬁcult to predict.
The extension of MDNetto complicated systems is worthy of further study.
A The Velocity Verlet Algorithm
The algorithm of the most commonly used integrator, velocity Verlet, is:
[Equation]
In practice, p i ( t + ∆ t ) may not need to be calculated explicitly.
Then, the coordinates and momentaare updated alternately, which is also known as leapfrog integration.
B Comparison of Time Cost
We demonstrate the time cost of roll-out of MDNet and Verlet algorithm here.forward time
[s] time cost per fs [s]
Verlet-
[Table]
MDNet-
[Table] [Table]