Toward Generalizable Surrogate Models for Molecular Dynamics via Graph Neural Networks

TL;DR

This work addresses the computational bottleneck of classical molecular dynamics by introducing a graph neural network surrogate that directly predicts atomic displacements over a finite time interval $\Delta t$, effectively learning the atomic evolution operator without explicit force evaluations. The model combines graph-based message passing with transformer-style self-attention over per-timestep graphs constructed from a large aluminum MD dataset generated with an EAM potential, and it uses rich node/edge features, static embeddings, and radial basis function encodings to capture local and many-body interactions. Across training and validation, the surrogate achieves sub-angstrom accuracy within the training horizon, remains stable for short-to-mid horizon extrapolations ($\Delta t$ up to 50), and preserves key physical signatures such as radial distribution functions and temperature-dependent mean squared displacement, while also maintaining thermodynamic fidelity (temperature and volume) relative to MD. Ablation studies show that the approach remains accurate with relatively shallow networks, while a minimum hidden dimension is required to properly represent high-dimensional atomic motion. Together, these results establish GNN-based surrogate integrators as a promising, computationally efficient complement to traditional MD for accelerated, physically faithful atomistic simulations within a validated regime.

Abstract

We present a graph neural network (GNN) based surrogate framework for molecular dynamics simulations that directly predicts atomic displacements and learns the underlying evolution operator of an atomistic system. Unlike conventional molecular dynamics, which relies on repeated force evaluations and numerical time integration, the proposed surrogate model propagates atomic configurations forward in time without explicit force computation. The approach represents atomic environments as graphs and combines message-passing layers with attention mechanisms to capture local coordination and many-body interactions in metallic systems. Trained on classical molecular dynamics trajectories of bulk aluminum, the surrogate achieves sub angstrom level accuracy within the training horizon and exhibits stable behavior during short- to mid-horizon temporal extrapolation. Structural and dynamical fidelity are validated through agreement with reference radial distribution functions and mean squared displacement trends, demonstrating that the model preserves key physical signatures beyond pointwise coordinate accuracy. These results establish GNN-based surrogate integrators as a promising and computationally efficient complement to traditional molecular dynamics for accelerated atomistic simulations within a validated regime.

Figures (6)

• Figure 1: (a–d) Atomic configurations of bulk aluminum at increasing temperatures showing the progression from a perfect FCC lattice to thermally disordered states. (e) Corresponding radial distribution function $g(r)$ at 300 K, 500 K, and 800 K.
• Figure 2: Training and validation loss as a function of epoch.
• Figure 3: Model displacement error (in Å) as a function of $\Delta t$ at 350 K.
• Figure 4: Actual versus model-derived radial distribution functions at 300--800 K for $\Delta t = 5$.
• Figure 5: Actual versus model-derived radial distribution functions at 300--800 K for $\Delta t = 50$.