Symmetry From Scratch:
Group Equivariance as a Supervised Learning Task

We present foundational work on a proof-of-concept training schematic that allows unconstrained, group-agnostic models to learn equivariance directly from equivariant architectures. While this result could shed new insight on topics of model distillation (Hinton, Vinyals, and Dean 2015), model extraction (Tramèr et al. 2016), or even further our understanding of neural network training dynamics, in this work, we investigate the efficacy of having learned symmetries as an initial weight condition for group-agnostic models to learn downstream tasks. Additionally, as a proof-of-concept work, we will only focus on feature extraction for images and, therefore, 2D signals over the discrete grid $\mathbb{Z}^{2}$ using group-equivariant convolutions.

For 2D planar signals, current works that model symmetry-breaking with equivariant models are either fixed MLP layers that comply with relaxed equivariance constraints (Kaba and Ravanbakhsh 2024) or constrained to group convolutions (Cohen and Welling 2016a) with steerable filters (Cohen and Welling 2016b) on arbitrarily chosen equivariant bases (Wang, Walters, and Yu 2022). We show that the types of equivariance enforced by group convolution architectures can actually be learned directly via supervised learning by a more general class of group-agnostic architectures (fig.1), i.e., general MLPs¹¹1though our methodology easily extends to transformers and other widely used architectures., and once trained, can handle tasks involving both symmetric and symmetry-breaking data samples.

In the more theoretical line of work, tuning higher-symmetry models to lower-symmetry models generally involves some parameter-sharing process. From this perspective, Ravanbakhsh, Schneider, and Poczos (2017) explored designing model parameters to reflect equivariance over discrete group actions. Shakerinava et al. (2024) also introduced a weight-sharing regularization scheme, defining a loss that directly encourages weight-sharing to encourage symmetries in parameter space for machine learning in low-data regimes.

In direct relation to our experimental results, which focus on MLP architectures, is work by Finzi, Welling, and Wilson (2021) on equivariant MLPs (EMLP). They showed that EMLP layers can be constructed by decomposing arbitrary matrix groups (i.e., discrete groups and Lie groups) into their generators. The structure of a single MLP layer under any such group is shown to be the same as the solution of the equivariance constraint over the finite set of generators. While their EMLP architecture is constricted to stacks of EMLP layers with equivariant non-linearities, we show that more general MLP-based architectures, such as the MLP-mixer (Tolstikhin et al. 2021) (Section 2), may also learn group symmetries under a supervised learning context.

Provided that the ground-up approach to equivariant architectures is laborious and computationally expensive, many have worked more broadly on allowing group-agnostic general-purpose architectures to be part of a pipeline that, as a whole and in some limit, can be considered group-equivariant. Notably, frame-averaging introduced by Puny et al. (2022) and probabilistic symmetrization by Kim et al. (2024) both leverage group averaging to convert group-agnostic universal approximators into group-equivariant approximators, while Kaba et al. (2023) focused on mapping all samples to their canonical orientation with a learnable canonicalization function. However, for theoretical guarantees, these works often constrain their scope to work within single groups and do not consider symmetry-breaking.

In other works more specific to symmetry breaking, the angle of attack has been engineering modifications or adding additional weights to insert dependence on previously equivariant group transformations. For example, Kaba and Ravanbakhsh (2024) defined a relaxed equivariance constraint that modified the original construction of an EMLP to handle symmetry breaking. For group convolutions, Elsayed et al. (2020) tested the practicality of relaxing spatial invariance with a linear combination of a basis set of filter banks; and Wang, Walters, and Yu (2022) generalized the idea to arbitrary groups with the construction of relaxed group convolution, reintroducing symmetry-breaking dependence on specific pairs of feature map signals and group transformations.

We propose to use supervised learning for learning group equivariance on a model $h_{\theta}$ that is not inherently equivariant. Let $\rho:G\rightarrow GL_{n}(\mathbb{R})$ be a group representation of $G$ and $\Phi_{\tau}:\mathcal{X}\rightarrow\mathcal{Y}$ be a $G$-equivariant neural network parameterized by $\tau\in\mathcal{T}$ such that $\forall g\in G,x\in\mathcal{X}$:

We train a model $h_{\theta}:\mathcal{X}\rightarrow\mathcal{Y}$ to become approximately equivariant $h_{\theta}(\rho_{in}(g)x)\approx\rho_{out}(g)h_{\theta}(x)$ through supervised learning on dataset $\mathcal{D}=\{(x_{i},\tau_{i}),\Phi_{\tau_{i}}(x_{i})\}_{i}$, where $x_{i}\sim\mathcal{N}(0,I_{\mathcal{X}})$ and $\tau_{i}\sim\mathcal{N}(0,I_{\mathcal{T}})$. We call this process symmetry-cloning (alg.1).

