LIT-640: Query-Key Normalization for Transformers

Henry, Dachapally, Pawar and Chen (2020) — ARXIV-2010.04245

Key takeaways

• The technique, stated exactly. Apply ℓ₂ normalization along the head dimension of each query and key matrix before multiplying them, then scale up by a learnable parameter instead of dividing by the square root of the embedding dimension. Both halves matter and the second is the one that gets dropped in retellings.

• The motivation is softmax saturation, not stability at scale. The stated aim is to make the softmax “less prone to arbitrary saturation without sacrificing expressivity” — an expressivity argument about the attention distribution, arrived at from low-resource translation rather than from large-model training divergence.

• The evidence is modest and clearly bounded, and the baseline is narrower than it first looks. +0.928 BLEU averaged over five low-resource translation pairs from the TED Talks corpus and IWSLT’15. The abstract says this is against state-of-the-art bilingual benchmarks, which is true and imprecise: the baseline row is Nguyen and Salazar’s own system (LIT-651), whose repository this work starts from and whose five pairs it reuses. Both arms use PreNorm and FixNorm; the comparison changes two things at once — ScaleNorm becomes vanilla LayerNorm, and scaled dot-product attention becomes QKNorm. So the number is not QK-norm against a plain Transformer, it is QKNorm + LayerNorm against ScaleNorm + 1/√d.

  Per pair: en→vi 32.79→33.24, ar→en 36.09→36.75, en→he 28.28→28.96, gl→en 22.01→24.21, sk→en 32.58→33.23, all at p < 0.01. Five language pairs at translation scale — not a claim about frontier pretraining, which is where the record’s own practice now applies it.

  Measured against LIT-651’s system, then, QKNorm + LayerNorm wins on all five pairs, by 0.45 to 2.20 BLEU, and the gain cannot be assigned to QKNorm alone because ScaleNorm was removed in the same step.

• QKNorm is ScaleNorm, moved onto the queries and keys. The source says so directly: ScaleNorm “replaces layer normalization with ℓ₂ normalization along the embedding dimension, multiplied by a learnable scalar parameter”, and QKNorm “complements LayerNorm rather than replacing it”. The technique SOTA-192 recommends is LIT-651’s normalization applied at a different point in the block.

• It replaces a constant with a learned scalar, which is the detail that makes LIT-641’s measurement possible at all: because the temperature is learned by backpropagation, it can be read off a trained model and compared against what a theory would predict. Bricken and Pehlevan (LIT-641) did that: they obtained the learned temperatures of this paper’s trained QKNorm heads, found β between 10 and 25, and set that range against the three optimal values sparse distributed memory gives, which it falls between.

• Removing it has since been ablated outside translation. In LIT-708’s 20-layer attention-only decoder, the one run with QK-norm removed diverges at the tuned learning rate, the only divergence in that study. It is a single ablation at one rate and at most 24M parameters, and it tests the technique this paper introduced rather than repeating its BLEU comparison.

Standing in the anthology

Filed as a defect repair. SOTA-192, Active and converged, recommends normalizing queries and keys before the dot product — and carries source: LIT-088 and introduced_by: LIT-088, which is ViT-22B.

LIT-088 says so itself: “nothing in the record reads a source establishing QK-norm. This is one of the earliest at scale.” One of the earliest at scale is an adoption claim. ViT-22B did not introduce QK-norm; this paper did, two and a half years earlier, for a different reason and at a different scale.

30 documents in this record name QK-norm. SOTA-192 recommends it, SOTA-131 rescales query and key weights under Muon, SOTA-050 is what it extends, and LIT-449 records an independent arrival at it in diffusion transformers. The technique is everywhere in the record and its origin was not held.

#290 promoted this identifier as candidate #7 with exactly that reasoning — “named origin of QK-norm; LIT-088 uses it with no source for the technique” — and its promoted queue has not been worked: 21 of its 25 promotions are unfiled, and this unit reached the paper through the re-read of the same issue’s declines rather than through the list that already named the defect.

introduced_by is corrected to this paper in the same contribution. The distinction DP-005 draws is between evidence and adoption; this is the same distinction applied to origin, and the record had the later, larger, better-known paper in the slot.