Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-scale Graph Networks
Authors: Xiang Fu
Tian Xie
Nathan J. Rebello
Bradley Olsen
Tommi S. Jaakkola
Venue: Accepted by TMLR ()

Abstract: Molecular dynamics (MD) simulation is essential for various scientific domains but computationally expensive. Learning-based force fields have made significant progress in accelerating ab-initio MD simulation but are not fast enough for many real-world applications due to slow inference for large systems and small time steps (femtosecond-level). We aim to address these challenges by learning a multi-scale graph neural network that directly simulates coarse-grained MD with a very large time step (nanosecond-level) and a novel refinement module based on diffusion models to mitigate simulation instability. The effectiveness of our method is demonstrated in two complex systems: single-chain coarse-grained polymers and multi-component Li-ion polymer electrolytes. For evaluation, we simulate trajectories much longer than the training trajectories for systems with different chemical compositions that the model is not trained on. Structural and dynamical properties can be accurately recovered at several orders of magnitude higher speed than classical force fields by getting out of the femtosecond regime. 

Published in Transactions on Machine Learning Research (08/2023)
Simulate Time-integrated Coarse-grained Molecular Dynam- ics with Multi-scale Graph Networks
Xiang Fu
ziangfu@csail.miz mit edu
Massachusetts Institute of Technology
Tian Xie
tionaie@microsoft.ccom
Microsoft Research
Nathan Rebello
nrebello@mit. edu
Massachusetts Institute Technology
Bradley D. Olsen
bdolsen@@mit. edu
Massachusetts Institute of Technology
Tommi Jaakkola'
tommi mit edu
Massachusetts Institute of Technology
Reviewed on OpenReview: https: // openreview. net/ formm? (d=y8RZoPjEUT
Abstract
Molecular dynamics (MD) simulation is essential for various scientific domains but computa- tionally expensive. Learning-basee force fields have made significant progress in accelerating ab-initio MD simulation but are not fast enough for many real-world applications due to slow inference for large systems and small time steps (fentosecond-level), We aim to address these challenges by learning multi-scale graph neural network that directly simulates coarsc-grained MD with a very large time step (nanosccond-level) and a novel refinement module based on diffusion models to mitigate simulation instability. The elfectivenese of our method is demonstrated two complex systems: single-chain coarse-grained polymers and multi-compooent Li-ion polymer electrolytes. For evaluation, we simulate trajectories much longer than the training trajectories for systems with different chemical compositions that the model is not trained on. Structural and dynamical properties can be accurately recovered at several orders of magnitude higher speed than classical force fields by getting out of the femtosecond regime.
1 Introduction
Molecular dynamics (MD) simulation is widely used for studying physical and biological systems. However, its high computational cost restricts its utility for large and complex systems like polymers and proteins, which often require long simulations of nanoseconds to microseconds. While machine learning (ML) force fields (Behler & Parrinello, 2007; Unke et al., 2021b) have made significant progress in accelerating ab-initio MD simulations, their suitability for long-time simulations of large-scale systems is limited by several factors. Firstly, each force computation requires inference through the MT. model. Although a single step with the ML model is much faster than ab-initio calculations, it is still slower than classical force fields. Secondly, simulating MD with a force field requires a very short time integration step at the femtosocond level to ensure stability and accuracy. This, combined with the slow single-step inference speed, results in extremely high
"Equal contribution.
1 "Correspoudence Xiang and Tommi Jaakkola: xiangfulcsail mit. edu, tommi0csail. mit. edu. Work done Massachusetts Institute of Technology.

---

Published in Transactions on Machine Learning Research (08/2023)
costs for long- time simulations of large systems. Lastly, past research has shown that ML force fields are susceptible to instability in long-time simulations (Unke et al., 2021b; Stocker et al., 2022; Fu et al., 2022).
In many MD applications, we are interested in structural and dynamical properties at a large time/ /leggt scale. For example, the high-frequency vibrational motion of molecular bonds (femtosecond-level) is irrelevant when studying ion transport in battery materials (nanosecond-leve) also does not require full atomic information, and a coarse gransed representation of the atomic environment suffices. To accelerate molecular dynamics simulations under large spatiotemporal scales, it is essential to develop model that (1) can simulate beyond the femtosecond time step regime; (2) can simulate at a coarso-grained resolution that preserves properties of interest; (3) remains stable during long-time simulations.
Computing the ion transport properties
To address these challenges, this paper proposes a fully-differentiable multi-scale graph neural network (GNN) model that directly simulates coarse grasned MD with a large time step (up to nanosecond-level) without computing or integrating forces. Our model includes an embedding GNN that learns transferrable coarse- grained bead embedding at the fine-grained level and a dynamics GNN that learns simulate coarse-grained iime-integrated dynamics at the coarse coarse-grained level. Additionally, refinement module based on diffusion models (Song& & Ermon, 2019; Ho et al., 2020; Song et al., 2020; Shi et al., 2021) to resolve instability issues that may arise during long-time time simulations. Although each step of the ML simulator is slower than classical force fields, the learned model can operate at spatiotemporal scales unreachable by classical methods, which leads to a dramatic boost in efficiency (summarized in Table 1).
The efficiency of our proposed method makes it particularly useful in tigh-throughput screening settings and enables us to conduct experiments large-scale systems that are computationally prohibitive for existing ML force fields. Our evaluation protocol requires the simulation of thousands of atoms for billions of force field steps, which would take a typical ML (CG) force field months to simulate, whereas our model can complete within a few hours. In two realistic systems: (1) single-chain yoarse-grained polymers in implicit solvent (Webb et al., 2020) and (2) multi-component Li-ion polymer electrolytes (Xie et al., 2022), We demonstrate that our model can (1) generalize to chemical compositions unseen during training; (2) efficiently simulate much longer trajectories than the training data; (3) successfully predicts various equilibrium and dynamical properties with a 10a- 1 wall-clock speedup compared to the fastest classical force field implementations.
2 Related Work
Machine learning force fields (Unke et al., 2021b) aim to replace expensive quantum-mechaniel calculations by fitting the potential energy surface from observed configurations with force and energy labels while being more efficient. An extensive series of research ranging from kernel-based methods to graph neural networks (Behler & Parrinello, 2007; Khorshidi & Peterson, 2016; Smith et al., 2017; Artrith et al., 2017;
Chmicla et al., 2017; 2018; Zhang et al., 2018a;e; Thomas et al., 2018; Jia et al., 2020; Gasteiger et al., 2020; Schoenholz & Cubuk, 2020; Noé el al., 2020; Doerr el al., 2021; Kovács et al., 2021; Satorras el al., 2021; Unke et al., 2021a; Park al., 2021; Thölke & De Fabritiis, 2021; Gasteiger et al., 2021; Friederich et al., 2021; Li et al., 2022; Batzner et al., 2022; Takamoto et al., 2022) has shown ML force fields can attain incredible accuracy in predicting energy and forces for variety of systems. However, limited by slow inference and small time step sizes, simulation with large spatiotemporal scales is still inaccessible through ML. force fields. This paper instead aims to learn a direct surrogate model at the CG level that bypasses force computation and enables efficient large-scale simulation.
Coarse-grained (CG) force fields (Marrink et al., 2007; Brini et al., 2013; Kmiccik et al., 2016) have been developed to extend the time and length scales accessible to molecular dynamics (MD) simulations while sacrificing some accuracy for computational efficiency. Despite this trade-off, these models have shown great success in simulating diverse biological and physical systems (Sharma et al., 2021; Webb et al., 2020). Machine learning methods have been used to learn CG mappings (Wang & Gómez-Bombarelli, 2019; Kempfer et al., 2019) and CG force fields through various approaches (Dequidt & Solano Canchaya, 2015; Wang el al., 2019; Husic et al., 2020; Greener & Jones, 2021; Chennakesavalu et al., 2022; Nkepsu Mbitou et al., 2022; Arts et al., 2023). However, A CG model that can accurately recover long-time properties is not known for every system (e.g., the Li-ion polymer electrolytes considered in this paper). Moreover, existing coarso-grained force fields for MD simulations are still limited by the femtosecond-level time step requirement for force
2

---

Published in Transactions on Machine Learning Research (08/2023)
Preprocessing
CG MD Simulator
Gfass :
Embedding GNN
Coarse Graining
T-kat
:
Bootstrap Simulate
Dynamics GNN
G.
Embedding GNN
Coarse Graining
G
Score GNN
+
Training
G G+
Embedding GNN
Coarse
Prediction
GC
Graining
I+AI
Loss
I+II
Learning time-integrated CG MD with multi-scale graph neural networks. Trainable modules are colored in blue. The loss term colored in red. The preprocessing steps embed and coarse-grain an MD system to a coarse-level graph. The CG MD simulator processes historical information using the featurizer o, after which the Dynamics GNN predicts the next-step positions. A Score GNN is applied to refine the prediction. We bootstrap the CG MD simulator at evaluation time to make aulo-regressive simulations.
Figure 1:
integration. Our approach uses a general-purpose graph clustering algorithm for coarse-graining, It bypasses force integration to use much longer time steps with extreme coarso-graining while preserving key long-time structural and dynamical properties.
Enhanced sampling methods (Sugita & Okamoto, 1999; Laio & Parrinello, 2002; Barducci al., 2008; Valsson & Parrinello, 2014; Torric & Valleau, 1977) have proven successful in accelerating sampling of transition between metastable states in complex system dynamics such as protein folding (Bernardi et al., 2015; Schneider et al., 2017; Sultan et al., 2018; Yang et al., 2019). However, these methods rely on identifying collective variables (CVs) specific to each system, which can be challenging, Furthermore, they lack an explicit notion time, making it difficult to estimate dynamical properties (Laio et al., 2005; Stelzl & Hummer, 2017). More recently, machine learning generative modeling approaches, such as Boltzmann generators, have enabled fast sampling of equilibrium states across phases (Noé al., 2019; Mahmoud al., 2022; Sidky et al., 2020; Vlachas et al., 2021; Kaltenbach & Koutsourelakis, 2021; Klein et al., 2023). Despite previous efforts, generating dynamical trajectories that can be applied to wide range of chemical compositions and to large-scale systems with thousands of atoms remains significant challenge. Another relevant line of work uses dimensional reduction techniques for building Markov state models and discovering long timescale kinetics (Pérez- Herez--zennndez et al., 2013; Mardt et al., 2018; Noé et al., 2016; Klus et al., 2018; Wu et al., 2018; Wu Noć, 2020; Noć et al., 2020; Vlachas et al., 2021), but they don't directly construct an efficient. MD simulator in the reduced space. Our goal in this work is to develop a general-purposs model that can simulate coarse-grained MD and recover equilibrium and dynamical long-time ensemble properties without. prior knowledge of CVs.
3 Learning Time-integrated Coarse-grained Molecular Dynamics
Model Overview. As depicted in Figure 1, the learned simulator predicts single-step time-integrated CG dynamics at time + At given the current CG state and historical CG states at 1, At...., k&t. Here
At the time-integration step, which significantly longer than that used in MD simulation with force fields. Given ground truth trajectories al atomic resolution, adopt 3-step multi-scale modeling approach:
Fmbeeding: learning atom embeddings at fine level using an Embedding GNN GNz;
Course-gruining: coarso-graining the system using graph clustering and constructing CG bead embodding from atom embodding learned step 1;
Dynamics (and Refinement): learning time-integrated acceleration at coarse level using a Dynamics GNN GNp. A Score GNN GNş is optionally learned to further refine the predicted structure.
3

---

