PaperlayerView exact revision on arXiv

Source-linked AI summary

How Do Transformers Learn to Represent Symmetries?

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

arXiv:2610.10305v1cs.LG

TL;DR

The paper asks how Transformer architecture and finite augmentation budgets shape the symmetries a model can learn, a question not resolved by the popularity of augmentation-based approaches. It studies a vanilla Transformer on point clouds and finds an ordering of learnability across symmetry groups, then analyzes invariance mechanisms and their use in equivariant prediction.

  • Problem

    It remains unclear which symmetries Transformers can efficiently learn from finite augmentation, where compute budgets limit the transformed views each sample receives.

    Source: §1 · pp. 1–2

  • Method

    The paper systematically studies a vanilla Transformer on point clouds, tests learned invariance and extrapolation, analyzes trained-model mechanisms, and extends the analysis to equivariant prediction.

    Source: §1 · p. 2

  • Results

    Learnability increases from non-angle-preserving to angle-preserving symmetries, with base angle-preserving groups easiest; the paper also identifies mechanisms underlying their learned approximate invariance.

    Source: §1 · p. 2, §3 · p. 4

  • Takeaways & Limitations

    The identified mechanisms for approximate invariance can provide building blocks for learning equivariance.

    Source: §1 · p. 2

    The study focuses on small vanilla Transformers and point-cloud datasets, and its extension to equivariance remains preliminary.

    Source: §7 · p. 10

Abstract

from arXiv · show

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/

1 Introduction

The paper asks how Transformers learn symmetries from finite augmentation, a practical setting where only limited transformed views are available. It studies a vanilla Transformer on point clouds, finding an affinity for angle-preserving and especially base symmetries, then examines extrapolation, mechanisms, and equivariance.

1 Introduction

  • Finite augmentation makes symmetry approximate, and which architectures efficiently learn which symmetries from limited transformed views remains an open question.The question is especially relevant because Transformer adoption has not established whether they learn all symmetry types equally well.Source: §1 · pp. 1–2
  • The study systematically tests a vanilla Transformer’s ability to learn different symmetries under limited augmentation and examines how trained models represent them.Source: §1 · p. 2
  • Transformers show a special affinity for angle-preserving symmetries, particularly base groups such as translation, rotation, and scale.Source: §1 · p. 2
  • For base angle-preserving groups, the paper tests generalization to unseen transformation ranges and identifies interpretable invariance mechanisms that can also support learned equivariance.Source: §1 · p. 2

2 Related work

Related work covers architectural constraints and data augmentation as routes to symmetry-aware models, alongside other approaches such as frame averaging and canonicalization. It also notes trade-offs of exact equivariance and newer efforts to encourage equivariance through losses.

2 Related work

  • Common approaches include hard architectural equivariance and approximate equivariance through data augmentation, alongside frame averaging, canonicalization, and multi-task learning.Source: §2 · pp. 2–3
  • The study’s Table 1 separates the considered groups into angle-preserving and non-angle-preserving transformations.Source: Table 1 · p. 3
  • Exact equivariance can improve data efficiency and generalization, but may increase implementation complexity, challenge scalability, and restrict the hypothesis space when data symmetries are approximate.The passage also describes equivariance-promoting loss terms as an increasingly explored alternative.Source: §2 · p. 3

3 Finite augmentation reveals a symmetry inductive bias in Transformers

The section evaluates whether finite augmentation lets a vanilla Transformer learn symmetry on point-cloud tasks. Learnability follows a hierarchy: base angle-preserving groups are easiest, followed by their compositions, while non-angle-preserving groups are hardest.

