Measuring Learned Equivariance: An Information Theoretic Approach

Claire Schlesinger; Wolfgang Gatterbauer; Robin Walters

cs.LGOct 12, 2026

Measuring Learned Equivariance: An Information Theoretic Approach

Authors: Claire Schlesinger, Wolfgang Gatterbauer, Robin Walters

Organizations: Khoury College of Computer Sciences Northeastern University Boston, MA 02115, USA

Abstract

Learning symmetries through data augmentation offers an alternative to architectural equivariance, but does not guarantee that a network respects those symmetries. Measuring learned equivariance is especially difficult for intermediate representations, whose group actions are not specified. In a finite discrete setting, we show that equivariance requires representations of original and transformed inputs to determine one another exactly, yielding a test based on normalized mutual information (NMI). For high-dimensional activations, we use centered kernel alignment (CKA) as a scalable test, with an exact guarantee for orthogonal feature transformations. Across architectures and tasks, augmentation usually increases transformation consistency, showing that learned equivariance can be probed directly from representations without recovering the hidden transformation rules that relate them.

Figures & tables

Figure 1: Comparison of equivariance error between

1-\text{CKA}

and the local equivariance error (LEE) using the assumptions of Gruver et al. (2023) . Each model is an exactly equivariant model made with ESCNN ( Weiler and Cesa, 2019 ) and equivariance error is computed at every half a degree in the range

[0,359.5]

. CKA is able to properly detect the zeros at the right intervals where the LEE is unable to detect equivariance for any model, even the SO(2).

Figure 2: Per-layer CKALE test for

C_{4}

: CKA

<1

at any rotation rules out exact equivariance at that layer.

	Filter
Rotation	Sobel-X	Laplacian
$90^{\circ}$	0.1954	1.0000
$180^{\circ}$	1.0000	1.0000
$270^{\circ}$	0.1954	1.0000

Table 1: NMI between filtered original and rotated images for the Sobel-X filter (not

C_{4}

-equivariant) and the Laplacian (

C_{4}

-equivariant) on a restricted, discretized image set. The Laplacian stays at

1

. Sobel-X is

C_{2}

-equivariant and anti-commutes, so it reads

1

180^{\circ}

; a metric assuming the output action matches the input action would report a large error here, but NMI detects the equivariance without being told the action is a sign flip.

		linear CKA		exact MI
readout	nature	value	null	bits	MI $/H$	null $/H$
identity	reference	1.0000	0.0004	7.3205	1.000	0.111
scramble $p{=}1.00$	bijective, lossless	0.0376	0.0009	7.3205	1.000	0.111
merge $q{=}0.10$	90% of alphabet destroyed	0.0431	0.0006	3.1574	0.606	0.005
Sobel-x control (not C4-equivariant)		0.0021	0.0004	1.2676	0.195	0.031

Table 2: CKA versus MI on a discrete set of data points. We apply the Laplacian to the input image, then either keep the output unchanged, scramble the alphabet (relabel outputs while keeping points with the same label together), or merge the alphabet (map multiple distinct outputs to the same value). CKA reads the lossless relabeling as non-equivariant, while MI correctly detects equivariance. A high null indicates the metric may be less informative due to a large alphabet.

readout	nature	MI	null	MI $/H$	linCKA	rbfCKA	SVCCA
identity	reference	7.3205	0.8113	1.000	1.0000	1.0000	1.0000
bit-permute	bijective, isometry	7.3205	0.8113	1.000	1.0000	1.0000	1.0000
ill-cond. linear	bijective, linear, not isometry	7.3205	0.8113	1.000	0.3584	0.2280	1.0000
alphabet-permute	bijective, not linear	7.3205	0.8113	1.000	0.1113	0.1543	0.2597
lossy-half	many-to-one, destroyed	1.8643	0.0065	0.505	0.4835	0.4528	0.5042

Table 3: Measuring equivariance after a Laplace operation is applied to the discretized images. The transformation is applied so we have a more interesting probability distribution under transformations. Under certain operations applied to the output we see how CKA, MI, and SVCCA either fails or succeeds to detect equivariance. SVCCA is able to detect one more group than CKA, but lacks an unbiased estimator, so we chose to utilize CKA.

		Laplace		Sobel-X
	metric	value	null	value	null	separation
