How Do Transformers Learn to Represent Symmetries?

Eduardo Santos-Escriche; Valerie Engelmayer; Ya-Wei Eileen Lin; Stefanie Jegelka

cs.LGOct 7, 2026

How Do Transformers Learn to Represent Symmetries?

Authors: Eduardo Santos-Escriche, Valerie Engelmayer, Ya-Wei Eileen Lin, Stefanie Jegelka

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

Abstract

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/

Figures & tables

Figure 1 : Overview of our study and main findings. Under finite data augmentation, vanilla Transformers learn invariance more efficiently for angle-preserving symmetries, especially their base groups, than for non-angle-preserving transformations. For these base symmetries, we identify interpretable mechanisms underlying the learned approximate invariance in the trained models.

Panel of Figure 2 — Figure 2 : ( Left ) Samples from different classes of the 3D Topology dataset (top), and legend for the plots on the right, summarizing our findings (bottom). ( Right ) Mean $\pm$ std for test accuracies (top) and invariance errors (bottom) across 5 random seeds. Angle-preserving symmetries are efficiently learned by the Transformer, while it struggles with non-angle-preserving ones.

Figure 3 : Accuracy ( top ) and invariance error ( bottom ) for different architectures on base angle-preserving symmetries under OOD evaluation. DeepSets exhibits the weakest performance, PointNet the strongest with Transformers as close second. Across architectures, scaling invariance appears easiest to learn, whereas translation invariance appears most difficult.

Panel of Figure 4 — Figure 4 : Attention maps of a full Transformer trained on ModelNet10 augmented with 16 rotational augmentations per sample remain highly stable across augmentations. ( Left ) Invariance error of $\operatorname{A}$ across the entire test set. ( Right ) Examples of the first $8\times 8$ attention scores. Rows correspond to different samples, and columns correspond to rotations by the indicated angles.

Figure 5 : A model with 2 blocks of

{\bm{X}}+\operatorname{SA}(\operatorname{LN}({\bm{X}}))

( Residual ) can learn to be translation invariant (dashed) and solve the 3D Topology task (solid lines), but

\operatorname{SA}(\operatorname{LN}({\bm{X}}))

( No Res. ) cannot.

Panel of Figure 6 — Figure 6 : Results for the first block of a full Transformer trained on ModelNet10 with 16 translation augmentations per sample. ( Left ) Cosine similarity between ${\bm{X}}$ and $\operatorname{SA}({\bm{X}})$ is strictly negative, indicating that the attention output has signs opposite to those of the residuals. ( Middle ) The attention maps have low invariance error across the full test set. ( Right ) Each row shows the first $8\times 8$ attention scores of a sample and four of its translated augmentations spanning the training range. A translation coefficient of, e.g., 2 means the points are shifted by 2 in the positive direction for each axis. Samples from the same orbit often focus their attention on the same landmark point.

Panel of Figure 7 — Figure 7 : Mechanism indicators for the second block of a Transformer trained on ModelNet10 with 16 scale augmentations. ( Left ) The feature vectors in the attention output have larger norms than those in the residual stream, reducing the residual non-invariance to approximately the noise level. ( Middle and Right ) Attention maps are mostly invariant under scalings. Rows correspond to different samples.

Figure 8 : A model consisting only of

\operatorname{SA}(\operatorname{LN}({\bm{X}}))

blocks can learn to be scale invariant and solve the 3D Topology task.

Figure 9 : Task losses and equivariance errors of Transformers, trained with 16 augmentations on ShapeNet tasks, are reasonably low. Note that

\operatorname{Scale}

and

\operatorname{T}

consider a different task than

\operatorname{O}

and

\operatorname{SO}

Figure 10 : Results for the first layer of a Transformer trained on ShapeNet with

16

\operatorname{O}(3)

augmentations per sample. On the test set, the value representations are approximately

\operatorname{O}(3)

-equivariant ( left ) and the attention maps are approximately

\operatorname{O}(3)

