Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

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 stateof-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.

Neural Scaling of Deep Chemical Models

“…5) was used for neural scaling experiments to improve convergence and avoid issues with systematic drift in predicted energies, which we identified during the course of this work and plan to address in future work. We use the SchNet [59], PaiNN [36], Allegro [10], and SpookyNet [37] models. Model implementations are from the NeuralForceField repository [34,60,61] and the Allegro repository [10].…”

“…That is, not only do the equivariant GNNs achieve better performance for a given data budget, they achieve increasingly greater performance gains given more training data. This is due to the models' equivariance, which is known to produce greater sample efficiency [9,10], but it is interesting to note that this trend persists to much larger and more chemically diverse datasets than were previously considered, which typically include only 10 2 − 10 3 molecular geometries from a single molecular species.…”

“…In this work we consider four flavors of NFFs: SchNet [35], PaiNN [36], Allegro [10], and SpookyNet [37]. This series of models represents increasingly physics-informed model architectures, from models with internal layers that manipulate only E(3) invariant quantities (SchNet) to those that use E(3) equivariant quantities (PaiNN, Allegro, SpookyNet), strictly local models with learned many-body functions and no message passing (Allegro), and physically-informed via empirical corrections (SpookyNet).…”

“…For efficiency and to isolate scaling behavior, we fix hyperparameters from TPE as m and d are varied, but strictly speaking the optimal hyperparameters will change as m and d vary [21]. Due to computational resource limitations, we train ChemGPT models for a fixed number of epochs (10) to determine the loss.…”

“…However, unlike in CV and NLP, the path to scaling deep chemical networks and the potential benefits are unclear. Chemical deep learning can incorporate physics-based priors that may ameliorate the steep resource requirements seen in other fields [9][10][11][12]. Moreover, because of the heterogeneity and complexity of chemical space [13] and molecular machine learning (ML) tasks [14,15], training general and robust models that perform well on a wide variety of downstream tasks remains a pressing challenge [8,16,17].…”