(a) identity readout — reference
assumes $\rho$	LEE, $\rho=I$ $\downarrow$	1.4933	1.4145	1.3959	1.4150	0.93 $\times$
fits $\rho$	affine fit $\downarrow$	0.0000	1.0001	1.0836	1.0001	$\infty$
fits $\rho$	Procrustes $\downarrow$	0.0000	1.4874	1.4884	1.4641	$\infty$
no $\rho$	linear CKA $\uparrow$	1.0000	0.0020	0.0023	0.0007	435 $\times$

Table 4: Comparison between different methods on detecting equivariance given an identity and an energy readout. All methods except LEE succeed on the identity readout. Under the energy readout, only mutual information stays exact, while CKA still clearly separates the equivariant kernel from the control.

no convolution — raw pixels
granularity	symbols	samp./sym	MI	$H$	MI $/H$	null	margin
full 16-bit	65,536	1.00	16.0000	16.0000	1.000	16.0000	0.0000
Laplacian — exactly C4-equivariant
full 16-bit	45,707	1.43	15.1834	15.1834	1.000	14.3668	0.8165
quadrant	486	134.8	7.3205	7.3205	1.000	0.8157	6.5048
Sobel-x — not C4-equivariant (control)

Table 5: Changing the granularity of what mutual information sees from 4 bits to the full 16 bits shows that the null reading rises with it due to the larger alphabet, making mutual information a much worse measure the more unique outputs there are. In a continuous space with infinite outputs, mutual information is not only intractable to compute, but also loses its ability to distinguish real equivariance from scrambles.

model	sym. var.	null dim	align	$\lVert\nabla\rVert$	LEE	$1-$ CKA
	LieGG
Equivariant SO(2)	0.00000	1.0	1.0000	5.9228	0.9999	0.000000
MLP, augmented	0.05760	0.0	0.0000	5.9187	0.3702	0.008651
MLP, no augmentation	0.06845	0.0	0.0000	6.1034	0.6413	0.042468
MLP, untrained	0.16950	0.0	0.0000	0.0083	0.2518	0.460116

Table 6: Three measures on exactly-SO(2) point clouds (three seeds, 300 epochs, 2048 samples for LieGG; lower is more equivariant). LEE ranks the provably equivariant model last and an untrained network first; LieGG and CKALE recover the correct order. The

\lVert\nabla\rVert

column is the mechanism: LEE

=\lVert L_{X}f\rVert/\lVert f\rVert

is invariant to uniformly rescaling the features but not to their being nearly constant, where the numerator vanishes and the denominator does not. The untrained network’s input gradient is over

700\times

smaller than the others’. LieGG normalizes by its polarization matrix’s largest singular value, where rescaling cancels, and is unaffected.

readout	$m$	sym. var. $\downarrow$	null dim	align $\uparrow$	$A_{\rm true}/$ rand $\downarrow$	$1-$ CKA $\downarrow$
		LieGG
invariant scalar, head(\|\|h\|\|)	1	0.00003	1.0	1.0000	0.00008	0.000000
invariant vector, \|\|h\|\|	64	0.00022	1.0	1.0000	0.00007	0.000033
equivariant , h	128	0.43667	0.0	0.0000	0.87357	0.000000

Table 7: LieGG makes an assumption that the input space has to be equivariant and the output space has to be invariant to rotation. Purposely building a model that has an equivariant output space shows that LieGG is unable to detect it.

transformation of the features	$\mathcal{S}_{\rm equiv}$	$\mathcal{S}_{\rm inv}$	SEIS reads it as
identity (reference)	1.0000	1.0000	—
spatial re-encoding rotate the grid	1.0000	0.1311	re-encoded ✓
channel re-encoding the true $\rho(g)$	0.2238	0.0941	destroyed $\times$
disjoint images (reference)	0.1361	0.0767	—

Table 8: Comparison on SEIS equivariance measure between spatial and channel re-encoding. SEIS is able to detect spatial re-encodings but not channel re-encodings and is therefore unable to detect equivariance of steerable models.

Figure 3: Final-layer CKA across models and datasets for

C_{4}

and

SO(2)

tasks (“naive” denotes the perceptron). Augmentation increases learned equivariance on most tasks but not colorization, likely because colorization relies more on texture than position. Error bars show the 25th and 75th percentiles.

Appendix figures & tables31 assets

Supplementary material from the paper’s appendix.

Appendix