-invariant ( right ).

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Symmetry group	Lie algebra basis
$\operatorname{SO}(3)$	$\left\{\begin{bmatrix}0&1&0\\ -1&0&0\\ 0&0&0\\ \end{bmatrix},\ \begin{bmatrix}0&0&1\\ 0&0&0\\ -1&0&0\\ \end{bmatrix},\ \begin{bmatrix}0&0&0\\ 0&0&1\\ 0&-1&0\\ \end{bmatrix}\right\}$
$\operatorname{T}(3)$	$\left\{\begin{bmatrix}0&0&0&1\\ 0&0&0&0\\ 0&0&0&0\\ 0&0&0&0\\ \end{bmatrix},\ \begin{bmatrix}0&0&0&0\\ 0&0&0&1\\ 0&0&0&0\\ 0&0&0&0\\ \end{bmatrix},\ \begin{bmatrix}0&0&0&0\\ 0&0&0&0\\ 0&0&0&1\\ 0&0&0&0\\ \end{bmatrix}\right\}$

Appendix

Table 2 : Example of Lie algebra bases for the rotation

\operatorname{SO}(3)

and translation groups

\operatorname{T}(3)

Figure 11 : OOD Topology invariance metrics correlation for base angle-preserving symmetries.

Figure 12 : Illustration of augmentations from the different angle-preserving and non-angle-preserving symmetry augmentations that we consider in our experiments for the 3D Topology dataset.

Model	Block	Layer sizes	#Params
DeepSets	Encoder $\phi$	$3\!\to\!156\!\to\!156$ , sum-pool	25,740
DeepSets	Decoder $\rho$	$156\!\to\!156\!\to\!C$ (dropout $0.1$ )	25,748
PointNet	Input T-Net ( $k{=}3$ )	$3\!\to\!9\!\to\!18\!\to\!140$ , maxpool, $140\!\to\!70\!\to\!35\!\to\!9$	16,099
	MLP 1	$3\!\to\!9\!\to\!9$	162
	Feature T-Net ( $k{=}9$ )	$9\!\to\!9\!\to\!18\!\to\!140$ , maxpool, $140\!\to\!70\!\to\!35\!\to\!81$	18,745
	MLP 2	$9\!\to\!18\!\to\!140$ , maxpool	3,156

Appendix

Table 3 : Parameter overview of compared architectures. Both operate on 3D point sets (

d_{\text{in}}=3

) and are parameter-matched to

\approx\!51

C

denotes the number of classes.

Transformation	Symmetry Group	Train / Validation / Test	Generalization
Rotation	$\operatorname{SO}(3)$	$[0^{\circ},90^{\circ})$	$[90^{\circ},180^{\circ})$
Rotation and reflection	$\operatorname{O}(3)$	$[0^{\circ},90^{\circ})$	$[90^{\circ},180^{\circ})$
Translation	$\operatorname{T}(3)$	$[0,1.5)$	$[1.5,3)$
Scaling	$\operatorname{Scale}(3)$	$[\times 1,\times 1.5)$	$[\times 1.5,\times 2)$

Appendix

Table 4 : Ranges of augmentations used for training, validation, and testing, as well as for measuring the generalization of the model for the considered based symmetries in § 4 .

Panel of Figure 13 — Figure 13 : Illustration of samples of 3D Tetris (bottom-left) and results measuring the mean and standard deviation for test accuracies (middle) and invariance losses (right) across 5 runs with different random seeds. Angle-preserving symmetries (blue and green tones) are efficiently learned by the Transformer, while it struggles with non-angle-preserving ones (orange tones).

Panel of Figure 14 — Figure 14 : Attention maps of a full Transformer trained on Tetris augmented by 32 rotations per sample are very stable across augmentations, as also indicated by their low invariance loss.

Panel of Figure 15 — Figure 15 : Results for the first block a full Transformer trained on Tetris with 32 translations per sample. Right: Each row shows the attention maps of a sample (Original) and four of its augmentations over the full training range. A coefficient of, e.g., 2 means the points are shifted by 2 in the positive direction for each axis. Samples from the same orbit focus their attention on the same landmark point. Middle: Invariance loss of the attention maps is low for the entire test set. Left: Cosine similarity between ${\bm{X}}$ and $\operatorname{SA}({\bm{X}})$ is strictly negative, i.e. the attention output has signs opposite to those of the residuals.

Figure 16 : Mechanism ablations for Tetris dataset. Left: A model with 2 blocks of

{\bm{X}}+\operatorname{SA}(\operatorname{LN}({\bm{X}}))

( Residual ) can learn to be translation invariant and solve the Tetris task, but

\operatorname{SA}(\operatorname{LN}({\bm{X}}))

( No Res. ) cannot. Right: A model consisting only of

\operatorname{SA}(\operatorname{LN}({\bm{X}}))

blocks can learn to be scale invariant and solve the Tetris task.