3 Finite augmentation reveals a symmetry inductive bias in Transformers

  • The controlled setup uses a two-block, single-head vanilla Transformer without positional encodings or masks, and defines invariance as unchanged output under group action.Source: §3 · p. 3
  • Evaluation combines classification accuracy on transformed test samples with KL-based invariance error, used only as a diagnostic, on 3D Topology and ModelNet10.Source: §3 · pp. 3–4
  • Angle-preserving groups including SO(3), O(3), T(3), SE(3), Scale(3), and Sim(3) need few augmentations for high accuracy and low invariance loss; Affine(3), SL(3), and GL(3) need more and still perform worse.Source: §3 · p. 4
  • On ModelNet10, T(3) and Scale(3) perform strongest on accuracy and invariance loss and remain stable as augmentation increases; SO(3) and O(3) need larger budgets, while SE(3) and Sim(3) are harder.Source: §3 · p. 4
  • Transformers learn base angle-preserving groups most easily, compositional angle-preserving groups next, and non-angle-preserving groups least efficiently.This ordering motivates focusing subsequent analysis on base groups to isolate how each symmetry is learned.Source: §3 · pp. 4–5, Figure 2 · p. 4

4 Learning base angle-preserving symmetries

On base angle-preserving groups, scaling is easiest, and Transformers approach PointNet on rotations and reflections while outperforming it on translation invariance error and generalization; learned invariance extends beyond training ranges.

4 Learning base angle-preserving symmetries

  • The study tests in- and out-of-distribution generalization on ModelNet10 for translation, rotation, and scale after training on smaller-magnitude transformations.The larger-magnitude test transformations form a generalization set disjoint from the training range.Source: §4 · p. 5
  • DeepSets, PointNet, and a vanilla Transformer are compared to isolate how well each architecture learns base angle-preserving invariances, rather than to maximize benchmark performance.PointNet adds T-nets that predict affine matrices and can canonicalize orientation.Source: §4 · p. 5
  • Scaling is easiest across architectures; PointNet is strongest overall, while Transformers come close on rotations and reflections and outperform PointNet on translation invariance error and generalization.DeepSets has the weakest performance beyond scaling, with translation predictions nearly reduced to random guessing.Source: Figure 3 · p. 5, §4 · pp. 5–6
  • For base angle-preserving symmetries, Transformers learn invariance beyond the limited augmentation ranges seen during training.The evaluation compares transformations from the training range with a disjoint, larger-magnitude generalization range.Source: §4 · pp. 5–6

5 Mechanisms that can induce invariance to angle-preserving symmetries

The analysis identifies attention-based mechanisms for rotation/reflection, translation, and scale invariance, then checks their signatures in trained Transformers and ablations.

5 Mechanisms that can induce invariance to angle-preserving symmetries

  • The analysis isolates mechanisms for rotation/reflection, translation, and scale invariance through targeted ablations and activation-level checks in full Transformers.The study focuses on features invariant to each base angle-preserving symmetry.Source: §5 · p. 6

5.1 Rotation and reflection: stable attention over scalar relations

  • For rotations and reflections, attention can expose invariant scalar relations through inner-product scores; row-wise softmax requires only within-row score differences to remain stable.For point clouds, the Gram matrix contains the scalar information preserved by rotations and reflections.Source: §5.1 · p. 6
  • With 16 rotated augmentations per sample, trained attention maps remain nearly unchanged across rotations, matching the predicted preservation of score differences.The analysis reports small attention-map invariance errors across held-out inputs.Source: §5.1 · p. 7, Figure 4 · p. 7

5.2 Translation: constructing relative coordinates

  • For translation, attention often focuses on one landmark, while its output opposes the residual stream, consistent with subtracting a reference point or mean.The observed focus rules out local-neighborhood attention for convolution-like behavior and uniform attention for mean computation.Source: §5.2 · pp. 7–8
  • With invariant row-stochastic attention, subtracting its weighted point aggregate yields translation-invariant features; residual connections support expressive landmark shifts, and removing them breaks expressive invariance.A two-block residual model learns translation invariance and solves the 3D Topology task, while the no-residual model cannot.Source: Figure 5 · p. 7, §5.2 · p. 8

5.3 Scale: LayerNorm and residual suppression

  • Stable attention maps make attention-only stacks positively homogeneous; without additive transformations, argmax predictions are invariant to positive scale changes.The passage identifies LayerNorm as a natural candidate for learning scale invariance and notes that input-projection bias is an error source in the tested setting.Source: §5.3 · p. 9
  • In a full Transformer, scale invariance is supported when attention-output features have larger norms than residual features, reducing residual non-invariance toward the noise level.The second-block attention maps are also mostly invariant under scalings.Source: Figure 7 · p. 8, §5.3 · p. 9

