PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics

Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, Boris Kozinsky

arXiv:2204.05249v1physics.comp-phcond-mat.mtrl-scics.LGphysics.chem-ph

TL;DR

Accurate, efficient atomistic potentials remain difficult to scale because message passing hinders parallel computation, while local descriptors have lagged in accuracy. Allegro uses strictly local equivariant tensor products without message passing, achieving strong benchmark accuracy, transferability, molecular-dynamics agreement, and large-system scalability.

  • Problem

    Accurate and computationally efficient energy and force models are needed for molecular and materials simulations, but message passing complicates parallelization while local descriptors have lower accuracy.

  • Method

    Allegro is a strictly local equivariant interatomic potential that learns many-body functions through tensor products of learned equivariant representations without message passing.

  • Results

    Allegro achieves state-of-the-art accuracy, strong out-of-distribution transferability, agreement with first-principles dynamics, and parallel scaling to hundreds of millions of atoms.

  • Takeaways & Limitations

    Allegro challenges message passing as the standard for atomistic machine learning while supporting accurate simulations of complex systems at large scale.

  • Takeaways & Limitations

    The strict-local formulation leaves the treatment of explicit long-range interactions outside the demonstrated scope, while message-passing models face extended-receptive-field parallelization burdens.

Abstract

from arXiv · show

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.

INTRODUCTION … Equivariant Neural Networks

The paper motivates machine-learning interatomic potentials that combine accurate energy and force predictions with computational efficiency, then introduces Allegro as a strictly local equivariant approach designed for scalable atomistic simulation. It also explains how message passing expands receptive fields and how E(3)-equivariant operations preserve the symmetries of atomic systems.

  • INTRODUCTION: Molecular dynamics and Monte Carlo simulations require accurate, computationally efficient predictions of potential energies and atomic forces for long-time studies of complex systems.
  • Related Work: Machine-learning interatomic potentials approximate high-fidelity energy and force labels while targeting computational efficiency that scales linearly with atom number.
  • INTRODUCTION: Allegro combines the accuracy of equivariant message-passing models with strict locality and parallelization to access systems containing hundreds of millions of atoms.
  • Message Passing Interatomic Potentials: Message-passing interatomic potentials represent atoms as graph nodes and connect pairs within a fixed cutoff distance, iteratively updating hidden atomic states using edge features.
  • Message Passing Interatomic Potentials: 6 message-passing layers with a 6 Å local cutoff produce a 36 Å effective cutoff in a 64-molecule liquid-water example at P = 1 bar and T = 300 K.
  • Equivariant Neural Networks: Atomic-system physics is unchanged by E(3) symmetries, with energies invariant and atomic forces equivariant under rotations, inversion, and translation.
  • Equivariant Neural Networks: Equivariant neural networks operate directly on geometric inputs such as displacement vectors using E(3)-equivariant operations, producing internal features that transform consistently.
  • Equivariant Neural Networks: Tensor products combine equivariant representations while respecting valid irreducible representations, underpinning operations such as scalar multiplication, dot products, and cross products.

Atom-density Representations

The Atomic Cluster Expansion represents local atomic environments through body-ordered expansions of projected local atomic densities, retaining invariant scalar products for energy prediction. Its computational cost grows polynomially with the number of radial-chemical basis functions as body order increases, motivating species embeddings and restricted tensor-product combinations.

  • Atomic Cluster Expansion: ACE represents local atomic environments with coefficients from a body-ordered expansion of the local atomic density, providing invariant environment descriptions.The local density is projected onto combinations of radial basis functions and species channels.
  • Atomic Cluster Expansion: Only scalar, rotation-invariant tensor products are retained, and total energy is fit linearly over these scalars through a chosen maximum body order.Scalar outputs can arise, for example, when combined basis functions have matching rotation orders ℓ1 = ℓ2.
  • Computational scaling: ACE has polynomial computational scaling in Nfull-basis as body order ν+1 increases, creating a core evaluation bottleneck.Nfull-basis denotes the number of two-body radial-chemical basis functions.
  • Computational scaling: In the basic ACE descriptor, Nfull-basis = Nbasis × S, directly coupling basis size to the number of radial basis functions and species.Species embeddings have been proposed to remove the direct dependence on S.
  • Computational scaling: NequIP and related equivariant neural networks avoid unfavorable embedded-basis scaling by computing tensor products for only a restricted set of input-tensor combinations.This contrasts with computing the broader combinations used in ACE.

RESULTS · Energy decomposition · The Allegro model

