A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups

Marc Finzi; Max Welling; Andrew Gordon Wilson

Source-linked AI summary

A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups

Marc Finzi, Max Welling, Andrew Gordon Wilson

arXiv:2104.09459v1Apr 19, 2021cs.LGmath.DSstat.ML

TL;DR

Existing equivariant-network methods lacked a completely general construction for arbitrary matrix groups and representations. The paper introduces a polynomial-time constraint-solving framework and builds EMLP from it, recovering established layers and improving generalization across symmetry-focused applications. The approach also supports groups previously considered infeasible, including O(5), O(1, 3), Sp(n), and the Rubik’s cube group.

Problem
Existing methods did not provide an algorithm completely general to the choice of symmetry group or representation, despite growing applications requiring broader symmetry handling.
Method
The paper reduces equivariance conditions to generator constraints, solves them efficiently, and combines the resulting linear layers with a bilinear layer to form EMLP.
Results
EMLP recovers established equivariant layers, applies to previously infeasible groups, and consistently improves generalization on Lorentz-invariant particle scattering and dynamical systems.
Takeaways & Limitations
The released library enables researchers to construct equivariant layers for arbitrary matrix groups by specifying group generators.
Takeaways & Limitations
Dense matrix multiplies make EMLP slow to train at the scale of convolutional networks or large graph networks with specialized implementations.

Abstract

from arXiv · show

Symmetries and equivariance are fundamental to the generalization of neural networks on domains such as images, graphs, and point clouds. Existing work has primarily focused on a small number of groups, such as the translation, rotation, and permutation groups. In this work we provide a completely general algorithm for solving for the equivariant layers of matrix groups. In addition to recovering solutions from other works as special cases, we construct multilayer perceptrons equivariant to multiple groups that have never been tackled before, including $\mathrm{O}(1,3)$, $\mathrm{O}(5)$, $\mathrm{Sp}(n)$, and the Rubik's cube group. Our approach outperforms non-equivariant baselines, with applications to particle physics and dynamical systems. We release our software library to enable researchers to construct equivariant layers for arbitrary matrix groups.

1. Introduction

The paper addresses the need for equivariant networks beyond commonly supported symmetries by introducing a general, efficient construction for equivariant multilayer perceptrons. It recovers established layers as special cases and applies the same approach to previously infeasible groups.

Core contribution: The paper proposes EMLP, which characterizes all linear maps between finite-dimensional group representations and computes them with a polynomial-time algorithm.The authors also release documentation and examples for the library.
Unification: Convolutional, permutation-equivariant, graph, and point-cloud layers arise as special cases of the general algorithm.
Core contribution: Equivariance conditions for matrix groups with arbitrary linear representations reduce to M + D constraints from discrete and continuous generators.
Core contribution: The constraint-solving approach can be accelerated by exploiting structure and recasting the computation as an optimization problem.
Architecture: Adding a bilinear layer yields an EMLP architecture that extends to new groups by specifying their generators.
Applications: The same underlying architecture is applied to O(5), O(1, 3), Sp(n), and the Rubik’s cube group, outperforming non-equivariant baselines.

2. Related Work

Prior work developed equivariant networks for selected discrete and continuous groups, but no construction was completely general to both the symmetry group and its representation. Existing alternatives also faced scope or computational limitations.

Existing approaches: Earlier equivariant networks addressed translations, discrete groups, rotations, permutations, sets, graphs, and related structures through specialized constructions.
Existing approaches: Group FFTs and regular group convolutions offered flexibility, but continuous-group regular representations are fundamentally challenging because they are infinite dimensional.
Limitations: Lang and Weiler provided a process for arbitrary compact groups, but it required considerable mathematical work and did not extend beyond compact groups.
Limitations: Orbit-based weight sharing was limited to regular representations of finite groups, while another explicit algorithm scaled with group size and became infeasible for large or continuous groups.

3. Background

The background represents symmetries through group and Lie algebra generators acting on finite-dimensional vector spaces. More complex representation spaces are built using duals, direct sums, and tensor products.

