Exact Attention Sensitivity and the Geometry of Transformer Stability

TL;DR

This paper develops a stability theory for transformers, deriving exact operator norms and geometric bounds to explain training stability and the roles of normalization techniques, validated on large models.

Contribution

It introduces a first-principles framework for understanding transformer stability, deriving exact norms, and explaining the effects of normalization and depth on gradients.

Findings

Pre-LN preserves identity gradient paths.

DeepNorm's scaling emerges from attention matrix structure.

Attention sensitivity remains high throughout training.

Abstract

Despite powering modern AI, transformers remain mysteriously brittle to train. We develop a stability theory that explains why pre-LayerNorm works, why DeepNorm uses $N^{-1/4}$ scaling, and why warmup is necessary, all from first principles. Our framework has two pillars: (1) We derive the \emph{exact} operator norm of the softmax Jacobian, $\|J_{softmax}(u/\tau)\|_{\infty\to 1} = \theta(p)/\tau$, where the balanced-mass factor $\theta(p)\in[0,1]$ quantifies attention sensitivity. (2) We introduce a block-$\infty$/RMS geometry aligned with tokenwise computation, yielding Lipschitz bounds independent of sequence length. Using this framework, we prove that pre-LN preserves identity gradient paths while post-LN compounds LayerNorm Jacobians exponentially with depth, and we show that DeepNorm's $N^{-1/4}$ emerges from the quartic structure of attention's four projection matrices. We validate our theory on 774M-parameter models and find that, contrary to the intuition that attention sharpens during training to reduce sensitivity, $\theta(p) \approx 1$ persists throughout. Transformer stability arises entirely from architectural gradient flow, not from attention dynamics. This finding changes how we reason about training: the architecture itself must handle sensitivity, not learned attention patterns.