The method learns high-dimensional potential energy surfaces with strictly local many-body equivariant representations. Allegro decomposes system energy into central-atom pair contributions and uses an arbitrarily deep equivariant architecture to predict those contributions while preserving energy conservation.

  • RESULTS: The proposed method targets high-dimensional potential energy surfaces using strictly local many-body equivariant representations.
  • Energy decomposition: System potential energy is decomposed into per-atom energies E_i, following previous approaches.
  • Energy decomposition: Each per-atom energy is further decomposed into pairwise energies indexed by a central atom and one neighbor, unlike standard MLIP practice.
  • Energy decomposition: Although indexed by atom i and neighbor j, E_ij may depend on all neighbors of i; E_ij and E_ji use their corresponding central-atom environments and are therefore unequal by design.
  • Energy decomposition: The resulting force field conserves energy.
  • The Allegro model: Allegro is an arbitrarily deep equivariant neural network with N_layer ≥1 layers that learns representations associated with neighboring-atom pairs through invariant and equivariant latent spaces.
  • The Allegro model: The invariant latent space contains scalar ℓ= 0 features, while the equivariant space processes tensors of arbitrary rank ℓ≥0, with both spaces interacting at every layer.
  • The Allegro model: The final pair energy E_ij is computed by a multi-layer perceptron acting on the final layer’s scalar features, whose equivariant representations are truncated at hyperparameter ℓ_max.

Two-body latent embedding

The two-body latent embedding initializes Allegro by encoding pairwise scalar properties with a nonlinear MLP, while geometric information is represented through radial and spherical-harmonic features. These initial pair features provide inputs for subsequent equivariant processing.

  • Scalar pair embedding: A nonlinear MLP embeds the scalar properties of each atom pair ij into initial scalar latent features x_ij,L=0.This embedding occurs before the first tensor product layer.
  • Scalar pair embedding: The pair embedding incorporates a one-hot encoding of the center and neighbor atom species, Z_i and Z_j.The species encoding is concatenated within the embedding expression.
  • Geometric features: Interatomic distance r_ij is projected onto a radial basis using Bessel functions with a polynomial envelope.The basis is normalized as described in Appendix A and plotted in Figure 6.
  • Geometric features: The equivariant features n,ℓ,p are computed by linearly embedding the spherical-harmonic projection of the unit bond direction r̂_ij.These features represent directional information for each center-neighbor pair.
  • Geometric features: For each center-neighbor pair ij, the equivariant features n,ℓ,p are computed from the initial two-body scalar latent features.The channel index n ranges over the equivariant channels.

Layer architecture

Each Allegro tensor product layer combines learned environment embedding, an equivariant tensor product, scalar-latent updating, and equivariant channel mixing. The tensor product couples center-neighbor features with other neighbors through learned environment representations and valid equivariant paths.

  • Layer architecture: Each layer has four components: an environment-weighting MLP, an equivariant tensor product, a scalar-latent update MLP, and an equivariant linear mixing layer.These components respectively embed the central atom’s environment, process equivariant features, update scalar information, and mix equivariant channels.
  • Tensor product: The tensor product combines the current center-neighbor equivariant state with geometry from other neighbor pairs in the central atom’s local environment.Bilinearity enables the density trick, expressing the update with one tensor product rather than one per neighbor and saving computational effort.
  • Tensor product: Valid tensor product paths satisfy |ℓ1 −ℓ2| ≤ℓout ≤|ℓ1 +ℓ2|, pout = p1p2, and ℓout ≤ℓmax; this work includes all allowable paths.Restricting output irreps to values among the inputs is another possible hyperparameter choice.
  • Environment embedding: Allegro learns the weights used to form its embedded environment, rather than relying on a predetermined radial-chemical basis as in descriptor methods such as ACE.The learned embedding uses the scalar featurization of each center-neighbor pair and produces a fixed number of channels nequivariant.
  • Latent update and feature mixing: Scalar tensor-product outputs are compressed into the fixed-dimensional invariant latent space, coupling scalar and equivariant representations, while same-irrep paths are linearly mixed into one output space.The mixing preserves the input channel count and compresses information from multiple paths sharing an output irrep.

Residual update · Output block

