PaperlayerView exact revision on arXiv

Source-linked AI summary

Frame Averaging for Invariant and Equivariant Network Design

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

arXiv:2110.03336v4cs.LGstat.ML

TL;DR

Invariant and equivariant neural networks can be difficult to make expressive and computationally efficient, especially for large or continuous symmetry groups. The paper introduces Frame Averaging, which replaces full group averaging with averaging over equivariant frames while preserving exact symmetry properties and backbone expressive power. FA yields universal graph and Euclidean-motion models and achieves state-of-the-art results across several benchmark tasks.

  • Problem

    Designing neural networks that are simultaneously invariant or equivariant, expressive, and computationally efficient remains difficult for important symmetry types.

  • Method

    Frame Averaging adapts backbone architectures by replacing full group averaging with averaging over a carefully selected finite subset of group elements called a frame.

  • Results

    FA preserves backbone expressive power, supports universal graph and Euclidean-motion invariant or equivariant models, and achieves state-of-the-art performance across point-cloud, graph-separation, and n-body tasks.

  • Takeaways & Limitations

    Frame Averaging provides a generic framework for constructing exact invariant or equivariant models across multiple symmetry groups and backbone architectures.

  • Takeaways & Limitations

    The paper considers only finite frames, and its point-cloud frame construction is undefined at rare repeated-eigenvalue points and can change sharply near them.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Frame Averaging (FA) addresses the difficulty of building expressive, efficient invariant or equivariant networks by replacing intractable group averaging with averaging over carefully selected frames. The framework supports multiple backbone architectures and symmetry types, while preserving expressive power and enabling universal models.

  • Motivation: The framework targets symmetry settings where invariant or equivariant architecture design is otherwise elusive, computationally expensive, or insufficiently expressive.Exact averaging becomes intractable for large combinatorial groups or infinite continuous groups, motivating the frame-based alternative.
  • Core idea: FA replaces full group averaging with averaging over a carefully selected subset called a frame, retaining exact invariance or equivariance and expressive power.The frame is required to satisfy a set equivariance property, and small frames can make computation more efficient.
  • Instantiations: FA is instantiated for permutation-invariant graph models, Euclidean-motion invariant or equivariant message-passing GNNs, and Euclidean-motion equivariant set and point-cloud networks.The backbones include MLPs, node-identification GNNs, Message-Passing GNNs, DeepSets, PointNet, and DGCNN.
  • Theory: Theoretically, FA preserves the expressive power of backbone architectures and yields universal graph and Euclidean-motion invariant or equivariant models.The paper also reports an E(d) invariant/equivariant GNN that maintains message-passing power.
  • Experiments: Experiments apply FA to point-cloud normal estimation, beyond 2-WL graph separation, and n-body dynamics prediction, achieving state-of-the-art performance in all reported benchmarks.These experiments evaluate FA across several invariant and equivariant learning tasks.

2 FRAME AVERAGING

Frame Averaging replaces intractable full group averaging with averaging over an equivariant frame, preserving exact invariance or equivariance while reducing computation. Under bounded, frame-finite conditions, it also preserves the backbone's expressive power and supports universal constructions.

  • Generic formulation: Frame Averaging replaces full group averaging with averaging over a nonempty, set-valued frame F(X) that is equivariant under the group action.A frame is G-equivariant when F(ρ1(g)X) = gF(X).
  • Function symmetrization: Frame averaging is computationally attractive when the frame is easy to compute and has manageable cardinality, unlike averaging over large or infinite groups.Taking the entire group as the frame recovers standard group averaging but can be intractable.
  • Function symmetrization: Averaging arbitrary scalar or vector-valued functions over an equivariant frame yields G-invariant or G-equivariant functions, respectively.The invariant case is recovered by using scalar outputs with the trivial representation.
  • Second symmetry: FA can combine a backbone's existing H symmetry with a new G symmetry when the representations commute and the frame is H-invariant.Under these conditions, frame-averaged invariant and equivariant backbones become G × H invariant or equivariant.
  • Efficient calculation: For invariant FA, quotienting the frame by the stabilizer reduces evaluations from |F(X)| to m_F = |F(X)|/|G_X|.Uniform random samples from F(X) induce uniform samples over the quotient, enabling approximation; smaller m_F is especially useful for highly symmetric inputs.
  • Expressive power: If the frame is bounded and frame-finite, FA preserves approximation of arbitrary continuous equivariant functions and therefore can inherit universal approximation from a universal backbone.The approximation domain includes the points sampled by the frame operator.

3 MODEL INSTANCES

