Force-Free Molecular Dynamics Through Autoregressive Equivariant Networks

TL;DR

This work introduces TrajCast, a transferable and data-efficient framework based on autoregressive equivariant message passing networks that directly updates atomic positions and velocities lifting the constraints imposed by traditional numerical integration.

Abstract

Molecular dynamics (MD) simulations play a crucial role in scientific research. Yet their computational cost often limits the timescales and system sizes that can be explored. Most data-driven efforts have been focused on reducing the computational cost of accurate interatomic forces required for solving the equations of motion. Despite their success, however, these machine learning interatomic potentials (MLIPs) are still bound to small time-steps. In this work, we introduce TrajCast, a transferable and data-efficient framework based on autoregressive equivariant message passing networks that directly updates atomic positions and velocities lifting the constraints imposed by traditional numerical integration. We benchmark our framework across various systems, including a small molecule, crystalline material, and bulk liquid, demonstrating excellent agreement with reference MD simulations for structural, dynamical, and energetic properties. Depending on the system, TrajCast allows for forecast intervals up to $30\times$ larger than traditional MD time-steps, generating over 15 ns of trajectory data per day for a solid with more than 4,000 atoms. By enabling efficient large-scale simulations over extended timescales, TrajCast can accelerate materials discovery and explore physical phenomena beyond the reach of traditional simulations and experiments. An open-source implementation of TrajCast is accessible under https://github.com/IBM/trajcast.

Figures (8)

• Figure 1: Overview of the TrajCast architecture.(A) Autoregressive workflow: An atomistic system at time $t_0$ is passed through an equivariant MPNN (grey box) to predict the new positions and velocities at time $t_1$. Atomic attributes (positions, velocities, chemical elements) are encoded into initial features, which are refined over $T$ message passing blocks. Estimates of the displacement and velocity vectors are generated based on the final features. These are then refined to ensure momentum conservation. The trajectory is built by rolling out predictions, where outputs from one step serve as inputs for the next. A thermostat ensures sampling from the canonical (NVT) ensemble at constant temperature $T$, with states following the Boltzmann distribution. (B) The embedding block encodes node and edge attributes and generates the initial features. (C) Messages are constructed by convolving latent features with filters derived from a learnable radial basis and the spherical harmonics expansion of edge vectors. (D) In the update block, messages from neighbors are pooled and combined via a tensor product with velocity vectors in a learnable radial and spherical harmonic basis. The result is passed through a non-linearity and added to the previous layer's features, weighted by the node's chemical element. (E) Conservation of total linear and angular momentum is enforced by adjusting the displacements and velocities.
• Figure 2: Comparison of properties of paracetamol from MD and TrajCast with a 7 fs prediction horizon. (A) Dynamical properties are compared based on the element-specific vibrational density of states (VDOS). (B) Energetics are validated by comparing the distribution of the potential energy, computed in post-processing for the Trajcast-generated trajectory by using the same force field as in the reference MD. (C) Free energy barriers drive the system’s kinetics and are computed from the free energy surface (FES), shown for both MD and TrajCast at the top. The FES is generated by sampling two dihedral angles, $\psi_1$ and $\psi_2$, visualized by the inset within the bottom panel. The center panel visualizes different realizations of paracetamol for various dihedral pairs, matching the colors of corresponding points in the FES. Barriers are calculated along the cuts shown on the FES and colored accordingly in the bottom panel.
• Figure 3: Evaluation of TrajCast on crystalline quartz with a 30 fs prediction horizon. (A) Illustration of an atomic configuration of quartz used to train and validate TrajCast. (B) Dynamical properties are compared based on the element-specific vibrational density of states (VDOS). (C) Energetics are validated by comparing the distribution of the potential energy, computed in post-processing for the Trajcast-generated trajectory by using the same force field as in the reference MD.
• Figure 4: Benchmarking TrajCast on liquid water with a 5 fs prediction horizon. (A) Snapshot of an instantaneous configuration of bulk water. (B) Comparison of the element-specific vibrational density of states (VDOS) between TrajCast and MD. (C) Mean squared displacement (MSD) of oxygen atoms over a correlation time $\tau$ of up to 10 ps, complementing the dynamic evaluation. (D) Validation of the molecular liquid structure based on the pairwise combinations of chemical elements in the radial distribution function (RDF). The shaded areas represent the threefold standard deviation, computed over 4 blocks. The diffusion coefficients are computed based on a fit to the slope in the range from 1 to 10 ps following the procedure outlined in marsalek_quantum_2017.
• Figure 5: Impact of prediction horizon and training set size on performance of TrajCast. (A) Parity plot comparing predicted and reference velocities for TrajCast models trained for water with prediction horizons of 0.5 and 15 fs, respectively. (B) Dependence of the relative velocity MAE on the MD time-step multiple. Error bars correspond to the threefold standard deviations computed across three models varying in the initial weights. (C) Exponents from power law fits to learning curves, quantifying the relationship between model accuracy and training set size across different chemical systems and MD time-step multiples. The inset shows that these exponents correspond to the slope $s$ of a linear fit on a log-log plot of MAE versus training set size. The larger the absolute value of $s$ the larger the improvement on the generalization error of the model when additional data is used. The dashed lines serve as guide to the eye.