6 From invariance to equivariance

The paper extends its invariance analysis to equivariant prediction, finding low equivariance errors and mechanisms consistent with approximate equivariance.

6 From invariance to equivariance

  • An SA(LN(X))-only model can learn scale invariance and solve the 3D Topology task.Source: Figure 8 · p. 9
  • The experiments use ShapeNet surface normals for O(3)- and SO(3)-equivariant targets, and part centroids for T(3)- and Scale(3)-equivariant targets.Source: §6 · p. 9
  • ShapeNet Transformers trained with 16 augmentations achieve reasonably low task losses and equivariance errors, suggesting approximately equivariant representations.The reported tasks differ: Scale and T use different tasks from O and SO.Source: Figure 9 · p. 9, §6 · p. 9
  • Attention maps are approximately invariant, while value projections appear approximately equivariant on the data support and augmentation orbits, despite lacking global isotropy.The value-projection assessment uses invariance of its Gram matrix as a proxy for equivariance.Source: §6 · p. 10
  • The analysis suggests that invariant components can serve as building blocks for learned equivariant computation.Source: §6 · p. 10

7 Conclusions

The paper finds that vanilla Transformers learn some symmetries more reliably than others, identifies mechanisms for base angle-preserving invariance, and frames the scope of those findings.

7 Conclusions

  • For O(3)-trained ShapeNet models, first-layer value representations are approximately equivariant and attention maps approximately invariant.Source: Figure 10 · p. 10
  • Transformers learn angle-preserving symmetries more reliably than non-angle-preserving ones, and learn base symmetries more efficiently than their compositions.The base groups studied are translation, rotation, and scale.Source: §7 · p. 10
  • Structural analyses identify mechanisms for invariance to each base angle-preserving symmetry, and suggest these mechanisms can also support learned equivariance.Source: §7 · p. 10
  • The study is limited to small vanilla Transformers and point-cloud datasets; its equivariance analysis remains preliminary.Whether the mechanisms persist in other data types, larger datasets, richer tasks, or settings with noise and interacting symmetries remains open.Source: §7 · p. 10

Supplementary Material

The supplement introduces additional discussion of related work covered in the paper’s opening sections.

Supplementary Material

A.1 The shift from hard equivariance towards data augmentation

The related work describes growing use of augmentation and approximate equivariance, while reviewing studies of learned symmetries and measurement approaches.

A.1 The shift from hard equivariance towards data augmentation

  • Recent large-scale generative work has shifted from hard equivariant architectures toward data augmentation, including Transformer-based scientific models operating over restricted symmetry sets.The cited examples include AlphaFold3, MCFs [9], TransIP [2], Zatom-1 [10], and TABASCO [11].Source: §A.1 · p. 18

A.2 Approximate equivariance

  • Approximate equivariance is presented as useful when noisy or biased measurements and symmetry-breaking effects make strict equivariance potentially harmful to training dynamics and predictive performance.The related work also reviews augmentation and other approximate-equivariance approaches.Source: §A.2 · p. 18

A.3 Learning symmetries with neural networks

  • Prior work reports that Transformers can learn SO(2) invariance while avoiding SL(4) invariance, and that they can become more equivariant than CNNs after augmentation training.The section also notes drawbacks in existing learned-equivariance metrics, including ad-hoc decisions and limited applicability across transformations and data types.Source: §A.3 · p. 18

B Experimental details

This appendix details the experimental protocols, datasets, and equivariant targets, and reports supplementary benchmark and OOD findings alongside analysis of jointly learning compositional symmetries. It also examines architectural variants and the mechanisms associated with translation, rotation, scale, and learned equivariance.

B.1.1 Preliminaries

  • Augmentations are sampled by drawing Lie-algebra coefficients from a uniform distribution, exponentiating to group elements, and applying point-wise actions across each point cloud.For connected matrix Lie groups, the parameterization samples group elements through the associated Lie algebra; each input also retains its unaugmented view.Source: §B.1.2 · p. 19, §B.1.1 · p. 19

