FlashMD: long-stride, universal prediction of molecular dynamics

TL;DR

FlashMD tackles the time-scale gap in molecular dynamics by learning direct, long-stride trajectories with a graph neural network that predicts updates to positions and momenta over strides Δτ that can exceed standard MD steps by 1–2 orders of magnitude. The approach enforces energy conservation at inference, accommodates arbitrary thermodynamic ensembles, and explicitly addresses challenges from chaos, time-reversibility, and symmetry breaking, enabling stable, long-time MD propagation. Across water systems and universal models, FlashMD demonstrates accurate reproduction of equilibrium observables and meaningful time-dependent behavior for diverse chemistries, achieving substantial speedups over conventional MD. This direct MD trajectory predictor has the potential to dramatically extend accessible atomistic time scales and offers a practical, universal companion to traditional ML interatomic potentials for large-scale, long-time simulations.

Abstract

Molecular dynamics (MD) provides insights into atomic-scale processes by integrating over time the equations that describe the motion of atoms under the action of interatomic forces. Machine learning models have substantially accelerated MD by providing inexpensive predictions of the forces, but they remain constrained to minuscule time integration steps, which are required by the fast time scale of atomic motion. In this work, we propose FlashMD, a method to predict the evolution of positions and momenta over strides that are between one and two orders of magnitude longer than typical MD time steps. We incorporate considerations on the mathematical and physical properties of Hamiltonian dynamics in the architecture, generalize the approach to allow the simulation of any thermodynamic ensemble, and carefully assess the possible failure modes of such a long-stride MD approach. We validate FlashMD's accuracy in reproducing equilibrium and time-dependent properties, using both system-specific and general-purpose models, extending the ability of MD simulation to reach the long time scales needed to model microscopic processes of high scientific and technological relevance.

Figures (11)

• Figure 1: Schematic overview of FlashMD. Atoms of the system at time step $\tau$ are taken as inputs, with atomic numbers $Z_i$ and momenta $\boldsymbol{p}_i(\tau)$ embedded into the node features $\boldsymbol{h}_i$, and relative coordinates $\boldsymbol{q}_{ij}(\tau)$ embedded into the edge features $\boldsymbol{e}_{ij}$ of a GNN for the system. The node outputs are used to predict the new configuration $\boldsymbol{p}_i(\tau+\Delta\tau)$ and $\Delta\boldsymbol{q}_i(\tau+\Delta\tau)$ in a multi-head manner. Center-of-mass constraints are also enforced. Uncertainty quantification can be enabled as shown in the navy inset. Optional filters for energy conservation enforcement, thermodynamic ensemble control, and random rotation are provided, as discussed below. Conventional MD would require $\Delta\tau$ explicit numerical integrations to reach the final configuration as opposed to 1 pass of FlashMD.
• Figure 2: Comparison of physical observables obtained from MD (black) and FlashMD (other colors). Left and Center: radial distribution functions for oxygen and hydrogen atoms, respectively, from simulations in the $NVT$ ensemble using the Langevin thermostat. Right: densities from simulations in the $Np T$ ensemble.
• Figure 3: Results of case studies conducted for the universal FlashMD models. (a) Ramachandran plots of the main backbone dihedrals for a simulation of solvated alanine dipeptide at 450 K. (b) Mean square displacement (MSD) of the Al (110) surface atoms at 500 K, at different layers from the surface (B indicates the limiting value for the bulk). The premelting and defect formation phenomena are also visualized as traces of atomic positions from a FlashMD simulation at 600 K, run with $\Delta\tau=$ 64 fs. The ideal atomic positions are also shown for reference. (c) Li conductivities of $\gamma-$Li$_3$PS$_4$ at varying $T$, along with the initial system configuration overlaid with traces of the Li atom positions from the FlashMD trajectory at 700 K. In both (b) and (c), traces are obtained with a moving average in time to remove thermal fluctuations and visualize more clearly the diffusive behavior.
• Figure 4: Implementation of the random rotation filter. A random rotation matrix $\boldsymbol{R}$ is sampled and applied on all coordinates and momenta before the rotated inputs are provided to FlashMD. After model inference, $\boldsymbol{R}^{-1}$ is applied to rotate the system back to the original coordinate reference. Random rotation filter is only relevant for rotationally unconstrained GNNs.
• Figure 5: Implementation details of the energy conservation enforcement filter in FlashMD. Energy model can be any model of the interatomic potential (e.g. MLIP, classical force field, etc.) that can be used to compute $V({\{\boldsymbol{q}_i\})}$.