Simulate Time-integrated Coarse-grained Molecular Dynamics with Geometric Machine Learning

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

Keywords: molecular dynamics, learning to simulate, coarse-graining, graph neural network, score-based generative models, polymer, material science

TL;DR: A graph neural network simulator that simulates molecular dynamics at a lower spatio-temporal resolution while being several orders of magnitude faster than traditional molecular dynamics.

Abstract: Molecular dynamics (MD) simulation is the workhorse of various scientific domains but is limited by high computational cost. Learning-based force fields have made major progress in accelerating ab-initio MD simulation but are still not fast enough for many real-world applications that require long-time MD simulation. In this paper, we adopt a different machine learning approach where we coarse-grain a physical system using graph clustering and model the system evolution with a very large time-integration step using graph neural networks. A novel score-based GNN refinement module resolves the long-standing challenge of long-time simulation instability. Despite only being trained with short MD trajectory data, our learned simulator can generalize to unseen novel systems, and simulate for much longer than the training trajectories. Properties requiring 10-100 ns level long-time dynamics can be accurately recovered at several orders of magnitude higher speed than classical force fields. We demonstrate the effectiveness of our method on two realistic complex systems: (1) single-chain coarse-grained polymers in implicit solvent; (2) multi-component Li-ion polymer electrolyte systems.