(a) at the C16 model’s own group elements
rotation	on lattice?	input floor	model	model $/$ floor
$0^{\circ},90^{\circ},180^{\circ},270^{\circ}$	yes	0.00000	0.00000	—
$22.5^{\circ}$ and its 3 repeats	no	0.00011	0.00032	2.83
$67.5^{\circ}$ and its 3 repeats	no	0.00011	0.00026	2.30
$45^{\circ}$ and its 3 repeats	no	0.00013	0.00054	4.10
off-lattice mean		0.00012	0.00037	3.08

Appendix

Table 9: The resampling floor. (a) A C16 steerable conv at its own sixteen group elements: the input floor is

1-\mathrm{CKA}

between raw pixels and their own rotation, the model column the same quantity on the conv’s features. (b) Both, by grid resolution. All entries are

1-\mathrm{CKA}

; 128 MNIST images,

29^{2}

grid in panel (a).

condition	equivariance error $\downarrow$	linear CKA $\uparrow$	CKA null
grid-aligned reference ( $\theta=\pi/2$ )	0.00000	1.0000	0.0729
full resolution	0.13429	0.9994	0.0722
subsample	0.21735	0.9995	0.0727
blur, then subsample	0.15250	0.9989	0.0432

Appendix

Table 10: Downsampling a feature map at export aliases it. Equivariance error is

\lVert f(gx)-\rho(g)f(x)\rVert/\lVert\rho(g)f(x)\rVert

with

\rho(g)

read off escnn’s declared FieldType , so it is ground truth rather than a competing estimate. The null is the shuffled control: the pairing between the clean and rotated branches is destroyed while marginals, dimensionality and sample count are held fixed.

\theta=2.0

rad (

114.6^{\circ}

), 256 MNIST images.

Figure 4: Full view of a CNN progression on learning equivariance on CIFAR-10 classification.

Figure 5: Full view of a ViT progression on learning equivariance on CIFAR-10 classification.

Figure 6: Full view of a perceptron progression on learning equivariance on CIFAR-10 classification.

Figure 7: Full view of a CNN progression on learning equivariance on MNIST-Font classification.

Figure 8: Full view of a perceptron progression on learning equivariance on MNIST-Font classification.

Figure 9: Full view of a CNN learning equivariance on CIFAR-10 colorization.

Figure 10: Full view of a UNet learning equivariance on CIFAR-10 colorization.

Figure 11: Full view of a ViT learning equivariance on CIFAR-10 colorization.

Figure 12: Full view of a perceptron learning equivariance on CIFAR-10 colorization.

Figure 13: Full view of a CNN learning equivariance on STL-10 colorization.

Figure 14: Full view of a UNet learning equivariance on STL-10 colorization.

Figure 15: Full view of a ViT learning equivariance on STL-10 colorization.

Figure 16: Full view of a perceptron learning equivariance on STL-10 colorization.

Figure 17: Full view of a CNN learning fluid dynamics on a grid-like buoyancy dataset.

Figure 18: Full view of a ViT learning fluid dynamics on a grid-like buoyancy dataset.

Figure 19: Full view of a UNet learning fluid dynamics on a grid-like buoyancy dataset.

Figure 20: Full view of a perceptron learning fluid dynamics on a grid-like buoyancy dataset.

Figure 21: Full view of a PointNet learning fluid dynamics on a particle-like buoyancy dataset.

Figure 22: Full view of a SetTransformer learning fluid dynamics on a particle-like buoyancy dataset.

Figure 23: Full view of an MLP learning fluid dynamics on a particle-like buoyancy dataset.

Figure 24: Full view of a perceptron learning fluid dynamics on a particle-like buoyancy dataset.

Figure 25: Full view of a CNN learning stress tensor prediction on a grid-like anisotropic dataset.

Figure 26: Full view of a ViT learning stress tensor prediction on a grid-like anisotropic dataset.

Figure 27: Full view of a UNet learning stress tensor prediction on a grid-like anisotropic dataset.

Figure 28: Full view of a perceptron learning stress tensor prediction on a grid-like anisotropic dataset.

Figure 29: Full view of a PointNet learning stress tensor prediction on a particle-like anisotropic dataset.

Figure 30: Full view of a SetTransformer learning stress tensor prediction on a particle-like anisotropic dataset.

Figure 31: Full view of an MLP learning stress tensor prediction on a particle-like anisotropic dataset.

Figure 32: Full view of a perceptron learning stress tensor prediction on a particle-like anisotropic dataset.

Explore similar work

Oct 7, 2026

How Do Transformers Learn to Represent Symmetries?

