Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Frame Averaging for Invariant and Equivariant Network Design

Omri Puny, Matan Atzmon, Heli Ben-Hamu, Ishan Misra, Aditya Grover, Edward J. Smith, Yaron Lipman

TL;DR

The paper addresses the challenge of designing invariant/equivariant neural networks without sacrificing expressiveness or incurring prohibitive cost. It proposes Frame Averaging (FA), a general framework that replaces full group averaging with averaging over a finite, group-equivariant frame, enabling exact symmetry properties while preserving backbone power. The authors prove expressiveness preservation and universality under FA and instantiate the framework across point clouds (E(d)/SE(d)) and graphs (S_n, E(d)), yielding universal FA-PointNet/FA-DGCNN and FA-GNN models. Empirically, FA achieves state-of-the-art results on point cloud normal estimation, beyond-2-WL graph separation, and n-body dynamics, and they demonstrate effective approximate FA for large symmetry groups. Overall, FA provides a versatile, principled approach to building highly expressive, symmetry-preserving models with broad applicability andclear open questions about frame selection and stability.

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.

Frame Averaging for Invariant and Equivariant Network Design

TL;DR

The paper addresses the challenge of designing invariant/equivariant neural networks without sacrificing expressiveness or incurring prohibitive cost. It proposes Frame Averaging (FA), a general framework that replaces full group averaging with averaging over a finite, group-equivariant frame, enabling exact symmetry properties while preserving backbone power. The authors prove expressiveness preservation and universality under FA and instantiate the framework across point clouds (E(d)/SE(d)) and graphs (S_n, E(d)), yielding universal FA-PointNet/FA-DGCNN and FA-GNN models. Empirically, FA achieves state-of-the-art results on point cloud normal estimation, beyond-2-WL graph separation, and n-body dynamics, and they demonstrate effective approximate FA for large symmetry groups. Overall, FA provides a versatile, principled approach to building highly expressive, symmetry-preserving models with broad applicability andclear open questions about frame selection and stability.

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 -WL graph separation, and -body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.

Paper Structure

This paper contains 44 sections, 10 theorems, 35 equations, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Frame Averaging
  3. Frame averaging for function symmetrization
  4. Incorporating $G$ as a second symmetry.
  5. Right actions.
  6. Efficient calculation of invariant frame averaging.
  7. Approximation of invariant frame averaging.
  8. Expressive power
  9. Model instances
  10. Point clouds, Euclidean motions.
  11. Graphs, permutations
  12. Graphs, Euclidean motions.
  13. Symmetry and frame.
  14. Previous works
  15. Rotation invariant and equivariant point networks.
  16. ...and 29 more sections

Key Result

Theorem 1

Let ${\mathcal{F}}$ be a $G$ equivariant frame, and $\phi:V\rightarrow\mathbb R$, $\Phi:V\rightarrow W$ some functions. Then, $\left \langle \phi \right \rangle _{\mathcal{F}}$ is $G$ invariant, while $\left \langle \Phi \right \rangle _{\mathcal{F}}$ is $G$ equivariant.

Figures (5)

  • Figure 1: Frame equivariance (sphere shape represents the group $G$; square represents $V$).
  • Figure 2: Invariance error (left); $m_{_{\mathcal{F}}}$-$m_{_G}$ (right).
  • Figure 3: Minimal spacing histogram over the training set of $n$-body dataset satorras2021n .
  • Figure 4: Distance between original and noisy frames as a function of $\sigma$ (we plot average and std). The result is reported over the training set of $n$-body dataset satorras2021n .
  • Figure 5: FA-Local-PointNet architecture.

Theorems & Definitions (19)

  • Definition 1
  • Theorem 1: Frame Averaging
  • Theorem 2: Frame Average second symmetry
  • Theorem 3
  • Corollary 1
  • Theorem 4: Expressive power of FA
  • Proposition 1
  • Corollary 2: FA-DeepSets/PointNet are universal
  • Proposition 2
  • Corollary 3: FA-MLP and FA-GNN+ID is graph universal
  • ...and 9 more