Symmetry Breaking in Transformers for Efficient and Interpretable Training

TL;DR

This paper introduces a simple symmetry-breaking method in transformer attention mechanisms that enhances training efficiency and interpretability by reducing redundant rotational degrees of freedom.

Contribution

The authors propose a straightforward protocol that inserts a preferred direction in the rotational space, improving optimizer performance and enabling interpretability of attention heads.

Findings

Improved performance of simple optimizers in transformer training.

Enhanced interpretability of attention mechanisms.

Achieved competitive results with minimal architectural changes.

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.