Rank-Aware Spectral Bounds on Attention Logits for Stable Low-Precision Training

Attention scores in transformers are bilinear forms $S_{ij}=x_i^{\mathrm{\top }}M{x}_j\mathrm{/}\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\mathrm{\top }}$ has rank $r\ll d$, tail probabilities for ${\max }_{i,j}\mathrm{\mid }{S}_{ij}\mathrm{\mid }$ decay as $\exp (-d^2{\alpha }^2\mathrm{/}(\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 $\mathrm{\parallel }{W}^Q{W}^{K\mathrm{\top }}{\mathrm{\parallel }}_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.