Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Symmetry Breaking in Transformers for Efficient and Interpretable Training

Eva Silverstein, Daniel Kunin, Vasudev Shyam

TL;DR

Transformers contain rotational symmetries in attention that carry no direct gradient signal yet shape learning dynamics. The authors propose batchwise, unlearned biases $b_Q$ and $b_V$ to break these symmetries, guided by a Hamiltonian view that links optimization dynamics to conserved angular momenta. Empirically, symmetry breaking enables memory-efficient optimizers, particularly Energy Conserving Descent, to approach or exceed the performance of adaptive methods on GPT-2 scale models and reveals an interpretable mechanism where attention weights are amplified for token classes aligned with $\, ext{E}[b_Q]$. The work combines theoretical and empirical analysis to show that minimal architectural modifications can improve both efficiency and interpretability in transformer training, while outlining scaling and analysis directions for future work.

Abstract

The attention mechanism in its standard implementation contains extraneous rotational degrees of freedom that are carried through computation but do not affect model activations or outputs. We introduce a simple symmetry-breaking protocol that inserts a preferred direction into this rotational space through batchwise-sampled, unlearned query and value biases. This modification has two theoretically motivated and empirically validated consequences. First, it can substantially improve the performance of simple, memory-efficient optimizers, narrowing -- and in some cases closing -- the gap to successful but more complex memory-intensive adaptive methods. We demonstrate this by pretraining 124M parameter transformer models with four optimization algorithms (AdamW, SOAP, SGDM, and Energy Conserving Descent(ECD)) and evaluating both validation loss and downstream logical reasoning. Second, it enables an interpretable use of otherwise redundant rotational degrees of freedom, selectively amplifying semantically meaningful token classes within individual attention heads. Overall, our results show that minimal, principled architectural changes can simultaneously improve performance and interpretability.

Symmetry Breaking in Transformers for Efficient and Interpretable Training

TL;DR

Transformers contain rotational symmetries in attention that carry no direct gradient signal yet shape learning dynamics. The authors propose batchwise, unlearned biases and to break these symmetries, guided by a Hamiltonian view that links optimization dynamics to conserved angular momenta. Empirically, symmetry breaking enables memory-efficient optimizers, particularly Energy Conserving Descent, to approach or exceed the performance of adaptive methods on GPT-2 scale models and reveals an interpretable mechanism where attention weights are amplified for token classes aligned with . The work combines theoretical and empirical analysis to show that minimal architectural modifications can improve both efficiency and interpretability in transformer training, while outlining scaling and analysis directions for future work.

Abstract

The attention mechanism in its standard implementation contains extraneous rotational degrees of freedom that are carried through computation but do not affect model activations or outputs. We introduce a simple symmetry-breaking protocol that inserts a preferred direction into this rotational space through batchwise-sampled, unlearned query and value biases. This modification has two theoretically motivated and empirically validated consequences. First, it can substantially improve the performance of simple, memory-efficient optimizers, narrowing -- and in some cases closing -- the gap to successful but more complex memory-intensive adaptive methods. We demonstrate this by pretraining 124M parameter transformer models with four optimization algorithms (AdamW, SOAP, SGDM, and Energy Conserving Descent(ECD)) and evaluating both validation loss and downstream logical reasoning. Second, it enables an interpretable use of otherwise redundant rotational degrees of freedom, selectively amplifying semantically meaningful token classes within individual attention heads. Overall, our results show that minimal, principled architectural changes can simultaneously improve performance and interpretability.
Paper Structure (37 sections, 24 equations, 2 figures, 26 tables, 1 algorithm)

This paper contains 37 sections, 24 equations, 2 figures, 26 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Symmetry, Conservation, and Optimization
  4. A Hamiltonian Framework for Optimization
  5. Continuous Rotational Symmetries in Attention
  6. Implications of Conserved Angular Momentum
  7. Biases to break attention symmetries
  8. Empirical Evaluation and interpretability
  9. Symmetry Breaking Enables Efficient Optimization
  10. Interpretable $b_Q$ Alignment
  11. Summary of empirical evaluation
  12. Discussion and Limitations
  13. Use of LLMs
  14. The $ECD_{q=1}$ version of ECD optimization
  15. Rotational Drift Analysis
  16. ...and 22 more sections

Figures (2)

  • Figure 1: Friction-based versus energy-conserving optimization. Here we illustrate the difference between a friction-based versus energy-conserving optimization strategy in a quadratic basin. (a) SGDM is like a ball rolling downhill with friction: momentum can cause it to overshoot the minimum, but kinetic energy is dissipated, damping oscillations, leading to convergence. (b) ECD is like as a snowball rolling in the same basin: here the total energy is conserved, but the mass increases as the loss decreases, naturally reducing velocity and slowing motion near the minimum.
  • Figure 2: Conserved angular momentum hinders ECD. Trajectories of five optimizers on the synthetic error surface $f(x,y) = (x^2 + y^2 - 1 + z)^2$, which exhibits a circular symmetry and serves as a toy model for the rotational symmetries in attention. Random perturbations $z \sim \mathcal{N}(0, 0.1^2)$ mimic mini-batch noise. SGDM is initialized with zero momentum and consequently conserves zero angular momentum, proceeding radially toward the minimum. Adam with its fixed, symmetry-breaking coordinate axes proceeds along a different angle than SOAP and Shampoo, whose preconditioning steps orient them in a radial-angular bases. In contrast, ECD (shown with an exaggerated chaos parameter $\nu$) introduces angular momentum that interferes with optimization as in §\ref{['sec-ECD-SGDM-conserved-angular-momenta']}.