Numerical Fragility in Transformers: A Layer-wise Theory for Explaining, Forecasting, and Mitigating Instability

TL;DR

The paper tackles forward-pass instability in Transformers trained with low-precision arithmetic and develops a layer-wise, first-order theory for predicting where and when errors amplify in self-attention. It presents interpretable diagnostics $(\kappa_{\mathrm{score}},\ \kappa_{\mathrm{softmax}},\ \kappa(V))$, a residual-relaxation bound using $(1+\rho_\ell)$, and a precision/width-aware LayerNorm indicator $\rho_{\mathrm{LN}}$ to produce a unified, online-estimable stability bound. Empirical validation on Tiny-ViT/CIFAR-10 shows the combined predictor tracks forward mismatch across precisions, that a high $\kappa_{\mathrm{softmax}}$ can forewarn instability several steps in advance, and that a minimal LayerNorm-$\varepsilon$ intervention yields stable training with negligible overhead. Together, these results provide actionable diagnostics and a practical mitigation pathway for robust mixed-precision training of Transformers, with potential applicability to larger models and longer sequences.

Abstract

Transformers trained in low precision can suffer forward-error amplification. We give a first-order, module-wise theory that predicts when and where errors grow. For self-attention we derive a per-layer bound that factorizes into three interpretable diagnostics: a score-scale ratio $κ_{\rm score}$, a rowwise softmax sensitivity $κ_{\rm softmax}$, and value conditioning $κ(V)$. We prove a residual relaxation inequality showing that residual blocks attenuate depth-wise accumulation, and we introduce a precision- and width-aware LayerNorm indicator $ρ_{\rm LN}$ with a matching first-order bound in the $ε$-dominated regime. These pieces yield a unified forward-stability bound whose right-hand side is directly estimable during training. On Tiny-ViT/CIFAR-10 we evaluate the bound and components. (1) The combined predictor $κ_{\rm softmax},(1+κ_{\rm score}),κ(V),|W_O|2+κ{\rm eff}+C_{\rm LN}$ tracks FP32$\leftrightarrow$LP mismatches across seeds, widths, and precisions; scaling by $ε_{\rm mach}$ collapses mixed-precision points. (2) The time-series maximum of $κ_{\rm softmax}$ acts as an early-warning signal, leading error spikes by 16-24 steps (corr. 0.65-0.82; permutation $p!\approx!10^{-3}$; Precision@K 0.89-1.00). (3) Guided by $ρ_{\rm LN}$, a small LayerNorm-$ε$ tweak targeting $ρ_\star$ gives consistent stabilization (mean tail-loss $\downarrow\ \approx0.010$ at $ρ_\star!=!0.6$, cap$=10^{-2}$) with negligible overhead. Overall, our theory supplies actionable, unitless diagnostics that (i) explain when self-attention is fragile, (ii) forecast instability, and (iii) motivate a minimally invasive mitigation.

Figures (6)

• Figure 1: Exp-1. Observed mismatch vs. predictor (per-config tails). Colors denote precision.
• Figure 2: Exp-1 raw scatter (per seed/width/precision).
• Figure 3: Exp-1 mixed-precision collapse after scaling by $\epsilon_{\mathrm{mach}}$.
• Figure 4: Exp-2 example (seed 0): $\kappa_{\mathrm{softmax}}$ leads forward mismatch by $\sim$16 steps.
• Figure 5: Exp-3 (loss). Tail loss improvements across $(\rho_*,\text{cap})$; positive is better.

Theorems & Definitions (17)

• Theorem 1: Self-Attention Forward Error
• Remark 1
• Theorem 2: Residual Stabilization
• Corollary 1: Depth-wise relaxation
• Remark 2
• Proposition 1: First-order LN forward error in the $\varepsilon$-dominated regime
• Theorem 3: Unified forward stability
• Lemma 1: GEMM first-order bound
• proof
• Lemma 2: Row-wise softmax Jacobian norm