Allegro uses scalar-space residual updates to carry information from earlier layers forward. A fully connected neural network then maps the final latent features to the pair energy Eij.

  • Residual update: Residual updates operate in the scalar latent space after each layer.They update the previous scalar features from layer L −1 by adding the new features.
  • Residual update: The residual update adds new scalar features to those from layer L −1.This additive update forms the layer-to-layer feature propagation mechanism.
  • Residual update: Residual updates allow scalar information from earlier layers to propagate forward easily.The passage identifies this as the principal role of the residual connection.
  • Residual update: The residual operation is applied after each layer.Thus, the scalar latent representation is updated repeatedly throughout the network.
  • Output block: The output block predicts the pair energy Eij from the final layer’s latent features.The prediction is made using a fully connected neural network.
  • Output block: The fully connected output network has output dimension 1.Its single output is used for the pair-energy prediction.

Normalization … Quantum-chemical properties of small molecules

The paper specifies normalization schemes for internal features and targets, including choices that preserve size extensivity. It then evaluates Allegro on revised MD17 and QM9, emphasizing force accuracy and transferability to new molecular structures.

  • Internal normalization: e3nn-style initialization and normalization make network outputs approximately zero mean and unit variance element-wise.The scheme chooses initial weight distributions and normalization constants accordingly.
  • Internal normalization: Sums over multiple features are divided by the square root of the number of terms, assuming uncorrelated terms whose variances add.This normalization applies explicitly to the embedded environment and atomic energy, which sum over neighbors.
  • Normalization of targets: For fixed chemical composition, target normalization sets µZ to average per-atom potential energy and σZ to the root-mean-square force components.The scheme uses final prediction scale and shift parameters and is reported as highly important for training.
  • Normalization of targets: Target normalization ensures size extensivity, whereas subtracting mean total potential energy across the training set violates size extensivity.Size extensivity is required when evaluating systems larger or smaller than the training systems.
  • Normalization of targets: For varying chemical composition, a linear pre-fitting scheme computes µZ from species counts and reference energies.The formula uses [Nconfig,Z]−1 [Econfig], with [Nconfig,Z] containing species atom counts and [Econfig] containing reference energies.
  • Small molecule dynamics: On revised MD17, Allegro achieves state-of-the-art mean absolute error in force components, while NequIP performs better on energies for some molecules.The benchmark contains 10 small organic molecules at DFT accuracy.
  • Transferability to higher temperatures: Allegro’s transferability is tested on new structures that may be visited during long molecular dynamics simulations, including conformations at higher temperatures.A stronger distance-decay exponent induces a stronger bias toward smaller pair energies from atoms farther from the central atom.
  • Quantum-chemical properties of small molecules: On QM9, Allegro is evaluated across chemical space using approximately 134k DFT structures and four energy-related molecular-property targets.The structures contain C, H, O, N, and F and up to 9 heavy atoms.

Li-ion dynamics in an Phosphate Electrolyte · Scaling

Allegro reproduces structural and kinetic properties of amorphous Li3PO4, including Li-ion mean-squared displacement, while its strictly local design enables favorable scaling and parallel computation to large systems. Its costs scale linearly with atom count and neighbor count, remain constant with species count, and support strong scaling across GPUs.

  • Li-ion dynamics in an Phosphate Electrolyte: Allegro trains on 50 ps AIMD trajectories at 3000 K in the melted state and 600 K in the quenched state to model amorphous Li3PO4.The model is evaluated in the quenched amorphous state and used for molecular dynamics at 600 K.
  • Li-ion dynamics in an Phosphate Electrolyte: Allegro’s small, fast model sufficiently captures complex structural and kinetic properties during the Li3PO4 phase change.Ten 50 ps simulations at T=600K compare Allegro against AIMD.
  • Li-ion dynamics in an Phosphate Electrolyte: Allegro shows excellent agreement with AIMD for Li mean-squaredisplacement in the quenched state.The Li MSD results are averaged over 10 Allegro runs, with shading indicating ± one standard deviation.
  • Scaling: O(N) scaling in atom count enables Allegro to avoid the O(N 2) scaling of some global descriptor methods such as sGDML.Allegro is strictly local within the chosen cutoff, making it easy to parallelize.
  • Scaling: O(M) scaling in neighbors per atom contrasts with the quadratic O(M 2) scaling of DimeNet and Equivariant Transformers.The comparison concerns the number of neighbors per atom M.
  • Scaling: O(1) scaling in species count contrasts with SOAP’s O(S2) and ACE’s O(Sbody order−1) dependence.This scaling is attributed to Allegro’s local equivariant architecture.
  • Scaling: Independent central-atom neighborhoods allow energy and force contributions to be computed in parallel and combined through a final sum reduction.Allegro’s total energy is a sum over scaled pairwise energies Eij, and each term depends only on a local neighborhood.
  • Scaling: 421,824 atoms remain fixed while strong-scaling experiments vary A100 GPUs from 1 to 64 in LAMMPS.Scalability across devices uses an Allegro extension to LAMMPS compatible with spatial decomposition.

