G-RepsNet: A Fast and General Construction of Equivariant Networks for Arbitrary Matrix Groups

TL;DR

G-RepsNet introduces a lightweight, universal framework for constructing equivariant networks with respect to arbitrary matrix groups by leveraging tensor polynomial representations and inexpensive tensor operations. By distinguishing regular (finite groups) and non-regular (continuous groups) representations and employing a three-part layer design—input representations, tensor conversion, and neural processing—the approach achieves equivariance with simple linear operations and invariant mixing. The authors prove universality for both regular and non-regular settings and demonstrate strong empirical performance across N-body dynamics, 2D/3D image tasks, and PDE problems, often outperforming or matching state-of-the-art equivariant architectures while offering superior scalability. This framework generalizes and subsumes several existing designs (e.g., vector neurons, harmonic networks, equitune) and shows practical impact for broad applications where symmetry-aware learning is beneficial.

Abstract

Group equivariance is a strong inductive bias useful in a wide range of deep learning tasks. However, constructing efficient equivariant networks for general groups and domains is difficult. Recent work by Finzi et al. (2021) directly solves the equivariance constraint for arbitrary matrix groups to obtain equivariant MLPs (EMLPs). But this method does not scale well and scaling is crucial in deep learning. Here, we introduce Group Representation Networks (G-RepsNets), a lightweight equivariant network for arbitrary matrix groups with features represented using tensor polynomials. The key intuition for our design is that using tensor representations in the hidden layers of a neural network along with simple inexpensive tensor operations can lead to expressive universal equivariant networks. We find G-RepsNet to be competitive to EMLP on several tasks with group symmetries such as O(5), O(1, 3), and O(3) with scalars, vectors, and second-order tensors as data types. On image classification tasks, we find that G-RepsNet using second-order representations is competitive and often even outperforms sophisticated state-of-the-art equivariant models such as GCNNs (Cohen & Welling, 2016a) and E(2)-CNNs (Weiler & Cesa, 2019). To further illustrate the generality of our approach, we show that G-RepsNet is competitive to G-FNO (Helwig et al., 2023) and EGNN (Satorras et al., 2021) on N-body predictions and solving PDEs, respectively, while being efficient.

Figures (6)

• Figure 1: A summary of G-RepsNet layer construction exemplified with inputs of types $T_0, T_1,$ and $T_2$, and outputs of the same types. Each layer consists of three subcomponents: i) input feature representation shown as $T_i$, ii) converting tensor types appropriately shown using arrows from $T_i$ to $T_j$, and iii) neural processing the converted tensors using appropriate neural networks, as discussed in §. \ref{['sec: G-RepsNet_Architecture']}.
• Figure 2: Comparison of GRepsNets with EMLPs finzi2021practical and MLPs for (a) O(5)-invariant synthetic regression task with input type $2T_1$ and output type $T_0$, (b) O(3)-equivariant regression with input as masses and positions of 5 point masses using representation of type $5T_0 + 5T_1$ and output as the inertia matrix of type $T_2$, (c) SO(1, 3)-invariant regression computing the matrix element in electron-muon particle scattering with input of type $4T_1$ and output of type $T_0$. Across all the tasks, we find that GRepsNets, despite its simple design, are competitive with the more sophisticated EMLPs and significantly outperform MLPs.
• Figure 3: (a) and (b) show layers from vector neurons deng2021vector and equitune basu2023equi, which are special cases of GRepsNet..
• Figure 4: (a) shows a simple way to add residual connections in GRepsNet. (b) shows the architecture used for $T_2$ CNNs and equituning, where the first $k$ layers are made of $T_1$-layers to extract features, then the extracted features are converted in $T_2$ tensors, which are then processed by $T_2$-layers. Finally $T_0$ tensors, i.e., scalars are obtained as the final output.
• Figure 5: Times per epoch (in seconds) for different MLPs, GRepsNets, and EMLPs for varying dataset sizes. Note that MLPs and GRepsNets have comparable time per epoch, whereas EMLPs take huge amount of time. Hence, EMLPs, despite its excellent performance on equivariant tasks, is not scalable to larger datasets of practical importance.

Theorems & Definitions (6)

• Lemma 1: weyl1946classical
• Theorem 1
• Theorem 2
• proof : Proof to Thm. \ref{['thm:GReps_scalar_universal_proof']}
• Lemma 2: villar2021scalars
• proof : Proof to Thm. \ref{['thm:GReps_vector_universal_proof']}