Equivariant Foundation Model for Exascale Molecular Dynamics Simulations

Foundation model (FM) is a paradigm shift in Artificial Intelligence (AI) and has transformed our way of model training, where a single universal model acquires sufficient robustness and generalizability to enable diverse, out-of-distribution downstream tasks. Recently, great strides have been made in terms of both model architecture and training datasets in chemical and material research domains. Massive materials datasets have become publicly available while advanced network architectures have been proposed to take advantage the geometrical tensor and equivariant embedding of molecular geometry backed by group theory. We present the first exa-deployable foundation model for molecular dynamics simulations on exaflop/s parallel supercomputers 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 framework for the first time.