Theoretical Analysis

The theoretical analysis relates Allegro’s iterated equivariant tensor products to ACE while identifying important differences in basis construction, scaling, and body order. Allegro avoids ACE’s basis-size growth through learnable channel mixing, while its body order depends on the nonlinearities used in embedding networks.

  • Recursive construction: Allegro’s first layer forms sums of tensor products between pair features and spherical-harmonics projections of neighboring atoms, with bilinearity yielding sums over symmetrically valid irrep-index paths.Repeated substitution exposes a recursive relationship across layers.
  • Allegro–ACE comparison: Allegro with Nlayer layers and ACE body order ν + 1 = Nlayer + 2 share the same core iterated equivariant tensor products, but their features are not equivalent.Allegro features are analogous to, but distinct from, full equivariant ACE basis functions.
  • Scaling: ACE carries an expanding set of radial-chemical basis indices, whereas Allegro features do not increase this way across layers.This difference produces O(N ν full-basis) scaling for ACE versus O(1) for Allegro in the radial-chemical basis size Nfull-basis; learnable channel mixing enables Allegro’s favorable scaling.
  • Body order: Allegro’s final features have finite body order when environment embedding weights have finite body order, which holds for linear latent and embedding MLPs.Infinite-Taylor-expansion nonlinearities make the embedding weights and entire model infinite-body-order, while two-body MLP nonlinearities do not affect body order.

DISCUSSION

The discussion presents Allegro as a deep-learning interatomic potential that combines accurate energy and force prediction with scalability to large atomistic systems and molecular-dynamics property prediction. It also raises questions about representation choices and suggests that learned layerwise environment-weight dependence may contribute to Allegro’s performance relative to ACE.

  • Contributions: Allegro predicts energies and forces accurately, scales to large atomistic systems, and predicts structural and kinetic properties in molecular-dynamics simulations of complex systems.The authors describe the method as challenging the standard of widely used message-passing neural networks.
  • Representation: The findings raise questions about the optimal choice of representation across atomistic machine-learning methods.This question is framed in the context of Allegro challenging message-passing neural networks.
  • Relation to ACE: The authors speculate that Allegro may outperform systematic ACE basis expansion partly because each layer learns environment-embedding weights from scalar latent features in all previous layers.This dependence may let the model learn an atom’s importance for higher body-order interactions from lower body-order descriptions of its environment.

METHODS … Training details

The methods evaluate Allegro and reference models across several quantum-chemistry datasets and molecular-dynamics systems, using specified simulation protocols and software implementations. Training was conducted on a single NVIDIA V100 GPU.

  • Software: Experiments used Allegro at commit a5128c2a86350762215dad6bd8bb42875ebb06cb, NequIP 0.5.3, e3nn 0.4.4, PyTorch 1.10.0, and Python 3.9.7.LAMMPS experiments used the LAMMPS code and the Allegro interface’s pair allegro implementation.
  • Reference Training Sets: 950 revised MD-17 structures were used for training and 50 for validation, from 10 organic molecules with 100,000 DFT-computed structures.The revised MD-17 calculations used PBE/def2-SVP accuracy, tight SCF convergence, and a dense DFT integration grid.
  • Reference Training Sets: The 3BPA dataset contains 500 training structures at T=300K and test sets at 300K, 600K, and 1200K.The test-set sizes are 1669, 2138, and 2139 structures, respectively; data used ωB97X/6-31G(d) calculations.
  • Reference Training Sets: QM9 contains 133,885 structures with up to 9 heavy elements and species H, C, N, O, and F in relaxed geometries.Properties were computed at the DFT/B3LYP/6-31G(2df,p) level; 3,054 structures failing geometry consistency checks were excluded.
  • Reference Training Sets: Li3PO4 reference data came from two 50 ps AIMD simulations in VASP using PBE, 400eV cutoff energy, Γ-point sampling, 2 fs steps, and NVT dynamics.The simulations used PAW pseudopotentials and a Nosé-Hoover thermostat.
  • Molecular Dynamics Simulations: Production and timing simulations used LAMMPS with pair allegro; Li3PO4 ran at T=600K with 2 fs steps, while Ag timing ran at T=300K with 5 fs steps.Li3PO4 used NVT, a Nosé-Hoover thermostat, and a temperature damping parameter of 40 time steps; Ag also used NVT.
  • Training details: Models were trained in single-GPU training on a NVIDIA V100 GPU.The training setup used one GPU rather than distributed multi-GPU training.