Group generators: Finite groups can be generated by a discrete set of elements, often far fewer than the total number of group elements.The Rubik’s cube group, for example, has six generators despite having size 4 × 10^19.
Group generators: Continuous groups are analyzed through Lie algebras, whose infinitesimal generators form the tangent space at the identity and map back to the group through the exponential map.
Lie algebras: For SO(n), the Lie algebra consists of antisymmetric matrices obtained by differentiating the defining orthogonality constraint at the identity.
Lie groups: The exponential map is surjective for some groups, including SO(n), SE(n), SU(n), and Sp(n), but not for O(n), E(n), and O(1, 3).
Representations: A finite-dimensional group representation assigns each group element an invertible matrix describing how objects transform.
Representations: Lie algebra representations describe infinitesimal transformations, while tensor representations combine base representations using dual, direct-sum, and tensor-product operations.

4. Equivariant Linear Maps

The paper converts equivariance of linear maps into finite generator constraints and solves them by finding a constraint nullspace. This framework reproduces standard equivariant layers for specific groups and representations.

Equivariance constraints: Equivariant linear maps must transform inputs and outputs consistently under the group action, producing constraints on the map parameters.
Equivariance constraints: Although the group may contain combinatorially many or infinitely many elements, the constraints can be reduced to a finite, small collection.
Equivariance constraints: The reduced system contains one constraint for each discrete generator and infinitesimal generator, totaling O(M + D), and these constraints are necessary and sufficient.
Solving the constraints: Stacking the constraints into a matrix C allows singular-value decomposition to produce a nullspace basis Q, with all symmetric solutions written as v = Qβ or projected as v = QQ^T v0.
Solving the constraints: For finite-dimensional representations, singular-value decomposition solves the constraints with runtime O((M + D)m^3).
Special cases: The framework recovers convolution, deep sets, equivariant graph networks, and GCNNs by specifying the relevant groups and representations.
Special cases: For GCNNs, solving constraints for the semidirect product of grid translations and 90-degree rotations yields the group-convolutional layer.
General applicability: The method computes bases for established and previously unexplored groups, including SO(n), O(n), Sp(n), Lorentz groups, SU(n), and the Rubik’s cube group.

5. Efficiently Solving the Constraint

The paper makes equivariant-basis computation scalable by decomposing constraints and using structure-aware Krylov nullspace optimization, enabling high-dimensional representations and efficient matrix-vector products.

Scalability: The practical bottleneck is computing equivariant bases, whose direct SVD becomes too costly for representations with dimension m ≥ 5000.The method addresses initialization-time basis computation rather than forward-pass cost alone.
Independent sub-problems: Block structure separates constraints by representation pairs, allowing independent solutions that are later reassembled into the full matrix.This exploits direct-sum structure in the group and algebra representations.
Krylov method: The nullspace basis is computed by minimizing a constraint objective with gradient updates related to QR power iteration and Oja’s rule.The updates seek smallest singular vectors while preserving nullspace components.
Krylov method: Gradients require only matrix-vector multiplies with the constraint matrix, avoiding explicit formation of representation matrices.Kronecker structure further enables efficient matrix-vector multiplication.

6. Network Architecture

The EMLP combines equivariant linear maps, gated nonlinearities, and a cheap bilinear layer to process mixed tensor features while preserving group equivariance and improving data efficiency.

Feature representations: EMLP layers use tensor representations with different ranks and multiplicities, allocating intermediate channels across tensor ranks.For example, an SO(3)-equivariant layer can combine scalar, vector, matrix, and higher-rank tensor channels.
Data efficiency: Across dataset sizes and tasks, EMLP models with the relevant symmetry consistently outperform standard MLPs and data-augmented MLPs, often by orders of magnitude.Figure 5 reports 95% confidence intervals over three runs.
Gated nonlinearities: Gated nonlinearities multiply each feature object by a scalar gate, using Swish for scalar objects and regular representations.The gate is produced separately for the different objects concatenated in the feature space.
Universality: Gated nonlinearities alone are not universal for general groups and representations because they cannot express contractions between different feature objects.The paper identifies operations such as dot products as unavailable under this limitation.
Bilinear layer: The cheap bilinear layer addresses this limitation by contracting pairs of input tensors into specified output types, with learnable weights for non-scalar contractions.It applies one tensor as a reshaped linear map to the other when the representation types satisfy the stated condition.
Architecture: EMLPs stack equivariant linear layers, bilinear transformations, shortcut connections, and gated nonlinearities into full equivariant networks.The architecture maps collections of geometric quantities between input, internal, and output representations.

