Diagnose attention logit explosion in a large run and defend QK-norm as the fix

Short answer·Hard·4.0 · 0·~3 min·Asked atFractal Analytics Mercor Wipro

Attempt it

Halfway through a long pretraining run on a 30B model, monitoring shows max attention logits trending from ~20 to over 100, attention entropy in several heads dropping toward zero, and intermittent loss spikes that recoveries only partially undo. Diagnose what is happening mechanically, explain why the 1/sqrt(head_dim) scaling did not prevent it, and argue for QK-norm versus logit soft capping as interventions.

Free · 2 AI evals / day

**What is happening mechanically.** The Q and K projection matrices have learnable weights. Over hundreds of thousands of steps, those weights drift upward in norm. Larger Q and K vectors mean larger dot products, which means larger attention logits. Once logits hit ~100, the softmax saturates: one token attracts almost all the attention mass and the rest get effectively zero. Attention entropy in those heads collapses toward zero. This matters because saturated softmax has bad gradient properties. Almost everywhere, the gradient is near zero (the output barely changes if logits change). When the ranking flips (a different token becomes the argmax), the gradient spikes hard. Low precision (bf16) also struggles with `exp(100)` and can overflow to infinity. **Why `1/sqrt(d_k)` does not help.** The scaling factor is a constant. It was designed assuming Q and K have unit-variance entries (true at init), so `QK^T` would have variance proportional to `d_k`, and `1/sqrt(d_k)` brings it back to O(1). Once Q and K weight magnitudes drift, the variance of `QK^T` is no longer proportional to `d_k`, and the scaling factor is no longer the right correction. $$\tilde{Q} = \text{RMSNorm}(Q), \quad \tilde{K} = \text{RMSNorm}(K)$$

**Diagnosis.** Three symptoms point to the same root cause: max logits trending from ~20 to >100, attention entropy collapsing toward zero in some heads, and intermittent loss spikes. The mechanism is that Q and K projection weight norms have drifted upward across the run, inflating QK dot products. Softmax saturates as logits grow; one token captures essentially all attention mass, attention entropy collapses, and the gradient becomes erratic. In bf16, `exp(>100)` risks overflow. **Why `1/sqrt(d_k)` did not prevent it.** The scaling factor handles init-time variance under the assumption that Q and K have approximately unit-variance entries. Then `QK^T` has variance proportional to `d_k`, and dividing by `sqrt(d_k)` restores O(1) scale. The factor is a constant; it knows nothing about the learned magnitudes of Q and K. Over hundreds of thousands of steps with no constraint on weight norm growth, Q and K can drift by 5-10x in scale, and the scaling factor cannot compensate. **QK-norm.** Normalize Q and K per head after their projections: $$\tilde{q}_i = \text{RMSNorm}(q_i), \quad \tilde{k}_j = \text{RMSNorm}(k_j)$$ With RMSNorm, `tilde_q` has unit RMS (modulo a learnable gain). The dot product `tilde_q . tilde_k` is bounded by `head_dim` in worst case, but in practice by a much smaller value because vectors are rarely aligned. Crucially, the bound is structural: no matter how the Q and K projection weights drift, post-norm Q and K stay in a fixed scale.

TL;DR

Q and K weight norms drift upward across long training, inflating QK logits, saturating softmax, and collapsing attention entropy. The `1/sqrt(d_k)` factor only corrects init-time variance, not learned weight growth.

Memory aid

Easy to grasp

Imagine a microphone with a volume knob and a hard clip. At the start of recording, the gain is tuned so voices sit nicely in the dynamic range. Over hours, the gain knob slowly drifts up on its own. Voices that started at conversational level are now shouting and the recorder clips. Two ways to fix it: install a volume governor at the input that holds the gain constant (QK-norm) or put a brick-wall limiter on the output that clamps anything above a threshold (soft-capping). The governor prevents the problem; the limiter masks it. You usually want the governor, but if the show is already live, you reach for the limiter.

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.

Attention logit explosion is one of the most cited stability failures in long pretraining runs at frontier scale. The pattern is consistent enough across reported incidents that the 2025-2026 open-weights community has converged on QK-norm as the structural fix, and the question's symptoms (entropy collapse, intermittent loss spikes, logits trending past 100) are diagnostic for exactly this pathology.

The value of the question is not in naming QK-norm. It is in walking the full chain from learned weight drift, through the softmax saturation mechanism, to the choice between structural fix and symptomatic patch, with an honest tradeoff between mid-flight feasibility and from-scratch architectural cleanliness.

The cascade: weight drift to entropy collapse

Start with the attention computation: s_ij = (q_i . k_j) / sqrt(d_k), then a_ij = softmax(s_ij), then output o_i = sum_j a_ij * v_j.

At initialization, Q and K projection weights have entries with standard deviation around 1/sqrt(d_model), so post-projection Q and K vectors have approximately unit-variance entries. The dot product q . k has variance d_k, and the 1/sqrt(d_k) factor brings logits back to O(1). Softmax distributes attention smoothly. Entropy per head is healthy (typically log(seq_len) minus a small amount).

Across hundreds of thousands of steps with Adam updating Q and K projections, weight norms drift. There is no explicit constraint preventing this; Adam normalizes step direction by gradient magnitude but does not bound parameter magnitudes. Weight decay slows growth but does not prevent it entirely, especially when the loss landscape rewards larger Q and K norms (which it does in the early-to-mid training regime, where sharper attention helps).

By 200k-500k steps, Q and K weight Frobenius norms can be 3-10x their initial values. Post-projection q and k vectors have norms 3-10x larger than at init. Dot products q . k grow by 9-100x. Logits s_ij = (q . k) / sqrt(d_k) follow.

Once logits reach magnitude 30-50, softmax begins saturating. By 80-100, several heads have one-hot attention. Entropy collapses. Loss spikes intermittently. In bf16, exp of large logits risks overflow.

This is the failure mode the monitoring is showing. The fix needs to break the link between Q/K weight drift and logit magnitude.

Why the canonical `1/sqrt(d_k)` is structurally insufficient

QK-norm: the structural fix

Soft-capping: the symptom clamp

The mid-flight versus from-scratch tradeoff

Situations where this technique stops working.

2–4 min · Everything important, quickly.

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

Gemma 2 (Google, 2024) shipped both QK-norm and logit soft-capping (`cap = 30` for self-attention, `cap = 50` for final logits) as belt-and-suspenders stability measures.
Gemma 3 (Google, 2025) kept QK-norm and refined soft-capping; the combination is the production recipe for the Gemma line.

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

QCould you replace QK-norm with simple weight decay on `W_Q` and `W_K`?

Partially. Weight decay slows the drift but does not bound the post-projection vector norms structurally. The relationship between decay strength and effective logit scale is indirect and depends on optimizer state; QK-norm gives a structural guarantee.

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 `1/sqrt(d_k)` permanently bounds logits. It is a constant scalar that handles init-time variance only; it does nothing when learned weight magnitudes drift.

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

Trace the mechanism: weight drift inflates QK, softmax saturates, entropy collapses, gradients spike
Explain why 1/sqrt(d_k) is an init-time correction, not a runtime bound

Primary sources. Browse if you want the original framing.