Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics (2204.05249v1)

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.

• The paper presents Allegro, a novel deep learning interatomic potential that employs strictly local equivariant representations to capture many-body correlations.
• The methodology bypasses traditional message passing networks, enabling efficient, linear scaling for large-scale atomistic simulations.
• Experimental results demonstrate state-of-the-art accuracy on benchmarks like QM9 and MD-17 while ensuring strong transferability and reliable dynamics reproduction.

Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

The paper "Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics" introduces Allegro, a novel deep learning interatomic potential designed for modeling the energy and atomic forces of molecular and material systems. Allegro presents a solution to the longstanding challenge in computational chemistry and materials science: achieving both high accuracy and computational scalability in simulations of atomistic dynamics.

At the core of this research is the development of a strictly local equivariant representation of atomic structures that bypasses the traditional reliance on message passing neural networks (MPNNs). Allegro instead utilizes tensor products of learned equivariant representations to model many-body correlations, thereby retaining the desirable properties of equivariance under Euclidean group operations. This architectural choice allows Allegro to overcome the typical computational bottlenecks associated with MPNNs due to their large receptive fields and iterative information propagation.

Numerical Performance and Claims

The paper highlights several significant results:

• Allegro demonstrates state-of-the-art accuracy on benchmark datasets such as QM9 and revised MD-17, covering diverse molecular systems. Notably, even a single layer of the Allegro model outperforms existing advanced models like deep message passing networks and transformers on the QM9 benchmark.
• Allegro exhibits remarkable transferability, effectively generalizing to out-of-distribution data and maintaining performance across systems with varying atomic configurations and environmental conditions.
• A practical demonstration is provided through molecular dynamics simulations of an amorphous phosphate electrolyte, where Allegro accurately reproduces structural and kinetic properties that are consistent with first-principles calculations.

Theoretical and Practical Implications

Allegro's architecture challenges the predominance of MPNNs in atomistic machine learning by offering a method that provides scalable parallel computation while maintaining or exceeding the accuracy of existing models. The architecture's strict locality ensures linear scaling with respect to the number of atoms, making it feasible to perform simulations involving over 100 million atoms—a significant milestone for large-scale atomistic simulations.

Theoretically, the work opens discussions on the necessity and efficiency of capturing long-range interactions and many-body correlations without extensive message passing. Allegro's design, which bypasses the cubic growth in receptive fields typical of MPNNs, might indicate a paradigm shift in how machine learning potentials are constructed, advocating for a reduction in complexity while preserving interpretability and accuracy through local interactions.

Future Developments

The formalism presented by Allegro suggests several avenues for future research. It invites deeper theoretical exploration into the nature of local versus non-local interactions in complex materials. Additionally, the ability to extend Allegro's strictly local framework to incorporate explicit long-range interactions, if deemed necessary, could further enhance its applicability in systems where such interactions play a critical role.

From a practical standpoint, Allegro sets a foundation for scaling AI-driven materials simulations to unprecedented atom counts and simulation times, providing a powerful tool for researchers in chemistry, materials science, and related fields.

In conclusion, Allegro represents a significant advancement in the field of machine learning for atomistic simulations, offering a viable path towards achieving the dual objectives of accuracy and scalability—an essential combination for pushing the boundaries of computational materials science and chemistry.