The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials

TL;DR

This paper unifies and analyzes various E(3)-equivariant interatomic potential models, introduces a simplified architecture called BOTNet, and provides insights into design choices affecting accuracy and extrapolation.

Contribution

It develops a mathematical framework unifying ACE and NequIP models, and introduces BOTNet, a simplified, interpretable model with competitive accuracy.

Findings

Systematic ablation study highlights critical design choices.

BOTNet achieves high accuracy on benchmark datasets.

Framework enables exploration of design space for equivariant potentials.

Abstract

The rapid progress of machine learning interatomic potentials over the past couple of years produced a number of new architectures. Particularly notable among these are the Atomic Cluster Expansion (ACE), which unified many of the earlier ideas around atom density-based descriptors, and Neural Equivariant Interatomic Potentials (NequIP), a message passing neural network with equivariant features that showed state of the art accuracy. In this work, we construct a mathematical framework that unifies these models: ACE is generalised so that it can be recast as one layer of a multi-layer architecture. From another point of view, the linearised version of NequIP is understood as a particular sparsification of a much larger polynomial model. Our framework also provides a practical tool for systematically probing different choices in the unified design space. We demonstrate this by an ablation study of NequIP via a set of experiments looking at in- and out-of-domain accuracy and smooth extrapolation very far from the training data, and shed some light on which design choices are critical for achieving high accuracy. Finally, we present BOTNet (Body-Ordered-Tensor-Network), a much-simplified version of NequIP, which has an interpretable architecture and maintains accuracy on benchmark datasets.