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

Foundation models (FM) represent a paradigm shift in AI, where a single universal model acquires sufficient robustness and generalizability to enable diverse, out-of-distribution downstream tasks. Leveraging ALCF supercomputing resources, a research team led by the University of Southern California has introduced Allegro-FM to perform exascale molecular dynamics simulations to accurately describe a wide variety of material properties and processes with a single pretrained model. Allegro-FM covers 89 elements in the periodic table and can be applied to a wide range of tasks, including structural correlations, reaction kinetics, mechanical strengths, fracture, and solid/liquid dissolution, exhibiting emergent capabilities for which the model was not trained.

Challenge

Fracture mechanics of silicate materials are studied for technological importance as well as to answer fundamental scientific questions. To examine the applicability of Allegro-FM to fracture behavior, the researchers performed a tensile test on a tobermorite 11Å (T11) crystal, a calcium silicate mineral and the key ingredient for fire-resilient concrete. The researchers also simulated the carbonation process—important for carbon sequestration—for a nanoparticle of T11 placed in a mixture of H₂O and CO₂ molecules.

Approach

The researchers have developed an exascalable universal machine learning interatomic potential by leveraging an E(3) equivariant network architecture and a set of large-scale organic and inorganic materials datasets merged by the Total Energy Alignment framework.

Results

In a paper featured on the cover of the The Journal of Physical Chemistry Letters, the researchers demonstrated Allegro-FM’s excellent agreement with high-level quantum chemistry theories for describing structural, mechanical, and thermodynamic properties, as well as its emergent capabilities for structural correlations, reaction kinetics, mechanical strengths, fracture, and solid/liquid dissolution, for which the model has not been trained. The researchers also demonstrated the robust predictability and generalizability of Allegro-FM for chemical reactions using Transition1x and confirmed the scalability of Allegro-FM on the ALCF’s exascale Aurora supercomputer, with a 97.5 percent perfect speedup for 4.08 billion atoms across 4,096 GPUs. The obtained molecular dynamics trajectory for the carbonation process simulation was robust, with no system failures or spurious events observed.

Impact

This work demonstrates the potential of using AI foundation models to accelerate and advance novel materials design and discovery via large-scale atomistic simulations.