Rotation Invariant Graph Neural Networks using Spin Convolutions

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.

arXiv comment13 pages

References in corpus (18)

• Neural Message Passing for Quantum Chemistry
• Crystal Graph Convolutional Neural Networks for an Accurate and Interpretable Prediction of Material Properties
• SchNet - a deep learning architecture for molecules and materials
• Quantum-Chemical Insights from Deep Tensor Neural Networks
• Machine Learning of Accurate Energy-Conserving Molecular Force Fields
• PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments and Partial Charges
• The Open Catalyst 2020 (OC20) Dataset and Community Challenges
• Searching for Activation Functions
• Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
• SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
• 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data
• Cormorant: Covariant Molecular Neural Networks
• E(n) Equivariant Graph Neural Networks
• Neural Message Passing with Edge Updates for Predicting Properties of Molecules and Materials
• Relevance of Rotationally Equivariant Convolutions for Predicting Molecular Properties
• Spherical Message Passing for 3D Graph Networks
• ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations
• An Introduction to Electrocatalyst Design using Machine Learning for Renewable Energy Storage

Cited by in corpus (7)

• The Open Catalyst 2022 (OC22) Dataset and Challenges for Oxide Electrocatalysts
• DPA-2: a large atomic model as a multi-task learner
• AdsorbML: A Leap in Efficiency for Adsorption Energy Calculations using Generalizable Machine Learning Potentials
• Crystal Diffusion Variational Autoencoder for Periodic Material Generation
• Boosting Heterogeneous Catalyst Discovery by Structurally Constrained Deep Learning Models
• Frame Averaging for Invariant and Equivariant Network Design
• Unifying O(3) Equivariant Neural Networks Design with Tensor-Network Formalism