Allegro-FM: Toward an Equivariant Foundation Model for Exascale Molecular Dynamics Simulations

We present a foundation model for exascale molecular dynamics simulations by leveraging an E(3) equivariant network architecture (Allegro) and a set of large-scale organic and inorganic materials data sets merged by the Total Energy Alignment framework. The obtained model (Allegro-FM) is versatile for various material simulations for diverse downstream tasks covering 89 elements in the training sets. Allegro-FM exhibits excellent agreement with high-level quantum chemistry theories in describing structural, mechanical, and thermodynamic properties, while exhibiting emergent capabilities for structural correlations, reaction kinetics, mechanical strengths, fracture, and solid/liquid dissolution, for which the model has not been trained. Furthermore, we demonstrate the robust predictability and generalizability of Allegro-FM for chemical reactions using Transition1x, which consists of tens of thousands of organic reactions and 9.6 million configurations including transition state data, in addition to reactive simulations using calcium silicate hydrates as a test bed. With its computationally efficient, strictly local network architecture, Allegro-FM scales up to multibillion-atom systems with a parallel efficiency of 0.975 on the exaflop/s Aurora supercomputer at Argonne Leadership Computing Facility. The approach presented in this work demonstrates the potential of the foundation model for novel materials design and discovery based on large-scale atomistic simulations.