[Paper Review] Rotation Invariant Graph Neural Networks using Spin Convolutions

Muhammed Shuaibi, Adeesh Kolluru+5 more|arXiv (Cornell University)|Jun 17, 2021

Machine Learning in Materials Science34 references51 citations

TL;DR

SpinConv introduces a rotation-invariant graph neural network for atomic structures by using per-edge local frames and a spin convolution over the remaining rotation degree of freedom, achieving state-of-the-art results on OC20 for energy and force estimation.

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.

Motivation & Objective

• Motivate efficient, accurate DFT-like predictions of atomic forces and energies with rotation-invariant GNNs.
• Develop a per-edge local coordinate frame enabling angular information to be captured without losing rotational invariance.
• Propose a spin convolution over the remaining rotational degree of freedom to achieve rotation-invariant edge messages.
• Present energy-centric and force-centric SpinConv variants to balance energy conservation and direct force estimation.
• Demonstrate state-of-the-art performance on OC20 and competitive results on MD17 and QM9.

Proposed method

• Represent atomic structures as graphs with atoms as nodes and neighbor edges within a cutoff.
• Compute edge messages h_st iteratively with a shared update that includes spin convolution, distance encoding, and atomic-number embeddings.
• Use a local frame defined by the source s and target t to project neighbors onto a sphere around s, then apply spin convolution across the roll degree of freedom about st axis.
• Employ either an energy-centric path (E via -∂E/∂x for forces) or a force-centric path (directly predicting forces) for training objectives.
• Adopt a grid-based or spherical-harmonics-based spherical representation for the spin convolution, followed by pooling to ensure rotation invariance.
• Incorporate an embedding block (mixture-of-experts style) to fuse atomic-number information into message updates.
• Optionally rotate the entire structure during training for robustness and to encourage approximate rotation equivariance in the force block.

Experimental results

Research questions

• RQ1 Can a spin-convolution-based, rotation-invariant GNN accurately learn energy and force fields for atomic systems?
• RQ2 How do energy-centric and force-centric SpinConv variants compare in terms of energy conservation and predictive accuracy?
• RQ3 What is the effect of representation choice (grid-based vs spherical harmonics) and model size on OC20, MD17, and QM9 performance?
• RQ4 To what extent does inducing rotation robustness (via multiple random rotations) improve force predictions?
• RQ5 How does SpinConv perform on in-domain vs out-of-domain (adsorbates/catalysts) splits in OC20?

Key findings

• SpinConv with force-centric training achieves the lowest energy MAE and force MAE among tested models on OC20 S2EF tasks.
• Grid-based spin convolution trains faster and achieves competitive accuracy with spherical harmonics for OC20; both representations yield similar accuracies.
• Energy-centric SpinConv improves energy accuracy but is outperformed by force-centric variants in overall S2EF metrics and IS2RS results.
• On OC20 IS2RE, SpinConv outperforms previous approaches, especially in the direct regression setup for relaxed energy predictions.
• On OC20 IS2RS, SpinConv achieves the best ADwT (53.62% averaged across splits) but DimeNet++-large attains the best AFbT (21.82%), with SpinConv still offering faster inference (~263.2h vs 814.6h).
• Across MD17 and QM9, SpinConv shows competitive or superior performance to state-of-the-art methods, particularly when leveraging the force-centric formulation and appropriate representation size.