Paper deep dive
Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training
Seyed Morteza Emadi
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 91%
Last extracted: 7/20/2026, 9:12:03 PM
Summary
This paper introduces a rank-aware spectral bound for attention logits to enable stable low-precision (FP8) training in transformers. By analyzing the bilinear structure of attention scores, the authors derive a tighter concentration inequality that accounts for the low-rank nature of the query-key interaction matrix. This allows for the computation of geometry-aware scale factors based on spectral norms rather than observed activation statistics, eliminating overflow risks during transient states (e.g., checkpoint loading) while maintaining compatibility with fused attention kernels.
Entities (10)
Relation Signals (6)
Attention Logits → governs → Overflow Risk
confidence 95% · whose maximum magnitude governs overflow risk in low-precision training.
Query-Key Interaction Matrix → determines → Attention Logits
confidence 92% · Attention scores in transformers are bilinear forms... depends on the interaction of two projected vectors... M=W^Q W^{K\top}
Geometry-Aware Scaling → eliminates → Overflow
confidence 90% · geometry-aware scaling eliminates overflows in transient scenarios where delayed scaling fails
Rank-Aware Concentration Inequality → provides → Tighter Bounds
confidence 90% · yields 8--28x tighter concentration than rank-agnostic bounds
Geometry-Aware Scaling → iscompatiblewith → Fused Attention Kernels
confidence 88% · remains compatible with fused attention kernels
Delayed Scaling → failsduring → Transient Scenarios
confidence 85% · delayed scaling fails in transient scenarios such as checkpoint loading, resumption, and learning rate changes
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Attention scores in transformers are bilinear forms $S_{ij} = x_i^\top M x_j / \sqrt{d_h}$ whose maximum magnitude governs overflow risk in low-precision training. We derive a \emph{rank-aware concentration inequality}: when the interaction matrix $M = W^Q W^{K\top}$ has rank $r \ll d$, tail probabilities for $\max_{i,j}|S_{ij}|$ decay as $\exp(-d^{2}\alpha^{2}/(\gamma r))$ rather than $\exp(-d\alpha^{2})$, where $\gamma > 1$ is a typicality parameter. For transformer attention where $r = d_h$, this yields $8$--$28\times$ tighter concentration than rank-agnostic bounds in modern architectures. We apply this result to FP8 training, deriving \emph{geometry-aware scale factors} that provide principled overflow guarantees without observing activations. The method computes per-layer scales from the spectral norm $\|W^Q W^{K\top}\|_2$ via implicit power iteration, includes a grouped query attention formulation that avoids key expansion, and remains compatible with fused attention kernels. Across GPT-2 XL to Llama-2-70B, geometry-aware scaling eliminates overflows in transient scenarios where delayed scaling fails, while achieving comparable downstream MMLU accuracy.
Tags
Links
- Source: https://arxiv.org/abs/2602.18851v1
- Canonical: https://arxiv.org/abs/2602.18851v1
Trouble viewing inline? Open PDF directly →
Full Text
79,621 characters extracted from source content.
Expand or collapse full text
Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training Seyed Morteza Emadi Abstract Attention scores in transformers are bilinear forms Sij=xi⊤Mxj/dhS_ij=x_i Mx_j/ d_h whose maximum magnitude governs overflow risk in low-precision training. We derive a rank-aware concentration inequality: when the interaction matrix M=WQWK⊤M=W^QW^K has rank r≪dr d, tail probabilities for maxi,j|Sij| _i,j|S_ij| decay as exp(−d2α2/(γr)) (-d^2α^2/(γ r)) rather than exp(−dα2) (-dα^2), where γ>1γ>1 is a typicality parameter. For transformer attention where r=dhr=d_h, this yields 88–28×28× tighter concentration than rank-agnostic bounds in modern architectures. We apply this result to FP8 training, deriving geometry-aware scale factors that provide principled overflow guarantees without observing activations. The method computes per-layer scales from the spectral norm ‖WQWK⊤‖2\|W^QW^K \|_2 via implicit power iteration, includes a grouped query attention formulation that avoids key expansion, and remains compatible with fused attention kernels. Across GPT-2 XL to Llama-2-70B, geometry-aware scaling eliminates overflows in transient scenarios where delayed scaling fails, while achieving comparable downstream MMLU accuracy. FP8 training, transformer stability, spectral norms, low-precision training 1 Introduction Attention logits in transformers exhibit a bilinear structure: each score Sij=xi⊤WQWK⊤xj/dhS_ij=x_i W^QW^K x_j/ d_h depends on the interaction of two projected vectors. This structure determines the maximum magnitude of attention logits, a quantity with direct consequences for low-precision training. FP8 formats reduce precision from 16 to 8 bits, but the E4M3 format spans only ±448± 448 compared to FP16’s ±65,504± 65,504.111FP8 training can yield substantial speedups on memory-bound workloads (Micikevicius et al., 2022); our focus is on the calibration problem that enables safe deployment. When attention logits exceed this range, quantization overflows produce NaN values that corrupt training. Problem formulation. Let RmaxR_ denote the representable range of the target format (Rmax=448R_ =448 for E4M3). Given a scale factor s>0s>0, quantization overflows when maxi,j|Sij|>s⋅Rmax _i,j|S_ij|>s· R_ . The calibration problem is to select s that: 1. (Safety) Ensures Pr(maxi,j|Sij|>s⋅Rmax)≤δ∗ ( _i,j|S_ij|>s· R_ )≤δ^* for a target failure probability δ∗δ^*; 2. (Utilization) Maximizes precision utilization [maxi,j|Sij|]/(s⋅Rmax)E[ _i,j|S_ij|]/(s· R_ ), minimizing quantization error. These objectives trade off: smaller s improves utilization but increases overflow risk. The challenge is to calibrate s without observing maxi,j|Sij| _i,j|S_ij|, which would require materializing the full attention matrix. Existing methods derive scale factors from observed activation statistics, either from a history buffer (delayed scaling) or per-iteration (current scaling), without analyzing the bilinear mechanism that governs logit magnitudes. Delayed scaling and its limitations. The standard approach for FP8 training is delayed scaling (Micikevicius et al., 2022), which maintains a history buffer of activation maxima (typically 16 steps) and derives scale factors from past observations: scalet=max(history)448⋅η,scale_t= (history)448·η, (1) where η<1η<1 provides a safety margin. Scale factors are computed before the current forward pass, avoiding the need to observe current activations. However, delayed scaling assumes activation magnitudes change slowly. This assumption breaks whenever weight dynamics outpace the history buffer. We call this failure mode history staleness. Common triggers include: (i) loading pretrained checkpoints, where FP8 history initializes to defaults while weights reflect extensive prior training; (i) resuming from saved states, since standard checkpointing omits scaling state; and (i) learning rate transitions, where warmup or cyclic schedules cause weights to shift faster than history adapts. An alternative is current scaling, which computes scale factors from max|St| |S_t| in each forward pass. This eliminates staleness but requires materializing the full L×L× L attention score matrix, incompatible with fused attention kernels like FlashAttention (Dao et al., 2022; Dao, 2024) that achieve O(L)O(L) memory precisely by avoiding this materialization. Table 1 summarizes this dilemma: delayed scaling is compatible with fused kernels but fails during transients; current scaling handles transients but requires materializing the score matrix. Prior to this work, no method achieved both. Table 1: The FP8 scaling dilemma. Method Transient-Safe Fused-Compat. Delayed ✗ ✓ Current ✓ ✗ Ours ✓ ✓ Our approach: predictive calibration from weight geometry. We propose a fundamentally different paradigm. Instead of reacting to observed activations (whether past or present), we predict their bounds from the singular value structure of the query-key interaction matrix. Attention scores form a bilinear structure: Sij=xi⊤Mxjdh,where M=WQWK⊤∈ℝd×dS_ij= x_i Mx_j d_h, M=W^QW^K ^d× d (2) is the query-key interaction matrix and xi,xjx_i,x_j are input embeddings. For pre-LN architectures (including GPT-2, Llama, and Mistral, which place normalization before attention), LayerNorm (or RMSNorm (Zhang and Sennrich, 2019)) constrains ‖xi‖≈d\|x_i\|≈ d, yielding: maxi,j|Sij|≤∥M∥2⋅ddh=:Bmax. _i,j|S_ij|≤\|M\|_2· d d_h=:B_ . (3) This bound depends only on current weights, not activations, and can be estimated via power iteration (Golub and Van Loan, 2013) without forming the d×d× d matrix M. When weights change, the bound updates immediately, enabling predictive calibration that eliminates the failure modes of delayed scaling during transients such as checkpoint loading, resumption, and learning rate changes. Since we never observe activations, compatibility with fused attention kernels is preserved. The worst-case bound BmaxB_ assumes inputs align perfectly with top singular vectors, which becomes exponentially unlikely as dimension increases. We therefore introduce a calibration factor α∈(0,1)α∈(0,1) to obtain a tighter bound Bα=α⋅BmaxB_α=α· B_ . To quantify the associated risk, we derive a rank-aware concentration inequality on Pr(maxi,j|Sij|≥Bα) ( _i,j|S_ij|≥ B_α) that exploits the low-rank structure of the interaction matrix (rank(M)=dh≪drank(M)=d_h d). The key insight is that although M acts on d-dimensional space, its range is only dhd_h-dimensional; inputs are exponentially unlikely to find this low-dimensional subspace of maximum amplification. The resulting tail exponent improves from dα2dα^2 to d2α2/(γdh)d^2α^2/(γ d_h), a factor of d/(γdh)=8d/(γ d_h)=8–28×28× in modern architectures (where γ>1γ>1 is a parameter controlling the typicality threshold; see Section 3.2). This analysis provides a principled selection rule: given a target overflow probability δ∗δ^*, we derive the minimum α guaranteeing Pr(maxi,j|Sij|≥Bα)≤δ∗ ( _i,j|S_ij|≥ B_α)≤δ^*. Contributions. 1. Tight spectral interaction bound. We prove maxi,j|Sij|≤‖WQWK⊤‖2⋅d/dh _i,j|S_ij|≤\|W^QW^K \|_2· d/ d_h under the distributional constraints induced by LayerNorm or RMSNorm. This bound is never looser than the naive submultiplicative bound ‖WQ‖2‖WK‖2\|W^Q\|_2\|W^K\|_2 and is strictly tighter unless the top right singular vectors align (Corollary 3.3). We show extension to RoPE architectures, with the tighter interaction bound validated empirically (Corollary 3.6). 2. Rank-aware concentration inequality. We prove a concentration inequality for bilinear forms xi⊤Mxjx_i Mx_j when M has rank r≪dr d and inputs are approximately uniform on the sphere (Proposition 3.4). The key insight is a two-stage conditioning argument: first bound the probability that projections onto M’s row space are atypical, then apply concentration with a tighter Lipschitz constant. This yields tail exponents of order d2α2/(γr)d^2α^2/(γ r) versus dα2dα^2 for rank-agnostic bounds, an improvement of d/(γr)d/(γ r) that becomes 88–28×28× in modern transformers where r=dhr=d_h. Applied to attention, this provides a principled selection rule: given a target overflow probability δ∗δ^*, we derive the minimum calibration factor α guaranteeing safety. 3. Implicit spectral norm estimator. We estimate ‖WQWK⊤‖2\|W^QW^K \|_2 via power iteration without forming the d×d× d interaction matrix, achieving O(nheads⋅dh⋅d)O(n_heads· d_h· d) cost per layer. For grouped query attention, we derive an implicit formulation that avoids key matrix expansion (Section 4.2). 4. Empirical validation. Zero overflows across GPT-2 XL through Llama-2-70B on all transient scenarios where delayed scaling fails. Downstream evaluation on MMLU confirms preserved capability: our auto-α calibration achieves accuracy comparable to delayed scaling (33.6% vs. 33.0%) while eliminating all overflows, whereas conservative spectral bounds degrade to 28.1% due to under-utilization (Section 5.4). Positioning. The spectral bound in Equation (3) builds on classical matrix analysis; our contribution lies in developing a complete calibration framework for low-precision training. This requires solving three problems that the bound alone does not address: (i) principled probabilistic calibration under the rank-deficient structure of attention, (i) efficient per-layer computation via implicit power iteration, and (i) an implicit GQA formulation that avoids key expansion. Together, these enable the first calibration method that is both robust to transients and compatible with fused attention kernels. 2 Related Work FP8 training and scaling strategies. Micikevicius et al. (2022) introduced FP8 formats and delayed (history-based) scaling, now standard in low-precision training frameworks. DeepSeek-V3 (DeepSeek-AI, 2024) employs per-tile scaling (e.g., 1×1281× 128 blocks) but retains attention in BF16/FP32. Microscaling (Rouhani et al., 2023) similarly uses block-wise factors. These approaches address spatial heterogeneity within tensors but still derive scales from observed values, inheriting history staleness during transients. FlashAttention-3 (Shah et al., 2024) introduces native FP8 support with internal block quantization, but requires inputs cast to FP8 before kernel entry. Incorrect global scale factors cause overflow before internal mechanisms can compensate. Other recent work addresses orthogonal aspects: architectural modifications (Hernández-Cano et al., 2025), inference fallback (Lee et al., 2025), and LoRA fine-tuning (Choi et al., 2025). Our method provides a safety envelope: geometry-aware bounds ensure inputs are within safe range, and are orthogonal to block-wise optimizations. Predictive scaling. Unit Scaling (Blake et al., 2023) eliminates scaling factors via careful initialization that maintains unit variance throughout the network, but requires training from scratch and cannot be applied to pretrained checkpoints. MOSS (Zhang et al., 2025) exploits the bounded update property of AdamW (|ΔWt|≤η| W_t|≤η) to predict weight magnitudes without runtime max-reduction. However, MOSS requires the MXFP8 microscaling format with hardware-specific support. More fundamentally, bounds derived from individual weight norms ‖WQ‖\|W^Q\| and ‖WK‖\|W^K\| are too loose for effective FP8 calibration: the naive bound ‖WQ‖2‖WK‖2\|W^Q\|_2\|W^K\|_2 can substantially overestimate when singular vectors of WQW^Q and WKW^K do not align, wasting dynamic range (Corollary 3.3). Our method addresses attention-specific overflow using only the standard E4M3 format, requires no custom kernels, and is compatible with NVIDIA’s Transformer Engine. Spectral methods for transformer stability. Zhai et al. (2023) show that entropy collapse (attention concentrating on single tokens) is governed by the spectral norm of the query-key interaction matrix. Our work reveals the same quantity governs numerical overflow risk, connecting optimization stability and numerical stability through shared geometric structure. Miyato et al. (2018) introduced spectral normalization for GANs as a regularizer; we use spectral norms for prediction rather than constraint. 3 Theoretical Foundation This section develops the theoretical basis for geometry-aware scaling. 3.1 Spectral Bound on Attention Logits Pre-softmax attention scores are computed as S=QK⊤/dhS=QK / d_h, where Q=XWQQ=XW^Q and K=XWKK=XW^K. Each score can be written as: Sij=xi⊤Mxjdh,where M=WQWK⊤∈ℝd×d.S_ij= x_i Mx_j d_h, M=W^QW^K ^d× d. (4) The matrix M is the query-key interaction matrix. Unlike FFN layers (additive structure) or embeddings (element-wise), attention’s bilinear structure compounds magnitudes: large queries meeting large keys produce scores that can exceed the E4M3 threshold of 448. This makes attention the primary overflow bottleneck in FP8 training. To bound attention logits, we exploit this bilinear structure. Proofs for this section are in Appendix A. A natural first approach tracks ‖WQ‖2\|W^Q\|_2 and ‖WK‖2\|W^K\|_2 separately: Proposition 3.1 (Naive bound). Let ‖xi‖2≤BX\|x_i\|_2≤ B_X for all tokens. Then: maxi,j|Sij|≤‖WQ‖2‖WK‖2⋅BX2dh. _i,j|S_ij|≤ \|W^Q\|_2\|W^K\|_2· B_X^2 d_h. (5) This bound is loose because it treats WQW^Q and WKW^K independently. Analyzing M=WQWK⊤M=W^QW^K directly yields a tighter result: Proposition 3.2 (Interaction bound). Under the same conditions: maxi,j|Sij|≤‖WQWK⊤‖2⋅BX2dh. _i,j|S_ij|≤ \|W^QW^K \|_2· B_X^2 d_h. (6) Corollary 3.3 (Interaction bound is tighter). ‖WQWK⊤‖2≤‖WQ‖2‖WK‖2\|W^QW^K \|_2≤\|W^Q\|_2\|W^K\|_2, with equality iff the top right singular vectors of WQW^Q and WKW^K coincide. For pre-LN architectures (GPT-2, Llama, Mistral), LayerNorm or RMSNorm (Zhang and Sennrich, 2019) precedes attention, constraining ‖xi‖22≈d\|x_i\|_2^2≈ d and giving BX=dB_X= d. Substituting yields the worst-case bound: maxi,j|Sij|≤∥WQWK⊤∥2⋅ddh=:Bmax. _i,j|S_ij|≤\|W^QW^K \|_2· d d_h=:B_ . (7) Calibrating the bound. The bound BmaxB_ assumes input vectors align perfectly with the top singular vectors of M. In practice, this is overly pessimistic: the worst case requires inputs to find the direction of maximum amplification in a d-dimensional space, but LayerNorm-normalized inputs behave like random directions on the sphere. In high dimensions, random vectors are nearly orthogonal to any fixed direction, making such alignment exponentially unlikely. This concentration phenomenon motivates introducing a calibration factor α∈(0,1)α∈(0,1). The calibrated bound is: Bα:=α⋅Bmax=α⋅‖M‖2⋅ddh.B_α:=α· B_ =α·\|M\|_2· d d_h. (8) Smaller α yields a tighter bound and higher FP8 dynamic range utilization, but increases the probability that actual logits exceed the bound. The natural question is how to select α: small enough to be practical, yet large enough to guarantee safety. In the next section, we develop a probabilistic framework that provides a principled answer. 3.2 Rank-Aware Probabilistic Guarantee We derive a probabilistic bound on Pr(maxi,j|Sij|≥Bα) ( _i,j|S_ij|≥ B_α) that enables principled selection of the calibration factor α. Proofs for this section are in Appendix B. Assumption (Spherical token directions). For Pre-LN transformers, we model post-LayerNorm (or post-RMSNorm) token vectors as xi=duix_i= d\,u_i, where directions uiu_i are approximately isotropic (idealized as i.i.d. uniform on d−1S^d-1). For self-attention, where queries and keys derive from the same tokens, we treat each role’s projection direction as an independent draw. This assumption captures the effect of normalization in high dimensions and is the only distributional assumption required for Proposition 3.4. Remark (Diagonal terms). For diagonal terms (i=ji=j), the attention score involves the quadratic form ui⊤Muiu_i Mu_i rather than a bilinear form with independent arguments. Our analysis treats these as bilinear by assuming independence. This is conservative for two reasons: (i) diagonal terms constitute only L of L2L^2 query-key pairs, a vanishing fraction as L grows; and (i) quadratic forms u⊤Mu Mu for u uniform on the sphere exhibit tighter concentration than bilinear forms u⊤Mvu Mv with independent u,vu,v (the variance is smaller when both vectors are identical). Thus, the independence assumption yields an upper bound on overflow probability. Under this model, and exploiting the low-rank structure of the query–key interaction matrix rank(M)≤dh≪drank(M)≤ d_h d, we obtain a tail bound whose concentration exponent improves by a factor of d/(γdh)d/(γ d_h) over rank-agnostic arguments. Table 2 quantifies this improvement for representative architectures (see Appendix B.3 for derivation details). Table 2: Concentration exponent improvement from rank-aware bounds: d/(γdh)d/(γ d_h) tighter tail probabilities, where γ is determined by Eq. (12). Model d dhd_h γ Improvement GPT-2 XL 1600 64 2.98 8×8× Mistral-7B 4096 128 2.26 14×14× Llama-2-13B 5120 128 2.28 18×18× Llama-2-70B 8192 128 2.32 28×28× Our goal is a safety guarantee covering all tokens, heads, and layers. Accordingly, the analysis applies union bounds over L2L^2 query–key pairs and over all attention heads. This yields a conservative bound that remains robust under transient conditions (e.g., checkpoint loading and learning-rate discontinuities), while leaving slack in typical steady-state regimes. Proposition 3.4 (Rank-aware overflow probability bound). Let M=WQWK⊤M=W^QW^K with WQ,WK∈ℝd×dhW^Q,W^K ^d× d_h for a single attention head and rank(M)=dhrank(M)=d_h. For a sequence of length L and any γ>1γ>1, Pr(maxi,j|Sij|≥Bα)≤T1+T2, ( _i,j|S_ij|≥ B_α )\;≤\;T_1+T_2, -17.07182pt (9) where T1 T_1 =Lexp(−dh2(γ−1−lnγ)), =L \! (- d_h2(γ-1- γ) ), (10) T2 T_2 =2L2exp(−d2α22γdh). =2L^2 \! (- d^2α^22γ d_h ). (11) The term T1T_1 bounds the probability that any key projection is atypically large, while T2T_2 bounds overflow conditioned on typical key projections. The parameter γ>1γ>1 controls the typicality threshold. Extension to full transformers. Proposition 3.4 bounds overflow probability for a single attention head. For a transformer with N=nlayers×nheadsN=n_layers× n_heads heads, overflow occurs if any head overflows. Applying a union bound, Pr(overflow in any head)≤N(T1+T2). (overflow in any head)≤ N(T_1+T_2). To ensure this probability is at most δ∗δ^*, it suffices to require N⋅T1≤δ∗/2N· T_1≤δ^*/2 and N⋅T2≤δ∗/2N· T_2≤δ^*/2.222The equal allocation is a standard simplification; optimizing the split yields negligible improvement in practice. Solving these inequalities yields a principled selection rule for γ and α. Principled α selection. Step 1: Select γ. From N⋅T1≤δ∗/2N· T_1≤δ^*/2 and (10), we require h(γ):=γ−1−lnγ≥2dhln(2NLδ∗).h(γ):=γ-1- γ\;≥\; 2d_h \! ( 2NLδ^* ). (12) Step 2: Select α. From N⋅T2≤δ∗/2N· T_2≤δ^*/2 and (11), we require α≥αmin:=2γdhdln(4NL2δ∗).α\;≥\; _ := 2γ d_hd \! ( 4NL^2δ^* ). (13) Table 3 reports αmin _ for our evaluated models with δ∗=10−6δ^*=10^-6 and L=1024L=1024. Table 3: Minimum calibration factor αmin _ for δ∗=10−6δ^*=10^-6 and L=1024L=1024. Model d dhd_h N αmin _ GPT-2 XL 1600 64 1200 0.074 Mistral-7B 4096 128 1024 0.035 Llama-2-13B 5120 128 1600 0.028 Llama-2-70B 8192 128 5120 0.018 Selecting α in practice. We set α slightly above αmin _ for safety margin. Larger models permit smaller α due to stronger concentration arising from higher d/dhd/ d_h ratios. In our experiments, we use α=0.08α=0.08 for GPT-2 XL, α=0.04α=0.04 for Mistral-7B, α=0.03α=0.03 for Llama-2-13B, and α=0.02α=0.02 for Llama-2-70B, each exceeding the corresponding αmin _ . Transient vs. steady-state regimes. Eq. (13) yields a conservative αmin _ designed for transient-prone workflows (pretrained loading, checkpoint resumption without FP8 state, and LR discontinuities) where a single overflow can corrupt training. For steady-state fine-tuning, we optionally tighten α using a one-time empirical calibration (Sec. 3.5), then freeze it and revert to fully predictive scaling. Modern architectures (Llama, Mistral) use Rotary Position Embeddings (RoPE), which apply position-dependent rotations to query and key vectors. We show our bounds extend without modification. 3.3 Extension to Rotary Position Embeddings Rotary Position Embeddings (Su et al., 2021) apply rotation matrices to query and key vectors before computing attention: Smn=(Rmqm)⊤(Rnkn)dh=qm⊤Rm⊤Rnkndh,S_mn= (R_mq_m) (R_nk_n) d_h= q_m R_m R_nk_n d_h, (14) where Rm,Rn∈ℝdh×dhR_m,R_n ^d_h× d_h are position-dependent block-diagonal rotation matrices. The following results, proved in Appendix C, show these rotations cannot amplify scores beyond our bounds. Proposition 3.5 (RoPE preserves norms). Let Rθ∈ℝdh×dhR_θ ^d_h× d_h be any RoPE rotation matrix. Then: 1. RθR_θ is orthogonal: Rθ⊤Rθ=IR_θ R_θ=I. 2. ‖Rθ‖2=1\|R_θ\|_2=1. 3. For any q,k∈ℝdhq,k ^d_h: |(Rmq)⊤(Rnk)|≤‖q‖⋅‖k‖|(R_mq) (R_nk)|≤\|q\|·\|k\|. Corollary 3.6 (Geometry-aware scaling extends to RoPE). For RoPE-based attention, the worst-case bound |Smn|≤‖WQ‖2‖WK‖2⋅d/dh|S_mn|≤\|W^Q\|_2\|W^K\|_2· d/ d_h holds rigorously for all positions via ‖Rm‖2=1\|R_m\|_2=1. The tighter interaction bound ‖WQWK⊤‖2\|W^QW^K \|_2 applies when RoPE rotations do not systematically align with the singular subspaces of WQW^Q and WKW^K, a condition we verified empirically across all layers in Mistral-7B, Llama-2-13B, and Llama-2-70B. The probabilistic bound (Proposition 3.4) also extends conservatively; see Appendix C.4. 3.4 Geometry-Aware Scale Factor Spectral norms vary by up to 10×10× across layers (Appendix D), so a global scale would severely underutilize FP8 precision. We therefore compute scale factors per layer. For layer ℓ , let σQK(ℓ)=‖W(ℓ)QW(ℓ)K⊤‖2 _QK^( )=\|W^Q_( )W^K _( )\|_2 denote the spectral norm of the interaction matrix. The calibrated bound for layer ℓ is: Bα(ℓ)=α⋅σQK(ℓ)⋅ddh.B_α^( )=α· _QK^( )· d d_h. To map this bound to a scale factor, we define Rsafe=ηfp8⋅448R_ safe= _ fp8· 448, where ηfp8<1 _ fp8<1 provides margin below the E4M3 maximum. The geometry-aware scale factor is: scale(ℓ)=Bα(ℓ)Rsafe=α⋅σQK(ℓ)⋅d/dhηfp8⋅448. scale^( )= B_α^( )R_ safe= α· _QK^( )· d/ d_h _ fp8· 448. (15) During the forward pass, pre-softmax attention scores S(ℓ)=Q(ℓ)K(ℓ)⊤/dhS^( )=Q^( )K^( ) / d_h are divided by scale(ℓ) scale^( ) before FP8 quantization, ensuring scaled logits lie within the representable range. In all experiments, we set ηfp8=0.8 _ fp8=0.8 and select α per model following Section 3.2. For architectures with grouped query attention (GQA), where multiple query heads share key-value heads, we compute σQK(ℓ) _QK^( ) without expanding the key matrix; this implicit formulation is described in Section 4.2. 3.5 Auto-α Calibration for Steady-State Training The conservative α from Proposition 3.4 guarantees overflow probability below δ∗=10−6δ^*=10^-6, appropriate for transient-prone workflows. However, during steady-state fine-tuning, actual logits occupy a small fraction of the theoretical envelope. Auto-α is an optional enhancement for practitioners who prioritize utilization over worst-case robustness. We define the slack ratio at step t as: rt=maxi,j|Sij(t)|Bmax,r_t= _i,j|S_ij^(t)|B_ , (16) where the numerator is the observed maximum logit. Algorithm. Auto-α calibration proceeds in two phases: 1. Burn-in (steps 11 to TcalibT_calib): Run with conservative α0 _0 from Proposition 3.4. Collect slack ratios rt\r_t\. 2. Calibration: Set αfinal=P99.99(rt)×κ _final=P_99.99(\r_t\)×κ, where κ≥1κ≥ 1 is a safety multiplier. The calibrated αfinal _final is then frozen for all subsequent training. This is distinct from delayed scaling: we measure slack to select a static bound, not to continuously adapt. After burn-in, the method is again fully predictive. Theoretical guarantee. When using auto-α, the probabilistic guarantee from Proposition 3.4 no longer applies, since the calibration factor is determined empirically rather than from the theoretical selection rule (Eq. 13). Auto-α trades the worst-case guarantee for higher utilization, relying on the assumption that the burn-in distribution is representative of subsequent training. This is appropriate for steady-state fine-tuning but not for workflows with anticipated distribution shifts. Memory-efficient attention compatibility. Auto-α calibration requires observing maxi,j|Sij| _i,j|S_ij| during the burn-in phase, which necessitates materializing the score matrix and is incompatible with FlashAttention. For a 100-step burn-in out of thousands of training steps, this overhead is negligible (<0.1%<0.1\% of total compute). After burn-in, the calibrated α is frozen and the method reverts to fully predictive scaling compatible with memory-efficient attention. 4 Efficient Estimation of Spectral Bounds This section presents efficient algorithms for computing the spectral bounds. 4.1 Spectral Norm Estimation via Power Iteration Computing ‖M‖2\|M\|_2 where M=WQWK⊤∈ℝd×dM=W^QW^K ^d× d via full SVD costs O(d3)O(d^3), which is prohibitive for large models. For Llama-2-70B with d=8192d=8192, this would require forming a 6767 million entry matrix per layer. We instead use power iteration (Golub and Van Loan, 2013), a classical algorithm that estimates the top singular value through repeated matrix-vector multiplication. The key insight is that we never form M explicitly. Starting from random unit vectors u,v∈ℝdu,v ^d, we alternate: u←Mv‖Mv‖,thenv←M⊤u‖M⊤u‖.u← Mv\|Mv\|, v← M u\|M u\|. (17) After these updates, ‖Mv‖\|Mv\| provides an estimate of ‖M‖2\|M\|_2. Each matrix-vector product is computed implicitly: Mv Mv =WQ(WK⊤v), =W^Q(W^K v), (18) M⊤u M u =WK(WQ⊤u), =W^K(W^Q u), (19) requiring O(nheads⋅dh⋅d)O(n_heads· d_h· d) operations per iteration, where nheadsn_heads is the total number of attention heads. This is far cheaper than explicitly forming M∈ℝd×dM ^d× d, which would require O(nheads⋅dh⋅d2)O(n_heads· d_h· d^2) operations and O(d2)O(d^2) memory. We maintain the persistent vectors u,vu,v across training steps. During steady-state training, weights change gradually (controlled by the learning rate), so singular vectors drift slowly and one iteration per forward pass suffices to track them. For cold-start scenarios (initialization, checkpoint loading), we run 5 iterations to ensure convergence from random vectors. Remark (Transient robustness). Even during rapid weight changes (e.g., learning rate spikes), any underestimation of the spectral norm is bounded. If the true spectral norm increases by factor ρ from step t to t+1t+1, power iteration with warm-start vectors underestimates by at most ρ until the next update. Since our calibration factor α already provides margin above αmin _ (Section 3.2), moderate underestimation does not cause overflow. Our experiments with 100×100× learning rate spikes (Section 5.2) confirm zero overflows under this regime. The detailed algorithm is provided in Appendix E. 4.2 Implicit Formulation for Grouped Query Attention Modern architectures use Grouped Query Attention (GQA) (Ainslie et al., 2023), where multiple query heads share key-value heads. For example, Mistral-7B has nq=32n_q=32 query heads but only nkv=8n_kv=8 key-value heads (a 4:1 ratio). This creates a dimension mismatch: WQ∈ℝd×(nq⋅dh)W^Q ^d×(n_q· d_h) while WK∈ℝd×(nkv⋅dh)W^K ^d×(n_kv· d_h). To compute ‖WQWK⊤‖2\|W^QW^K \|_2 with matching dimensions, a naive approach would expand WKW^K by replicating each block of dhd_h columns g=nq/nkvg=n_q/n_kv times to form WexpK∈ℝd×(nq⋅dh)W^K_exp ^d×(n_q· d_h). This expansion consumes substantial memory: 32MB per layer on Mistral-7B. We observe that power iteration requires only matrix-vector products, not the matrices themselves. For the forward product MvMv where M=WQWexpK⊤M=W^QW^K _exp: Mv=WQ⋅RepeatBlocks(WK⊤v,g),Mv=W^Q· RepeatBlocks(W^K v,g), where RepeatBlocks(z,g) RepeatBlocks(z,g) replicates each dhd_h-block of z∈ℝnkv⋅dhz ^n_kv· d_h exactly g times to produce a vector in ℝnq⋅dhR^n_q· d_h. The reverse product M⊤uM u uses a dual operation: M⊤u=WK⋅SumGroups(WQ⊤u,g)M u=W^K· SumGroups(W^Q u,g), where SumGroups sums each group of g blocks. These operations require replicating only small intermediate vectors (nkv⋅dhn_kv· d_h elements) rather than the full weight matrix. Proposition 4.1 (Implicit GQA power iteration). Power iteration using the unexpanded WKW^K converges to the same spectral norm ‖WQWexpK⊤‖2\|W^QW^K _ exp\|_2 as explicit expansion. The proof and detailed algorithm are in Appendix F. By avoiding the expanded matrix, we reduce memory transactions by a factor of g (4–8× on modern GQA architectures). This saving can offset the cost of power iteration, yielding negligible or even negative overhead in practice (Table 9). 4.3 Complete Algorithm Algorithm 1 presents the complete forward pass. The algorithm is predictive (scale depends only on weights), fused-compatible (no activation observation), and per-layer. Algorithm 1 Geometry-Aware Attention Forward Pass 1: Input: Token embeddings X∈ℝL×dX ^L× d, layer weights WQ,WK,WV,WOW^Q,W^K,W^V,W^O 2: Parameters: Calibration factor α, FP8 margin ηfp8 _fp8 3: State: Persistent vectors u,vu,v for power iteration 4: 5: Stage 1: Estimate spectral norm from weights 6: σQK←PowerIteration(WQ,WK,u,v) _QK← PowerIteration(W^Q,W^K,u,v) 7: 8: Stage 2: Compute predictive scale factor 9: Bα←α⋅σQK⋅d/dhB_α←α· _QK· d/ d_h 10: scale←Bα/(ηfp8⋅448)scale← B_α\;/\;( _fp8· 448) 11: 12: Stage 3: Attention with pre-computed scale 13: Q,K,V←XWQ,XWK,XWVQ,K,V← XW^Q,XW^K,XW^V 14: S←QK⊤/dhS← QK / d_h 15: S~←S/scale S← S\;/\;scale 16: return softmax(S~)⋅V⋅WOsoftmax( S)· V· W^O 5 Experiments We evaluate geometry-aware scaling across four transformer architectures spanning 1.5B to 70B parameters. Our experiments address three questions: (1) Does geometry-aware scaling eliminate overflows in transient scenarios where delayed scaling fails? (2) Does the safety–utilization tradeoff affect downstream task quality? (3) What is the computational cost? The key property enabling transient safety is that geometry-aware scaling responds instantaneously to weight changes: the scale factor is computed from current weights, not historical activations. Appendix H validates this with a controlled experiment where we spike attention weights by 4×4×; delayed scaling overflows catastrophically while geometry-aware scaling adapts in the same forward pass. The following sections demonstrate that this instantaneous response eliminates overflow in realistic training scenarios. 5.1 Experimental Setup Models. We evaluate on four models: GPT-2 XL (1.5B parameters, 48 layers), Mistral-7B (32 layers), Llama-2-13B (40 layers), and Llama-2-70B (80 layers). The models include both standard multi-head attention (GPT-2 XL, Llama-2-13B) and grouped query attention (Mistral-7B with 4:1 ratio, Llama-2-70B with 8:1 ratio), enabling evaluation of the implicit GQA formulation. Architectural specifications and hardware details are in Appendix G.1. Implementation details. All experiments use FP8 E4M3 format with margin ηfp8=0.8 _fp8=0.8, giving a safe range of 0.8×448=358.40.8× 448=358.4. Calibration factors are set per model following the procedure in Section 3.2: α=0.08α=0.08 for GPT-2 XL, 0.040.04 for Mistral-7B, 0.030.03 for Llama-2-13B, and 0.020.02 for Llama-2-70B. Power iteration uses persistent vectors with one update per forward pass; we initialize with 5 iterations at the start of training or after checkpoint loading to ensure convergence from random initialization. 5.2 Transient Scenarios We evaluate three transient scenarios that expose the history staleness problem in delayed scaling. Loading pretrained models. The most common scenario is loading pretrained weights into an FP8 training pipeline. The history buffer initializes to default values while weights reflect extensive prior training, creating immediate mismatch. Table 4 shows overflow statistics on the first forward pass after loading pretrained weights. Delayed scaling overflows on 100% of layers across all four models, with maximum scaled logits exceeding 5000. Geometry-aware scaling achieves zero overflows because the scale factor is computed directly from the loaded weights in the same forward pass. Table 4: First forward pass after loading pretrained weights. Overflow: layers exceeding 448. Max Scaled: maximum scaled logit. Delayed Ours Model Overfl. Max Scaled Overfl. Max Scaled GPT-2 XL 48/48 8761 0/48 120 Mistral-7B 32/32 7123 0/32 196 Llama-2-13B 40/40 5600 0/40 2.5 Llama-2-70B 80/80 9498 0/80 1.0 Checkpoint resumption. Training large models involves interruptions due to job preemption, hardware failures, or scheduled maintenance. Upon resumption, the history buffer resets to default values while model weights reflect the training state at the checkpoint. We simulate this by training for 300 steps, saving a checkpoint, then resuming with a fresh history buffer. In the first 10 steps after resumption, delayed scaling overflows on 4 steps for GPT-2 XL, 2 steps for Llama-2-70B, and 1 step each for Mistral-7B and Llama-2-13B. Geometry-aware scaling maintains zero overflows across all models because the scale factor depends only on current weights, which are restored correctly from the checkpoint. While checkpointing scaling state is possible, standard frameworks do not include it by default, and loading pretrained weights from external sources provides no history at all. Learning rate transitions. Learning rate schedules involve sudden changes: warmup phases ramp the rate over initial steps, and cyclic schedules spike periodically. These transitions cause weights to change faster than the history buffer can track. We evaluate with a 100×100× spike: learning rate 10−510^-5 for 100 steps, then 10−310^-3. In the 10 steps following the transition, delayed scaling overflows on 3 to 5 steps depending on the model (5 for Mistral-7B, 4 for Llama-2-70B, 3 for Llama-2-13B). On GPT-2 XL, NaN gradients appear immediately, terminating training. Geometry-aware scaling handles the transition with zero overflows across all models. 5.3 Computational Overhead Geometry-aware scaling requires computing spectral norms via power iteration, adding computational cost. We measure forward pass timing averaged over 100 iterations (detailed methodology in Appendix I). On MHA architectures, overhead is minimal: +1.0%+1.0\% for GPT-2 XL and +1.9%+1.9\% for Llama-2-13B. On Mistral-7B with GQA (4:1 ratio), we observe negative overhead of −5.3%-5.3\%: geometry-aware scaling is faster than delayed scaling in our implementation. We attribute this to our implicit GQA formulation (Section 4.2), which has favorable memory access patterns, combined with synchronization overhead in the baseline’s history buffer updates. However, this result is implementation-specific; the key finding is that overhead is negligible across all architectures. Llama-2-70B shows +4.3%+4.3\% due to its 80 layers; full details in Appendix I. 5.4 Downstream Task Evaluation Training loss measures average next-token prediction but may mask degradation in reasoning. We evaluate on MMLU (Hendrycks et al., 2021) to verify that geometry-aware scaling preserves downstream capability. We fine-tune Llama-2-13B for 3000 steps on MMLU STEM subjects (17 subjects, 295 training examples) using three methods: delayed scaling, conservative spectral scaling (α=0.03α=0.03), and auto-α spectral scaling (100-step burn-in, κ=1.0κ=1.0). All methods use identical hyperparameters: learning rate 10−510^-5, batch size 4, sequence length 1024. We evaluate on held-out test sets (2783 examples). Table 5 shows training metrics and downstream accuracy. All three methods converge to similar training loss (≈0.012≈ 0.012), yet downstream accuracy differs substantially. Conservative spectral scaling achieves zero overflows but degrades accuracy by 5 percentage points (28.1% vs. 33.0%), indicating that 0.5% FP8 utilization introduces excessive quantization noise, losing fine-grained distinctions between attention logits even when average loss remains low. Auto-α calibration resolves this tradeoff: by tightening α from 0.03 to 0.00036 based on the 99.99th percentile of observed slack ratios, utilization increases to 31.2%, achieving comparable accuracy (33.6%) while maintaining zero overflows. Caveats. The primary contribution of this evaluation is demonstrating that geometry-aware scaling eliminates overflows without catastrophic accuracy degradation, not that it improves accuracy over the baseline. The 0.6 percentage point difference between auto-α (33.6%) and delayed scaling (33.0%) is within the variance expected from stochastic optimization on a small training set (295 examples); we do not claim statistical significance for this difference. For the delayed scaling baseline, overflows (68 occurrences) were handled by clamping to the representable range, allowing training to continue; without clamping, NaN propagation would terminate training entirely. Table 5: Training metrics and MMLU STEM accuracy on Llama-2-13B. Auto-α achieves comparable accuracy to the baseline while eliminating all overflows. Method Loss Overfl. Util. MMLU Delayed 0.0116 68 70.7% 33.0% Ours (conservative) 0.0119 0 0.5% 28.1% Ours + auto-α 0.0116 0 31.2% 33.6% These results suggest a practical workflow: conservative α for transient-prone scenarios (checkpoint loading, learning rate warmup) where safety is paramount, and auto-α for steady-state fine-tuning where utilization matters. Training loss curves are provided in Appendix K, per-subject accuracy in Appendix L, and FP8 utilization statistics in Appendix J. 6 Conclusion We established spectral bounds on transformer attention logits and derived a rank-aware probabilistic calibration framework for low-precision training. The key theoretical result is maxi,j|Sij|≤‖WQWK⊤‖2⋅d/dh _i,j|S_ij|≤\|W^QW^K \|_2· d/ d_h, which is never looser than the naive submultiplicative bound and is strictly tighter unless the top right singular vectors of WQW^Q and WKW^K align. This bound, combined with our rank-aware probabilistic analysis, provides principled calibration given a target failure probability δ∗δ^*. The bound is efficiently computable via power iteration and extends conservatively to RoPE architectures. Experiments demonstrate zero overflows across GPT-2 XL through Llama-2-70B on all transient scenarios where delayed scaling fails, with negligible computational overhead. Auto-α calibration resolves the utilization-accuracy trade-off: by learning a data-driven bound during burn-in, we achieve comparable MMLU accuracy to delayed scaling while eliminating all overflows. Impact Statement This work improves reliability and efficiency of FP8 training; we do not anticipate ethical risks beyond those inherent to large-scale ML training. References J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai (2023) GQA: training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, p. 4895–4901. Cited by: §4.2. C. Blake, D. Orr, and C. Luschi (2023) Unit scaling: out-of-the-box low-precision training. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, p. 2548–2576. Cited by: §2. K. Choi, H. Lee, S. Park, D. Kwon, and J. Lee (2025) FALQON: accelerating LoRA fine-tuning with low-bit floating-point arithmetic. In Advances in Neural Information Processing Systems, Vol. 38. Cited by: §2. T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré (2022) FlashAttention: fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, Vol. 35, p. 16344–16359. Cited by: §1. T. Dao (2024) FlashAttention-2: faster attention with better parallelism and work partitioning. In International Conference on Learning Representations, Cited by: §1. DeepSeek-AI (2024) DeepSeek-v3 technical report. arXiv preprint arXiv:2412.19437. Cited by: §2. G. H. Golub and C. F. Van Loan (2013) Matrix computations. 4th edition, Johns Hopkins University Press. Cited by: §1, §4.1. D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt (2021) Measuring massive multitask language understanding. In International Conference on Learning Representations, Cited by: §5.4. A. Hernández-Cano, D. Garbaya, I. Schlag, and M. Jaggi (2025) Towards fully FP8 GEMM LLM training at scale. In Advances in Neural Information Processing Systems, Vol. 38. Cited by: §2. B. Laurent and P. Massart (2000) Adaptive estimation of a quadratic functional by model selection. Annals of Statistics 28 (5), p. 1302–1338. Cited by: §B.1. H. Lee, O. Kwon, Y. Park, and J. W. Lee (2025) NestedFP: high-performance, memory-efficient dual-precision floating point support for LLMs. In Advances in Neural Information Processing Systems, Vol. 38. Cited by: §2. P. Micikevicius, D. Stosic, N. Burgess, M. Cornea, P. Dubey, R. Grisenthwaite, S. Ha, A. Heinecke, P. Judd, J. Kamalu, et al. (2022) FP8 formats for deep learning. arXiv preprint arXiv:2209.05433. Cited by: §G.1, §1, §2, footnote 1. T. Miyato, T. Kataoka, M. Koyama, and Y. Yoshida (2018) Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, External Links: Link Cited by: §2. B. D. Rouhani, R. Zhao, V. Elango, R. Shafipour, M. D. Rouhani, et al. (2023) Microscaling data formats for deep learning. arXiv preprint arXiv:2310.10537. Cited by: §2. J. Shah, G. Bikshandi, Y. Zhang, V. Thakkar, P. Ramani, and T. Dao (2024) FlashAttention-3: fast and accurate attention with asynchrony and low-precision. In Advances in Neural Information Processing Systems, Vol. 37. Cited by: §2. J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu (2021) RoFormer: enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864. Cited by: §3.3. S. Zhai, T. Likhomanenko, E. Littwin, D. Busbridge, J. Ramapuram, Y. Zhang, J. Gu, and J. M. Susskind (2023) Stabilizing Transformer training by preventing attention entropy collapse. In Proceedings of the 40th International Conference on Machine Learning, p. 40770–40803. Cited by: §2. B. Zhang and R. Sennrich (2019) Root mean square layer normalization. In Advances in Neural Information Processing Systems, Vol. 32. Cited by: §1, §3.1. Y. Zhang, H. Zhen, M. Yuan, and B. Yu (2025) MOSS: efficient and accurate FP8 LLM training with microscaling and automatic scaling. arXiv preprint arXiv:2511.05811. Cited by: §2. Appendix A Proofs for Section 3.1 This appendix contains the proofs for propositions and corollaries stated in Section 3.1. A.1 Proof of Proposition 3.1 (Naive Bound) Proof. By the Cauchy-Schwarz inequality: |Sij|=|qi⊤kj|dh≤‖qi‖2‖kj‖2dh.|S_ij|= |q_i k_j| d_h≤ \|q_i\|_2\|k_j\|_2 d_h. Since qi=WQ⊤xiq_i=W^Q x_i and kj=WK⊤xjk_j=W^K x_j, submultiplicativity of the spectral norm gives: ‖qi‖2=‖WQ⊤xi‖2≤‖WQ⊤‖2‖xi‖2=‖WQ‖2BX,\|q_i\|_2=\|W^Q x_i\|_2≤\|W^Q \|_2\|x_i\|_2=\|W^Q\|_2B_X, and similarly ‖kj‖2≤‖WK‖2BX\|k_j\|_2≤\|W^K\|_2B_X. Combining these: |Sij|≤‖WQ‖2‖WK‖2⋅BX2dh.∎|S_ij|≤ \|W^Q\|_2\|W^K\|_2· B_X^2 d_h. A.2 Proof of Proposition 3.2 (Interaction Bound) Proof. Define M=WQWK⊤∈ℝd×dM=W^QW^K ^d× d. The attention score can be written as a bilinear form: Sij=xi⊤WQWK⊤xjdh=xi⊤Mxjdh.S_ij= x_i W^QW^K x_j d_h= x_i Mx_j d_h. By the variational characterization of the spectral norm: ‖M‖2=max‖u‖2=‖v‖2=1|u⊤Mv|.\|M\|_2= _\|u\|_2=\|v\|_2=1|u Mv|. Therefore, for any vectors xi,xjx_i,x_j: |xi⊤Mxj|≤‖M‖2⋅‖xi‖2⋅‖xj‖2≤‖M‖2⋅BX2,|x_i Mx_j|≤\|M\|_2·\|x_i\|_2·\|x_j\|_2≤\|M\|_2· B_X^2, which gives |Sij|≤‖M‖2⋅BX2/dh|S_ij|≤\|M\|_2· B_X^2/ d_h. ∎ A.3 Proof of Corollary 3.3 (Interaction Bound is Tighter) Proof. The inequality ‖WQWK⊤‖2≤‖WQ‖2‖WK‖2\|W^QW^K \|_2≤\|W^Q\|_2\|W^K\|_2 follows directly from submultiplicativity of the spectral norm: for any matrices A,BA,B of compatible dimensions, ‖AB‖2≤‖A‖2‖B‖2\|AB\|_2≤\|A\|_2\|B\|_2. For the equality condition, let WQ=UQΣQVQ⊤W^Q=U_Q _QV_Q and WK=UKΣKVK⊤W^K=U_K _KV_K be the singular value decompositions. Then: WQWK⊤=UQΣQVQ⊤VKΣKUK⊤.W^QW^K =U_Q _QV_Q V_K _KU_K . The spectral norm of this product equals σ1(WQ)⋅σ1(WK) _1(W^Q)· _1(W^K) if and only if the (1,1)(1,1) entry of VQ⊤VKV_Q V_K has absolute value 11, which requires the top right singular vectors to coincide: v1Q=±v1Kv_1^Q=± v_1^K. Under random initialization, singular vectors are uniformly distributed on the sphere, making exact alignment a measure-zero event. ∎ Appendix B Proofs for Section 3.2 This appendix contains the proofs for Proposition 3.4, as well as supporting lemmas. B.1 Supporting Lemmas We first establish a tail bound for the projection of a uniform random vector onto a subspace. Lemma B.1 (Projection norm distribution). Let u∼Unif(d−1)u (S^d-1) and let V∈ℝd×kV ^d× k have orthonormal columns spanning a k-dimensional subspace. Then: ‖V⊤u‖22∼Beta(k2,d−k2),[‖V⊤u‖22]=kd.\|V u\|_2^2 ( k2, d-k2 ), [\|V u\|_2^2]= kd. Proof. The uniform distribution on d−1S^d-1 can be generated as u=g/‖g‖2u=g/\|g\|_2 where g∼(0,Id)g (0,I_d). Since V has orthonormal columns, V⊤g∼(0,Ik)V g (0,I_k), and let V⟂∈ℝd×(d−k)V_ ^d×(d-k) complete an orthonormal basis, so V⟂⊤g∼(0,Id−k)V_ g (0,I_d-k) independently. Then ‖V⊤g‖22∼χk2\|V g\|_2^2 χ^2_k and ‖V⟂⊤g‖22∼χd−k2\|V_ g\|_2^2 χ^2_d-k are independent, and: ‖V⊤u‖22=‖V⊤g‖22‖g‖22=‖V⊤g‖22‖V⊤g‖22+‖V⟂⊤g‖22.\|V u\|_2^2= \|V g\|_2^2\|g\|_2^2= \|V g\|_2^2\|V g\|_2^2+\|V_ g\|_2^2. The ratio of independent chi-squared variables χk2/(χk2+χd−k2)χ^2_k/(χ^2_k+χ^2_d-k) follows Beta(k/2,(d−k)/2)Beta(k/2,(d-k)/2). ∎ Lemma B.2 (Chi-squared ratio tail bound). Let X∼Beta(k/2,(d−k)/2)X (k/2,(d-k)/2). For any γ>1γ>1: Pr(X≥γ⋅kd)≤exp(−k2(γ−1−lnγ)). (X≥γ· kd )≤ (- k2(γ-1- γ) ). Proof. By Lemma B.1, X∼Beta(k/2,(d−k)/2)X (k/2,(d-k)/2) with mean μ=k/dμ=k/d. The stated bound follows from the multiplicative Chernoff bound for Beta distributions. For X∼Beta(a,b)X (a,b) and γ>1γ>1, the moment generating function yields: Pr(X≥γ⋅a+b)≤exp(−a⋅h(γ)), (X≥γ· aa+b )≤ (-a· h(γ) ), where h(γ)=γ−1−lnγ>0h(γ)=γ-1- γ>0 for γ>1γ>1. This is the Beta analogue of the chi-squared Chernoff bound (Laurent and Massart, 2000, Lemma 1) and follows from the representation of Beta as a ratio of independent Gamma variables. Setting a=k/2a=k/2 and noting that a/(a+b)=k/da/(a+b)=k/d: Pr(X≥γ⋅kd)≤exp(−k2(γ−1−lnγ)).∎ (X≥γ· kd )≤ (- k2(γ-1- γ) ). B.2 Proof of Proposition 3.4 (Rank-Aware Overflow Bound) Proof. Under the spherical assumption, token vectors satisfy xi=d⋅uix_i= d· u_i where ui∼Unif(d−1)u_i (S^d-1) independently. The attention score becomes: Sij=ddh⋅ui⊤Muj,S_ij= d d_h· u_i Mu_j, and the calibrated bound is Bα=α⋅‖M‖2⋅d/dhB_α=α·\|M\|_2· d/ d_h. We must bound: Pr(maxi,j|Sij|≥Bα)=Pr(maxi,j|ui⊤Muj|≥α‖M‖2). ( _i,j|S_ij|≥ B_α )= ( _i,j|u_i Mu_j|≥α\|M\|_2 ). Step 1: SVD decomposition. Let M=UΣV⊤M=U V be the SVD with rank(M)=dhrank(M)=d_h, where V∈ℝd×dhV ^d× d_h contains the right singular vectors. For any unit vector uju_j: Muj=UΣV⊤uj,Mu_j=U V u_j, so ‖Muj‖2=‖ΣV⊤uj‖2≤σ1‖V⊤uj‖2=‖M‖2‖V⊤uj‖2\|Mu_j\|_2=\| V u_j\|_2≤ _1\|V u_j\|_2=\|M\|_2\|V u_j\|_2. Step 2: Define typical keys. For γ>1γ>1, define the typical key event: ℰj:=‖V⊤uj‖22≤γ⋅dhd,E_j:= \\|V u_j\|_2^2≤γ· d_hd \, which implies ∥Muj∥2≤∥M∥2γdh/d=:β∥M∥2\|Mu_j\|_2≤\|M\|_2 γ d_h/d=:β\|M\|_2 where β=γdh/dβ= γ d_h/d. Let ℰ=⋂j=1LℰjE= _j=1^LE_j be the event that all keys are typical. Step 3: Bound probability of atypical keys. By Lemmas B.1 and B.2: Pr(ℰjc) (E_j^c) =Pr(‖V⊤uj‖22≥γ⋅dhd) = (\|V u_j\|_2^2≥γ· d_hd ) ≤exp(−dh2(γ−1−lnγ)). ≤ (- d_h2(γ-1- γ) ). By union bound over L keys: Pr(ℰc)≤Lexp(−dh2(γ−1−lnγ))=T1. (E^c)≤ L (- d_h2(γ-1- γ) )=T_1. Step 4: Concentration given typical keys. On event ℰE, for each j we have zj:=Mujz_j:=Mu_j with ‖zj‖2≤β‖M‖2\|z_j\|_2≤β\|M\|_2. For fixed zjz_j, the function f(ui)=ui⊤zjf(u_i)=u_i z_j on d−1S^d-1 is ‖zj‖2\|z_j\|_2-Lipschitz. By Lévy’s lemma, for ui∼Unif(d−1)u_i (S^d-1): Pr(|ui⊤zj|≥t) (|u_i z_j|≥ t) ≤2exp(−dt22‖zj‖22) ≤ 2 (- dt^22\|z_j\|_2^2 ) ≤2exp(−dt22β2‖M‖22). ≤ 2 (- dt^22β^2\|M\|_2^2 ). Setting t=α‖M‖2t=α\|M\|_2: Pr(|ui⊤Muj|≥α‖M‖2∣ℰ) (|u_i Mu_j|≥α\|M\|_2 ) ≤2exp(−dα22β2) ≤ 2 (- dα^22β^2 ) =2exp(−d2α22γdh). =2 (- d^2α^22γ d_h ). Step 5: Union bound over all pairs. By union bound over L2L^2 query-key pairs: Pr(maxi,j|ui⊤Muj|≥α‖M‖2∣ℰ) ( _i,j|u_i Mu_j|≥α\|M\|_2 ) ≤2L2exp(−d2α22γdh) ≤ 2L^2 (- d^2α^22γ d_h ) =T2. =T_2. Step 6: Combine bounds. By the law of total probability: Pr(maxi,j|ui⊤Muj|≥α‖M‖2) ( _i,j|u_i Mu_j|≥α\|M\|_2 ) ≤Pr(ℰc)+Pr(maxi,j|ui⊤Muj|≥α‖M‖2∣ℰ) ≤ (E^c)+ ( _i,j|u_i Mu_j|≥α\|M\|_2 ) ≤T1+T2.∎ ≤ T_1+T_2. B.3 Comparison with Rank-Agnostic Concentration This section compares the rank-aware concentration bound in Proposition 3.4 with a baseline bound obtained by applying standard concentration inequalities without exploiting the low-rank structure of the query–key interaction matrix. Rank-agnostic baseline. Without exploiting the rank constraint rank(M)=dhrank(M)=d_h, we apply Lévy’s lemma directly. For fixed uju_j, the function f(ui)=ui⊤Mujf(u_i)=u_i Mu_j on the unit sphere d−1S^d-1 is ‖Muj‖2\|Mu_j\|_2-Lipschitz. Since ‖Muj‖2≤‖M‖2\|Mu_j\|_2≤\|M\|_2, Lévy’s lemma yields Pr(|ui⊤Muj|≥α‖M‖2)≤2exp(−dα22). (|u_i Mu_j|≥α\|M\|_2 )≤ 2 \! (- dα^22 ). Applying a union bound over all L2L^2 query–key pairs gives Pr(maxi,j|ui⊤Muj|≥α‖M‖2)≤2L2exp(−dα22). \! ( _i,j|u_i Mu_j|≥α\|M\|_2 )≤ 2L^2 \! (- dα^22 ). Rank-aware improvement. In contrast, Proposition 3.4 exploits the low-rank structure rank(M)=dhrank(M)=d_h to obtain the term T2=2L2exp(−d2α22γdh).T_2=2L^2 \! (- d^2α^22γ d_h ). The concentration exponent in the rank-aware bound is therefore larger by a factor of d2α2/(2γdh)dα2/2=dγdh, d^2α^2/(2γ d_h)dα^2/2= dγ d_h, relative to the rank-agnostic baseline. This comparison isolates the quantitative benefit of exploiting rank(M)=dh≪drank(M)=d_h d: for fixed α, the overflow probability decays exponentially faster in dimension when the low-rank structure of the query–key interaction is taken into account. Appendix C Proofs for Section 3.3 This appendix contains the proofs for Proposition 3.5 and Corollary 3.6, establishing that geometry-aware scaling extends to architectures using Rotary Position Embeddings. C.1 Background on RoPE RoPE applies position-dependent rotations to query and key vectors. For head dimension dhd_h, the rotation matrix at position θ is block-diagonal: Rθ=(Rθ1⋱Rθdh/2),R_θ= pmatrixR_ _1&&\\ & &\\ &&R_ _d_h/2 pmatrix, where each 2×22× 2 block is a rotation matrix: Rθi=(cosθi−sinθisinθicosθi).R_ _i= pmatrix _i&- _i\\ _i& _i pmatrix. The angles θi _i depend on both the position index m and the dimension index i, typically as θi=m⋅ωi _i=m· _i where ωi _i are fixed frequencies. C.2 Proof of Proposition 3.5 (RoPE Preserves Norms) Proof. We prove each property in turn. (1) Orthogonality. Each 2×22× 2 block RθiR_ _i is a rotation matrix, hence orthogonal: Rθi⊤Rθi R_ _i R_ _i =(cosθisinθi−sinθicosθi)(cosθi−sinθisinθicosθi) = pmatrix _i& _i\\ - _i& _i pmatrix pmatrix _i&- _i\\ _i& _i pmatrix =(1001)=I2. = pmatrix1&0\\ 0&1 pmatrix=I_2. Since RθR_θ is block-diagonal with orthogonal blocks, Rθ⊤Rθ=IdhR_θ R_θ=I_d_h. (2) Unit spectral norm. For any orthogonal matrix Q, we have ‖Qx‖2=‖x‖2\|Qx\|_2=\|x\|_2 for all x. Taking the supremum over unit vectors: ‖Q‖2=sup‖x‖2=1‖Qx‖2=sup‖x‖2=1‖x‖2=1.\|Q\|_2= _\|x\|_2=1\|Qx\|_2= _\|x\|_2=1\|x\|_2=1. Since RθR_θ is orthogonal, ‖Rθ‖2=1\|R_θ\|_2=1. (3) Inner product bound. For any q,k∈ℝdhq,k ^d_h: |(Rmq)⊤(Rnk)|=|q⊤Rm⊤Rnk|.|(R_mq) (R_nk)|=|q R_m R_nk|. The matrix Rm⊤RnR_m R_n is orthogonal (product of orthogonal matrices), so ‖Rm⊤Rnk‖2=‖k‖2\|R_m R_nk\|_2=\|k\|_2. By Cauchy-Schwarz: |q⊤Rm⊤Rnk|≤‖q‖2⋅‖Rm⊤Rnk‖2=‖q‖2⋅‖k‖2.∎|q R_m R_nk|≤\|q\|_2·\|R_m R_nk\|_2=\|q\|_2·\|k\|_2. C.3 Proof of Corollary 3.6 (Extension to RoPE) Proof. With RoPE, the attention score between positions m and n is: Smn=(Rmqm)⊤(Rnkn)dh=xm⊤WQRm⊤RnWK⊤xndh,S_mn= (R_mq_m) (R_nk_n) d_h= x_m W^QR_m R_nW^K x_n d_h, where qm=WQ⊤xmq_m=W^Q x_m and kn=WK⊤xnk_n=W^K x_n. The effective interaction matrix becomes Mm,n=WQRm⊤RnWK⊤M_m,n=W^QR_m R_nW^K , which varies with positions. Worst-case bound. The deterministic worst-case bound follows from submultiplicativity: ‖Mm,n‖2=‖WQRm⊤RnWK⊤‖2≤‖WQ‖2‖Rm⊤Rn‖2‖WK‖2=‖WQ‖2‖WK‖2,\|M_m,n\|_2=\|W^QR_m R_nW^K \|_2≤\|W^Q\|_2\|R_m R_n\|_2\|W^K\|_2=\|W^Q\|_2\|W^K\|_2, since ‖Rm⊤Rn‖2=1\|R_m R_n\|_2=1 for orthogonal Rm,RnR_m,R_n. Combined with ‖xm‖2,‖xn‖2≤d\|x_m\|_2,\|x_n\|_2≤ d under LayerNorm/RMSNorm: |Smn|≤‖WQ‖2‖WK‖2⋅ddh.|S_mn|≤ \|W^Q\|_2\|W^K\|_2· d d_h. Tighter interaction bound. The tighter bound ‖WQWK⊤‖2\|W^QW^K \|_2 can be violated only if ‖Mm,n‖2>‖WQWK⊤‖2\|M_m,n\|_2>\|W^QW^K \|_2 for some positions, which requires RoPE rotations to align the singular subspaces of WQW^Q and WKW^K more favorably than the identity. Since RoPE rotation angles depend only on position indices and fixed frequency bands (independent of weight geometry), such systematic alignment does not occur in practice. We verified empirically that maxm,n‖Mm,n‖2≤‖WQWK⊤‖2 _m,n\|M_m,n\|_2≤\|W^QW^K \|_2 holds across all layers in Mistral-7B, Llama-2-13B, and Llama-2-70B. Probabilistic extension. The probabilistic bound (Proposition 3.4) extends conservatively to RoPE by applying it to the position-independent matrix M=WQWK⊤M=W^QW^K , which empirically upper-bounds all ‖Mm,n‖2\|M_m,n\|_2; see Appendix C.4 for details. ∎ C.4 Extension of Probabilistic Bounds to RoPE This section justifies the application of Proposition 3.4 to RoPE architectures. The challenge. Proposition 3.4 bounds overflow probability for a fixed interaction matrix M=WQWK⊤M=W^QW^K . With RoPE, the effective interaction matrix becomes position-dependent: Mm,n=WQRm⊤RnWK⊤M_m,n=W^QR_m R_nW^K , where Rm,RnR_m,R_n are position-dependent rotation matrices. Why the bound remains conservative. We apply Proposition 3.4 using ‖M‖2=‖WQWK⊤‖2\|M\|_2=\|W^QW^K \|_2 rather than ‖Mm,n‖2\|M_m,n\|_2. This is justified by two observations: 1. Empirical validation: We verified that maxm,n‖Mm,n‖2≤‖WQWK⊤‖2 _m,n\|M_m,n\|_2≤\|W^QW^K \|_2 holds across all layers in Mistral-7B, Llama-2-13B, and Llama-2-70B. RoPE rotations, which depend only on position indices and fixed frequency bands, do not systematically align with the singular subspaces of WQW^Q and WKW^K to amplify the spectral norm. 2. Conservative application: By using the unrotated spectral norm ‖WQWK⊤‖2\|W^QW^K \|_2, which empirically upper-bounds all position-dependent variants, our calibration factor α provides at least the claimed overflow probability guarantee. Remark on the spherical assumption. Proposition 3.4 assumes input directions are approximately uniform on d−1S^d-1. This models post-LayerNorm tokens, which have controlled norm and near-isotropic directions in high dimensions. The projection WQ⊤xW^Q x does not preserve uniformity (it concentrates along WQW^Q’s singular vectors), but this makes the actual overflow probability lower than our bound predicts: inputs are unlikely to align with both WQW^Q’s and WKW^K’s top singular directions simultaneously. Thus, the spherical assumption yields a conservative bound in practice. In summary, geometry-aware scaling works identically for RoPE architectures: we compute the spectral norm of the position-independent matrix WQWK⊤W^QW^K and apply the same calibration procedure. Appendix D Spectral Norm Distribution Across Layers The spectral norm σQK(ℓ)=‖W(ℓ)QW(ℓ)K⊤‖2 _QK^( )=\|W^Q_( )W^K _( )\|_2 varies substantially across layers within each model, motivating the per-layer scale computation in Eq. (15). Figure 1 shows the layer-by-layer spectral norms for all four architectures, and Table 6 summarizes the distribution statistics. Figure 1: Spectral norm σQK(ℓ) _QK^( ) by layer for all four models, computed on pretrained weights. Early layers consistently exhibit larger spectral norms, with layer 0 being the maximum in three of four models. Table 6: Spectral norm statistics across layers for pretrained weights. Max Layer indicates which layer has the largest spectral norm. Model Mean Max Min Max Layer GPT-2 XL 83.1 483.9 55.8 0 Mistral-7B 4.9 46.8 2.4 0 Llama-2-13B 198.4 463.5 134.4 0 Llama-2-70B 584.2 1786.1 264.6 67 The variation is substantial: GPT-2 XL exhibits a 8.7×8.7× range (55.8 to 483.9), Mistral-7B a 19.5×19.5× range (2.4 to 46.8), Llama-2-13B a 3.5×3.5× range (134.4 to 463.5), and Llama-2-70B a 6.7×6.7× range (264.6 to 1786.1). This heterogeneity confirms that per-layer scaling is essential: a global scale factor calibrated for the maximum spectral norm would waste FP8 dynamic range on most layers, while one calibrated for the average would risk overflow on high-norm layers. Appendix E Algorithm Details This appendix provides detailed algorithms for the spectral norm estimation procedures described in Section 4. E.1 Power Iteration for Standard Multi-Head Attention Algorithm 2 computes the spectral norm ‖M‖2\|M\|_2 where M=WQWK⊤M=W^QW^K for standard multi-head attention (where nq=nkvn_q=n_kv). The algorithm maintains persistent vectors u and v that track the top left and right singular vectors across training steps. Algorithm 2 Power Iteration for Spectral Norm Estimation 1: Input: Weight matrices WQ,WK∈ℝd×(nq⋅dh)W^Q,W^K ^d×(n_q· d_h) 2: State: Persistent vectors u,v∈ℝdu,v ^d 3: Output: Spectral norm estimate σ 4: 5: if not initialized then 6: u←random unit vector in ℝdu unit vector in R^d 7: v←random unit vector in ℝdv unit vector in R^d 8: end if 9: 10: One iteration of power method 11: z←WK⊤vz← W^K v z∈ℝnq⋅dhz ^n_q· d_h 12: u′←WQzu ← W^Qz u′=Mv∈ℝdu =Mv ^d 13: σ←‖u′‖2σ←\|u \|_2 Spectral norm estimate 14: u←u′/σu← u /σ Update persistent left singular vector 15: 16: y←WQ⊤uy← W^Q u y∈ℝnq⋅dhy ^n_q· d_h 17: v′←WKyv ← W^Ky v′=M⊤u∈ℝdv =M u ^d 18: v←v′/‖v′‖2v← v /\|v \|_2 Update persistent right singular vector 19: 20: return σ Complexity. Each iteration requires four matrix-vector products: WK⊤vW^K v, WQzW^Qz, WQ⊤uW^Q u, and WKyW^Ky. Each product costs O(nq⋅dh⋅d)O(n_q· d_h· d) operations, for a total of O(nq⋅dh⋅d)O(n_q· d_h· d) per iteration. This is far cheaper than forming M∈ℝd×dM ^d× d explicitly, which would cost O(nq⋅dh⋅d2)O(n_q· d_h· d^2). Convergence. Power iteration converges to the top singular value at rate (σ2/σ1)k( _2/ _1)^k after k iterations, where σ1≥σ2 _1≥ _2 are the two largest singular values. Since we maintain persistent vectors across forward passes and weights change gradually during training, one iteration per forward pass suffices to track the slowly drifting singular vectors. E.2 Implicit GQA Power Iteration For Grouped Query Attention with nq>nkvn_q>n_kv, Algorithm 3 computes the spectral norm without expanding the key matrix. Algorithm 3 Implicit GQA Power Iteration 1: Input: WQ∈ℝd×(nq⋅dh)W^Q ^d×(n_q· d_h), WK∈ℝd×(nkv⋅dh)W^K ^d×(n_kv· d_h) 2: Parameters: Number of query heads nqn_q, KV heads nkvn_kv, head dimension dhd_h 3: State: Persistent vectors u,v∈ℝdu,v ^d 4: Output: Spectral norm estimate σ≈‖WQWexpK⊤‖2σ≈\|W^QW^K _exp\|_2 5: 6: g←nq/nkvg← n_q/n_kv Group size 7: 8: if not initialized then 9: u←random unit vector in ℝdu unit vector in R^d 10: v←random unit vector in ℝdv unit vector in R^d 11: end if 12: 13: Forward: compute MvMv where M=WQWexpK⊤M=W^QW^K _exp 14: zkv←WK⊤vz_kv← W^K v zkv∈ℝnkv⋅dhz_kv ^n_kv· d_h 15: z←RepeatBlocks(zkv,g)z← RepeatBlocks(z_kv,g) z∈ℝnq⋅dhz ^n_q· d_h 16: u′←WQzu ← W^Qz u′∈ℝdu ^d 17: σ←‖u′‖2σ←\|u \|_2 18: u←u′/σu← u /σ 19: 20: Backward: compute M⊤uM u where M⊤=WexpKWQ⊤M =W^K_expW^Q 21: y←WQ⊤uy← W^Q u y∈ℝnq⋅dhy ^n_q· d_h 22: ykv←SumGroups(y,g)y_kv← SumGroups(y,g) ykv∈ℝnkv⋅dhy_kv ^n_kv· d_h 23: v′←WKykvv ← W^Ky_kv v′∈ℝdv ^d 24: v←v′/‖v′‖2v← v /\|v \|_2 25: 26: return σ Memory savings. The explicit approach would require expanding WKW^K from (nkv⋅dh)×d(n_kv· d_h)× d to (nq⋅dh)×d(n_q· d_h)× d by replicating columns. For Mistral-7B with nq=32n_q=32, nkv=8n_kv=8, dh=128d_h=128, and d=4096d=4096, this expansion would require 32×128×4096×2=3232× 128× 4096× 2=32MB per layer (in FP16). The implicit formulation avoids this entirely by operating on the unexpanded WKW^K and only replicating/summing small intermediate vectors of size nkv⋅dhn_kv· d_h. Appendix F Proof of Implicit GQA Equivalence This appendix proves Proposition 4.1, establishing that the implicit GQA formulation computes the same spectral norm as explicit expansion. Proof. Let g=nq/nkvg=n_q/n_kv be the group size. We have WQ∈ℝd×(nq⋅dh)W^Q ^d×(n_q· d_h) and WK∈ℝd×(nkv⋅dh)W^K ^d×(n_kv· d_h). Define WexpK∈ℝd×(nq⋅dh)W^K_exp ^d×(n_q· d_h) as the matrix obtained by replicating each block of dhd_h columns of WKW^K exactly g times. The interaction matrix is M=WQWexpK⊤∈ℝd×dM=W^QW^K _exp ^d× d. Forward direction. For any v∈ℝdv ^d, we show that Mv=WQ⋅RepeatBlocks(WK⊤v,g)Mv=W^Q· RepeatBlocks(W^K v,g). Since WexpK⊤∈ℝ(nq⋅dh)×dW^K _exp ^(n_q· d_h)× d has repeated row-blocks, WexpK⊤v∈ℝnq⋅dhW^K _expv ^n_q· d_h consists of nqn_q blocks of size dhd_h, where blocks (i−1)g+1,…,ig\(i-1)g+1,…,ig\ for i=1,…,nkvi=1,…,n_kv all equal the i-th block of WK⊤v∈ℝnkv⋅dhW^K v ^n_kv· d_h. This is exactly the output of RepeatBlocks(WK⊤v,g) RepeatBlocks(W^K v,g). Since WQ∈ℝd×(nq⋅dh)W^Q ^d×(n_q· d_h), multiplying by this repeated vector yields Mv∈ℝdMv ^d. Backward direction. For any u∈ℝdu ^d, we show that M⊤u=WK⋅SumGroups(WQ⊤u,g)M u=W^K· SumGroups(W^Q u,g). We have M⊤=WexpKWQ⊤M =W^K_expW^Q . First, WQ⊤u∈ℝnq⋅dhW^Q u ^n_q· d_h. Since WexpK∈ℝd×(nq⋅dh)W^K_exp ^d×(n_q· d_h) has g identical copies of each column block of WKW^K, multiplying WexpKW^K_exp by a vector y∈ℝnq⋅dhy ^n_q· d_h is equivalent to multiplying WKW^K by the vector that sums each group of g blocks: WexpKy=WK⋅SumGroups(y,g)W^K_expy=W^K· SumGroups(y,g). Thus M⊤u=WK⋅SumGroups(WQ⊤u,g)∈ℝdM u=W^K· SumGroups(W^Q u,g) ^d. Conclusion. Power iteration maintains u,v∈ℝdu,v ^d and computes MvMv and M⊤uM u via the implicit formulation. Since these equal the explicit products, the iteration converges to the same singular vectors and value: σimplicit=σexplicit=‖WQWexpK⊤‖2 _implicit= _explicit=\|W^QW^K _exp\|_2. ∎ Appendix G Experimental Details This appendix provides additional experimental details and extended results for Section 5. G.1 Training Configuration Table 7 summarizes the architectural specifications for all four models evaluated. Table 7: Model architectures. MHA: multi-head attention. GQA: grouped query attention. GPT-2 XL Mistral-7B Llama-2-13B Llama-2-70B Params 1.5B 7B 13B 70B Layers 48 32 40 80 Attention MHA GQA MHA GQA nqn_q / nkvn_kv 25/25 32/8 40/40 64/8 d 1600 4096 5120 8192 dhd_h 64 128 128 128 Hardware H100 H200 H200 B200 Table 8 provides the full training configuration for each model. Table 8: Training configuration for all experiments. GPT-2 XL Mistral-7B Llama-2-13B Llama-2-70B Batch size 8 4 4 2 Seq. length 1024 1024 1024 1024 Learning rate 10−410^-4 10−510^-5 10−510^-5 10−510^-5 Optimizer AdamW AdamW AdamW AdamW Weight decay 0.01 0.01 0.01 0.01 Grad. clip 1.0 1.0 1.0 1.0 FP8 format E4M3 E4M3 E4M3 E4M3 α 0.08 0.04 0.03 0.02 ηfp8 _fp8 0.8 0.8 0.8 0.8 Delayed scaling baseline configuration. For the delayed scaling baseline, we use the standard configuration from Micikevicius et al. (2022): history buffer length of 16 steps and safety margin η=0.9η=0.9 (see Eq. 1). The history buffer is initialized to 1.0 at the start of training or after checkpoint loading, which is the source of history staleness in transient scenarios. Appendix H Instantaneous Response Validation This appendix validates the instantaneous response property of geometry-aware scaling with a controlled stress test. Setup. We train GPT-2 XL normally for 10 steps, then multiply all attention weights by 4×4× at step 10, simulating an extreme transient. This artificial perturbation isolates the mechanism: delayed scaling relies on stale statistics and cannot adapt, while geometry-aware scaling recomputes the bound from current weights. Results. Figure 2 shows the results. Delayed scaling, relying on statistics that predate the weight change, produces scaled logits exceeding 6500 and overflows catastrophically. Geometry-aware scaling responds in the same forward pass: the scale factor jumps from 5.8 to 93.8 at step 10, keeping scaled logits below 85 throughout. Figure 2: Response to 4×4× weight spike at step 10 (GPT-2 XL). (a) Maximum scaled attention logit. (b) Scale factor over time. Delayed scaling overflows because its history buffer contains no information about the weight change. Geometry-aware scaling adapts instantaneously because the scale factor is computed from current weights. While the 4×4× spike is a stress test, the underlying mechanism explains why geometry-aware scaling succeeds in the realistic transient scenarios evaluated in Section 5.2: loading pretrained checkpoints, checkpoint resumption, and learning rate transitions all involve weight configurations for which the history buffer has no relevant statistics. Appendix I Overhead Measurement Table 9 reports forward pass timing for all models. Table 9: Computational overhead (forward pass time, averaged over 100 iterations). Model Attention Delayed Ours Overhead GPT-2 XL MHA 136.3 ms 137.6 ms +1.0%+1.0\% Mistral-7B GQA 4:1 58.2 ms 55.1 ms −5.3%-5.3\% Llama-2-13B MHA 96.2 ms 98.1 ms +1.9%+1.9\% Llama-2-70B GQA 8:1 206.3 ms 215.0 ms +4.3%+4.3\% Measurements follow this protocol: 1. Warm up GPU with 10 forward passes (discarded) 2. Run 100 forward passes with CUDA synchronization after each 3. Report mean time per forward pass 4. Repeat 3 times and report median across repetitions For fair comparison, both methods use identical model configurations, batch sizes, and sequence lengths. The only difference is the scale factor computation method. Appendix J FP8 Utilization Statistics Table 10 shows FP8 dynamic range utilization during training with different methods on GPT-2 XL. Table 10: FP8 dynamic range utilization (%) during training (GPT-2 XL). Method Median P10 P90 Delayed 90.5% 84.0% 98.0% Ours (conservative) 14.3% 13.6% 15.3% Ours + auto-α 31.2% 28.5% 33.7% Appendix K MMLU Training Curves Figure 3 shows training loss curves for the three scaling methods on Llama-2-13B fine-tuned on MMLU STEM subjects. Figure 3: Training loss comparison for delayed scaling, geometry-aware scaling (conservative), and geometry-aware scaling with auto-α on Llama-2-13B (MMLU STEM, 3000 steps). All methods converge to similar final loss (≈0.012≈ 0.012), yet downstream MMLU accuracy differs substantially (Table 5). The vertical dashed line indicates when auto-α calibration completes (step 100). The conservative variant shows slightly higher loss throughout training due to reduced FP8 utilization (0.5% vs. 31.2% for auto-α). The key observation is that training loss alone does not predict downstream performance: the conservative variant achieves comparable loss (0.0119 vs. 0.0116) but 5 percentage points lower MMLU accuracy (28.1% vs. 33.0%). This gap arises because excessive quantization noise from 0.5% FP8 utilization degrades the fine-grained attention patterns required for reasoning tasks, even when average loss remains low. Appendix L Detailed MMLU Results by Subject Table 11 shows accuracy on all 17 MMLU STEM subjects. Table 11: MMLU accuracy (%) by subject on Llama-2-13B after 3000 fine-tuning steps. Subject Delayed Cons. Auto-α Abstract Algebra 25.0 25.0 27.0 College Math 36.0 31.0 35.0 Elementary Math 27.0 25.9 27.8 HS Math 26.7 22.6 27.0 HS Statistics 30.1 31.0 30.6 Astronomy 41.4 29.6 37.5 College Physics 39.2 37.3 38.2 HS Physics 27.8 26.5 29.1 College CS 29.0 30.0 34.0 Computer Security 47.0 34.0 42.0 HS CS 37.0 26.0 34.0 College Chemistry 35.0 22.0 34.0 HS Chemistry 27.1 29.6 32.0 College Biology 36.8 24.3 41.7 HS Biology 45.8 31.6 46.8 Electrical Eng. 35.2 31.7 31.0 Machine Learning 21.4 22.3 25.9 Average 33.0 28.1 33.6 The largest accuracy gaps between conservative spectral and auto-α appear in subjects requiring fine-grained numerical reasoning: College Biology (+17.4 p), HS Biology (+15.2 p), College Chemistry (+12.0 p), and Computer Security (+8.0 p). This pattern suggests that excessive quantization noise particularly affects tasks requiring precise attention to numerical or logical relationships. Appendix M Auto-α Calibration Details M.1 Algorithm Algorithm 4 provides the complete auto-α calibration procedure. Algorithm 4 Auto-α Calibration 0: Conservative initial α0 _0, burn-in steps TcalibT_calib, quantile q, safety multiplier κ 1: Initialize slack ratio buffer R←∅R← 2: for t=1t=1 to TcalibT_calib do 3: Compute Bmax(t)=‖WQWK⊤‖2⋅d/dhB_ ^(t)=\|W^QW^K \|_2· d/ d_h 4: Run forward pass with scale =α0⋅Bmax(t)/(0.8×448)= _0· B_ ^(t)/(0.8× 448) 5: Record Mt=maxi,j|Sij(t)|M_t= _i,j|S_ij^(t)| (actual maximum logit) 6: Append rt=Mt/Bmax(t)r_t=M_t/B_ ^(t) to R 7: end for 8: αemp←Quantileq(R) _emp _q(R) 9: αfinal←αemp×κ _final← _emp×κ 10: Freeze αfinal _final for all subsequent training M.2 Calibration Statistics During the 100-step burn-in on Llama-2-13B: • Slack ratio range: [0.000073,0.000366][0.000073,0.000366] • Slack ratio mean: 0.0001770.000177 • P99.99P_99.99 quantile: 0.0003600.000360 • Calibrated α: 0.000360.00036 (with κ=1.0κ=1.0) This represents an 83×83× tightening compared to the conservative α=0.03α=0.03, resulting in proportionally higher FP8 utilization (31.2% vs. 0.5%). M.3 Safety Multiplier Selection The safety multiplier κ controls the trade-off between utilization and robustness: • κ=1.0κ=1.0: Maximum utilization, suitable for stable fine-tuning • κ=2.0κ=2.0: Moderate headroom for distribution shift • κ=3.0κ=3.0: Conservative, suitable for heterogeneous data In our experiments, κ=1.0κ=1.0 achieved zero overflows across 3000 steps, suggesting the P99.99P_99.99 quantile provides sufficient margin for MMLU STEM fine-tuning.