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 \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$. 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.

Modern language models are expensive to train: the computations involve trillions of numbers that have to be stored and updated as training progresses. One way to lower the cost is to use smaller numbers, for instance 8-bit floating point (FP8) instead of the standard 16-bit. The trouble is that 8-bit numbers can only represent a narrow range of values, and if any value grows too large during training, the result is "Not a Number" outputs that corrupt the entire model. The standard solution looks at how large the numbers have been over the most recent training steps and adjusts the scaling accordingly. This is essentially driving while watching only the rearview mirror: it works while training is steady, but fails predictably whenever something changes suddenly. When practitioners load a previously trained model into the 8-bit pipeline, resume after a job interruption, or ramp up the learning rate, the recent history does not reflect what is about to happen, and every model we tested overflows in 100% of its attention layers. We propose a different approach: predict what is about to happen, instead of reacting to what just did. The mathematical structure of transformer attention lets us compute, directly from the model's current weights, a safe upper bound on the largest value any attention computation will produce. We then use that bound to set the scaling factor before the computation runs, so the result is guaranteed to fit in the 8-bit range. Because we read the weights themselves rather than past activity, the prediction never goes stale. Across models ranging from 1.5 billion to 70 billion parameters, our method eliminates the overflow failures that break standard methods while matching downstream accuracy. The code is publicly released for the community to use.