arXiv:2602.18851 · v2
Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training
摘要
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 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}α^{2}/(γr))$ rather than $\exp(-dα^{2})$, where $γ> 1$. For transformer attention where $r = d_h$, this yields $8$ to $28\times$ 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 $\|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.
作者备注：ICML 2026 camera-ready version. 19 pages, 3 figures. Code: https://github.com/SeyedMEmadi/fp8-geometry-aware-scaling
期刊信息：Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:28021-28039, 2026