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

Abstract: 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 datasets merged by Total Energy Alignment (TEA) framework. Thanks to the large-scale training sets, the obtained model (Allegro-FM) is versatile for various materials simulations for diverse downstream tasks covering 89 elements in the periodic table. Allegro-FM exhibits excellent agreements 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 the Transition1x that consists of 10k organic reactions and 9.6 million configurations including transition state data, as well as calcium silicate hydrates as a testbed. With its computationally efficient, strictly-local network architecture, Allegro-FM scales up to multi-billion-atom systems with a parallel efficiency of 0.964 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 a novel materials design and discovery based on large-scale atomistic simulations.