Simulate Time-integrated Coarse-grained Molecular Dynamics with Multi-scale Graph Networks

Xiang Fu, Tian Xie, Nathan J. Rebello, Bradley Olsen, Tommi S. Jaakkola

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.

Submission Length: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

We thank all reviewers for their helpful feedback. All changes are colored brown in the updated draft.

- We added further details on the coarse-graining procedure and score-based correction procedure.
- We replaced the TICA analysis in Fig 5 with a contact map analysis and moved the TICA analysis to Appendix C.
- We extended and moved discussion on SE(3) equivariance to Appendix A.
- We further explain the intuition on the squared radius of gyration in Appendix C and Fig 12.
- We added baseline results using reference MD of various lengths for ACF and relaxation time of squared radius of gyration to Appendix C, Fig 13, and Fig 14.
- We added supervised learning baseline results on recovering the distribution of the squared radius of gyration to Appendix C, Fig 15.

Video: https://www.youtube.com/watch?v=l3aGVjQezsc

Code: https://github.com/kyonofx/mlcgmd

Assigned Action Editor: ~Jasper_Snoek1

License: Creative Commons Attribution 4.0 International (CC BY 4.0)

Submission Number: 1110