Learning local equivariant representations for large-scale atomistic dynamics

A Musaelian and S Batzner and A Johansson and LX Sun and CJ Owen and M Kornbluth and B Kozinsky, NATURE COMMUNICATIONS, 14, 579 (2023).

DOI: 10.1038/s41467-023-36329-y

A simultaneously accurate and computationally efficient parametrization of the potential energy surface of molecules and materials is a long- standing goal in the natural sciences. While atom-centered message passing neural networks (MPNNs) have shown remarkable accuracy, their information propagation has limited the accessible length-scales. Local methods, conversely, scale to large simulations but have suffered from inferior accuracy. This work introduces Allegro, a strictly local equivariant deep neural network interatomic potential architecture that simultaneously exhibits excellent accuracy and scalability. Allegro represents a many-body potential using iterated tensor products of learned equivariant representations without atom-centered message passing. Allegro obtains improvements over state-of-the-art methods on QM9 and revMD17. A single tensor product layer outperforms existing deep MPNNs and transformers on QM9. Furthermore, Allegro displays remarkable generalization to out-of-distribution data. Molecular simulations using Allegro recover structural and kinetic properties of an amorphous electrolyte in excellent agreement with ab-initio simulations. Finally, we demonstrate parallelization with a simulation of 100 million atoms. The paper presents a method that allows scaling machine learning interatomic potentials to extremely large systems, while at the same time retaining the remarkable accuracy and learning efficiency of deep equivariant models. This is obtained with an E(3)- equivariant neural network architecture that combines the high accuracy of equivariant neural networks with the scalability of local methods.

Return to Publications page