Why does normalizing Q and K before the dot product help training at scale?

**Attention** computes a similarity between a query q and every key k via the dot product `q . k`. That score, scaled by sqrt(d_k), goes through softmax to become a weight. The output is a weighted sum of the value vectors. Problem: q . k can get very large. Softmax of large differences saturates, one weight goes to nearly 1 and the rest go to nearly 0. Once that happens, the gradient signal to all non-peak keys vanishes, the model stops learning from them, and training can stall or destabilize. The original transformer added the sqrt(d_k) divisor to fix the initialization-time variance of the dot product. That helps at step 0, when weights are small. But over a long pretraining run, the Q and K projection weights drift larger, so q . k drifts larger too, and the sqrt(d_k) constant does nothing to compensate. **QK-norm** is a fix: apply L2 normalization to q and k per token before the dot product. After normalization, q . k is bounded between -1 and 1 (it is a cosine similarity scaled by their original norms, which are now 1). A learnable scalar is usually multiplied in afterward so the model can still control sharpness. The result: stable softmax inputs across the whole training run, no matter how much the weights grow. Used in ViT-22B, Idefics, and modern multimodal transformers.

QK-norm L2-normalizes the query and key vectors of every token before the attention dot product, then multiplies by a learnable scalar. The fix targets a specific failure mode of long pretraining runs: **logit explosion**. During training, the norms of W_Q and W_K drift upward as gradient updates accumulate. The dot product q . k scales with those norms, so the pre-softmax logits drift upward too. Once a logit grows large relative to its neighbors, softmax saturates: one weight goes near 1, the rest go near 0, and gradients to all non-peak keys vanish. Loss spikes, attention entropy collapses, the training run stalls. The original sqrt(d_k) divisor only stabilizes variance at initialization. It does not adapt to growing parameter norms. QK-norm does: after L2 normalization, |q . k| <= 1 regardless of the underlying W_Q, W_K magnitudes. The learnable scalar lets the model recover controllable sharpness without giving back the boundedness. The cost is tiny: two norm operations and one scalar multiply per attention call. The win is large: long pretraining runs no longer crash from logit explosion. ViT-22B (Dehghani et al. 2023) introduced it for stability; Idefics, several frontier vision-language models, and increasingly modern LLMs have adopted the pattern.