Define QK-norm and the training failure it prevents

Flashcard·Medium·4.0 · 0·~30s·Asked atCoreweave Tencent

Attempt it

QK-norm applies RMSNorm (or LayerNorm) to the query and key vectors per head, right after their projections and before the dot product. Without it, Q and K magnitudes can drift upward during long training runs, so attention logits explode, softmax saturates into near one-hot distributions (attention entropy collapse), and the loss spikes or diverges. Normalizing Q and K bounds the logit scale by construction. Shipped in ViT-22B, Gemma 2/3, OLMo 2, and many 2025-2026 frontier recipes, often alongside or instead of logit soft-capping.

Your answer

**Where it sits.** Inside the attention computation, the query and key projections produce Q and K tensors of shape `(batch, seq, n_heads, d_head)`. QK-norm inserts a RMSNorm (or LayerNorm) on these tensors per head, right after the projections and before the dot product: ``` q = norm(Q @ W_Q) k = norm(K @ W_K) scores = (q @ k.transpose(-2, -1)) / sqrt(d_head) ``` **What it prevents.** Across long pretraining runs, Q and K projection weights tend to grow in norm. Larger Q and K means larger dot products, which means larger attention logits. Once logits get big enough (~30-100), softmax saturates: one position captures essentially all attention mass and the rest get zero. Attention entropy in those heads collapses. The loss spikes intermittently because gradients through saturated softmax are erratic, and in bf16 the exponentials risk overflow. **Why QK-norm fixes it.** RMSNorm forces `q` and `k` to have a fixed RMS regardless of how their projection weights drift. The dot product is therefore bounded by `d_head` in the worst case and much smaller in practice. The `1/sqrt(d_head)` scaling now keeps logits in a stable range across the entire run, not just at init.

**The placement.** Inside the attention block, after the linear projections that produce Q and K from the residual stream, but before the dot product. The norm is applied per head (independently across the `d_head` dimension for each of the `n_heads` heads). Both LayerNorm and RMSNorm are used in practice; ViT-22B and Gemma use RMSNorm. **The failure mode it prevents.** Attention logit explosion driven by Q and K weight norm drift. Mechanically: the standard `(QK^T) / sqrt(d_k)` scaling factor is an init-time variance correction. It assumes Q and K have approximately unit-variance entries, true at initialization, but not after Adam has updated the projections for hundreds of thousands of steps. Once `||W_Q||` and `||W_K||` grow, post-projection q and k have larger entries, dot products grow quadratically with weight norms, and the `1/sqrt(d_k)` scaling cannot compensate. Logits drift from O(1) at init to magnitudes of 50-100+ by mid-training, saturating softmax. Saturated softmax produces one-hot attention and degenerate gradients; loss spikes intermittently and bf16 exp overflows. **Why QK-norm fixes it structurally.** RMSNorm enforces `mean(q_i^2) = 1` along the head dimension (modulo a learnable gain near 1). So `||q||_2 = sqrt(d_head)` regardless of how `W_Q` drifts. Same for k. The dot product `q . k` is bounded by `d_head` via Cauchy-Schwarz. After the `1/sqrt(d_head)` scaling, logits are bounded by `sqrt(d_head)`, typically around 11 for `d_head = 128`. This bound holds across the entire run; no drift can break it.

TL;DR

QK-norm normalizes Q and K per head, after projection and before the dot product, bounding logit scale by construction so softmax cannot saturate as Q and K weight norms drift across long training.

Memory aid

Easy to grasp

Imagine a tug-of-war where the rope length is what gets multiplied to score the pull. Without QK-norm, the team captains can sneak in longer ropes during the match, and pretty soon any pull from one side counts as a million points. QK-norm cuts every rope to a fixed length before each round. The teams can still angle the rope, point it sharply, or pull harder, but raw length cannot blow up the score. Softmax stays sane, attention stays diverse, and the training run does not catch fire halfway through.

Key concepts

Transformer architecture

Concept explanation~2 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

QK-norm became one of the standard stability tools in the open-weights transformer toolkit between 2023 and 2026. The technique is small (one normalization layer per attention call) but the failure mode it prevents (attention logit explosion driven by Q and K weight drift) is one of the most commonly cited stability incidents in long pretraining runs at frontier scale.

Knowing QK-norm cold means knowing three things: where it sits in the attention block, what specifically it prevents, and how it compares with the adjacent technique of logit soft-capping. This is interview-grade material for any 2026 LLM engineering role.

The placement inside attention

QK-norm operates on the query and key tensors inside the attention computation. The standard order of operations in a multi-head attention block:

Read residual stream x of shape (batch, seq, d_model).
Apply input LayerNorm or RMSNorm (this is the pre-norm step, separate from QK-norm).
Project to Q, K, V: Q = x @ W_Q, shape (batch, seq, n_heads, d_head). Same for K and V.
QK-norm: apply RMSNorm to Q and K per head. This step is what we are defining.
Apply RoPE rotation to Q and K (if RoPE is used).
Dot product: scores = (Q @ K.transpose) / sqrt(d_head).
Softmax to get attention weights.
Weighted sum with V; output projection.

The norm is per-head, meaning each of the n_heads heads has its own RMSNorm operating along the d_head axis. The learnable gain is typically a (n_heads, d_head) or (n_heads,) tensor; some implementations omit the gain entirely for maximum bounding.

The placement matters. Norming the input to the projection (step 2 in the pre-norm step) controls input variance but not post-projection variance. Norming after projection but after RoPE is also viable. Norming inside the dot product or after softmax is not QK-norm; those are different techniques (logit soft-capping is a post-softmax intervention).

The failure mode it prevents

Why QK-norm fixes it structurally

Shipping models and the soft-capping comparison

Situations where this technique stops working.

2–4 min · Everything important, quickly.

Real products, models, and research that use this idea.

ViT-22B (Google, 2023) introduced QK-norm to enable stable scaling of vision transformers to 22B parameters; the technique was identified as essential for stability at that scale.
Gemma 2 (Google, 2024) shipped QK-norm with RMSNorm together with logit soft-capping (`cap = 30` for attention, `cap = 50` for output) as a belt-and-suspenders stability stack.

What an interviewer would ask next. Try answering before peeking at the approach.

QDoes QK-norm work with RoPE?

Yes; apply QK-norm to Q and K post-projection, then apply RoPE rotation, then compute the dot product. The norm fixes magnitudes; the rotation changes direction. ViT-22B uses learned position; Gemma 2 and OLMo 2 combine QK-norm with RoPE successfully.

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases that signal junior thinking. Click to expand.

Most common mistake

Believing the `1/sqrt(d_k)` scaling already bounds logits. It corrects init-time variance only and cannot constrain weight growth across hundreds of thousands of steps.

60 second bullets to scan on the way to the call.

Identify where QK-norm sits: after Q and K projections, before the dot product, per head
State which norm types are used (RMSNorm and LayerNorm both valid; RMSNorm is the frontier choice)

Primary sources. Browse if you want the original framing.