Training Transformer-based architectures with finite data augmentation has become an increasingly popular approach in geometric machine learning. Despite its empirical success, the interplay between the Transformer architecture, invariance to different symmetries, and augmentation budgets remains underexplored. In this paper, we study the ability of a vanilla Transformer to learn various symmetries through finite data augmentation for point cloud datasets. We identify an ordering of increasing learnability across the following symmetry groups: (i) non-angle-preserving symmetries, (ii) angle-preserving symmetries, and (iii) base angle-preserving subgroups, such as translation, rotation, and scale. For the base angle-preserving groups, we further investigate the Transformer's extrapolation behavior and conduct a structural analysis of the trained models, allowing us to identify interpretable mechanisms that induce invariance. Finally, we extend our analysis to equivariant functions and show that the detected mechanisms for approximate invariance can also provide a key building block for learned equivariance. Our project page is available at https://transformers-learn-symmetries.github.io/

Eduardo Santos-Escriche, Valerie Engelmayer, Ya-Wei Eileen Lin +1

Technical University of Munich Munich Center for Machine Learning · Massachusetts Institute of Technology

Invariant Representation Learning Transformer Equivariant Representation Learning Data Augmentation

Ask PDF

Jun 24, 2026

Equivariance and Augmentation for Bayesian Neural Networks

Symmetries are important for many deep learning tasks, ranging from applications in the sciences to medical imaging. However, there is an ongoing debate about whether to impose symmetry constraints on the neural network architecture (yielding equivariant neural networks) or learn them from augmented training data. Although equivariant networks are well-studied theoretically, much less is known about data augmentation, since analyzing augmentation requires control over the training dynamics. Inspired by recent results that show that augmented infinite deep ensembles are exactly equivariant, we study data augmentation for Bayesian neural networks (BNNs) trained with variational inference. We focus on variational distributions in the exponential family and derive conditions under which exact equivariance is reached. We furthermore obtain bounds on the equivariance error and introduce three novel symmetrization techniques which boost the effect of data augmentation in this setting. We conduct extensive numerical experiments which show that one of our symmetrization methods (orbit expansion) outperforms the baseline in both equivariance and overall performance. Our code is available at github.com/dmw1998/augment-BNNs

Miaowen Dong, Axel Flinth, Jan E. Gerken

Department of Mathematical Sciences, Chalmers University of Technology and the University of Gothenburg, SE-412 96 Gothenburg, Sweden. · Department of Mathematics and Mathematical Statistics, Umeå University, Linnéus väg 49, 901 87 Umeå, Sweden.

Bayesian Neural Networks Data Augmentation Equivariant Neural Networks

Ask PDF

Jun 23, 2026

Data Augmentation: A Fourier Analysis Perspective

Data augmentation is a simple and model-agnostic approach for exploiting known invariances in learning problems. Given a group acting on the input space, one augments the training set with transformed copies of each sample. Because it exploits symmetries without modifying the underlying learning algorithm, data augmentation can be applied broadly across learning methods. However, this universality comes at a computational cost: when the group is large, full group-sized augmentation quickly becomes computationally infeasible. This raises a fundamental question: Can partial data augmentation achieve the same statistical benefits as full augmentation in terms of generalization and sample complexity? We develop a general framework for investigating this question using Fourier analysis and the representation theory of finite groups. We show that, for a broad class of classical learning problems, partial data augmentation based on a randomly sampled subset of group elements achieves the same minimax rates as full augmentation, up to an approximation error that vanishes as the subset size increases. Our results provide a theoretical explanation for why partial augmentation can retain the statistical benefits of full augmentation despite enforcing symmetry only approximately, and shed light on a recently raised question in learning with symmetries: whether statistically optimal learning under general group invariances can be achieved using computationally scalable methods. Moreover, we prove a complementary impossibility result: enforcing exact invariance via data augmentation requires averaging over the entire group, and cannot be achieved by any strict subset when the hypothesis space is sufficiently expressive. Together, these results provide a unified perspective on full and partial data augmentation, as well as exact and approximate symmetry enforcement.

Behrooz Tahmasebi, Melanie Weber, Stefanie Jegelka

Harvard John A. Paulson School of Engineering and Applied Sciences, Harvard University, Cambridge, MA 02138, USA. · Technical University of Munich (CIT, MCML, MDSI) · MIT Computer Science and Artificial Intelligence Laboratory (CSAIL)

Data Augmentation Statistical Learning Theory

Ask PDF