QK-Normalisation

Attention's dot product Q . K has no built-in ceiling. The 1/sqrt(d_k) scale factor in standard attention corrects for the expected growth of a dot product with dimension, but it does nothing to stop the query and key vectors themselves from growing during training, and at large enough scale, they do. When they grow enough, softmax's exponential turns modest logit differences into near-total dominance by a single key, a failure mode with a name, attention entropy collapse, and a measurable correlation with the loss spikes that plague very large training runs.

Why unbounded Q and K vectors are a problem

Softmax attention only cares about relative logit magnitude, but relative magnitude is exactly what gets distorted when Q and K vectors are free to grow. Two vectors that are only mildly more aligned than another pair can, once their norms are large, produce a raw dot product large enough to push softmax into its saturated regime, where the output is essentially one-hot and the gradient through it vanishes (see softmax-logits for the saturation mechanism, and attention-mechanism for why the fixed scale factor alone is not sufficient). At small scale this rarely bites. At the depth and width of frontier models, training dynamics can push Q and K norms into regimes where it does, and once it does, training can spike or diverge with little warning.

The fix: normalise before the dot product

QK-normalisation applies a normalisation, LayerNorm or RMSNorm, to the query and key vectors of each attention head independently, before the dot product, typically paired with a learned per-head scale that replaces the fixed 1/sqrt(d_k) factor. Because normalised vectors have bounded norm, their dot product is bounded regardless of how large the underlying linear projections grow during training; the raw magnitude of Q and K no longer determines how peaked the resulting attention distribution can become, only their direction does, modulated by a scale the model can still learn to adjust.

Dehghani et al., 2023 found QK-norm necessary to train ViT-22B, a 22-billion-parameter vision transformer, without the loss spikes that appeared without it, at a scale where smaller ViT variants had trained fine without any such normalisation. That result is the reason QK-norm went from a niche NLP proposal, Henry et al., 2020 had already proposed a query-key normalisation scheme for transformers years earlier without wide adoption, to a standard recommendation for training very large transformers of any modality. Several recent open frontier-adjacent language models have since adopted some form of QK-norm specifically to stabilise their largest training runs.

When it falls down

• It costs a normalisation op inside the hottest loop in the model. QK-norm runs per head, per layer, on every forward and backward pass, a real, if usually modest, throughput cost paid continuously rather than once.
• Bounding logits also bounds legitimate sharpness. Some attention patterns genuinely need to be close to one-hot, exact copying of a single token being the clearest example, and an aggressively bounded normalisation can make that harder to represent; implementations pair QK-norm with a learnable temperature specifically to recover the ability to sharpen when the task calls for it.
• It fixes one symptom, not every cause, of large-run instability. Loss spikes at frontier scale also arise from data issues, optimiser state corruption, and low-precision rounding; QK-norm is one stabiliser typically deployed alongside careful initialisation (see weight-initialisation-transformers), gradient clipping, and auxiliary losses on the output logits, not a substitute for them.