QK-Norm

Gemma 3

Easy

RMSNorm on Q and K before attention

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Gemma 3 drops the attention-logit soft-cap used in Gemma 2 and instead RMS-normalizes the query and key tensors after the linear projections. The technical report states this swap explicitly and points to ViT-22B (Dehghani et al., 2023), the small-scale instability study of Wortsman et al. (2023), and Chameleon (2024). Those works treat QK-norm as a stabilizer for attention logits when residual streams get large.

The paper does not write the RMSNorm formula, the axis, or where the op sits relative to RoPE. Those details come from the official library. Two independent RMSNorm modules, _query_norm and _key_norm, run on the projected heads. Values are left unnormalized. Soft-capping is configured to None on every released Gemma 3 size.

The point of the op is to keep (QK^\top) from drifting with sequence length or residual growth. After per-head RMS, each query and key vector has a controlled second moment, so the subsequent (1/\sqrt{d}) scale (or the 27B variant of that scale) sees a more predictable logit range.

How it works

Let a projected query or key head be (x \in \mathbb{R}^{d_h}), with (d_h) equal to head_dim (256 on 1B–12B, 128 on 27B). Official RMSNorm in gemma/gm/nn/_layers.py is

\mathrm{RMSNorm}(x) = \frac{x}{\sqrt{\mathrm{mean}(x^2)+\varepsilon}}\,(1+\gamma),\qquad \varepsilon=10^{-6}.

The learned (\gamma) is initialized to zeros, so at init the multiplier is exactly 1. The mean is taken over the last axis only; batch, time, and head axes are independent. After the Q and K projections produce tensors of shape [B, T, H, d_h] (or the GQA key shape [B, T, H_{kv}, d_h]), Gemma 3 applies

Q \leftarrow \mathrm{RMSNorm}_Q(Q),\qquad K \leftarrow \mathrm{RMSNorm}_K(K)

and only then applies RoPE. Query scaling by query_pre_attn_scalar happens after RoPE. That order is specified by Attention.__call__ in gemma/gm/nn/_modules.py, not by the report.

QK-norm is therefore not a substitute for the usual (1/\sqrt{d}) factor. It is an extra, learned per-channel rescale of the head vectors. Because (\gamma) is shared across positions but not across the Q and K modules, the two sides of the dot product can drift independently during training.

What the paper specifies: replace Gemma 2 soft-capping with QK-norm; keep GQA and the pre/post RMSNorm sandwich around the block. What is inferred from code: last-axis RMS, ((1+\gamma)) parameterization, (\varepsilon=10^{-6}), application after the Q/K einsums and before RoPE, and no V-norm.

Official code

• gemma/gm/nn/_layers.py — RMSNorm (scale init zeros, rsqrt(var + 1e-06), multiply by 1 + scale).
• gemma/gm/nn/_modules.py — Attention.setup constructs _query_norm / _key_norm when use_qk_norm=True; __call__ applies them immediately after the Q/K projections.
• gemma/gm/nn/_gemma.py — every Gemma3_* config sets use_qk_norm=True and attn_logits_soft_cap=None.

Watch-outs

• Norm before RoPE, not after. Rotating first, then normalizing, mixes the phase that RoPE just wrote into the vector and changes the intended geometry.
• Do not reuse one RMSNorm for Q and K. The official modules have separate parameters; tying them is a different model.
• Gemma RMSNorm is ((1+\gamma)\hat{x}), not (\gamma\hat{x}). A zeros-init (\gamma) then behaves as identity. Implementing T5-style (\gamma\hat{x}) with zeros init zeroes the heads.
• QK-norm does not replace the query scale. Omitting query_pre_attn_scalar after the norms over-sharpens softmax.

Sources

• Paper: Gemma 3 Technical Report, §2 (architecture; QK-norm replaces soft-capping)
• Code: google-deepmind/gemma gemma/gm/nn/_layers.py, gemma/gm/nn/_modules.py, gemma/gm/nn/_gemma.py