Exact Attention Sensitivity and the Geometry of Transformer Stability

arXiv:2602.18849v2 Announce Type: replace Abstract: We develop a sensitivity analysis for transformer attention in a geometry aligned with tokenwise computation. Our main result is the exact identity $\|J_\tau(u)\|_{\infty\to1}=\theta(p)/\tau$ for the Jacobian $J_\tau(u)$ of the tempered softmax $u\mapsto\mathrm{softmax}(u/\tau)$, where $\theta(p)=4\max_{S\subseteq[L]}p(S)(1-p(S))$ measures how evenly the attention distribution can be bisected rather than how concentrated it is. We combine this identity with a block-$\infty$/RMS norm under which row-stochastic attention mixing is nonexpansive. This yields a distribution-aware local Jacobian bound for multi-head attention and a sequence-length-independent Lipschitz bound on bounded input sets, with explicit dependence on width, input magnitude, temperature, and projection norms. We also identify a structural distinction between normalization placements: a pre-LN residual-sublayer Jacobian contains an additive identity term, whereas a post-LN residual-sublayer Jacobian does not. A LayerNorm projection lemma gives a sufficient condition under which the LayerNorm-only term in the post-LN expansion contracts geometrically; the condition is not tested by our experiments. Across three Pre-LN early-training runs of $774$M-parameter models, attention becomes substantially more concentrated while the median lower-bound certificate for $\theta(p)$ remains near one at every sampled layer and checkpoint. This certifies near-maximal exact sensitivity for at least half of the sampled rows within each layer. A minority of rows enters a dominant-atom regime with lower exact sensitivity, consistent with the deterministic relationship between $p_{\max}$ and $\theta(p)$.

Abstract:We develop a sensitivity analysis for transformer attention in a geometry aligned with tokenwise computation. Our main result is the exact identity $\|J_\tau(u)\|_{\infty\to1}=\theta(p)/\tau$ for the Jacobian $J_\tau(u)$ of the tempered softmax $u\mapsto\mathrm{softmax}(u/\tau)$, where $\theta(p)=4\max_{S\subseteq[L]}p(S)(1-p(S))$ measures how evenly the attention distribution can be bisected rather than how concentrated it is. We combine this identity with a block-$\infty$/RMS norm under which row-stochastic attention mixing is nonexpansive. This yields a distribution-aware local Jacobian bound for multi-head attention and a sequence-length-independent Lipschitz bound on bounded input sets, with explicit dependence on width, input magnitude, temperature, and projection norms. We also identify a structural distinction between normalization placements: a pre-LN residual-sublayer Jacobian contains an additive identity term, whereas a post-LN residual-sublayer Jacobian does not. A LayerNorm projection lemma gives a sufficient condition under which the LayerNorm-only term in the post-LN expansion contracts geometrically; the condition is not tested by our experiments. Across three Pre-LN early-training runs of $774$M-parameter models, attention becomes substantially more concentrated while the median lower-bound certificate for $\theta(p)$ remains near one at every sampled layer and checkpoint. This certifies near-maximal exact sensitivity for at least half of the sampled rows within each layer. A minority of rows enters a dominant-atom regime with lower exact sensitivity, consistent with the deterministic relationship between $p_{\max}$ and $\theta(p)$.

27 pages, 4 figures. v2: Major revision. Exact softmax sensitivity theorem unchanged. Removed an unsupported path-length scaling theorem and its DeepNorm derivation; corrected the uniform-attention parity case and the loss normalization/update accounting; reran the headline experiments; clarified the local vs. bounded-domain bounds and the conditional pre-LN/post-LN analysis

Machine Learning (cs.LG); Artificial Intelligence (cs.AI)

arXiv:2602.18849 [cs.LG]

(or arXiv:2602.18849v2 [cs.LG] for this version)

aiqbench analysis

Exact Attention Sensitivity and the Geometry of Transformer Stability

Related coverage