7. Experiments

The experiments evaluate EMLP on synthetic invariant and equivariant tasks, Lorentz-equivariant particle scattering, and symmetric dynamical systems. Across these settings, EMLP improves performance over standard MLP baselines and can yield conservation properties in Hamiltonian dynamics.

Synthetic and particle-physics experiments: EMLP with the given equivariance consistently outperforms baseline MLPs trained with and without data augmentation, often by orders of magnitude.The comparison spans different dataset sizes and tasks.
Dynamical systems: The dynamical-systems experiment models a 3D double spring pendulum with O(2) rotational and reflectional symmetry and Hamiltonian structure.The system is not a homogeneous space because group elements alone cannot traverse its state space.
Dynamical systems: Exploiting O(2), SO(2), or D6 symmetry with EMLP improves performance for both Neural ODE and HNN models.The comparison is reported in Table 1 and Figure 6 against ordinary MLP HNNs and Neural ODEs.
Dynamical systems: Continuous rotation symmetry in EMLP-HNN models yields conservation of angular momentum about the z-axis, whereas Neural ODE models do not conserve it.D6 also gives approximate angular-momentum conservation, but D2 does not; Noether’s theorem applies here to Hamiltonians rather than general ODEs.

8. Discussion

The paper presents a general construction for equivariant layers and connects it to established architectures and theoretical guarantees. It also identifies limitations involving nonlinear universality and computational scalability.

Conclusions: The construction is general to the choice of representation and matrix group, with convolutions, deep sets, graph networks, and GCNNs arising as special cases.The same framework supports equivariant bases for very large representations through iterative matrix-vector-multiply methods.
Conclusions: EMLP applies these general equivariant layers to Lorentz-invariant particle scattering and dynamical systems, showing consistently improved generalization.The discussion frames these applications as practical demonstrations of the construction’s capabilities.
Limitations: Dense matrix multiplications make EMLP and standard MLPs slow to train at convnet or large-graph-network scales.Specialized implementations may be needed to overcome this generality-specialization tradeoff.
Limitations: Gated nonlinearities and Norm-ReLUs alone are insufficient for universality in the general case.The limitation arises because these nonlinearities cannot express contractions between different feature objects, such as dot products.

C. Krylov Method for Efficiently Finding the Equivariant Subspace

The Krylov method computes equivariant subspaces without explicitly forming large representation matrices, using gradient descent followed by a small SVD. The method has exponential convergence and supports large representations through efficient matrix-vector products.

Algorithm: Matrix-vector multiplies with the constraint matrix avoid explicitly forming representation matrices.This enables computation for massively large constraint matrices, including one sized (4 × 10^5) × (7 × 10^4).
Algorithm: The method finds the nullspace of the constraint matrix using simple gradient descent followed by a small SVD.It is closely related to power iteration, Oja’s rule, and PCA-style optimization methods.
Convergence: The nullspace iteration converges exponentially, with O(log(1/ϵ)) iterations for an ϵ-accurate solution with probability 1.With an optimal learning rate, the convergence time is T = O(κ log(1/ϵ)), where κ = (σ_max/σ_min)^2.
Convergence: For rmax = 20, X reaches floating-point precision within 300 iterations across a range of groups and tensor representations.The empirical convergence is shown in Figure 7.
Complexity: The runtime depends on the number of discrete and continuous generators, representation-specific matrix-vector multiplication costs, and the target accuracy.The analysis gives specialized bounds for subgroups of symmetric groups, orthogonal groups, Lorentz groups, complex groups, symplectic groups, SO(3), and SO(2).

D. Linear Layers and Gated Nonlinearities are Not Universal

For groups with the parity property, rank-preserving linear layers and nonlinearities can disconnect even- and odd-order features, preventing universal approximation of some simple equivariant functions.