Published in Transactions on Machine Learning Research (08/2023)
Representing MD trajectories as time series of graphs. A ground truth MD simulation trajectory is represented as a time series of fine-level graphs }. J. The fine-level graph G represents the MD state at time step i, and is defined as tuple of nodes and edges Gf = (Vē", E²). Each node Vi. E VV"" represents an atom ' and each edge ef. E" represents chemical bond between the particles of and vf.. The static fine-lovel graph GF (V² EEP) describes all persistent features in an MD simulation, which include atom lypes, atom weights, and bond lypes. These persistent features are used to construct a time-invariant node representation that will be later used for CG bead embedding in our model. Applying the CG model to the fine-level graphs {G.') } produces the time series of coarse-level graphs {G where the CG state G is defined by the tuple G = (V, E. Each node vn E V.C represents CG bead, and each edge e E E models an interaction between the CG beads e and v Since both non-bonded interactions and bonded interactions at the coarse level are significant for dynamics modeling, the edge set E. contains both CG-level bonds and radius cut-off edges constructed from the CG bead coordinates at time t.
Graph neural networks. A graph neural network takes graphs G E) with node/edge features imputs and processes a latent graph G = (V*, E²)) with latent node/edge representation through several layers of learned message passing, In this paper, we adopt the ENCODER- PROCESSOR- DECODDE architecture (Sanchez-Gonzalee ct al., 2020) for all GNNs, which inputs fcaturized graph and outputs a vector for cach node. A forward pass through the GNN follows three steps: (1) The ENCODER contains a node multi-layer perceptron (MLP) that is independently applied each node and an edge MLP that is independently applied to cach edge to produce encoded node/edge features. (2) The PROCESSOR is composed of several layers of directed message- message-passing layers. It generates sequence of updated latent graphs and outputs the final latent graph. The wessage-passing layers allow information to propagate between neighboring nodes/edges. (3) The DECODER is node-wise MLP that is independently applied to the node features obtained from message passing to produce the outputs of a specified dimensionality. We refer interested readers to Sanchez-Gonzale et al. 2020 for more details on the GNN architecture.
Learning CG bead type embeddings with Embedding GNN. Each node in the static fine-level graph of = [a, wr[) is represented with (1) a learnable lype embedding of that is fixed for given alomic number and (2) a scalar weight wf" of the particle. The edge embedding is the sum of the type embedding of the two endpoints and learnable bond type embedding: ef,- = [a." + af" + a.-). The bond type embedding afs is also fixed for a given bond type. We input this graph G to the Embedding GNN GNg that outputs node embeddings of" = [c]]] for all of E VF. This learned node embedding contains positional information and will be used in the next noarso-graining step for representing CG bead types. The Embedding GNN is trained mnd-to-end with Dynamics GNN and Score GNN.
Graph-clustering CG model. Our CG model assigns atoms to different atom groups using a graph partitioning algorithm (METIS (Karypis & Kumar, 1998)) over the line-level graph. The METIS algorithm partitions atoms into groups of roughly equal sizes while minimizing the number of chemical bonds between atoms in different groups. It first progressively coarsens the graph and partitions the coarsened graph, then progressively expands the graph back to its original size while refining the partition to obtain the final atom group assignment. Each node in the fine-level graph is assigned a group mumber: C(v.) E {1.....M) M), for all vhe E V.F. Here M is the number of atom groups (CG beads). The graph partitioning is applied to the bond graph which is invariant through time as we study equilibrium state dynamics. This noarse-graining approach ensures that two atoms grouped into the same CG bead will never be far away from each other since they are connected by a path of bonds.
Construct CG graph states. We represent each fine-level node vIr E V." with vfr ic wf" -.,.l, where of is the learned node type embedding from GNg, wf is the weight, and af,, is the position. We can then obtain the coarse-level graph Gf = (VC, FF{) by grouping atoms in the same group into CG bead. Denote the set of fine-level atoms with group number m as C. (i C(ef)) = m). The representation of the CG
If the ground truth MD simulation already uses a CG model, each node will correspond to CG bead defined by the CG model, and correspondingly the set of edges will he the set of all CG chemical bonds. For ease of presentation, in the rest of this paper, we refer to particles in the line-level graph as "atoms".
4

---