The paper instantiates FA for Euclidean-symmetric point clouds and permutation-symmetric graphs using PCA-, sorting-, and feature-based frames. These constructions yield universal architectures, including FA-PointNet, FA-DGCNN, FA-MLP, FA-GNN+ID, and FA-GNN.

  • Point clouds: For point clouds, FA targets E(d) or SE(d) symmetry acting on coordinates through rotations or reflections and translations.The input space is V = R^(n×d), with Euclidean motions represented by ρ1(g)X = X R^T + 1 t^T.
  • Point clouds: The point-cloud frame uses covariance eigenvectors and centroid translations, producing 2^d elements, typically 4 in d = 2 and 8 in d = 3.The frame is E(d)-equivariant and bounded when the covariance has a simple spectrum.
  • Point clouds: The PCA-based point-cloud frame is defined only when the covariance matrix has non-repeating eigenvalues and can change sharply near repeated eigenvalues.Repeated-eigenvalue cases are described as rare singular points, while near-degeneracies can cause large frame changes under small perturbations.
  • Point-cloud architectures: Frame averaging turns PointNet and DGCNN backbones into universal E(3)-equivariant models, named FA-PointNet and FA-DGCNN.A composed local construction is called FA-Local-PointNet.
  • Graphs: For graphs, the permutation frame contains permutations that lexicographically sort rows of a feature matrix S(X), and it is S_n-equivariant and bounded.The graph input combines node features and an adjacency or edge-attribute matrix.
  • Graphs: Using MLP or GNN+ID backbones with the sorting frame yields universal S_n-equivariant graph models over bounded domains.The graph stabilizer is the automorphism group, which can make the frame large for highly symmetric graphs.
  • Euclidean graph dynamics: The Euclidean graph instantiation applies PCA-based frames to node features and composes FA with Message Passing GNN layers to form FA-GNN.The n-body input uses velocities alongside initial positions.

4 PREVIOUS WORKS

Prior work includes rotation-invariant point networks and E(3)-symmetric GNNs, but the paper positions FA as a broader framework with theoretical expressive-power analysis and universal constructions.

  • Rotation-invariant point networks: Earlier point-network approaches obtain rotation invariance through modified convolutions, transformed input representations, geometric relations, or local and global frames.The paper identifies PCA-based methods as instances of the FA framework.
  • Rotation-invariant point networks: FA extends prior point-network work to rotation/reflection and translation equivariance, broader architectures and symmetries, and theoretical expressive-power analysis.The cited prior approaches primarily address rotation invariance or related frame constructions.
  • Equivariant GNNs: Many SE(3)-equivariant GNNs inherit expensive higher-order intermediate representations from Tensor Field Networks.A recent efficient message-passing construction is described as E(d)-equivariant but not yet shown universal.

5 EXPERIMENTS

The experiments apply Frame Averaging to point-cloud normal estimation, graph separation, and n-body dynamics, evaluating invariant or equivariant learning across these tasks. FA provides tractable symmetry-aware models and shows stronger practical behavior in the reported comparisons.

  • Experimental tasks: FA is evaluated on point-cloud normal estimation, graph separation, and E(3)-equivariant n-body particle-position prediction.The tasks cover O(3) equivariance with translation invariance, S_n invariance, and E(3) equivariance.
  • Point clouds: normal estimation: Normal estimation uses the ABC CAD-model dataset and compares original-data and randomly rotated train/test settings using normal-prediction error.The reported error is 1 − (n^T n̂)^2, with PointNet, DGCNN, VN-PointNet, and VN-DGCNN among the baselines.
  • Graphs: expressive power: Both FA and full-group averaging achieve perfect graph separation, while approximate FA has lower invariance error than approximate GA even with k = 1 samples.The experiment uses GRAPH8c and EXP, and the figure reports invariance error versus sample size together with m_F − m_G.
  • Graphs: expressive power: With k = 1 samples, FA learns the graph-classification task whereas GA fails, consistent with approximate FA's stronger invariance in the reported comparison.The classification experiment is the second graph-separation task in Table 2.
  • n-body problem: FA-GNN improves over the state of the art by more than 20% on n-body particle-location prediction, using roughly the same number of parameters as EGNN.Performance is measured by mean squared error between predicted and ground-truth locations.

6 CONCLUSIONS

The paper concludes that Frame Averaging adapts neural backbones to desired symmetries while preserving expressive power and offering efficient constructions. It validates this approach across universal GNNs, point-cloud networks, and Euclidean-motion-aware GNNs, while identifying open questions about frame design.

  • Contributions: Frame Averaging is presented as a generic methodology for adapting existing neural architectures to desired invariant or equivariant symmetries.The framework is intended for symmetries appearing in the data.
  • Theory: The paper proves that FA preserves the expressive power of backbone models and is efficient to compute in several cases of interest.The conclusion states that these properties support the proposed model constructions.
  • Model families: FA yields universal GNNs, universal Euclidean-motion-invariant or equivariant point-cloud networks, and Euclidean-motion-invariant or equivariant GNNs.These are the principal architecture families highlighted in the conclusion.
  • Open questions: The paper leaves systematic efficient-frame discovery, frame-choice effects on learning, and useful FA architectures or modules as open questions.It also points to extending invariance or equivariance from a subgroup H to a larger group G.

A.1 PROOF OF THEOREM 1

The proof uses frame equivariance: transforming the input transforms the frame by the corresponding group action while preserving its cardinality.

  • Frame equivariance: Frame equivariance is defined by F(ρ1(g)X) = gF(X), which implies that transformed inputs yield frames with equal cardinality.This cardinality preservation is used in the proof of Theorem 1.
  • Proof role: The proof therefore treats the frame as transforming consistently with the input under the group action.The supplied passage states the defining relation and its cardinality consequence.

