Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

TL;DR

Allegro introduces a strictly local, $E(3)$-equivariant interatomic potential that captures many-body physics without message passing, achieving state-of-the-art accuracy on QM9 and revMD-17 while enabling parallel scaling to hundreds of millions of atoms. The model uses a two-space latent framework (invariant and equivariant) and relies on tensor-product interactions with an embedded environment that is learned from data, yielding linear $\mathcal{O}(N)$ scaling and strong GPU scalability. Extensive benchmarks demonstrate strong transferability to high-temperature conformations and complex materials like Li3PO4, with accurate structural and kinetic properties recovered in MD simulations. Theoretical analysis connects Allegro to ACE, clarifying how learned environment weights enable finite effective body-order and efficient learning, while empirical results show Allegro can outperform both message-passing networks and descriptor-based baselines in key settings. Overall, Allegro offers a scalable, accurate, and transferable framework for large-scale atomistic dynamics with practical impact for materials design and molecular simulations.

Abstract

A simultaneously accurate and computationally efficient parametrization of the energy and atomic forces of molecules and materials is a long-standing goal in the natural sciences. In pursuit of this goal, neural message passing has lead to a paradigm shift by describing many-body correlations of atoms through iteratively passing messages along an atomistic graph. This propagation of information, however, makes parallel computation difficult and limits the length scales that can be studied. Strictly local descriptor-based methods, on the other hand, can scale to large systems but do not currently match the high accuracy observed with message passing approaches. This work introduces Allegro, a strictly local equivariant deep learning interatomic potential that simultaneously exhibits excellent accuracy and scalability of parallel computation. Allegro learns many-body functions of atomic coordinates using a series of tensor products of learned equivariant representations, but without relying on message passing. Allegro obtains improvements over state-of-the-art methods on the QM9 and revised MD-17 data sets. A single tensor product layer is shown to outperform existing deep message passing neural networks and transformers on the QM9 benchmark. Furthermore, Allegro displays remarkable generalization to out-of-distribution data. Molecular dynamics simulations based on Allegro recover structural and kinetic properties of an amorphous phosphate electrolyte in excellent agreement with first principles calculations. Finally, we demonstrate the parallel scaling of Allegro with a dynamics simulation of 100 million atoms.

Figures (6)

• Figure 1: The Allegro network. Left: the Allegro model architecture. Right: a tensor product layer. Blue and red arrows represent scalar and tensor information, respectively, $\otimes$ denotes the tensor product, and $\oplus$ is concatenation.
• Figure 2: Structural Properties. Left: Radial Distribution Function, right: Angular Distribution Function of tetrahedral bond angle. All defined as probability density functions.
• Figure 3: Li Dynamics. Comparison of the Li MSD of AIMD vs. Allegro. Results are averaged over 10 runs of Allegro, shading indicates +/- one standard deviation.
• Figure 4: The quenched Li3PO4 structure at T=600K
• Figure 5: Strong Scaling results on a Li3PO4 structure of 421,824 atoms, performed in LAMMPS.