QK Normalisation

The $1/\sqrt{d_k}$ in scaled dot-product attention assumes $q$ and $k$ have unit variance per component. At initialisation they do. After fifty thousand steps they do not — $W_Q$ and $W_K$ have grown, the logits have grown with them, and nothing in the architecture pushes back.

The end state is a softmax that has saturated: one weight is 1, the rest are 0, the gradient through the row is zero, and the head is dead. In a large run this shows up as a loss spike that does not recover.

If both vectors are normalised, their dot product is bounded by $d_k$ in the worst case and concentrated near zero in the typical one, regardless of what the projection weights have become. The logit cannot run away because it is no longer a function of the weight norms at all.

The evidence

Dehghani et al. hit this training ViT-22B: attention logits grew until the entropy of the attention distribution collapsed to near zero, always in the first few thousand steps, always fatal. QK-norm fixed it outright and was the change that made the 22-billion-parameter run possible.

Wortsman et al. then studied it at small scale deliberately, and found the same instability reproduces in models of a few hundred million parameters if you push the learning rate — which makes it a cheap thing to test. Their measurements show QK-norm widening the range of usable learning rates by roughly an order of magnitude, with the divergence threshold moving out rather than the optimum moving.

The variants, and what they cost

LayerNorm on $q$ and $k$ per head — the ViT-22B form, with the mean subtraction retained.

RMSNorm on $q$ and $k$ — what Gemma 2 and most recent models use, on the same reasoning that dropped the mean elsewhere.

$\ell_2$ normalisation with an explicit learned scalar — the original Henry et al. formulation, which makes the cosine interpretation exact.

All three cost two extra reductions over $[B, H, T, d_k]$ per layer. Against a $d_k$-deep matmul that is a small overhead in FLOPs, but it is bandwidth-bound and it sits between the projection and the attention kernel, where fusion is awkward — figure on 3–5% of step time in a naïve implementation, and less if the norm is folded into the projection epilogue.

import torch
from torch import Tensor, nn
from torch.nn import functional as F


class Attention(nn.Module):
    def __init__(self, dim: int, heads: int, qk_norm: bool = True):
        super().__init__()
        self.h, self.dk = heads, dim // heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)

        # Per-head, over the head dimension only.
        self.qn = nn.RMSNorm(self.dk) if qk_norm else nn.Identity()
        self.kn = nn.RMSNorm(self.dk) if qk_norm else nn.Identity()

    def forward(self, x: Tensor) -> Tensor:                   # [B, T, D]
        b, t, _ = x.shape
        q, k, v = self.qkv(x).view(b, t, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)

        q, k = self.qn(q), self.kn(k)                         # bound the logits

        z = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(z.transpose(1, 2).reshape(b, t, -1))

Applying the norm before the rotary rotation rather than after is the convention, and it matters: RoPE is orthogonal, so it preserves the norm, but the two orders produce different results once a learned per-channel gain is involved.

QK-norm bounds the attention logit. It does nothing for the output logits, which have their own growth problem — the usual answer there is a soft cap, $c \cdot \tanh(z/c)$, which Gemma 2 applies to both the attention scores and the final projection. Nor does it help with the loss spikes that come from the optimiser state rather than the forward pass; those need AdamW epsilon tuning or gradient clipping, and misdiagnosing one as the other wastes a lot of compute.