Panel of Figure 17 — Figure 17 : Mechanism indicators for the second block of a Transformer trained on Tetris with 32 scale augmentations. Left: points in the attention output have larger norms than points in the residual. The non-invariance of the residuals is therefore reduced to the noise level. Middle and Right: Attention maps are mostly invariant under scalings. Different rows correspond to different samples.

Panel of Figure 18 — Figure 18 : Test (in-distribution) and OOD generalization (disjoint distribution) task (solid line) and invariance losses (dashed line) for base angle-preserving symmetries in the QM9 dataset across 3 random seeds reported as MAE in eVs (lower is better). The gray lines represent the performance of a Transformer trained without data augmentation and tested in each augmented testing set as a reference to evaluate the improvement derived from data augmentation. All symmetries are learned robustly with enough augmentations, with translations being the most sensitive to distribution shifts.

Figure 19 : Test accuracy for different positional encoding approaches as measured for each symmetry group considered in the 3D Tetris dataset.

Figure 20 : Test invariant error for different positional encoding approaches as measured for each symmetry group considered in the 3D Tetris dataset.

Figure 21 : Attention softmax for first layer of a Transformer trained on $32$ augmentations using each considered geometric positional encoding and for each base angle-preserving symmetry. The attention maps remain approximately invariant across augmentations for all positional encodings.

Figure 22 : Invariance of attention heads to

\operatorname{T}(3)

(left) and

\operatorname{SO}(3)

(right) for first layer of a Transformer trained with

32

\operatorname{SE}(3)

-augmentations per sample. Head 0 is more invariant than head 1 to both base subgroups, so they do not seem to specialize in either of these groups.

Panel of Figure 23 — Figure 23 : In- and out-of-distribution results for 3D Tetris comparing 1, 2, and 4 attention heads. Increasing the number of heads slightly increases performance on compositional groups. The only visible qualitative difference seems to be extrapolation on $\operatorname{T}(3)$ , which seems sensitive to hyperparameter settings with no clear trend.

Panel of Figure 24 — Figure 24 : Mechanism indicators for first layer of all models trained on ModelNet10 with translations.

Panel of Figure 25 — Figure 25 : Mechanism indicators for second layer of all models trained on ModelNet10 with translations.

Figure 26 : Invariance loss per block of the attention maps of all models trained with rotations on ModelNet10.

Figure 27 : Mechanism indicators of first layer (layer 0) of all models trained on ModelNet10 with scale augmentations.

Figure 28 : Mechanism indicators of second layer (layer 1) of all models trained on ModelNet10 with scale augmentations.

Figure 29 : Comparison of equivariance of

{\bm{V}}

and

\operatorname{MLP}({\bm{X}})

in the first block of a Transformer trained on ShapeNet normals prediction with

16

\operatorname{O}(3)

augmentations. As a proxy, we measure the invariance of inner products, divided by a random baseline. Latent representations after the MLP appear less equivariant than those after the value projection, but more equivariant than the random baseline.

Panel of Figure 30 — Figure 30 : Results for first block of Transformers trained on ShapeNet per-part centroid prediction with $16$ $\operatorname{Scale}(3)$ (left) or $\operatorname{T}(3)$ augmentations (right) per sample. Both attention maps are approximately invariant to the respective group.

Explore similar work

Data Augmentation: A Fourier Analysis Perspective
Jun 23, 2026Behrooz Tahmasebi, Melanie Weber, Stefanie JegelkaData Augmentation Equivariant Learning
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)
A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of $\mathrm{O}(2)$
Jun 26, 2026Tīkun Ông, Georg BökmanSelf-Supervised Vision Transformers Vision Transformer
Independent Researcher · University of Amsterdam
Rotational Equivariance in Machine Learning: A Comprehensive Tutorial
Aug 31, 2026Peter Lippmann, Fred A. HamprechtRotation Equivariance Equivariant Neural Networks
Interdisciplinary Center for Scientific Computing (IWR), Heidelberg University, 69120 Heidelberg, Germany

Jun 23, 2026, cs.LGJ/K move · Enter open · S save

Data Augmentation: A Fourier Analysis Perspective

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 Equivariant Learning Fourier Optimal Learner Perspective

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Data Augmentation: A Fourier Analysis Perspective

A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of O(2)\mathrm{O}(2)O(2)

Rotational Equivariance in Machine Learning: A Comprehensive Tutorial

A Unified Framework for Vision Transformers Equivariant to Discrete Subgroups of $\mathrm{O}(2)$