B.1.4 Additional remarks

  • The augmentation scheme may not cover all of SL(n) or GL(n), but covers the connected components of the principal angle-preserving groups; O(n) additionally uses sampled reflections.Restricting to local continuous transformations provides a consistent, efficient comparison while avoiding transformations that could destroy semantic information.Source: §B.1.4 · pp. 19–20

B.1.5 Hardness of compositional groups

  • Exact invariance to constituent groups implies invariance to their composition, but finite joint augmentations cover larger transformation spaces sparsely and may compound approximation errors.SE(3) has six continuous degrees of freedom and Sim(3) seven, versus three each for translation and rotation and one for scale.Source: §B.1.5 · pp. 20–21
  • Compositional invariance requires compatible conditions in shared attention and residual pathways; adding attention heads did not materially improve compositional-group performance or yield clear head specialization.Base-group mechanisms include orthogonally stable attention-score differences for rotation/reflection, reference subtraction for translation, and stable attention with positive homogeneity and LayerNorm for scale.Source: §B.1.5 · p. 21

B.2 Measuring learned equivariance

  • Equivariance is measured by prediction consistency under sampled group transformations; classification uses KL divergence, and alternative metrics preserve the qualitative conclusions across base symmetries.The alternatives are symmetric KL, Jensen–Shannon divergence, and a logit-based L2 metric; their correlations with KL range from moderate to near-perfect, weakest for translation.Source: §B.2 · p. 21, Figure 11 · p. 22

B.3 Finite augmentation reveals a symmetry inductive bias in Transformers

  • The 3D Topology benchmark uses five topologically distinct shapes preserved by the studied transformations; equivariant tasks use ShapeNet normals for SO(3)/O(3) and part centroids for translation and scale.For equivariant training, the vanilla Transformer produces per-point outputs without final sum aggregation and is trained with mean squared error.Source: §B.3 · p. 22, §B.5 · p. 25

C.1 3D Tetris

  • On 3D Tetris, performance again improves in learnability from non-angle-preserving to compositional angle-preserving to base angle-preserving groups, though labels may not be invariant to strong non-angle-preserving transforms.Tetris classes differ by angles and distances, which non-angle-preserving transformations distort; reflections can also collapse opposite-chirality classes.Source: §C.1 · p. 25, Figure 13 · p. 26

C.2 Molecular property prediction

  • On QM9, the Transformer largely preserves task performance and invariance on disjoint, unseen transformation ranges; more training augmentations help, but translation generalizes less well than rotation and scale.The translation gap is more sensitive to the number of training augmentations, consistent with its more difficult OOD case.Source: §C.2 · p. 26

C.3.1 Positional encodings

  • Test performance remained similar to the no-encoding baseline across geometric positional encodings, while invariant computations and activation signatures remained qualitatively present.The experiments suggest that auxiliary geometric positional information does not eliminate the tendency to learn the reported invariant computations.Source: §C.3.1 · p. 29, Figure 19 · p. 30, Figure 20 · p. 31

RoPE

  • Geometric RoPE uses coordinate-based phases whose query-key contribution depends on relative displacement; attention maps remained approximately invariant across augmentations for all tested geometric encodings.The phases are defined from spatial coordinates, but planar rotations make each rotary pair's query-key contribution depend on the difference between point coordinates.Source: §C.3.1 · p. 29, Figure 21 · p. 32

C.4.1 Translation

  • The translation mechanism persisted across seeds and augmentation counts: first-block attention subtracts approximately equivariant features from residuals, while second-block attention maps are highly invariant.The reported diagnostics cover Transformers trained on ModelNet10 with translation augmentations.Source: §C.4.1 · p. 33, Figure 24 · p. 35, Figure 25 · p. 36

C.4.2 Rotation

  • For ModelNet10 rotation augmentations, both Transformer blocks exhibited strong rotation invariance, which increased with the number of augmentations.The passage reports the trend for both blocks, without specifying a numerical increase.Source: §C.4.2 · p. 33, Figure 26 · p. 37

