Approximate equivariance via projection-based regularisation

TL;DR

This work approaches approximate equivariance via a projection-based regulariser which leverages the orthogonal decomposition of linear layers into equivariant and non-equivariant components and penalises non-equivariance at an operator level across the full group orbit, rather than point-wise.

Abstract

Equivariance is a powerful inductive bias in neural networks, improving generalisation and physical consistency. Recently, however, non-equivariant models have regained attention, due to their better runtime performance and imperfect symmetries that might arise in real-world applications. This has motivated the development of approximately equivariant models that strike a middle ground between respecting symmetries and fitting the data distribution. Existing approaches in this field usually apply sample-based regularisers which depend on data augmentation at training time, incurring a high sample complexity, in particular for continuous groups such as $SO(3)$. This work instead approaches approximate equivariance via a projection-based regulariser which leverages the orthogonal decomposition of linear layers into equivariant and non-equivariant components. In contrast to existing methods, this penalises non-equivariance at an operator level across the full group orbit, rather than point-wise. We present a mathematical framework for computing the non-equivariance penalty exactly and efficiently in both the spatial and spectral domain. In our experiments, our method consistently outperforms prior approximate equivariance approaches in both model performance and efficiency, achieving substantial runtime gains over sample-based regularisers.

Figures (9)

• Figure 1: Pseudo-code for the equivariant projection for finite (left) and continuous groups (right).
• Figure 2: Commutative diagrams showing how to apply the projection operator in Fourier space.
• Figure 3: Controlling the degree of learned $SO(2)$ invariance by tuning the parameters $\lambda_G$ and $\lambda_\perp$, which penalise the projections of the equivariant and non-equivariant components, respectively.
• Figure 4: Effect of increasing angular perturbation at fixed projection strength. Each panel shows the decision boundary and level sets of the approximately $\mathrm{SO}(2)$-invariant network (blue) and an MLP (orange) trained with a fixed non-equivariant penalty $\lambda_\perp = 1.0$ on datasets with growing angular "wave" amplitude $\sigma_\perp$ (left to right). As $\sigma_\perp$ increases, the decision boundary becomes more angle-dependent and the learned classifier departs from perfect radial symmetry only where required to fit the data, while remaining nearly circular elsewhere. The empirical invariance defect $\mathcal{E}(T)$ for each setting is reported beneath the corresponding panel.
• Figure 5: Effect of angular perturbations and projection strength. Columns vary the angular wave amplitude $\sigma_\perp$, rows vary the non-equivariant penalty weight $\lambda_\perp$. Blue contours show level sets of the approximately $\mathrm{SO}(2)$-invariant network and points denote training samples. Orange dashed lines are the decision boundary of a non-equivariant MLP. The value $\mathcal{E}(T)$ underneath each panel is the empirical invariance defect, demonstrating that larger $\lambda_\perp$ keeps the classifier close to invariant even as the Bayes decision boundary becomes increasingly angle-dependent.

Theorems & Definitions (26)

• Lemma 2.1: elesedy2021provably, Lemma 1
• Corollary 2.2
• Definition 3.1: Equivariance defect
• Lemma 3.2
• proof
• Lemma 3.3
• proof
• Corollary 3.4: Kim2023
• proof
• Theorem 3.5: Informal