Symmetry Breaking in Transformers for Efficient and Interpretable Training

1. Problem Statement

The paper addresses two primary challenges in training Transformer models:

1. Optimizer Inefficiency: While adaptive optimizers like AdamW and SOAP are highly effective for training Transformers, they are memory-intensive (requiring $\sim3N$ auxiliary variables). Conversely, memory-efficient, physics-inspired optimizers like Energy Conserving Descent (ECD) have shown promise in scientific applications but fail to match the performance of adaptive methods when applied to Transformers.
2. Theoretical Gap: The authors hypothesize that the failure of ECD in Transformers is due to the rotational symmetries inherent in the attention mechanism. These symmetries create "degenerate" directions in the parameter space that do not affect the model's output but shape the optimization dynamics.

Key Insight: In the attention mechanism, joint rotations of Query ($W_Q$) and Key ($W_K$) matrices, or Value ($W_V$) and Output ($W_O$) matrices, preserve attention scores (inner products). According to Noether's theorem, these continuous symmetries induce conserved angular momenta. In ECD, which relies on chaotic Hamiltonian dynamics with conserved total energy, this conservation traps kinetic energy in rotational motion, preventing the optimizer from effectively exploring the loss landscape in descent directions.

2. Methodology

The authors propose a simple, principled architectural modification to break these symmetries without sacrificing memory efficiency.

A. Symmetry-Breaking Protocol

They introduce unlearned, batchwise-sampled biases ($b_Q$ and $b_V$) into the Query and Value projections of the attention heads:

• Mechanism: During training, for each batch, random biases are sampled from normal distributions $N(\mu, \sigma^2)$ and added to the query and value vectors:
  $$q \leftarrow W_Q x + b_Q(\text{batch})$$
  $$v \leftarrow W_V x + b_V(\text{batch})$$
• Inference: The mean biases ($\mu_Q, \mu_V$) are applied.
• Symmetry Breaking: The random variance in the batchwise sampling breaks the continuous $O(d)$ rotational symmetry down to a discrete set, eliminating the conserved angular momenta that hinder ECD. The fixed mean $\mu_Q$ introduces a "preferred direction" in the rotational space.

B. Theoretical Framework (Hamiltonian Dynamics)

The paper models optimization using Hamiltonian mechanics:

• SGDM/Adam: Dissipative systems where friction removes kinetic energy, allowing convergence regardless of conserved quantities.
• ECD: A conservative system where total energy is constant. If angular momentum is conserved (due to symmetry), the system cannot efficiently convert kinetic energy into loss reduction. Breaking the symmetry allows the "wasted" rotational energy to be redirected into productive descent.

C. Interpretability Mechanism

The fixed mean bias $\mu_Q$ creates a preferred axis. The model learns to align the Key vectors ($W_K x$) of specific token classes with this bias. The attention score is modulated by a factor $e^{k \cdot b_Q}$, allowing the model to exponentially amplify or suppress specific token classes based on their alignment with the bias.

3. Key Contributions

1. Hamiltonian Explanation for ECD Failure: The authors provide a theoretical proof that rotational symmetries in attention heads induce conserved angular momenta, which obstruct the chaotic mixing required for ECD to function effectively.
2. Symmetry-Breaking Intervention: They propose a minimal architectural change (unlearned batchwise biases) that restores ECD's performance while maintaining its memory efficiency ($2N$ variables vs. $3N$ for adaptive methods).
3. Empirical Validation: They demonstrate that symmetry-broken ECD matches or exceeds the performance of adaptive optimizers (AdamW, SOAP) on GPT-2 (124M) models in terms of validation loss.
4. Interpretability: They show that the symmetry-breaking mechanism allows for a direct analysis of how models learn to amplify semantically meaningful tokens (e.g., sentence starters, punctuation) and suppress noise (e.g., encoding artifacts) via the alignment of Key vectors with the bias.

4. Experimental Results

The authors pretrained GPT-2 (124M) models on 500M tokens using four optimizers: ECD, SGDM, AdamW, and SOAP, under symmetric and symmetry-broken conditions.

• Validation Loss:
  • ECD: Without symmetry breaking, ECD performed poorly (Val Loss $\approx 3.93$). With symmetry breaking ($b_Q + b_V$), ECD improved significantly to 3.35, nearly matching SOAP (3.33) and outperforming its baseline by a large margin.
  • SGDM: Also improved with symmetry breaking (3.84 $\to$ 3.67).
  • AdamW: Did not benefit (performance slightly degraded), likely because Adam's adaptive preconditioning already breaks the symmetry implicitly.
• Downstream Reasoning (Logic Puzzles):
  • Performance on 14 logic tasks was heterogeneous. While validation loss improved consistently, logic gains varied by seed.
  • Predictor of Success: The study found that semantic alignment (how well the model aligns structural tokens like punctuation and sentence starters with the bias) was a better predictor of logic puzzle success than validation loss alone.
  • Models that successfully learned to amplify structural markers and suppress unicode/noise artifacts showed improved reasoning.
• Activation Functions: The effect was more pronounced with PReLU activations than GELU, suggesting GELU's inherent asymmetry partially mitigates the need for explicit symmetry breaking.

5. Significance and Implications

• Bridging the Gap: The work successfully bridges the gap between memory-efficient, physics-inspired optimizers and the high-performance adaptive optimizers currently dominating the field. This makes ECD a viable candidate for large-scale training where memory is a constraint.
• Principled Architecture Design: It demonstrates that understanding the geometric and symmetry structure of neural networks can lead to simple, effective architectural modifications.
• New Interpretability Lens: The method provides a novel way to interpret Transformer attention. By analyzing the alignment between learned Key vectors and the fixed bias, researchers can directly observe which token classes the model deems "important" for reasoning, offering a window into the model's internal logic.
• Scalability: The approach requires no additional trainable parameters (the biases are unlearned), making it highly scalable and compatible with existing open-source models (e.g., Llama, Gemma) that lack attention biases.

In conclusion, the paper argues that breaking rotational symmetries is not just a theoretical necessity for certain optimizers but a practical tool that enhances both the efficiency of training and the interpretability of the resulting models.