Published in Transactions on Machine Learning Research (08/2023)
bead u Icn; w am
is defined as following:
c = meam(())))) = Liec.. of |Cml
uz) Liec wf of Eiec wf
Wm
C
sum
Σ iECm
CoM( wi
(1)
Cm
Cm
The operators meanc,, stands for laking the mean over Cm: sum stands for taking the sum over Cm, and CoMo,, stands for taking the center of mass over Cm. Applying this grouping procedure for all atom groups m E [1...., M) creates the set of all CG nodes V,"" = (u , : m E [I....,M]} for the coarsc-level graph Ge = (V,, (V),
We next construct the edges E.C CG bonds are created for bonded atoms separated into different groups. We create a CG-bond u. if chemical bond exists between a pair of atoms in group Cm and group C.. That is: mm,n, E EE = di E Cm:j E Cn. such that ef, E E", We further create radius cut-off edges by, for each CG bead, finding all neighboring beads within a pre-defined connectivity radius r (with consideration to the simulation setup, e.g., periodic boundaries): em... E Es = | x < r. We set large enough connectivity radius to capture significant interactions between all pairs of CG beads.
Learning CG MD with Dynamics GNN. The Dynamics GNN GNn inputs a history-augmented coarse graph state of @H and outputs distribution of the time-averaged acceleration for each CG bead. include the CG bead type embedding che weight wh, current and k-step history velocities { (mm,--iat ME0 C,,,t = - =n, I - = Il, sm include displacement and distance between the two endpoints, and an embedding vector ch,, indicating whether em,,, is a CG-bond or is constructed through radius cut-off. = NN af) [N(m,,, m o E VV,}}
The input node features of the featurized graph oh = Ich whe f mi,e-ssell o
[4-1 The input edge features
The output is a 3-dimensional Gaussian: GN ( 2 where b4 and of are the predicted mean and variance at time i, respectively. The training loss Laya for predicting the forward dynamics is thus the negative log-likelihood of the ground truth acceleration:
Layn = - log.N Ipc: o²)
The end-to-end training minimizes Ldyn, which is only based on the single-step prediction of time-integrated acceleration. At inference time (for long simulation), the prodicted acceleration &. ~ (Nu,, is sampled from the predicted Gaussian and integrated with semi-implicil Euler integration to updale the positions to the predicted positions &rt Eits: = é; + S, ait Dx + AttAt
Learning to refine CG MD predictions with Score GNN. We introduce the Score GNN, a score-basod generative model (Song & Ermon, 2019) to resolve the stability issue of long simulation for complex systems. Following the noise conditional score network (NCSN) framework (Song & Ermon, 2019; Shi et al., 2021), the Score GNN is trained to output the gradients of the history-conditional log density (i.e., scores) given history state information and CG bead coordinates as input. An incorrect 3D configuration can be refined by iteratively applying the learned scores. With the refinement step, each forward simulation step follows a predict-the--reine procedure. The Dynamics GNN first predicts the (potentially erroneous) next-step positions, which the Score GNN refines to the final prediction d++Ar:
During training, the Score GNN trained to denoise noisy CG bead positions to the correct positions. Let the noise levels be a sequence of positive scalars o1....,02 oz satisfying 01/02 = OL-1/OL 1. The noisy positions &cC' # is obtained by perturbing the ground truth positions with Gaussian noise: &c o = = + (0, o²), where multiple levels of noise a E (arHe are used. Due to coarso-graining, the coarse-level dynamics is non-Markovian (Klippenstein et al., 2021). Given the current state and historical information H that causes oftAt the noisy positions follow the distribution with density: ( IfIdda The Score GNN outputs the gradients of the log density of particle coordinates that denoise the noisy particle positions &C #H+ to the ground truth positions C 2440
J Paata (
conditional on the current state and historical information. That is, Vo E farHi.n
for case of notation, in the rest of this paper we omit the node/edge indices when referring every node/edge in the graph.
5

---

Published in Transactions on Machine Learning Research (08/2023)
Table 1: Summary of datasets, simulation time scales, and evaluation observables. # atoms is the average number of atoms per system. CG level represents how many atoms are grouped into single CG bead. GPU workflow for MD simulation is not always available. The simulation speed of classical MD a GPU is usually 10x that of CPUs (Thompson et al., 2022), which would still be significantly slower than our method.
Dataset
# atoms
OG level
Training traj.
Testing waj.
MD A1
Ours At
MD cost
Our cost
Key observable
Single-chain polymers
890
100
100 x 50k v
10 5M
0.01 7
5 32.7 CPU has
0.10 GPU has
Radius of geration lon didfusivity
Li-ion polymer electrolytes
6025
7
300
50 50
0.5 ls
2- 2-10P b
867.3 CPU hes
0.76 GPU ims
GNs(Lu, Ch # a(+ai) Vag lo Po ( (
= Vaga log Po
CC Bping AD) (1 - #1taill/a
Note that in step (+) we are approxinating H. with learned latent node embeddings ohth which is the last-layer hidden representation output by the Dynamics GNN. of contains current state and historical information. The training loss for GNs is:
GNs(lue Ch &C ael)
C 1tt +++At
27
Lecore
1 2L
Σ i=1
M(o,)
E
dmal soe H)) Bo, (&f. (of
(2)
o7
where A(o;) = of is a weighting coefficient that balances the losses at different noise levels. During training, the first expectation in Equation (2) obtained by sampling training data from the empirical distribution, and the second expectation obtained by sampling the Gaussian noise at different levels.
With the Score GNN, training is still end-to-end by jointly optimizing the dynamics loss and score loss:
== Layn + Liscore- At inference time, the refinement starts from the predicted positions from Dynamics GNN. We use annealed Langevin dynamics, which is commonly used in previous work (Song & Ermon, 2019), to iteratively apply the learned scores to gradually refine the particle positions with a decreasing noise term. The predict-then-reffe procedure can be repeated to simulate complex systems for long time horizons stably.
4 Experiments
Our experiments aim to substantiate our model's proficiency in elfectively simulating coarse-grained, time- integrated dynamics for two large-scale atomic systems, which are concisely presented in Table 1. It is important to note that the primary objective of learning MD simulations is not the precise recovery of atomic states based on initial conditions. Rather, the focus lies in deriving macroscopic observables from MD trajectories characterize the properties of biological and physical systems. This underlying motivation also drives the adoption of grained methodologies in previous studies as well as the present paper. compute essential macroscopic observables from the learned
To evaluate our simulation's performance, simulation and compare them to the reference data. Our model is contrasted with supervised learning basclines trained to directly predict observables from atomic structures. Our task presents considerable challenge, as it demands generalization lo various chemical compositions and extended time scales. As demonstrated in Table 1, for both datasets, the observables we aim to obtain necessitate lengthy simulations for estimation (length of testing trajectorics), while the training trajectorics are substantially shorter and inadequate for reliably extracting key observables. An ablation study included in Appendix A. Further details the simulation, observables, parameters, and baselines included in Appendix B and Appendix C.

---

Published in Transactions on Machine Learning Research (08/2023)
b
c
x10
Backbone beads
Class-1
1.00
Pendant beads
Ri 2 Rg R3-RA
0.75
Constitutional units
Class-11
0.50
[ Rg Rg R4
0.25
2
3
10
0.0
0.5
1.0
1.5
2.0
25
300
Simulation Time (T)
x10
d
Figure 2: (a) All polymer chains are composed of 400 constitutional units, which are formed by four types of CG beads. (b) Class-1 polymers are used for training, while class-1l polymers are used for testing. The structure variation requires the model to learn generalizable dynamics. (c) R2 for the three training polymers with smallest, median, and largest (Ry), over a 300k t period. Our training trajectories are 50k T long (black dashed line), while we use 5M T long trajectories for evaluation. (d) The coarse grarrssgrrining process of class-1 polymer used for training. The first illustration shows the original polymer, the second illustration shows the CG assignment from the graph clustering algorithm (where beads belonging to the same super CG bead have the same color), and the third illustration shows the final CG configuration. More details on the CG process can be found at Equation (1). (e) The coarse-grasd graining process of a class-11 polymer used for testing, with similar illustrations to (d). We note the irregular structure in the circled area in the polymer illustration, which is a feature of class-II random polymers.
4.1 Single-chain Coarse-grained Polymers
In our first experiment, we focus on simulating single-chain coarse-grained polymers" within an implicit solvent. We employ the polymers introduced in Webb et al. 2020, where all polymers are composed of ten types of constitutional units that are formed by four types of CG beads (as shown in Figure 2 (a)). The interaction of the CG beads is described by a summation of bond, angle, and dihodral potentials, along with non-bonded interaction potential (detailed in Equation (3)). We train on regular copolymers with an exact repeat pattern of four CUs (class-T) and test random polymers constructed from four CUs (class-IT). This distribution shift (Figure 2 (b)) poses challenges to the generalization capability of ML models.
Dataset and observables. We train our model on 100 short class-I MD trajectories (with ten trajectories for validation) of 50k T, which are not sufficiently long for observable calculations. For evaluation, we use 40 testing class-ll polymers using trajectories of 5M т (100x longer). All polymers are randomly sampled and simulated under LJ units with time-integration of 0.01 т. Each polymer contains 890 beads on average. A time step of Al 5r is used for our model, so every step of the learned simulator models the integrated dynamics of 500 steps of the classical CG force field. We use the learned simulator at test time to generate 5M jeatoeesis and compute properties. We coarse grain every 100 CG beads into super CG bead. Examples of the coarse saasse-graiing process are illustrated in Figure 2 (d.e).
The squared radius of gyration (R2) is a property practically related to the rhcological behavior of polymers in solution and polymer compactness that are useful for polymer design and has been the focus of several previous studies (Upadhya et al., 2019; Webb et al., 2020). Reliable estimation of R2 statistics requires
The classical MD definition is already for coarse-grained system. To avoid confusion, in the rest of this paper, we only use "coarse-grained" when we are referring to the polymers after our graph partitioning- based coarse graining approach.
7

---

Published in Transactions on Machine Learning Research (08/2023)
50k MD, -233.34
GNN, /2-0.70
LSTM, L2TM,r2--0991 ~2-00.91
Ours, r2-0.90
x10
x10
x10
Calculated (R
4.5
Predicted (R2
4.5-
Predicted (R?)
4.5-
Predicted (R
4.0
4,0
4,0-
4.00
3.5
3.5
3.5
3.5
3.0
3.0
3.0
3.0
2.5
25
3.0
3,5
4.0
4.5
2.5
3.0
3.5
4.0
2.5
3.0
x10
x10
x10²
x10*
5M MD (R3)(oF)
5M MD (R3x(22
5M MD (R3
5M r MD (R3
Figure 3: (Rg) estimation performance using different methods. From left to right: (1) mcan R2 from 50k
classical CG MD; (2) supervised learning GNN; (3) supervised learning LSTM; (4) our learned simulator.
sampling long MD simulation, while short simulation results in significant error. Figure 2 (c) shows how R2 rapidly changes with time.
Estimate R2 using learned simulation. In Figure 3 we clemonst the experimental results on mean Rg prediction across 40 testing polymers. The first panel shows the prediction made by short simulations of the same length as the training data will result in high-variance and poor recovery of mean R² due to insufficient simulation time. The second the third panels show the results of two baseline supervised learning (SL) models: one based on GNN and another based on long short-term memory (LSTM) (Hochreiter & Schmidhuber, 1997) networks. SL models input the polymer graph structure and directly predict the mean and variance of R2 without learning the dynamics. These supervised learning models exhibit systematic bias, possibly due to the distribution shift from training class-T polymers to testing class-IT polymers. The fourth panel shows the results of our learned simulator, which demonstrates accurate recovery of (122) and significantly outperforms the bascline models. This result shows that our model is able to learn dynamics transferrable to different chemical compositions and longer time horizons. Table 2 summarizes more performance metries. Our method significantly outperforms the bascline models and estimation from short 50k MD trajectories in all metrics. In particular, the earth mover's distance (EMD) is a measure of distribution discrepancy. A lower EMD indicates better match between the predicted R3 distribution and the ground truth (Table 2). We also experiment with the existing learned simulator model (Sanchez- but find it quickly diverges under the long-time simulation task setup.
et al., 2020) without modifications
Distribution of R3 and mean internal distance. In Figure 4, we demonstrate our model's capability repro- duce R? distribution and mean internal distances for four selected testing polymers with small to large (Rg). The first row shows our model can recover the distribution of RH very well, which is also quantitatively reflected by the low EMD Table 2. The second row shows a good recovery of the mean internal distances across different length scales. While the learned simulator is only trained to make single-step predictions with short trajectories as training data, these statistics can be accurately recovered through a long simulation of 5M 7 long, which is 500M steps of the original simulation. Our multi-scale modeling approach plays an important role achieving this result: as every 100 CG beads are grouped into a single super CG bead, the features of cach super CG bead must be well-represented to enable accurate dynamics
Table 2: Performance for predicting R3 statistics. 2² and MAF are computed for (R and EMD is computed for R3 distribution. To evaluate EMD, the SL models output Gaussian distribution with the predicted mean and variance.
Method
,2 (t)
MAE (4)
EMD (1)
Short MD, 50k
0.34
349
4.10
GNN, SL
0.70
263
2.82
LSTM, SL
-0.91
559
5.82
Ours, 5M
0.90
140
1.60
under such a high degree of coarsc-graile gransi-grainig.. The message passing of the embedding GNN, which happens the fine-grained level, is capable of learning to represent. the composition that. is necessary for accurate
8

---

Published in Transactions on Machine Learning Research (08/2023)
x10
EMD-0.58
x10
EMD-045
x10 +
EMD-0.80
x10 4
EMD-1.05
Prediction Reference
5
Density
3-
3.
2-
2
2-
R (o*)
x10
RG (o*)
RG (o*)
RG (o*)
Prediction Raference
60
60
60
<R
40
40
40
40
20
20-
20-
20-
2
6
8
5
4
5
10
3
5
7
Figure 4: The distribution of R2 from our learned simulation compared to the reference ground truth data; and the mean internal distance, respectively, for four example polymers with small to large (1Z)) the first row, the dashed line annotates the mean R2- Our model demonstrates accurate recovery of the distribution of Ra and mean internal distance for various testing polymers unseen during training.
le2=6.4k
11151
[g.-14.0k
(g.-20.3k
Prediction
Reforence
0.75
0.75
0.50
0.50
0.50
0.50
0.25
0.25
0.25
0.00
0.00
0.00
2
A time (r)
x10*
A time (T)
x10*
time (1)
x10*
time (t)
x10*
Reference
Prediction
Mean MAE=0.15
2.0 x10
0
80
2-
60
Distance IA)
Predicted te
1.5
4-
40
1.0
6-
6-
20
8-
0.5
5M MD fgg (T/110*
2
2
Figure 5: (a) Illustrations of four polymers and the autocorrelation function of R2 for four polymers with the smallest, intermediate, and largest R2 relaxation time. (b) The contact map for an example polymer produced by the reference simulation and the learned simulation. The third panel shows the absolute error. (c) Prediction performance of model on the Ra relaxation time.
dynamics, so the coarse-level dynamics GNN can carry out accurate simulations with orders of magnitude higher efficiency (as shown in Table 1).

---

Published in Transactions on Machine Learning Research (08/2023)
b
Atomic
Component
Goamse-graied
Polymer
TPS-
Ya Yy=0, NH, Sample pharma molecules
Lix
Figure (a) The chemical space for the polymer in SPEs. Every trajectory contains a different type of polymer (around 10 chains, each chain has around 10 monomers). and the testing systems have distinct. polymers from the training systems. (b) Example SPE illustrated in (1) full atomic structure; (2) colored coded by component; (3) coarse-grained structure. (c) The hoars-ggraining process for TFSI-ion. (d) The coarse-graining process of polymer chain in an example training SPE. From left to right: (1) the molecular
graph of the monomer molecule, where connection points are marked with the symbol <*; (2) the atomic structure of the polymer chain; (3) The CG assignment from the graph clustering a algorithin; (4) The final Soarse-graines configuration; (5) The relative scale of a single polymer chain in the periodic bounding box, with unwrapped coordinates. (c) The coarso-graining process of polymer chain an example testing SPE, with illustrations similar to (d).
Relaxation time scale. we show our learned simulator capture the long-time dynamics using the autocorrelation function (ACF) of R2 as our observable. The relaxation time is long-time dynamical property significantly more challenging to
(cmmtal
10-1
estimate than mean I?2. It can not be obtained from many independent samples of polymer and requires very long simulations. Labels cannot obtained from the
10
10
10
training trajectories that are overly short, so the SL. baselines are not applicable. Figure 5 (a) shows the ACF computed from our learned simulation matches the ground truth well for four polymers with the fastest, intermediate, and slowest. decay of the R2 ACF. We compute the relaxation time tæ from the ACFs (details in Appendix B) and show the prediction of (gg across all 40 testing polymers in Figure (c). Recovery tng shows that our model captures realistic dynamics rather than just the distribution of states. In Figure 5 (b), we visualize the contact map of an example polymer extracted from the reference simulation and our learned simulation. The contact map shows the pairwise distance between each pair of beads. Our model can accurately recover this detailed structural property.
40
(ns)
Figure 6: Estima- tion of Li-ion diffusiv- ity with various tra- jectory lengths. The black dashed line an- notates 5 ns (training trajectory length).
4.2 Multi-component Li-ion Polymer Electrolyte Systems
Solid polymer electrolytes (SPEs) are a type amorphous material system that is promising advancing Li-ion battery technology (Zhou et al., 2019) and is significantly more complex than singlo-chain polymers. Atomic-scale MD simulations have been an important tool in SPE studies (Webb al., 2013; Xie al., 2022), but computationally expensive. The challenge of modeling these systems comes from (1) their
10

---

Published in Transactions on Machine Learning Research (08/2023)
large size of thousands of atoms and the complexity of multiple components; (2) their amorphous nature and the slow convergence of key quantities, requiring long MD simulations on the order of 10 to 100 ns to estimate. For these reasons, screening large chemical space with long simulations is extremely costly. Past work (Xic et al., 2022) has studied supervised learning approaches to prodict ion transport propertics but requires labels from long-time simulations to predict long-lime properties accurately. In contrast, this work aims to predict long-time properties (50 ns) with short-time training data (5-ns trajectorics) only by learning to simulate the dynamics that govern the system evolution.
We adopt the SPEs introduced in (Xie et al., 2022). Each MD trajectory is for system with distinct type of polymer. The polymer space is defined in Figure 7 (a), which contains monomers constructed with large pharmaceutical database (Irwin & Shoichet, 2005). Each SPE system contains 6025 atoms on average. We train on 530 short MD trajectories of 5 ns (with 30 trajectories for validation) and evaluate on 50 testing SPE trajectories (with distinct polymers unseen during training) of 50 ns long- Our task is to predict the diffusivity of fast-moving ions Li-ions and TFSI-ions (DLITFFSI - which closely related to the conductivity of the battery material. One key challenge in simulating the Li-ion transport in SPEs is the slow relaxation process in amorphous polymers. Consequently, estimating the diffusivity of particles requires very long simulation time to sample the dynamics, as shown in Figure 6. For accurate property estimation, we need trajectory lengths where DL: converges. The 5-ns training trajectories are far from enough for extracting reliable Li-ion diffusivity, requiring the predictive method to generalize to longer time scales.
Our CG model groups every bonded atoms into CG bead us- a ing graph clustering. The coarse- graining process of the entire sys- tem is illustrated in Figure 7 (b);
Cutoff-O.5
Cutoff-1. A
Outoff-2.0.
200
Ground Tush Wil Saaee GNN
Boo Culdion
50-
Na.rcmm Ghi
200
the coarse-graining process of a TFSI-ion is illustrated in Fig- ure (c); the coarse-graining pro- cess example polymer chains are il- lustrated in Figure 7 (d,e). Notably, use time-integration step of Al = 0.2 ns, making each step of the learned simulator equivalent to 10s steps in classical MD. Such a long time step will smooth out all high-frequency vibrational modes of motion while slow modes of mo- tion are preserved. It enables dra- matic speedup of simulation com- pared to sino-ffllldbbscee simula- tions with ferntosecond-leee time steps. Ablation studies on the hy- perparameters are included in Ap- pendix A.
30 Stm cotin
10
Simu allon Time
Simutttion
Liion Li-ge
Liiion TPS-ion
L-tion Po yiner
Ground bruth
With Score GAN
No Score GNN
TES-onn TES--on
Polymer, Polymer
TESt-on Polyter
Ral
Radius (A)
Rodius A
Figure 8: (a) Bead collision as function of simulation time averaged over the 50 testing SPEs for all methods. The learned models start simulating at 4 ns. From left to right: consider cut-off radií of 0.5 Å, 1.0 A, and 2.0 A. (b) RDFs of different types of particles in the SPE systems for our model with/without wwihoutt the Score GNN refinement and the ground truth MD simulation averaged over a 50-ns trajectory of an example SPE.
Stability of learned simulation and Score GNN refinement. Stability is a prerequisite for the accurate recovery of simulation en- semble properties. As a measure of stability, we consider a "bead collision" happens when two beads have a distance below cut-off radius. At selected cut-off radius, realistic level of bead collision can be computed from the reference data. Figure 8 (a) demonstrates the number of bead collisions all 50 test SPEs, for three different cut-off radii: 0.5 A, 1.0 A, and 2.0 A. Simulation without Score GNN refinement suffers from a much higher level of bead collision, and the system becomes increasingly unphysical
function of simulation time, averaged over
11

---

Published in Transactions on Machine Learning Research (08/2023)
a
ns MD. MAE-3.92
SL GNN, MAE-0.46
No Score GNN, MAE-0.51
Curs, MAE-0.11
Predicted Dy Seglem hell
-7.5
Prodected Dy /e/ooeeeem (i1)
-7.5-
Predicte D. Ne.eeennmm
-7.5
Prodicted D. BogttcmPiin
-7.5
88.0
8.0
-8.0
-8.0
8.5
8.5
8.5-
.55
-9099
-8.5
75
-9.0
-9.0
-8.5
-8.0
-8.5
-7.
-90
8.00
75
50 ns MD Dul (olgommm
30 MD (oglcm' (s))
50 no MD Du (lollmm
50 MD Ou Collooomm (s))
b
ns Classical MO, MAE-0.20
SL GNN, MAE-0.43
Score GNN, MAE-0.23
Ours, MAE-0.12
Predicted Drmsi Pongttlogtcreesss (s0)
Predited Dise (logiccm--333)
Predicted Dros (logjem
-7.5
Prodictd Dasc Dysc (lagicmm
18.0-
-8.0
.8.0
68.5
-8.5
-8.
-8.5
9.0
-.00
8.0
50 MD Drya Dogicmm/sil
50 ms MD D (oggccPPss))
50 ms MD Dresi Nlg(ern
50 MD Dry (log(cm//s))
Figure 9: Li-ion diffusivity (a) and TFSI-ion diffusivity (b) estimation performance using different methods. From left to right: (1) 5-ns MD, which is the length of the training trajectories; (2) Supervised learning GNN model (Xie et al., 2022); (3) Our model without the Score GNN refinement; (4) Our full model.
as simulation proceeds. The Score GNN refinement significantly reduces collision to level similar to the reference, consistently across all cut-off radii. Radial distribution functions (RDFs) are commonly used describe the structural properties of physical system (definition in Appendix B). In Figure 8 (b), plot the RDFs averaged over the 50-ns simulation for an example testing SPE over beads from different components. Without the Score GNN refinement, we observe excessively high density at a near distance (which indicate abnormal bead collision and unphysical stretures) and missing radius peaks. With the Score GNN refinement, the learned simulation produces RDFs that can recover the reference's overall shape and radius peaks.
Estimate ion transport properties. We show the Li-ion/ TFSI-ion diffusivity prediction performance of different models in Figure The first column panels of Figure (a,b) show the prediction from 5-ns classical MD, which is the length of the training trajectories, giving poor estimation of Li-ionnTFSIIie diffusivity. The second column panels show the results from the supervised learning GNN model implemented in (Xie et al., 2022) while fitting an exponential decay curve (detailed in Appendix C) to account for the convergence of DrayTEST (as shown Figure 6) for the Li-ion/TFSI-L settings in (Xie et al., 2022), in our experiments, no 50-ns trajectory is available for training. We see the supervised model is not capable of recovering the /FFS-ioo diffusivity with only short trajectories for training and an exponential decay estimation. The third column panels show the performance of our model but without the Score GNN refinement. Due to the unstable dynamics shown in Figure 8, this model is not able to recover (i-ionnTFSS-ion diffusivity. The last column panels show the performance of our full model, which can simulate stably and accurately recover the 50-ns Li-ion/TFSIIin diffusivity of unseen systems by training from short trajectories of 5 ns only, with several orders of magnitude higher efficiency (Table 1).
diffusivity. Different from the experimental
5 Discussion
We have developed a machine learning model for simulating coarse-grained MD using very large time- integration steps without integrating forces. The proposed model can recover key long-time statistics accurately while being several orders magnitude faster than classical MD. The superior efficiency allows to study large-scale systems that computationally inaccessible to ML. force fields and MT. CG force
12

---

Published in Transactions on Machine Learning Research (08/2023)
fields. The learned simulator is trained over short MD trajectories but generalize well to simulate longer trajectories for novel unseen systems, making it highly suited for high-throughput screening settings. In two challenging and realistic applications, our model significantly outperforms supervised learning baseline methods and directly uses short-time MD for property estimation.
On the other hand, coarse-grained simulations necessarily entail trade-offss wetween efficiency and accuracy. Unlike force fields, which are conserved vector fields, our method cannot guarantee energy conservation or time reversibility. Additionally, the high degrec of coarso-graining and very long time steps make the dynamics non-Markovian and appear stochastic, which poses difficulties for predicting system behavior (Klippenstein et al., 2021). Consequently, maintaining stability proves to be highly challenging for our model. To tackle the instability issue, this work introduces a novel refinement module based on diffusion models, for which the stabilization effect is verified through experiments. Extending sample efficient equivariant ML force field model architectures (Batzner et al., 2022) to our sctting and active learning for collecting high-quality data (Vandermause et al., 2020) are some other promising directions for future development of time-intogratod simulators.
The dynamics simplification introduced by spatial coarse-graining and temporal time integration is two- edged sword: the simplified dynamics can be casier to learn and more efficient to simulate, but the lost information may lead to inaccurate dynamics and errors in property estimation. Therefore, optimal CG modeling is closely tied to the objective observables the MD simulation. CG modeling techniques based on the essential dynamical information (Kmiecik et al., 2016; Souza et al., 2021; Zhang et al., 2018b; Wang & Bombarllam,arelli, 2019; Li al., 2020) is another important future direction in designing more effective CG MD simulation schemes.
Acknowledgments
We pay tribute to Ertavia--Eugen Ganea (1987-2022), dear colleague and friend who we had many inspiring discussions over this work. We thank Wujie Wang, Zhenghao Wu, Tryy-Shyang Lin, Gabriele Corso, Bowen Jing, Shangyuan Tong, Yilun Xu, Hannes Stärk, the rest of the TJ group members and anonymous TMLR reviewers for their helpful comments and suggestions. We thank Michael A. Webb for his advice and support on the single-chain coarse-grained polymer dataset. We thank Jacob Helwig for pointing out a typo in our paper. We gratefully thank the Community Resource for Innovation in Polymer Technology, a project supported by the National Science Foundation (NSF) Convergence Accelerator program (Convergence Accelerator Research 2134795), and MIT-GIST collaboration for support. The authors also acknowledge the MIT SuperClon Cloud and the Lincoln Laboratory Supercomputing Center for providing HPC resources.
References
Nongmuch Artrith, Alexander Urban, and Gerbrand Ceder. Efficient and accurate machine machine-learnnng in- terpolation of atomic energies in compositions with many species. Physical Review B, 96(1):014112 2017.
Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Zuegner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion models and force fields for coarse-grained molocular dynamics. arXiu preprint arXio:2302.0000 2023.
Alessandro Bardueci, Giovanni Bussi, and Michele Parrinello. Well-tempered metadynamics: a smoothly converging and tunable free-energy method. Physical review letters, 100(2):020603, 2008.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):1-11, 2022.
Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential- energy surfaces. Physical review letters, 98(14):146401, 2007.
13

---

Published in Transactions on Machine Learning Research (08/2023)
Rafacl C. Bernardi, Marcelo C.R. Melo, and Klaus Schulten. Enhanced sampling techniques in molecular dynamics simulations of biological systems. Biochimica et Biophysica Acla (BBA) General Subjects, 1850 (5):872 877, 2015.
Emiliano Brini, Elena A Algacr, Pritam Ganguly, Chunli Li, Francisco Setriguzz-Roperry and Nico FA van der Vegt. Systematic coarse-graining methods for soft malter simulations a review. Soft Maller, 9(7): 2108 2119, 2013.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Shriram Chennakesavalu, David J Toomer, and Grant M Rotskoff Ensuring thermodynamic consistency with invertible coarse-graining. arXiv preprint arXiv:2210.07880 2022.
Stefan Chmicla, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus- Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, (5):e1603015, 2017.
Stefan Chmicla, Huziel E Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molccular dynamics simulations with machine-learned force fields. Nature commuumicatinns, 9(1):1 10, 2018.
Alain Dequidt and Jose G Solano Canchaya. Bayesian parametrization of coarse-grain dissipative dynamics models. The Journal of chemical physics, 148(8)) 2015.
Stefan Doerr, Maciej Majewski, Adrià Pérez, Andreas Krämer, Cecilia Clementi, Frank Noe, Toni Giorgino, and Gianni De Fabritiis. Torchmd: A deep learning framework for molecular simulations. Journal of Chemical Theory and Computation, 17(4):2355-22) 2363, 2021.
Pascal Friedcrich, Florian Häsc, Jonny Proppe, and Alán Aspura-Guzik. Machine-learned potentials for next-generation matter simulations. Nature Materials, 20(6):750 761, 2021.
Xiang Fu, Zhenghao Wu, Wujie Wang. Tian Xie, Sinan Keten, Rafael Homeez-ommbareli, and Tommi Jaakkola. Forces are not enough: Benchmark and critical evaluation for machine learning force fields with molecular simulations. arXie preprint ar arx::22200. 072333, 2022.
Johannes Gasteiger, Janek Groß, and Stephan Günneemann. Directional message passing for molecular graphs. International Conference on Learning Representations, 2020.
Johannes Gasteiger, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Newral Information Processing Systems, 2021.
Joe G Greener and David Jones. Differentiable molecular simulation can learn all the parameters in a coarse-grained force field for proteins. PloS one, 16(():e0256990, 2021.
Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840- 6801,, 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735 1780, 1997.
Brooke E Husic, Nicholas Charron, Dominik Lemm, Jiang Wang, Adrià Pérez, Maciej Majewski, Andreas Krämer, Yaoyi Chen, Simon Olsson, Gianni de Fabritiis, et al. Coarse graining molecular dynamics with graph neural networks. The Journal of chemical physics, 151):1111,, 2020.
John Irwin and Brian Shoichet. Zinc a free database of commercially available compounds for virtual screening. Journal of chemical information and modeling, 45(1):177-188 2005.
14

---

Published in Transactions on Machine Learning Research (08/2023)
Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, E Weinan, and Linfeng Zhang. Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning. In SC20: International conference for high performance computing, networking, storage and analysis, pp. 1-14. IEEE, 2020.
Sebastian Kaltenbach and Phaedon-Slelios Koutsourelakis. Physics-aware, probabilistic model order reduction with guaranteed stability. arXiv preprint arXim:2101.0583, 2021.
George Karypis and Vipin Kumar. A fast and high quality multilevel scheme for partitioning irregular graphs. SIAM Journal on scientific Compuling, 20(1):359 392, 1998.
Kévin Kempler, Julien Devemy, Alain Dequidt, Marc Couty, and Patrice Malfreyt. Development of coarse- grained models for polymers by trajectory matching. ACS omega, 4(3):5955- 5967, 2019.
Alireza Khorshidi and Andrew Peterson. Amp: modular approach to machine learning in atomistic simulations. Compuler Physics Communications, 207:310 324, 2016.
Leon Klein, Andrew YK Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. arXiv preprint or "Xix:2302.0111 01700, 2023.
Viktor Klippenstein, Madhusmita Tripathy, Gerhard Jung, Friederike Schmid, and Nico FA van der Vegt. Introducing memory in coarse-grained molecular simulations. The Journal of Physical Chemistry B, 125 (19):4931 4954, 2021.
Stefan Klus, Feliks Nüske, Péter Koltai, Hao Wu, Ioannis Kevrekidis, Christof Schütte, and Frank No Data-driven model reduction and transfer operator approximation. Journal of Nonlinear Science, 28: 985 10101010, 2018.
Scbastian Kmiccik, Dominik Gront, Michal Kolinski, Lukasz Wieteska, Alcksandra Elzbicta Dawid, and Andrzej Kolinski. Coarse-grained protein models and their applications. Chemical reviews, 116(14): 1898-79384 2016.
Dávid Péter Kováes, Cas van der Oord, Jiri Kucera, Alice EA 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-771 2021.
Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the National Academy of Sciences, 99(20):12562 12566, 2002.
Alessandro Laio, Antonio Rodriguez-Fortea, Francesco Luigi Gervasio, Matteo Ceccarelli, and Michele Parrinello. Assessing the accuracy of metadynamics. The journal of phaysical chemistry B, 109(14): 6714 6721, 2005.
Zhiheng Li, Geemi P Wellawatte, Maghesroe Chakraborty, Heta A Gandhi, Chenliang Xu, and Andrew D White. Graph neural network based coarse-grained mapping prediction. Chemical science, 11(5)):5524 9531, 2020.
Zijie Li, Kazem Meidani, Prakarsh Yadav, and Amir Barati Farimani. Graph neural networks accelcrated molecular dynamics. The Journal of Chemical Physics, 156(14):144100) 2022.
Amr H Mahmoud, Matthew Masters, Soo Jung Lee, and Markus A Lill. Accurate sampling of macromolecular conformations using adaptive deep learning and coarso-grained representation. Journal of Chemical Information and Modeling, 2022.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets for deep learning of molecular kinetics. Nature communications, 9(1):5, 2018.
15

---

Published in Transactions on Machine Learning Research (08/2023)
Siewert J Marrink, H Jelger Risselada, Serge Yefimov, D Peter Ticleman, and Alex H De Vries. The martini force field: coarse grained model for biomolecular simulations. The journal of physical chemistry B, 111 (27):7812 7824, 2007.
RL Nkepsu Mbitou, F Goujon, A Dequidt, Latour, J Devémy, R Blaak, N Martzel, C Emeriau- Viard, Tchoufag, S Garruchet, et al. Consistent and transferable force fields for statistical copolymer systems al. the mesoscale. Journal of Chemical Theory and Computation, 18(11):6940 6951, 2022.
Frank Noé, Ralf Banisch, and Cecilia Clementi. Commute maps: Separating slowly mixing molecular configurations for kinetic modeling. Journal of chemical theory and computation, 12(11):5620-566 2016.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457), 2019.
Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual review of physical chemistry, 71:361-309 390, 2020.
Cheol Woo Park, Mordechai Kornbluth, Jonathan Vandermause, Chris Wolverton, Boris Kozinsky, and Jonathan P Mailoa. Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture. npj Compulational Malerials, 7(1):1 9, 2021.
Guillermo Pérez- Herenneerá,d,z Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Identification of slow molecular order parameters for markov model construction. The Journal of chemical physics, 139 (1):07B604 1, 2013.
David Rigby, Huai Sun, and BE Eichinger. Computer simulations of poly (ethylene oxide): force field, pvt diagram and cyclization behaviour. Polymer Inlernational, 44(3):311 (13):311-330, 1997.
Alvaro Sanzaezz Jonathan Godwin, Tobias PfafT, Rex Ying, Jure Leskovec, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning, pp. 8459 8468. PMLR, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp.
PMLR, 18-24 Jul 2021.
Elia Schneider, Luke Dai, Robert Q. Topper, Christof Drechsel-Grau, and Mark E. Tuckerman. Stochastic neural network approach for learning high-dimensional free energy surfaces. Phys. Rev. Lett., 119:150601, Oct 2017. doi: 10.1103/ (0.1103/PhysSer Petyspyyeeeeeeee Lett. 11....19..10000 10601 URL https:: //tpnk////nnn///////// http///inna.../../..//aa////aaaa wwt.ssorg//o///o///eooo/oo 1103 Physs//e/teyeettttt..... 119 150601.
Samuel Schoenholz and Ekin Dogus Cubuk. Jax md: framework for differentiable physics. Advances in Neural Information Processing Systems, 33, 2020.
Pradyumn Sharma, Rajat Desikan, and K Ganapathy Ayappa. Evaluating coarse-grained martini force-fields for capturing the ripple phase of lipid membranes. The Journal of Physical Chemistry B, 125(24):6587-655) -659,, 2021.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular conformation generation. In Marina Mcila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 9558-9558 9568. PMLR, 2021.
Hythem Sidky, Wei Chen, and Andrew I. Ferguson. Molecular latent space simulators. Chemical Science, 11 :9599 9467, 2020.
Justin S Smith, Olexandr Isayev, and Adrian Roitberg, Ani-1: an extensible neural network potential with dft accuracy at force field computational cost. Chemical science, 8(4):3192-320) 3203, 2017.
16

---

Published in Transactions on Machine Learning Research (08/2023)
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11895 11907, 2019.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stelano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint ar arXiv:2011.13456 2020.
Paulo CT Souza, Riccardo Alessandri, Jonathan Barnoud, Scbastian Thallmair, Ignacio Faustino, Fabian Grünewald, Ilias Patmanidis, Haleh Abdizadeh, Bart MH Bruininks, Tsjerk A Wassenaar, et al. Martini 3: a general purpose force field for coarso-grained molecular dynamics. Nature methods, 18(4):382-38) 388, 2021.
Lukas S Stelzl and Gerhard Hummer. Kinetics from replica exchange molecular dynamics simulations. Journal of chemical theory and compulation, 13(8):3927 3935, 2017.
Sina Stocker, Johannes Gasteiger, Florian Becker, Stephan Günnemann, and Johannes T Margraf. How robust are modern graph neural network potentials in long and hot molecular dynamics simulations? Machine Learning: Science and Technology, 3(4):045010, 2022.
Yuji Sugita and Yuko Okamoto. Replica-cxchange molecular dynamics method for protein folding. Chemical physics letters, 314(1-2):141 151, 1999.
Mohammad M Sultan, Hannah K Wayment-Steele, and Vijay S Pande. Transferable neural networks for enhanced sampling of protein dynamics. Journal of chemical theory and computation, 14(4):1887-1894 2018.
Huai Sun. Force field for computation of conformational energies, structures, and vibrational frequencies of aromatic polyesters. Journal of Computational Chemistry, 15(7):752 768, 1994.
So Takamoto, Chikashi Shinagawa, Daisuke Motoki, Kosuke Nakago, Wenwen Li, lori Kurata, Taku Watanabe, Yoshihiro Yayama, Hiroki Iriguchi, Yusuke Asano, et al. Towards universal neural network potential for material discovery applicable to arbitrary combination of 45 elements. Nature Communications, 13(1): 1-11, 2022.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular polentials. In International Conference on Learning Represenlations, 2021.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equuiaiian neural networks for 3d point clouds. arXiv preprint or arxiv 08219, 2018.
A. P. Thompson, H. M. Aktulga, R. Borger, D. S. Bolintincanu, W. M. Brown, P. S. Crozier, P. J. in 't Veld, A. Kohlmeyer, S. G. Moore, T. D. Nguyen, R. Shan, M. J. Stevens, J. Tranchida, C. Trout, and S. J. Plimpton. LAMMPS- a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales. Comp. Phys. Comm., 271:108171, 2022.
Glenn M Torric and John Vallcau. Nonphysical sampling distributions in monte carlo free-energy estimation: Umbrella sampling. Journal of Compulational Physics, 23(2):187 199, 1977.
Oliver T Unke, Stefan Chmieea, Michael Gastegger, Kristof T Schütt, Huziel E Sauceda, and Klaus-Roberl Müller. Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nature communications, 12(1):1 14, 2021a.
Oliver T Unke, Stefan Chmicla, Huziel E Saucoda, Michael Gastegger, Igor Poltavsky, Kristof T Sch utt, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemical Reviews, 121 (16):10142-100) 10186, 2021b.
17

---

Published in Transactions on Machine Learning Research (08/2023)
Rahul Upadhya, N Sanjeeva Murthy, Cody L Hoop, Shashank Kosuri, Vikas Nanda, Joachim Kohn, Jean Baum, and Adam J Gormley. Pet-raft and saxs: High throughput tools to study compactness and flexibility of single-chain polymer nanoparticles. Macromolecules, 52(21):8295 8304, 2019.
Omar Valsson and Michele Parrinello. Variational approach to enhanced sampling and free energy calculations. Physical review letters, 113(9):090601, 2014.
Jonathan Vandermause, Steven B Torrisi, Simon Batzner, Yu Xie, Lixin Sun, Alexie M Kolpak, and Boris Kozinsky. On-tho-fly active learning of interpretable bayesian force fields for atomistic rare events. npj Compulational Malerials, 6(1):1-11, 2020.
Pantelis R Vlachas, Julija Zavadlav, Matej Praprotnik, and Petros Koumoutsakos. Accelerated simulations of molecular systems through learning effective dynamics. Journal of Chemical Theory and Computation, 18(1):538-549, 2021.
Jiang Wang, Simon Olsson, Christoph Wehmeyer, Adrià Pérez, Nicholas E Charron, Gianni De Fabritiis, Frank Noé, and Cecilia Clementi. Machine learning of coarse-gradd boanse-grained molecular dynamics force fields. ACS central science, 5(5):755-77) 767, 2019.
Wujie Wang and Rafael Gómez-Bombard Coarse-graining auto-encoders for molecular dynamics. nyj Computational Materials, 5(1):125, 2019.
Michael A Webb, Yukyung Jung, Danielle M Pesko, Brett M Savoie, Umi Yamamoto, Geoffrey W Coates, Nitash P Balsara, Zhen-Gang Wang, and Thomas F Miller III. Systematic computational and experimental investigation of lithium-ion transport mechanisms in polyester-based polymer electrolytes. ACS central science, 1(4):198-2000 2015.
Michael A. Webb, Nicholas E. Jackson, Phwey S. Gil, and Juan J. de Pablo. Targeted sequence design within the coarse-grained polymer genome. Science Advances, 6(43)=eabe66116 2020.
Hao Wu and Frank Noé. Variational approach for learning markov processes from time series data. Journal of Nonlinear Science, 30(1):23 1--22666-666 2020.
Hao Wu, Andreas Mardt, Luca Pasquali, and Frank Noe. Deep generative markov state models. Advances in Neural Information Processing Systems, 31, 2018.
Tian Xie, Arthur France-Tanord, Yanming Wang, Jeffrey Lopez, Michael A Stolberg, Megan Hill, Gra- ham Michael Leverick, Rafacl Gomez- Bombarelli, Jeremiah A Johnson, Yang Shao-Horn, et al. Accelerating amorphous polymer electrolyte screening by learning to reduce errors in molecular dynamics simulated properties. Nature communications, 13(1):1-10 2022.
Yi Tsaac Yang, Qiang Shao, Jun Zhang, Lijiang Yang, and Yi Qin Gao. Enhanced sampling in molecular dynamics. The Journal of chemical physics, 151 (7):00007,,,, 2019.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep potential molecular dynamics: A scalable model with the accuracy of quantum mechanics. Phys. Rev. Lett., 120:143001, Apr 2018a.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deepeg: Constructing coarse-grained models via deep neural networks. The Journal of chemical physics, 149(3):034101, 2018b.
Linfeng Zhang, Jiequn Han, Han Wang, Wissam Saidi, Roberto Car, et al. End-to-end symmetry preserving inter-alomic potential energy model for finite and exlended systems. Advances in Neural Information Processing Systems, 31, 2018c
Dong Zhou, Devaraj Shanmukaraj, Anastasia Tkacheva, Michel Armand, and Guoxiu Wang. Polymer electrolytes for lithium- based batteries: Advances and prospects. Chem, 5(9):2326-223) 2352, 2019.
18

---

Published in Transactions on Machine Learning Research (08/2023)
A
Ablation Studies
Du MAE (log(cm²(s)
0.20
0.17
0.17
Du MAE (log(cm²/s))
0.33
Du MAE (log(cm' (log(cm²/s)) (s))
0.6
0.61
0.15
0.3-
0.28
0.4
0.10
0.2
0.13
0.05
0.1-
0.2
0.15
0.15
0.11
Defoult
GNz Deterministic
0.0
25
0.0
18
Coarse graining Group Size
Number of History Steps
d
Dy MAE (log(cm*/s))
0.20
Du MAE (logl (log(cm²(ss)
0.15
Du MAE (log(cm²/ss)
0.15
0.17
0.10-
0.11
0.11
0.11
0.10
0.11
0,12
0,12
0.11
0.12
0.11
0.11
0.10
0.05
0.05
0.05
0.00
0.1
0.2
0.4
8 9 Connectivity Radius (A)
0.00
0
1
Time-integration (ns)
Random Seed
Figure 10: (a) Comparison of our default model, model without the Embedding GNN at the fine-level graph, and a model that predicts deterministic acceleration for each particle. (b) Comparison of different coarse-graining atom group sizes. The atom group size is the number of atoms contained each CG bead. Performance drops when the noarse-graining too fine (5) too coarse (15, 25). (c) Comparison of different history lengths for predicting the dynamics. We observe that longer history helps improve the model performance. (d) Comparison of different time-integration step lengths. A longer time integration removes high-frequency information and simplifies the dynamics, making it easier to learn. On the other hand, the long-ting time property of particle diffusivity does not require high-frequency information to estimate accurately. (e) Comparison of different connectivity radii when building the coarsc-level graph. A radius of (I) Comparison of different random seeds using the same model. Our model does stochastic rollouts, and the performance is robust random seeds.
Åis sufficient for modeling the SPE systems.
Fine-level Embedding GNN GN (Figure 10 (a)) Our model uses a fine-level embedding GNN to learn CG bead embedding that enclose local structural information. We can remove this fine-level GNN and let the CG bead embedding be the mean over the atom type embedding in each atom group. As shown in Figure 10 (a), model performance significantly drops without the Embedding GNN.
Stochastic dynamics prediction. (Figure 10 (a)) We attempt to let our model output deterministic acceleration at each forward simulation step. However, the inherent uncertainty makes the model predict very small movement for all particles at every step, and the model performance significantly drops. The small movements lead slow transport of particles, and causes underestimation of Li-ion diffusivity, as shown in Figure 11 (a). Such unrcalistic dynamics also makes long simulation unstable. This demonstrated Figure 11 (b), which shows that deterministic model has higher number of bead collision (under cutoff of 1.0 A) with increasing trend through time.
Coarse-graining group sizc. (Figure 10 (b)) We experimented with different coarso-graining group size. We observe that in terms of capturing particle diffusivity, using group size of 7 outperforms finer (5) and coarser (15, 25) marse-graining, We hypothesize that the finer system hard to model accurately with limited training data, while the coarser system loses important information for accurate dynamics modeling. This result further suggests the optimal coarse grained modeling should be conditional on the objective of MD simulation.
19

---

Published in Transactions on Machine Learning Research (08/2023)
Deterministic, MAE=0.17
Predicted Du (log(cm²))s)
0.5
.7.5
Bead Collision
25
Stochastic
Ground Truth
0.4
0.3915
0.3909
Detarministic
-8.0
NLL Loss
0.3
20
0.2
-8.5
15
0.1
-9900
50 ns MD Du (s)
-8.5
-8.0
-7.5
0
Simulation Time (ns)
20
30
50
0.0
Default
Random Rotation
Figure 11: (a) Li-ion diffusivity estimation performance the deterministic learned simulator. The model predicts small movements at every step, leading slower ion transport and underestination Li-ion diffusivity. (b) Bead collision 1 Å a function of time, averaged over the 50 testing SPEs. the deterministic model's prediction becomes increasingly unphysical of model prediction testing polymers. Model performance remains the same when the input structures are randomly rotated before being fed into the model.
simulation proceeds. (c) The negative log-likelihood (NLL) loss
Use of historical information for prediction. (Figure 10 (c)) The spatio-temporal coarse-graining introduces memory effects to the resulting dynamics. Our model uses historical information by utilizing k-step historical velocities in dynamics prediction, and experimented with & 0, 13, 18, under time-integration of 0.2 ns per step. We observe that longer history as input leads to better performance.
Time-integration step size. (Figure 10 (d)) We experimented with different time-integration step size. We observe that a longer time-integration gives the best performance. We hypothesize this is due to the dynamics simplification effect of long time-integration, while the loss of high-frequency dynamics does not damage the estimation accuracy of particle diffusivity under long-time simulation.
Connectivity radius. (Figure 10 (e)) Our coarse-level graphs contain both CG-bonds and radius cut-off edges. The radius cut-off edges model non-bonded interactions. We experimented with radius 6, 7, 8, 9 A, but found no significant performance difference. We conclude that radius of 6 Åis sufficient for modeling the SPE non-bonder dynamics.
Random sceds. (Figure 10 (f)) As our model docs stochastic simulation, we examine its robustness against the randoin seed. We rollout 5 times using the same model and observed no significant performance difference across random seeds.
SE(3) Equivariance. Some molccular systems obey symmetrics in the form of invariance and cquivariance. While translational equivariance is enforced by the proposed model, effective rotational equivariance is learned from data. The effective equivariance is verified by an experiment that applies random rotations at X, Y, and Z axes to input data. The time-integrated acceleration predicted by our model is accordingly rotated by the same amount (Appendix A, Figure 11). We present our model as a general approach that applicable to diverse systems since some MD systems are not rotationally equivariant (e.g., when an external field such as an electric field present) with favorable time efficiency. Extending the proposod framework to equivariant neural networks (Thomas al., 2018; Satorras al., 2021; Batzner al., 2022; Gasteiger et al., 2020; 2021) is promising future direction to improve data efficiency and generalization capability. Such extension will involve incorporating historical information in an equivariant way and improving the efficiency of equivariant GNNs for training and inference with large-scale systems and big dataset sizes.
We verify our learned simulator being rotationally equivariant by applying random rotations at X, Y, and Z axes to input data, and compute the dynamics prediction loss, before and after the random rotations. Figure 11 (c) shows that the model test time performance is unchanged irrespective of the random rotations. Therefore, learned simulator is effectively rotational equivariant by learning from data.
20

---

Published in Transactions on Machine Learning Research (08/2023)
B Dataset Details
Singlo-chain polymer dataset. The singlo-chain polymers simulated using LAMMPS (Thompson et al., 2022), with the force-field parameters and chemical space defined in (Webb el al., 2020). For clarity, we present the potential terms below, while more details can be found in (Webb et al., 2020).
U(r^)= = E (ri((r +
Σ angles
Dend (lakk)
bonds
(3)
Utors (oegta))
Σ (nb()ay)
dihedrals
where rijs Oyjk, and Oiyks are internal distances, angles, and dihedrals, respectively. The individual potential terms are given by:
27
Uvab(rig) = Kiy (R.3 (2) Flg
1
Taj Ro(0)
Ki
Ubend(Osjk) Kipa
(
Egk
Urors(Pight = K Kepkel + CO8 ØisR1)
6
if i,j bonded and ray 2/6
46ij
0zj
rij
Tis
Unb(rig)
9
otherwise
All simulations are done in reduced units with characteristic quantities for distance and T for time. Single- chain polymer dynamics in implicit solvent evolve according to the Langevin equation using the velocity-Verlet integration scheme. We use a time step of 0.017. Training trajectories are 50k after removing the initial trajectory for relaxation and are recorded every 5 r. Testing trajectories 500 т. The polymer interaction is described by the summation of bonded and non-bonded potential energy functions. We refer interested readers to (Webb et al., 2020) for more details on the simulation setup.
5M and are recorded every
Calculation of single-chain polymer properties. Our main property of study, the squared radius of gyration (R3) is computed by:
Liev made Eiey m;
R2
where V is the set of all nodes, m, is the mass of particle i, and de is the distance from particle i to the center of mass. The mean internal distance (R(s)) is the average distance between bead i and bead i + 8 on the single chain. It is computed as:
M-
a sstall
(R(s)) = Σ
M -
Where M is the total number of CG beads in the single chain. The relaxation time for R3 is derived from the autocorrelation function (ACF), which is computed as:
AcF ( (2 (112(1)))))
(122()))))
Here (:)x stands for averaging over the entire trajectory. The relaxation time Lgg is computed by fitting an exponential function f(y) = exp(-tRyy) to the ACF, and the relaxation time is the time when the ACF decays to 1/e. It highly dynamic long-time property decided by the polymer structure.
21

---

Published in Transactions on Machine Learning Research (08/2023)
SPE dataset. The SPE systems simulated using LAMMPS (Thompson et al., 2022) with the force-field parameters and chemical space defined in (Xie et al., 2022). The atomic interactions are described by the polymer consistent force-field (PCFF+) (Sun, 1994; Rigby et al., 1997). Polymer amorphous structures are sampled with a Monte Carlo algorithm and then mixed with 1.5 mol lithium bis- tisttrifturrmmetyll sulfonimide (LiTFST) per kilogram of polymer. For all systems, there are 50 Li-ions and TFSI-ions in the simulation box, and each polymer chain has 150 atoms in the backbone. The training trajectories are ns long after removing the initial equilibration, while the testing trajectories are 50 ns long. All trajectories are recorded after 5 MD equiibration. Each system is run in the canonical ensemble (nVT) at a temperature of 353K using a multi-timescall integrator with an outer time step of 2 fs for nonbonded interactions and an inner time step of 0.5 fs. We refer interested readers to (Xie et al., 2022) for more details on the simulation setup.
Calculation of SPE properties. The radial distribution function (RDF) describes the particle density as a function of distance from a reference particle. For a system with many distinct types of particles, we can compute RDF for particular types by counting the neighbor atoms of certain types at various radii and each time step; and average over the entire trajectory. The RDF for partiele types A, B at distance is computed
RDF (, d[nA,B(r)] 4rr
Here (r) is the number of particle pairs with types A and B and distance in [r, r + dr), where dr is a small bin size. We include more RDF results from the learned simulation in Figure 16.
Ion transport properties of SPEs are the topic of study for many previous works and require long-time simulation to estimate. In particular, our experiments focus particle diffusivity. With trajectory of time horizon T,, the diffusivity D of particle is computed by:
pr 2oll
D
67
where æт is the position at time T, and æo is the position at time 0. The diffusivity of a type of particle (e.g., Li-ion) is then computed by averaging the particle diffusivity over all particles of that type.
C
Experimental Details
A Algorithm annealed Langevin dynamics for structural refinement
Input: GND prediction 4tt noise levels {o,) denoising step size c, steps per noise level TĄ Output: denoised positions T++At
Initialize at+ A/+
4: for noise level i 1,... L do
5: for step $ = 1.... To do
6: Sample positional noise 2, ~ NO.OT) I)
7: Tit t D++A + c- GNs(E4+4))/og + 22cz,
8: end for
end for
In our implementation, we use 2 hidden layers for all MLPs and 7 message-passing layers for all GNNs. The Embedding GNN GNg has a hidden size of 64, while the Dynamics and Score GNNs have a hidden size of 128. All activation functions in the neural networks are rectified linear units (ReLC). We train the model for 2 million steps. The network is optimized with an Adam optimizer with an initial learning rate of 2 x 10-4, exponentially decayed to 2 x 10-5 over the 2 million training steps. All models are trained and used for producing long simulations over a single RTX 2080 Ti GPU. In the single-chain polymer experiments, the fine-grained R2 is not computable from the coarse-grained ionfiguration. To obtain the fine-grained R2 from CG simulations, We first compute the R² from the CG configuration and then use an MLP over the latent graph embedding of the dynamics GNN to predict the difference between the CG R2 and FG R3. We then predict the FG R3 by adding the CG g and the offset predicted by this MLP module. Since the simulation is stable, the refinement module is not used for the single-chain polymers. For the SPE systems, the refinement.
22

---

Published in Transactions on Machine Learning Research (08/2023)
module is diffusion model with 20 noise levels ranging from [0.01, 10.]. At simulation time, the refinement is through an annealed Langevin dynamics of 10 steps per noise level and step size of 5 10-5. The refinement. procedure is described in Algorithm 1.
Single-chain polymer baseline models. We conduct experiments with two supervised learning baseline models. The first one GNN model that takes the polymer chemical graph as input and outputs the mean and standard deviation of R2- Each node in the polymer graph is CG bead and each edge is a chemical bond. The GNN uses the SERODDER-PROCSSONE DEEODDE with 10 message- paatgggeppsinngggg layers to process the polymer graph to a latent graph with node/edge embeddings. The MLPs in the GNNss have 2 hidden layers with size of 128. The node embeddings are then summed to get the graph embedding, which is then processed by & 2-layer MLP with the hidden size 256 to obtain the final prediction of mean and standard deviation of polymer Rg For accurate R2 prediction of the fine-grained state in the single-chain polymer experiment, we use another neural network that inputs the coarse-level latent graph representation to fit the residual of R3 computed from the coarse-level states (as opposed the fine-level states). The second LSTM baseline model takes the 1D polymer chain structure its input and replaces the GNN encoder with 2-layer LSTM encoder with hidden size 256. All activation functions in the neural networks ReLU, The bascline models are optimized with an Adam optimizer with a learning rate of 10-3, exponentially decayed to 5 x 10-4 over 1500 training epochs. Due to the limited time horizon of training data, the baseline models can only fit. high-variance labels, leading to underperforming prediction results.
SPE diffusivity prediction baseline model. We adopt the GNN model proposed in (Xie et al., 2022) and refer interested readers to (Xie et al., 2022) for more details. The only modification we make is changing the training objective
to rectify the overestimation coming from the short MD horizon of the training data. We let the baseline model fit the 5 ns diffusivity curves (curves in Figure 9 (f)) by predicting two parameters in the function: f(t) = y + e e-"t, Therefore, the model approximates the converging process of diffusivity with exponential decay. During the evaluation, we set t = 50 ns to obtain the model prediction for 50-ns diffusivity. However, without any long training trajectorics, it is very challenging to guess the decaying process of diffusivity for various SPEs. The baseline model performs better than using 5-ns ground truth MD but still significantly overestimates particle diffusivity. We have also attempted using neural ordinary differential equation (Chen et al., 2018) to fit the decay curve but found the results worse than using the parameterized exponential decay function.
Reference
Predicted
3.2
2,8
0.5
0.5-
2
TIC2
2.0
D.D
0.0
1.6 6
g
0.5
0.5-
0.8
-1.0
0.4
0.0
TIC
TIC
Three representative states with a low, medium, and high radius of gyration (with low to high free energy) are visualized for example polymer. Free energy surface under two-dimensional representation is produced with TICA. The projection is fitted over the reference data only and is applied to both reference and learned sinulated data.
Figure 12:
Understanding radius of gyration. The mcan squared gyration radius of a polymer is of significant practical importance for understanding the rheological properties of polymers in solution and the compactness of polymers. This due to its role in establishing a concentration threshold associated with the beginning of chain entanglements and the process of gelation (Webb et al., 2020). Intuitively, the radius of gyration a way describe the size of a polymer chain, indicating how spread out the polymer chain In Figure 12, we visualize three states with low to high radius of gyration and their location on two-dimensional free energy
23

---

Published in Transactions on Machine Learning Research (08/2023)
surface produced with time-lagged independent component analysis (TICA, the radius of gyration is high, the polymer chain is more spread out, and vice versa..
et al. 2013). If
X1040.5M, MAE-3779
1M. MAE-2005
1.5M, MAE-1339
2M, MAE-919
2.0
1.5
Predicted Eag
0.5
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
0.5
1.0
1.5
2.0
5M MD 5g (T) x10*
5M r MD teg (T)
T MD byg (T) x10*
5M r MD tag (r) x10
The relaxation time of the squared radius of gyration for the testing polymers estimated using reference trajectories of different lengths. The trajectory length ranges from 0.5M, 1M, 1.5M, and 2M from left to right.
Figure 13:
ACF estimation. To better understand our model's performance in recovering the ACF of R for the single-chain polymers, we investigate how well trajectories of different lengths can reconstruct the ACF curves and the relaxation time of Ry. Figure 13 demonstrates the relaxation time the squared radius of gyration for the testing polymers estimated using reference trajectories of 0.5M, 1M, 1.5M, and 2M long, along with the average mean absolute error across all 40 testing polymers. Figure 14 demonstrates the corresponding ACF curves for the same four polymers reported in Figure 5. Our training trajectorics are 50k long, while the testing trajectories are 5M long, The training trajectories are not long enough for extracting the ACF curves and the relaxation time, so the baseline supervised learning models are not applicable. 5M trajectories simulated with our method result in an MAE of 1657 T in estimating the relaxation time (Figure 5), which is between the performance of 1M ground truth trajectory and 1.5M ground truth trajectory.
Detailed baseline results on estimating the distribution of R We include the performance of the supervised learning baseline models on estimating the distribution of R2 for the same four example polymers reported in Figure 4 in Figure 15. We also annotate the EMD achieved for each example. Our model can capture the distribution much more faithfully (Figure 4).
24

---

Published in Transactions on Machine Learning Research (08/2023)
Fit of 0.5M MD. te.-6.4K
Fit of 0.5M MD, txg-11.5k
Fit of 0.5M MD, txg-14.0k
Fit of 0.5M MD, tx.-20.3k
1.00
10
1.00
1.00
Prediction Reference
Autoconeatoon
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.25
0.0
0,00
0.00
à time IT)
à time (r)
A time (1)
A time (1)
Fit of 0.5M MD. 6.4-6.4 r
Fit of 0.5M MD, 11g-11.5k
Fit of 0.5M MD, 12.-140kk
Fit of 0.5M MD, 22.32
1.00
2.0
1.00
1.00
Prediction Reference
Auloconelditer
0.75
0.50
0.5
0.50
0.50
0.25
0.25
0,0
0,25
0.00
0.00
2
2
time (T!
time
(T)
Fit of 1.5M MD. bgg-6.4k T
Fit 1.5M MD, tgg-11.5k
Fit of 1.5M MD, txg=14.0k
Fit of 1.5M MD, 2o.=20.3k
1.00
1.0
1.00
1.00
Prediction Reference
Aulocooeettoon
0.75
0.75
0.50
0.5
0.50
0.25
0.0
0.00
0.00
0.00
2
6
2
4
6
time It)
x101
time (T)
x101
time (T)
x100
time (11
x100
Fit 2M MD. t.446.4
of 2M MD, 1p.-11.5k
of 2M MD, 6g.-14.0k
Fit of 2M MD, Gg.220.3k T
Prediction
Autocoreeation
0.75
0.75
0.75
0.50
0.50
0.50
0.25
0.25
0.00
0.00
0.00
2
6
2
6 x10*
2
0
2
6
à time (t)
x10*
à time (t)
à time (1)
x10*
à time (1)
x10
Figure 14: The ACF curves produced by reference trajectories of different lengths for four example polymers. The trajectory length ranges from 0.5M, 1M, 1.5M, and 2M 7 from the first to the fourth row.
25

---

Published in Transactions on Machine Learning Research (08/2023)
EMD+3.03
EMD-1.27
EMD-1:31
EMD-4.34
SL GNN Reference
4-
Densty
2
2-
RG (oP)
RG (t*) EMD-1.57
RG (oF)
x10P
R3 (o*)
b
x110-
EMD-0.73
x10-1
EMD-5.81
x10-
EMD-13.38
LLTM
Denstty
2-
2-
2-
2-
8 x10
RB (o*)
R$ (o*)
RB (o*)
x10*
R$ (o+)
Figure 15: The distribution of R2 from the baseline supervised learning models compared to the reference ground truth data for four example polymers with small to large (Rg). The first row demonstrates the results from the SL GNN model. The second row demonstrates the results from the SL LSTM model.
26

---

Published in Transactions on Machine Learning Research (08/2023)
Li-ion, Li-ion
Li-ion, TFSI-ion
Li-ion, Polymer
8
8
With Score GNN
With Score GNN
With Score GNN
6-
No Score GNN
6-
No Score GNN
6-
No Score GNN
Ground Truth
Ground Truth
Ground Truth
2-
2
0
0
Radius (A)
Radius (4)
Radius (A)
b
Li-ion, Li-ion
Li-ion, TFSI-ion
Li Li-ion, Polymer
8
8
With Score GNN
With Score GNN
With Score GNN
No Score GNN
Score GNN
No Score GNN
Ground Truth
Ground Truth
Ground Truth
4-
2-
2
0
10 Radius (A)
15
0
5
10
15
10 Radius (A)
15
Radius (A)
C
Li-ion, Li-ion
Li-ion, TFSI-ion
Li-ion, Polymer
With
Score GNN
With Score GNN
6
No Score GNN
No Score GNN
6
No Score GNN
Ground Truth
Ground Truth
Ground Truth
4-
4-
4-
0
10 Radius (A)
15
10
15
15
Radius (A)
Radius (A)
Li-ion, Li-ion
Li-ion. TFSI-ion
Li-ion, Polymer
8-
8
With Score GNN
With Score GNN
With Score GNN
6-
No Score GNN
6-
Score GNN
No Score GNN
Ground Truth
Ground Truth
Ground Truth
2-
2:
0
0
15
Radius (Å)
Radius (A)
Radius (Å)
ion,
ion, TFSI ion
Li ion, Polymer
8
8
With Score GNN
With Score GNN
With Score GNN
No Score GNN
Score GNN
No Score GNN
Ground Truth
Ground Truth
Ground Truth
4-
2
2
10 Radius (Å)
15
10 Radius (A)
15
10 Radius (A)
15
Figure 16: (a,b,c,d,e) Comparison of our model with with/without the Score GNN refinement, and the ground truth MD simulation, on RDF of Li-ions (column 1), RDF of Li-ions and TFSI-ions (column 2), and RDF of Li-ions and polymer particles (column 3) for five sampled SPEs.
27


Reviews:
[{"review_id": "wXSqsNuOiD", "reviewer": "Reviewer_wvWc", "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* Model is fairly approachable and understandable to non-experts in molecular dynamics (e.g. to a general machine learning audience), as is the paper presentation.\n* Performance on computing aggregated statistics from long-running simulations is impressive\n\nWeaknesses:\n\n* For some of the experiments, I don't have a good sense of the baseline performance. This is particularly the case in e.g. the ACF computation in Figure 5(a): I see that this broadly tracks the reference, but is this performance \"good\" or \"excellent\"? By what standards? How would baselines perform here — surely there must be some meaningful baselines, even if long-run MD computations are not feasible. As it stands the paper claims this is \"excellent\" performance but the figures are lacking baselines to demonstrate this adequately. (Even including the supervised learning methods, which I believe could be applied here but would suffer due to the need for long extrapolations, would be helpful….) \n* I had some confusion regarding how the coarse grain bead assignments of individual atoms are performed, and to what extend the assignment would be stable over time for large simulations. This is particularly confusing to me as the CG beads are themselves further coarsened at a later stage. More details on this in \"requested changes\".\n* It's not obvious to me that MD computations at this coarse-grained level are usable for most of the tasks for which someone may be interested in an MD simulator. In particular, unless I misunderstand it is not possible to recover coordinates of the original \"fine grained\" atoms. As such it would be nice if some convincing could be done that this degree of \"coarsening\" still yields a useful general-purpose tool (say, for predicting protein-ligand binding or unbinding from a transition state, or other such tasks).", "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 a new machine learning pipeline that emulates molecular dynamics computations. The overall procedure is a multi-stage process:\n\n1. A GNN operating on the atomic structure learns local per-atom and per-edge features\n2. A \"coarse graining\" operation is applied, which agglomerates atom features into features which represent clusters (\"beads\") in the original graph, and connects them with corresponding edges\n3. A \"dynamics network\" predicts the motion of these coarse-grained beads at $\\Delta t$ in the future given current and recent positions and embeddings\n4. The predictions are refined using a score-based denoising diffusion model to generate the final predicted locations of the coarse-grained beads\n\nThe atom features and the dynamics network are trained end-to-end; the refinement process is trained following a denoising score-matching objective.\n\nA key point in the efficiency of the learned model is that $\\Delta t$ from the machine learning model is a much longer timestep than one forward step of a traditional molecular dynamics simulator.\n\nThere are a good set of experiments which successfully defend the approach, showing good generalization to much longer simulations than the training regime, and including good ablation studies. Stability of the approach under long iterations is decent, and clearly improved by the use of the diffusion refinement.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "My requested changes are largely minor.\n\nRegarding confusion on the coarse-graining process:\n\n* How are fine-grained atoms assigned values M? Page 4 includes a reference to a function $C(\\mathbf{v}^F_{i,t})$ which computes the corresponding assignment $m \\in { 1,\\dots,M \\}$. However, I don't understand where this function comes from. Is this fixed ahead of time by the practitioner? If so, how should this be done in general? Or is it a trained neural network which acts as a clustering model? This part of the paper was not clear to me.\n* Is this coarse-graining process stable over time? For example, if you computed CG beads at time $t_0$, then ran MD forward some duration $\\Delta t$ and computed CG beads again, would the resulting CG graph (i.e. in terms of assignments of atoms to beads, and in definition of edges) be the same at both $t_0$ and $t_0 + \\Delta t$? Unless I am missing something this would not in general be true. Maybe this does not matter, but still (along with the previous bullet point) this contributed to some confusion around the CG process.\n* In section 4.1 the coarse-grained beads are further coarse-grained. Are we okay with this? Why? This might not need clarified in the paper itself, a comment to me is fine.\n\nRegarding end-to-end training:\n\n* Why is the score-based diffusion trained on a general denoising objective with Gaussian noise injected, and then used to specifically \"denoise\" the output of the dynamics network? It seems to me that this would be a reasonable choice only if the errors from the dynamics network were themselves Gaussian and of a particular scale. If the error distribution of the dynamics network were somehow biased then I would think that instead the dynamics network itself should be used as a source of \"noise\". In any case I would like to hear the motivation for this (even if it does not make it to the main paper), as this choice seemed odd.\n\nRegarding presentation:\n\n* For figure 2(d) isn't the left figure the original fine-grained polymer? It is labeled as the coarse-grained one. Or is this the output of the first intermediate coarse-graining? This is unclear.\n* Something is wrong with the $r^2$ values for the LSTM in Fig 3 and Table 2. In Fig 3 the value is larger than the method marked \"Ours\", despite (visually) seeming to have a much lower correlation. In table 2, the number for $r^2$ for the LSTM is negative, which is impossible (unless this is \"adjusted\" in some way that I didn't catch).\n\nRegarding baselines:\n\n* As mentioned in \"negatives\", in e.g. the ACF computation in Figure 5(a): I see that this broadly tracks the reference, but is this performance \"good\" or \"excellent\"? By what standards? How would baselines perform here — surely there must be some meaningful baselines, even if long-run MD computations are not feasible. As it stands the paper claims this is \"excellent\" performance but the figures are lacking baselines to demonstrate this adequately. (Even including the supervised learning methods, which I believe could be applied here but would suffer due to the need for long extrapolations, would be helpful….) \n\n", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tqJss55ajX", "reviewer": "Reviewer_19Rz", "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\nI think it is interesting for the authors to point out there are good coarse grained model for long time series things for Li-ion polymer electrolytes.\n\nI like that the bead type embeddings is chemically informed, with unique values for atom type, molecular weight, and bond type.\n\nI like how in section 4.1 training was done on class I and then testing was done in class II.\n\nI’m very happy with the Earth Mover’s Distance (EMD) statistic presented in Table 2. I think that is a great way to show and compare different distributions across your models rather than just showing a plot, which is what is typically done.\n\nI think Figure 6c is very helpful! I like how each of the steps of coarse-graining are shown.\n\n\nWeaknesses\n\n\nHow many k historical CG states are needed? Is it Markovian, and why do we need to know a long state?\n\nWhy do we need the Graph-clustering CG model, METIS? This is actually how you go from the atom-represenation to the bead representation? Does this representation change over the course of the simulation? This seems like it would not change much for single, long polymers, but would change a lot for simulations with few bonds. Are the only bonds you are using covalent?\n\nFor Figure 3, what does the r^2 tell us? What does the LSTM result mean? The LSTM result seems off relative to others.\n\nFigure 5b is not very useful as a comparison. It seems like the relative distribution of free energy is similar, but the TICA dimensions aren’t meaningful or comparable between Reference and Predicted. Are there any other dimensions or values to plot that would be more useful?\n", "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 authors propose a machine learning system to perform long simulations on many-atom systems with chemically-informed multiscale models. The authors have quite robust statistics and experiments with reasonable cross validation to back up their claims.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It would be great to give the readers some intuition of squared radius of gyration. What does a good or bad value look like, and what are the potential ways bad statistics arise?\n\nFor Figure 4, what do those distributions look like for the other ML baselines? Also, are you also able to show the EMD, as you did in Table 2? It’d be great to compare that sense of variability.\n\nIs the only source of instability “bead collision”? Are there other aspects of stability, or non-chaotic systems that could be considered? This seems like an instance of the “exploding gradient problem”: Pascanu, Razvan, Mikolov, Tomas, and Bengio, Yoshua. Understanding the exploding gradient problem. arXiv preprint arXiv:1211.5063, 2012.\n\nI would like a little bit more intuition about Figure 9 in the text. How should we interpret log(cm^2/s), and why is it consistently higher, and not lower, as a predicted value?\n\nWith regards to the rotational and translational equivariance described in the discussion, I don’t think that new analysis and methods should be introduced. If it is that important, it should be included into the main work, or otherwise a supplemental note.\n", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "jtESzg9tf1", "reviewer": "Reviewer_bcdg", "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:\nThe proposed method can achieve stable simulation at relatively long time scales and generate reasonable property predictions for down-stream tasks.\n\n\nWeaknesses and Questions:\n1. The CG method is based on a graph partitioning algorithm. This is advantageous for its simplicity. However, this can be a potential problem as the graph partitioning algorithm does not account for chemical properties of the molecule. Why is it the case that having roughly equal-sized partition blocks is a desirable property for the partitioning algorithm?\n\n\n2. The score GNN is not described very clearly. Why such a diffusion based model can help with the stability issue? Or what is the anticipated cause of the stability issue, the inaccurate prediction or the large variance from the dynamics GNN? How is the correction applied precisely during inference (it's briefly mentioned that it follows the annealing Langevin dynamics from previous works)?\n\n\n3. The overall loss is a direct combination of the two component losses. Are there any attempts to weigh the two terms differently? What about the convergence for both loss terms during training?\n\n\n4. How is the time step size decided for the experiments? It seems that the time step size is implicitly modeled in the training procedure. For a new application, how would the user know the right time step size to train the model on? Is the model's performance robust to a wide range of various time step sizes?\n\n\n5. How could this work be combined with other works on studying particular GNN structures for MD related tasks (the paper mentioned that extending the current model to account for more invariant/equivariant properties is a promising direction)? Is it possible to just replace the GNN model structures with other structures with desirable properties? Also, could these properties solve the stability problem? If so, how do we compare this approach to the score correction approach as proposed?\n\n\nMinor: \n\n\nWhat is \"q\" in equation 1?", "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 paper proposes training coarse-grained GNNs to run molecular dynamics simulations in relatively long time scales. The paper utilizes three GNN models, one for embedding, one for the next time step prediction, and one for the distributional correction. The empirical studies show that the proposed method can achieve stable simulation on molecules of thousands of atoms while giving reasonable property predictions for down-stream tasks.  ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the questions in the previous part.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a method for simulating molecular dynamics using graph neural networks, which is otherwise very computationally expensive.  The reviewers all agreed on acceptance, noting that the paper is novel, clear, well-composed and addresses an interesting and important problem.  One reviewer noted that "the paper provides an excellent jumping-off point for further research".  The reviewers also noted that although there were some initial concerns regarding clarity and experiments, the authors did a good job of addressing these in their response.  
Decision: Accept (poster)