revMD-17 and 3BPA · 3BPA, NequIP · Li3PO4

The paper specifies separate training configurations for revMD-17 and 3BPA, including dataset budgets, architectures, cutoffs, bases, and loss normalization. It also describes the NequIP baseline on 3BPA and a compact Li3PO4 configuration.

  • revMD-17 and 3BPA: revMD-17 used 1,000 structures, split into 950 training and 50 validation, with 3 layers, 128 even and odd features, and ℓmax = 3.The radial cutoff was 7.0 ˚A for all molecules except Naphthalene, which used 9.0 ˚A.
  • revMD-17 and 3BPA: 3BPA used 500 structures, split into 450 training and 50 validation, with a 5.0 ˚A cutoff and an 8-function Bessel basis.Its polynomial envelope used p = 2.
  • revMD-17 and 3BPA: revMD-17 training used a joint energy-and-force loss with force weight 1,000 and total potential-energy weight 1.For 3BPA, the energy term was divided by N^2 and, after normalization, both energy and force terms received weight 1.
  • 3BPA, NequIP: NequIP on 3BPA used 5 layers, 64 even and odd features, and ℓmax = 3.The model used a radial network with 3 layers and 64 hidden neurons, plus equivariant SiLU-based gate nonlinearities.
  • 3BPA, NequIP: NequIP used a 5.0 ˚A radial cutoff, an 8-function non-trainable Bessel basis, and a polynomial envelope with p = 2.Its 3BPA data budget was 500 molecules, split into 450 training and 50 validation.
  • Li3PO4: Li3PO4 used 11,000 structures, split into 10,000 training and 1,000 validation, with 1 layer, 1 even-parity feature, and ℓmax = 1.Its 2-body latent MLP had hidden dimensions [32, 64], while the later latent MLP had one -dimensional hidden layer.
  • Li3PO4: Li3PO4 training used Adam in PyTorch with learning rate 0.001, batch size 1, and an on-plateau scheduler with patience 25 and decay factor 0.5.An exponential moving average with weight 0.99 was used for validation and the final model.
  • Li3PO4: The Ag model used 1,000 structures split into 950 training and 50 validation, while QM9 used 110,000 training and 10,000 validation structures.The Ag configuration used 1 layer, 1 even-parity feature, and ℓmax = 1; QM9 compared 3-layer ℓmax = 2 and 1-layer ℓmax = 3 models with 256 even and odd features.

APPENDIX … C. Variable per-layer cutoffs

The appendix describes normalization procedures for radial bases and scalar residual updates, then explains how Allegro implements decreasing per-layer cutoffs. Variable cutoffs restrict neighborhoods, edges, and cutoff envelopes by layer while preserving residual feature accumulation under an ordering condition.

  • A. Normalization of radial basis: Allegro components are designed and normalized to produce zero-mean, unit-variance outputs for inputs with matching statistics.This normalization target motivates adjustments to the radial basis.
  • A. Normalization of radial basis: The radial Bessel basis does not inherently satisfy the desired normalization, so its basis functions are numerically shifted and scaled.The adjustment makes the basis functions more closely satisfy the network’s normalization property.
  • A. Normalization of radial basis: The radial-basis normalization equations assume a simple uniform distribution of r between 0 and rmax.The construction uses the nth radial Bessel basis function eRBF,n and smooth polynomial cutoff u(r).
  • B. Normalization of scalar residual update: The scalar residual update at layer L is computed as a weighted sum with α = 1 2.The coefficient forms are selected to enforce normalization.
  • B. Normalization of scalar residual update: Assuming successive latent features are negligibly correlated and each has approximately zero mean and unit variance, the residual sum retains those statistics approximately.The assumption applies to xij,L−1 and xij,L at initialization.
  • C. Variable per-layer cutoffs: Variable cutoffs restrict neighborhood sums and embedded-environment sums to neighbors within the current layer’s cutoff rLmax.Each layer therefore uses only the neighbors allowed by its own cutoff.
  • C. Variable per-layer cutoffs: Each layer processes only edges within its cutoff, so residual scalar pair features sum contributions from layers whose cutoffs include the pair.The scalar-latent cutoff envelope u(rij) is also replaced by one using the current layer’s cutoff.
  • C. Variable per-layer cutoffs: The variable-cutoff scheme requires rL+1max; otherwise, previous-layer equivariant latent features may be missing for some edges.The supplied passage states this ordering condition is necessary for the scheme to make sense.
Loading 2204.05249v1…