C.4.3 Scale

  • With ModelNet10 scale augmentations, only the second block showed the signature of scaling self-attention outputs relative to residuals.The comparison of first- and second-block measurements supports this block-specific finding.Source: §C.4.3 · p. 33, Figure 27 · p. 38, Figure 28 · p. 39

C.3.2 Multi-head attention

  • Across head counts, results were broadly consistent; more heads slightly improved compositional-group performance, but showed no major overall impact.The proposed specialization of heads to base subgroups was not supported empirically; one head was often more invariant than the other across both subgroups.Source: §C.3.2 · p. 33, Figure 22 · p. 33, Figure 23 · p. 34

C.5 Equivariance

  • The equivariance analysis traces approximate O(3)-equivariance through a forward pass and identifies initial components relevant to learned equivariant mechanisms.For ShapeNet normals prediction with 16 O(3) augmentations, value-projected representations appeared more equivariant than post-MLP representations, which remained more equivariant than the random baseline.Source: §C.5 · p. 40, Figure 29 · p. 40

C.5.1 O(3)-equivariance of remaining architecture

Approximate O(3)-equivariance is maintained across Transformer submodules, while attention maps for translation and scale are also approximately invariant. The mechanisms and role of those attention maps remain unresolved.

C.5.1 O(3)-equivariance of remaining architecture

  • Unlike invariance, equivariance must be maintained through subsequent modules; residuals preserve it conditionally, while MLPs raise proxy error but remain below the random baseline.Strictly equivariant MLPs would require zero biases and isotropic projections, which are not observed in trained models.Source: §C.5.1 · p. 40
  • Approximate equivariance appears maintained by each submodule, rather than emerging only at the end; data-dependent approximations are used where full-space constraints are too restrictive.The authors propose O-equivariance in self-attention, but the remaining components must also support it.Source: §C.5.1 · p. 40

C.5.2 Other base groups: attention maps are invariant

  • Translation- and scale-trained Transformers also show approximately invariant attention maps, but their mechanisms and functional role remain unidentified.Figure 30 shows attention maps for models trained with 16 Scale(3) or T(3) augmentations per sample.Source: Figure 30 · p. 40, §C.5.2 · p. 40

D Proofs

  • Section D provides the proof of the theoretical result and follows the paper’s statement numbering.Source: D Proofs · p. 41

D.1 Proof of Lemma 1

The proofs characterize attention-map invariance through rowwise score differences and show that subtracting an attention-weighted representation yields translation invariance. They also establish that applying any function to this centered representation preserves that invariance.

D.1 Proof of Lemma 1

  • Row-wise softmax is unchanged by adding a constant to every logit, so attention-map invariance reduces to preserving score differences across keys.The proof expresses this condition as invariance of pairwise score differences within each row.Source: §D.1 · p. 41
  • Attention maps agree exactly when each row of the score-matrix change is constant across keys, equivalently when its centered difference vanishes.The centering condition is equivalent to every row of the score difference being constant across columns.Source: §D.1 · p. 41

D.2 Proof of Lemma 2

  • For translation-invariant, row-stochastic Λ, subtracting Λ⊤X from X yields a translation-invariant representation, so any downstream function g preserves that invariance.The proof verifies that translating every point leaves X − Λ⊤X unchanged.Source: §D.2 · p. 41

D.3 Proof of Lemma 3

When attention maps remain stable over a scale-augmentation range, attention-only layers and stacks scale their outputs by the same positive factor as their inputs. Consequently, positively scaled logits preserve the predicted class under argmax.

D.3 Proof of Lemma 3

  • If attention maps are stable over the scale range and logits scale by the same positive factor, the predicted class is invariant to scale.A positive factor preserves the ordering of the logits used by argmax.Source: §D.3 · p. 42
  • For a single attention-only layer, an unchanged attention map makes scaling the input scale the layer output by the same factor.This follows from SA(X) = A(X)XWV when A(cX) = A(X).Source: §D.3 · p. 42
  • If every layer's attention map is stable over the scale range, an attention-only stack is positively homogeneous, including with residual connections.The induction establishes Φ(cX) = cΦ(X) for c > 0.Source: §D.3 · p. 42

Related papers

Loading 2610.10305v1…