[Paper Review] Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

Albert Musaelian, Simon Batzner+5 more|arXiv (Cornell University)|Apr 11, 2022

Machine Learning in Materials Science21 citations

TL;DR

Allegro is a strictly local, E(3)-equivariant neural network interatomic potential that achieves state-of-the-art accuracy on small-molecule benchmarks while enabling scalable, parallel simulations and strong transferability to out-of-distribution data.

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.

Motivation & Objective

• Develop a strictly local, equivariant interatomic potential that matches or exceeds the accuracy of message-passing approaches.
• Demonstrate scalability and parallelizability for large systems without sacrificing fidelity.
• Show strong transferability to out-of-distribution and higher-temperature data.
• Validate physical realism by reproducing structural and kinetic properties in MD simulations of complex electrolytes.

Proposed method

• Introduce Allegro, an arbitrarily deep E(3)-equivariant network with scalar (invariant) and tensor (equivariant) latent spaces.
• Decompose total energy into per-species and per-center-pair energies to ensure size-extensive, local energy contributions.
• Use tensor-product layers to propagate information via an embedded environment (weighted sum of neighbor spherical harmonics) with a density-trick to reduce computational cost.
• Embed two-body scalar features and then iteratively update invariant and equivariant latent states across layers, followed by an MLP to predict pair energies.
• Normalize internal features and targets to improve training stability and size extensivity across varying system sizes and compositions.

Experimental results

Research questions

• RQ1 Can a strictly local, equivariant neural network achieve state-of-the-art accuracy comparable to or better than message-passing interatomic potentials?
• RQ2 How well does a local Allegro model generalize to out-of-distribution configurations and higher temperatures?
• RQ3 Is Allegro scalable to very large systems while preserving accuracy and energy conservation during dynamics?
• RQ4 Can Allegro faithfully reproduce both structural and kinetic properties in MD simulations of complex electrolytes?

Key findings

• Allegro achieves state-of-the-art force-component MAE on revised MD-17 benchmarks, and shows competitive or superior performance on several small-molecule targets.
• A single tensor-product layer in Allegro can outperform existing deep MPNNs and transformers on the QM9 benchmark.
• Allegro demonstrates strong transferability to out-of-distribution data, outperforming other local MLIPs.
• Molecular dynamics with Allegro recovers structural and kinetic properties of Li3PO4 with excellent agreement to first-principles calculations.
• The model scales in parallel to simulate systems with up to 100 million atoms, illustrating practical large-scale applicability.