Parity constraints: For groups with the parity property, every equivariant tensor of odd rank is zero, and maps from even-rank tensors to odd-rank tensors vanish.This follows because the representation can contain paired elements ρ(g′) = −ρ(g), imposing opposing equivariance constraints for odd ranks.
Architectural limitation: Rank-preserving nonlinearities make even- and odd-order quantities computationally disconnected when even-to-odd linear maps are zero.Gated nonlinearities permit odd features to depend on both ranks, but still provide no path from odd inputs to later even outputs.
Concrete example: For O(3), a network receiving a vector cannot learn the scalar norm ∥v∥ using Norm-ReLU or gated nonlinearities; its only valid output is an input-independent constant.The vector is rank 1 while the norm is rank 0, so no nonzero computational path connects them.
Scope: This limitation can prevent inner products, matrix-vector multiplications, and other valid equivariant functions regardless of network size.The paper identifies the same parity property for standard representations of O(n) and SO(2n).

E.1. Discrete Translation Group Zn

The algorithm recovers known equivariant bases for discrete and continuous groups while supporting large or previously unexplored groups through generator-based computation and tensor representations.

Discrete translation group Zn: For Zn, convolutional maps V → V have r = n basis matrices, while average pooling V → T0 has a one-dimensional basis.The recovered convolution basis is represented by circulant matrices.
Permutation group Sn: For Sn, the computed bases approach the kth Bell number Bk as n →∞, while small n can yield fewer independent solutions because some equivalence classes are empty.The paper reports these small-n dimensions as precise values and notes that prior Bell-number counting can overcount dependent solutions.
Rubik’s Cube group: The Rubik’s Cube group has more than 4 × 10^19 elements but is specified by only six face-turn generators, allowing equivariant bases to be computed without enumerating the group.Its 48-dimensional regular representation assigns one component to each non-center facet.
Continuous rotation groups: For SO(n), the Lie algebra has dimension D = n(n − 1)/2, while O(n) adds a discrete generator and excludes orientation-sensitive Levi-Civita basis elements.The additional diagonal basis element for SO(n) does not respect orientation-reversing isometries.
Symplectic group: For Sp(n), basis elements differ from those of O(n), and distributing rank between a base space and its dual produces different equivariant solutions.This remains true even though the basis dimensions become similar for large n.
Adding groups and representations: New groups require discrete generators, Lie-algebra generators when present, and their base representations; derived representations can be obtained automatically with Jacobian-vector products.The implementation also supports tensor constructions through duals, direct sums, and tensor products.

I. Implementation Details

The experiments use matched multilayer architectures and standardized optimization settings, with a smaller channel width for the double spring pendulum.

Synthetic experiments: Synthetic experiments use three EMLP layers with c = 384 channels, followed by an equivariant output layer; baseline MLPs use three hidden layers of size 384.All models use batch size 500, learning rate 3 × 10^-3, and Adam optimization.
Dynamical system: The double spring pendulum uses the same hyperparameters but reduces the channel width to c = 128 for all models.Trajectories are integrated with Jax’s adaptive Runge–Kutta integrator at tolerance 2 × 10^-6.

J. Datasets

The datasets cover invariant and equivariant functions for orthogonal and Lorentz groups, plus trajectories from a Hamiltonian double spring pendulum system.

Synthetic invariant and equivariant tasks: The O(5) invariant dataset samples x1 and x2 from N(0, 1)^5 and uses their normalized inner-product target.The O(3) inertia dataset computes inertia tensors from Gaussian positions and positive softplus-transformed masses.
Lorentz task: The Lorentz dataset defines particle-interaction targets from four sampled momenta pμ, p̃μ, qμ, and q̃μ drawn from N(0, 1/4^2).The target is constructed from Lorentz-covariant momentum contractions and the metric ημν.
Data splits: Each synthetic dataset reserves 5,000 examples for testing and 1,000 for validation with early stopping.These held-out splits are applied across the listed synthetic tasks.
Double spring pendulum: The double spring pendulum trajectories are generated from Hamiltonian dynamics with kinetic and potential energy terms for two masses, springs, and gravity.The system is integrated for T = 30s, and randomly selected 1s trajectory chunks provide training examples.