To demonstrate the merits of symmetry-cloning, we limit $\rho_{out}$ to be the trivial representation (i.e., group-invariant) for classification tasks and introduce a benchmark consisting of a symmetric task and a symmetry-breaking task. By comparing the performance of group-agnostic models, $G$-cloned models, and group-equivariant models on both tasks, we demonstrate the effectiveness of our method in clonning the symmetries of the groups $T(2)$ and $C_{4}$, the groups of 2D translations and cyclic rotations by 90 degrees (clockwise), respectively:

We run the benchmarking tasks for both $T(2)$ and $C_{4}$ group transformations on the well-studied MNIST handwritten dataset (fig.2), respectively. In the symmetric task, we test for symmetries baked into the model by evaluating models on transformed samples not encountered by the model during training, while in the symmetry-breaking task, we evaluate the model’s ability to differentiate between certain symmetries.

More specifically, we specify the tasks as follows:

• •

  $T(2)$ Symmetric Task: we train the model on the original dataset (fig.2a) and evaluate its performance on a test set where every sample is randomly translated within 25% of the maximum image size (fig.2b).

• •

  $T(2)$ Symmetry-breaking Task: we train the model on the translated dataset (fig.2b) but with labels modified as follows:

  $$\hat{y}=\begin{cases}y&\text{ if shifted left}\\ y+10&\text{ otherwise}\end{cases}$$

• •

  $C_{4}$ Symmetric Task: we train the model on the original dataset (fig.2a) and evaluate its performance on a test set where a random rotation $r\in\{0,\pi/4,\pi/2,3\pi/4\}$ is randomly applied to each sample (fig.2c).

• •

  $C_{4}$ Symmetry-breaking Task: we train the model on the rotated dataset (fig.2c) with the labels modified as follows:

  $$\hat{y}=\begin{cases}y&\text{ if $r\in\{0,\pi/4\}$}\\ y+10&\text{ otherwise}\end{cases}$$

Building towards increasingly general architectures, we start with a simple case of symmetry-cloning: $T(2)$-cloning a single-channel convolutional layer with a 3x3 kernel and $C_{4}$-cloning a single-channel group convolutional layer with a 4x3x3 kernel.

• •

  9-block mlp2cnn: We observe that the convolution operation, when unrolled as a single left matrix multiplication with the image, displays a block Toeplitz pattern. Therefore, the most straightforward group-agnostic architecture would be one that learns the permutation matrix, which, along with kernel parameters, would combine to reconstruct the Toeplitz matrix (fig.3). Note that in this case, the MLP layer is constrained to have as many linear components as there are kernel parameters.

• •

  approx-mlp2cnn: We relax the constraint on requiring the architecture to have as many linear components as there are kernels (i.e., having 7, 8, or 10 layers); instead, we add an additional embedding layer followed by a projection to the number of linear components involved. In effect, this allows kernel parameters to act more as input to the group-agnostic model as depicted in (fig.1).

• •

  9-block mlp2gcnn, approx-mlp2gcnn: In very much the fashion as mlp2cnn, we use four-stacked mlp2cnns to clone a single $C_{4}$ group-equivariant group convolution layer with a lifting filter of four channels; the approx-mlp2gcnn layers now having four additional projection heads.

In both the mlp2cnn and mlp2gcnn cases, we may use the symmetry-cloned MLP layers analogous to how one would stack a full CNN, apply the appropriate pooling, and add a classification head to produce an MNIST classifier. Abusing notation, we also refer to the $T(2)$-cloned classifiers by their layer names.

Eventually, we wish to apply symmetry-cloning to less constricting architectures on groups with infinite elements. The current procedure for building entire classifiers via stacking symmetry-cloned layers becomes prohibitively expensive as the cloning architectures become more generalized. Nevertheless, without being applied to the benchmarking tasks, we still demonstrate that a more practical architecture, MLP-Mixer, can also be symmetry-cloned.

• •

  mlpmixer2cnn: We apply a first MLP-Mixer layer to encode the concatenated input and kernel parameters, followed by a second to decode the convolutional output. As the architecture is much more general, large enough MLP-Mixers should be able to learn entire CNNs directly, but we leave that for future work.

• •

  mlpmixer2scnn: Instead of regular group convolutions, we extend symmetry-cloning to a more general class of steerable CNN that is $SO(2)$-equivariant, albeit on a small input size and show that it can roughly capture the equivariant nature of a steerable CNN.

When training symmetry-cloned classifiers on the benchmarking tasks, to better leverage the learned symmetries of symmetry-cloned models, some regularization is needed to ensure that weights do not deviate too quickly from the learned initialization $\theta_{\text{init}}$. Therefore, following (Jaques et al. 2017), we use a KL constraint to prevent the weight distribution from drifting away from the equivariant initialization too far too quickly:

The work empirically presents most results, and as such, there are no theoretical bounds on the approximate clonability of any symmetry with any unconstrained model. Most tunable hyperparameters were chosen arbitrarily to demonstrate our approach’s general applicability. However, this also leaves much room for improvement on finding more rigorous approaches and optimal architectures.