A.2 PROOF OF THEOREM 2

The proof of Theorem 2 is presented as a generalization of the preceding proof, with the earlier argument recovered when the second symmetry is trivial.

  • Proof relationship: The proof above Theorem 1 is identified as a special case of the more general argument used for Theorem 2.The relationship is established by choosing the second symmetry to be trivial.
  • Proof structure: Theorem 2's proof and the preceding proof are described as being quite similar in principle.The supplied passage does not provide further proof steps.

A.3 PROOF OF THEOREM 3 AND COROLLARY 1

The proof establishes structural properties of equivariant frames and analyzes how frame sampling approximates full frame averaging. The approximation guarantee improves when the frame has fewer orbit classes.

  • Frame structure: Equivariance makes F(X) closed under the stabilizer group GX, partitioning it into disjoint orbits of equal cardinality.This equal-orbit structure implies uniform sampling over F(X) selects each orbit with equal probability.
  • Approximation: Frame averaging can be expressed as a weighted sum over orbit representatives, with empirical weights determined by sampled frame elements.The empirical distribution assigns each orbit the fraction of samples that fall into it.
  • Approximation: Theorem 5 lower-bounds the probability that the empirical distribution yields approximation error at most ϵ over bounded continuous functions.The bound concerns a particular good empirical distribution and its approximation to frame averaging.
  • Approximation: For fixed sample size k and tolerance ϵ, the probability bound becomes better as the number of frame orbits mF decreases.The passage attributes the counter-intuitive dependence on k and ϵ to the growing set of good empirical distributions.
  • Frame construction: The proof constructs an O(d) frame from covariance eigenvectors and verifies equivariance through the transformation rule (R,t)(O,s) = (RO,Rs + t).Boundedness follows on compact domains because centered translations remain bounded while orthogonal matrices have spectral norm one.
  • Frame construction: The resulting frame is shown to satisfy both inclusions required for equivariance, completing the argument that F((R,t)X)=(R,t)F(X).The reverse inclusion is obtained by applying the forward inclusion to the inverse transformation.

B EMPIRICAL FRAME ANALYSIS

The empirical analysis examines eigenvalue spacing and frame stability on the n-body training set. The results support finite frames when spectra are simple and frames vary continuously under noise.

  • Repeating eigenvalues: The covariance-spectrum experiment measures normalized eigenvalue spacing, with minimal spacing indicating proximity to repeated eigenvalues.The n-body particle sets are centered and scaled before computing the covariance eigenvalues.
  • Repeating eigenvalues: 10^-2 was the smallest observed minimal spacing across 3000 n-body particle sets, empirically supporting finite frames for E(d) equivariance.The spacing histogram summarizes the training set distribution.
  • Frame stability: Frame stability is evaluated by measuring how much the constructed O(d) frame changes when the input point set is perturbed with noise.The comparison uses original and noisy frames from the n-body dataset.
  • Frame stability: Figure 4 plots the average distance and standard deviation between original and noisy frames as functions of noise level σ.The evaluation is performed over the n-body training set.
  • Frame stability: The stability plot validates continuity of the constructed frames when the covariance spectrum is simple.In the simple-spectrum case, the distance is invariant to the choice of frame representatives.

C IMPLEMENTATION DETAILS

The implementation adapts MLP, GNN, PointNet, DGCNN, and message-passing backbones to invariant or equivariant settings. Experiments specify architectures, invariance evaluation, and training protocols.

  • Point-cloud architectures: PointNet uses learnable layers with matrix weights, biases, concatenation, and ReLU activation, while baseline models retain T-Net components.The FA-PointNet backbone excludes the T-Net networks used by the baseline implementation.
  • Point-cloud architectures: FA-Local-PointNet combines equivariant features from FA-applied MLPs on k-nearest-neighbor patches with an FA-PointNet normal-estimation stage.Each patch contains k=20 points, and the second stage operates on features in R^(n×3d).
  • Point-cloud architectures: DGCNN uses EdgeConv, fully connected, and MaxPool layers, with a T-Net applied to the input.The architecture is based on an object-part segmentation network.
  • Training: Training uses ADAM with fixed learning rate 0.001, while epoch counts and batch sizes vary across experiments and hardware.The n-body model is trained for 10000 epochs with batch size 100 on a single Nvidia RTX-6000 GPU.
  • Graph architectures: Graph experiments use MLP and GIN-with-identifiers backbones, with graph-size-specific padding and flattened representations for selected datasets.The GIN+ID model has three layers, feature dimension 64, and ReLU activation.
  • Evaluation: Permutation invariance is evaluated by comparing outputs on 50 randomly permuted graphs using an invariance-error metric.FA-MLP, GA-MLP, and regular MLP share the same backbone in the GRAPH8c evaluation, whose errors are normalized by regular MLP error.
  • n-body architecture: The n-body FA-GNN uses a message-passing update with neighbor aggregation, edge attributes, and node features containing positions and velocities.The edge attributes encode charge difference and distance, while node updates use two-layer SiLU MLPs.
Loading 2110.03336v4…