Rotation Invariant Graph Neural Networks using Spin Convolutions

TL;DR

This work tackles the high computational cost of first-principles atomic simulations by introducing SpinConv, a rotation-invariant graph convolution that encodes angular information via a per-edge local reference frame and a spin convolution over the remaining roll degree of freedom. It presents two variants—an energy-centric model enforcing energy conservation through differentiating energy with respect to positions, and a force-centric model that directly regresses atomic forces—achieving state-of-the-art results on the OC20 dataset and strong performance on MD17 and QM9. The approach captures rich angular relations beyond triplets, scales to large datasets, and supports structure relaxation and molecular dynamics tasks, offering a practical path toward faster catalyst discovery and molecular simulations. While not yet reaching practical accuracy across all tasks, SpinConv demonstrates the viability of rotation-invariant angular modeling in atomic GNNs and provides a foundation for further efficiency and domain-specific enhancements.

Abstract

Progress towards the energy breakthroughs needed to combat climate change can be significantly accelerated through the efficient simulation of atomic systems. Simulation techniques based on first principles, such as Density Functional Theory (DFT), are limited in their practical use due to their high computational expense. Machine learning approaches have the potential to approximate DFT in a computationally efficient manner, which could dramatically increase the impact of computational simulations on real-world problems. Approximating DFT poses several challenges. These include accurately modeling the subtle changes in the relative positions and angles between atoms, and enforcing constraints such as rotation invariance or energy conservation. We introduce a novel approach to modeling angular information between sets of neighboring atoms in a graph neural network. Rotation invariance is achieved for the network's edge messages through the use of a per-edge local coordinate frame and a novel spin convolution over the remaining degree of freedom. Two model variants are proposed for the applications of structure relaxation and molecular dynamics. State-of-the-art results are demonstrated on the large-scale Open Catalyst 2020 dataset. Comparisons are also performed on the MD17 and QM9 datasets.

Figures (4)

• Figure 1: Illustration of projecting an atom $\acute{s}$ in the neighborhood of $s$ onto a sphere in a local coordinate frame defined by atom $s$ and $t$ (left). For each projected atom, a corresponding latitude $\phi$ (inclination) and longitude $\theta$ (azimuth) is computed for its projection onto a 2D reference frame (middle). The spin convolution is done in the longitudinal direction, corresponding to a roll is 3D space. (right) Example channel filters that are learned using the grid-based approach for the first through third message blocks and the force block.
• Figure 2: (left) Overall model diagram for energy-centric model taking atom positions $\bm{x}$ and atomic numbers $a$ as input and estimating the energy $E$. (right) Diagram of the embedding and force blocks. The force block is only used in the force-centric model to estimate the per-atom forces after the message blocks.
• Figure 3: Illustration of learned embeddings (weights on the one-hot embeddings) for the source $a_s$ and target $a_t$ atomic numbers plotted on a periodic table. A random sample of 12 values from each embedding are shown. Embeddings are from the first embedding block in the first message update. Note that neighboring atoms in the periodic table with similar properties have similar weights. Elements not in the OC20 dataset are marked with a light grey checkerboard pattern.
• Figure 4: Performance of SpinConv ablations on OC20 Val ID $30k$ (Table \ref{['tab:comp-ablation']}). All models trained for $560k$ steps and plotted against wall-clock training time. Note force-centric models and grid-based approaches converge more quickly than energy-centric models and those using spherical harmonics.