Exact Attention Sensitivity and the Geometry of Transformer Stability

arXiv:2602.18849v2 Announce Type: replace-cross Abstract: We develop a sensitivity analysis for transformer attention in a geometry aligned with tokenwise computation. Our main result is the exact identity |J_au(u)|{inftyo1}=heta(p)/au for the Jacobian J_au(u) of the tempered softmax umapstosoftmax(u/au), where heta(p)=4max{Ssubseteq[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 774M-parameter models, attention becomes substantially more concentrated while the median lower-bound certificate for heta(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 heta(p).

Source: arXiv cs.AI | 2026-08-18