Frame Averaging for Invariant and Equivariant Network Design

Abstract

Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data. However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efficient. For example, Euclidean motion invariant/equivariant graph or point cloud neural networks. We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable. In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame. We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group. Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures. Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs. We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond $2$-WL graph separation, and $n$-body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.

References in corpus (23)

• Adam: A Method for Stochastic Optimization
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Semi-Supervised Classification with Graph Convolutional Networks
• Neural Message Passing for Quantum Chemistry
• PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space
• Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
• How Powerful are Graph Neural Networks?
• Improving Graph Neural Network Expressivity via Subgraph Isomorphism Counting
• SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks
• 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data
• Spherical CNNs
• Cormorant: Covariant Molecular Neural Networks
• E(n) Equivariant Graph Neural Networks
• Covariant Compositional Networks For Learning Graphs
• N-body Networks: a Covariant Hierarchical Neural Network Architecture for Learning Atomic Potentials
• Janossy Pooling: Learning Deep Permutation-Invariant Functions for Variable-Size Inputs
• Rotation Invariant Graph Neural Networks using Spin Convolutions
• Equivariant Flows: sampling configurations for multi-body systems with symmetric energies
• Breaking the Limits of Message Passing Graph Neural Networks
• ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations
• On the Universality of Rotation Equivariant Point Cloud Networks
• Global Attention Improves Graph Networks Generalization
• On Universal Equivariant Set Networks

Cited by in corpus (5)

• Overcoming the Barrier of Orbital-Free Density Functional Theory for Molecular Systems Using Deep Learning
• On the design space between molecular mechanics and machine learning force fields
• Completeness of Atomic Structure Representations
• Reflection-Equivariant Diffusion for 3D Structure Determination from Isotopologue Rotational Spectra in Natural Abundance
• Complete Neural Networks for Complete Euclidean Graphs