QK-norm applies L2 normalization (or a learned scale) to the query and key vectors of every token before computing the attention score q · k. Why does this stabilize large-scale training? What problem in vanilla attention is it fixing, and what is the cost?
Free · 2 AI evals / day
**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.
Quick reference
TL;DR
At scale, Q and K magnitudes drift up, softmax saturates, gradients die. L2-normalizing Q and K per token bounds the dot product to a stable range no matter how big the weights grow.
Picture turning up the volume on a microphone over months. At first the speakers sound balanced. After a year, the loud channel completely drowns out everything else, you only hear one voice, the rest are silent. That is what happens inside a long training run: the numbers the model uses to compare words slowly get bigger and bigger, until one word's vote completely buries everyone else's. The model stops learning from the silent voices and training stalls. QK-norm is a volume limiter that clamps every voice to a fixed loudness before they hit the mixer. The model can still pick winners, but never loud enough to silence the rest of the room.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.
Short training runs where Q and K magnitudes do not have time to drift QK-norm targets long-run instabilities. For short fine-tunes (hundreds of millions of tokens), sqrt(d_k) plus careful learning-rate schedules is usually enough; adding QK-norm at fine-tune time conflicts with the pretrained model's attention sharpness distribution.
Models where attention sharpness needs to be very high (e.g., retrieval-style heads) QK-norm bounds |q . k| <= alpha, capping the maximum sharpness softmax can produce. Tasks that benefit from near-hard argmax retrieval can underperform unless alpha is tuned high enough to recover that sharpness.
Retrofitting an already-trained model that did not use QK-norm Inserting QK-norm into a pretrained model shifts the attention distribution and usually requires a full re-pretraining (or long fine-tuning) to recover quality. It is a pretrain-time decision, not a post-hoc patch.
Inference-only optimization scenarios QK-norm provides no inference speedup; it adds two L2-norm operations per attention call. Its value is entirely in training stability. For inference perf, look at FlashAttention, MQA, GQA, or KV-quantization instead.
Quick scan · 2–4 min
Everything important, quickly.
Name the failure mode (softmax saturation), explain why sqrt(d_k) does not fix training-time drift, state the L2-norm + learnable scalar recipe, cite ViT-22B as the originating model.
Real-world examples
Real products, models, and research that use this idea.
1ViT-22B (Dehghani et al. 2023) introduced QK-norm specifically to stabilize 22B-parameter vision pretraining runs that were diverging.
2Idefics-3 multimodal model uses QK-norm to keep vision-language pretraining stable at scale.
3Gemma 2 / Gemma 3 use a softer variant: per-head logit soft-cap that bounds pre-softmax values without strict L2 normalization.
4Several frontier LLMs in 2026 (including parts of the Gemini 3 line) adopt QK-norm or close variants to support stable >1T-token training runs.
5PaLM-540B's loss spikes during pretraining were widely attributed to softmax saturation and motivated subsequent adoption of variants like QK-norm.
Interviewer follow-ups
What an interviewer would ask next. Try answering before peeking at the approach.
QHow does logit soft-cap (Gemma) compare to QK-norm as a stability mechanism?
Sign in or go Pro to see the approach for this follow-up.
QIf QK-norm is so effective, why doesn't every model use it?
Sign in or go Pro to see the approach for this follow-up.
Don't say this
Red flags and common mistakes that signal junior thinking. Click to expand.
Most common mistake
Claiming sqrt(d_k) already prevents this. It fixes initialization-time variance only; it does not track Q and K growth during training, which is the source of the drift.