Masters_Thesis/02-literature/qknorm-research-findings.md

Research Findings — QK-Norm at Small Scale (Option 4)

Compiled 24/07/2026. Working file for later review — not final notebook text. Updated 24/07/2026 after a second prior-art pass — see NOVELTY VERDICT.

Study in one line: At <100M parameters, characterise — with a dense learning-rate sweep — whether isolated Query–Key Normalization (QK-norm) widens the band of stable learning rates and speeds convergence, on commodity hardware.

Design axis: independent variable = QK-norm {on, off}; moderator = learning rate (dense sweep, safe → divergent); dependent = convergence speed (steps-to-threshold), max stable LR / LR-sensitivity, final val perplexity, attention entropy + activation kurtosis (mechanism). Base = plain nanoGPT, muP explicitly OFF, M4 Pro, ~10–15M params, TinyStories / FineWeb-Edu.

---

⚠ NOVELTY VERDICT (read first) — updated after 2nd search pass

Blunt assessment: this is a replication / verification study, not a novel contribution. Two rounds of searching show the effect and even the nuanced reconciliation are already established:

• QK-norm's LR-stabilisation is extensively documented: rigorously by Wortsman et al. 2023 (formal LR-sensitivity, small→4.8B), at 70M by IMU-1 (2026) where it's the single largest architectural win, at 830M by 2410.16682 ("~1.5× higher LR"), and in 181M–1.3B ablations (SimpleGPT / "Controlling changes to attention logits").
• The LR-conditional reconciliation I was going to claim as the "gap" is already the community view: a 360M analysis shows QK-norm gives no benefit (or slightly worse) at low LR but prevents divergence at high LR — i.e. "shallow small models don't need it unless you push the LR." That is now the accepted mechanism, not an open question.
• A blog (aflah02) already isolated QK-norm and ran a (coarse, 3-point) LR sweep at 1B / 100B tokens.

What genuinely remains un-done (thin, incremental, but real):

1. A dense LR-sensitivity curve for isolated, muP-free QK-norm specifically in the 10–100M band. Wortsman is larger & partly bundled; IMU-1 uses kurtosis not a sweep & bundles muP; the blog is 1B & coarse. Nobody has published the clean small-scale swept curve.
2. Laptop-reproducible artifact (open code/logs/seeds runnable on an M4). Soft, but true.

Honest recommendation: for a research-methods module, a well-executed reproduction/verification study is legitimate and low-risk — it demonstrates method, which is what's being graded. But the pitch must be framed as "a rigorous, isolated, commodity-hardware verification of a known effect," NOT as a discovery. Do not claim novelty of technique, "first at small scale," or "resolves a contradiction."

If more novelty is wanted (optional pivot for the parent to decide):

• Cross QK-norm with a less-saturated axis — e.g. its interaction with positional encoding (NoPE/RoPE) or with depth/width at fixed budget, where the LR story is less mapped; or
• Study the precise LR threshold at which QK-norm flips from "slightly harmful" to "necessary" as a function of model depth (a small, clean, and not-yet-charted boundary) — this is the most defensible sharpening that stays within the current setup.

---

A. CORE ACADEMIC PAPERS

1. Henry et al. (2020) — "Query-Key Normalization for Transformers" (origin of QKNorm)

• EMNLP 2020 Findings. arXiv:2010.04245 — https://arxiv.org/abs/2010.04245
• ℓ2-normalise Q and K before the dot product; replace 1/√d with a learnable g. For low-resource NMT.
• Relevance: the technique's origin & definition. Also proves QK-norm was born on small models.

2. Dehghani et al. (2023) — "Scaling Vision Transformers to 22B Parameters" (the "why" citation)

• arXiv:2302.05442 — https://arxiv.org/abs/2302.05442
• Popularised QK-norm to stop attention-logit growth / divergence at scale.
• Relevance: canonical large-scale motivation.

3. Zhai et al. (2023, ICML) — "Preventing Attention Entropy Collapse" (mechanism)

• arXiv:2303.06296 — https://arxiv.org/abs/2303.06296
• "Entropy collapse" (spiky attention) ↔ instability. Proposes σReparam.
• Relevance: the why; supplies attention entropy as a metric to log.

4. Wortsman et al. (2023) — "Small-scale proxies for large-scale Transformer training instabilities" (THE prior work)

• ICLR 2024 oral. arXiv:2309.14322 — https://arxiv.org/abs/2309.14322
• Instabilities reproduce in small models at high LR; qk-layernorm + z-loss are "equally effective" small. Introduces the LR-sensitivity metric. Confirms with a 4.8B run.
• Relevance: the methodological template and the reason this study is verification, not discovery. Our only deltas vs. this: smaller (10–100M), isolated QK-norm, muP-off, laptop-reproducible.

5. IMU-1 (Grigorev et al., 2026) — "Sample-Efficient Pre-training of Small Language Models" (closest prior art)

• arXiv:2602.02522 — https://arxiv.org/abs/2602.02522
• 430M model + 70M ablations. QK-norm is the single largest individual architectural win (−0.63% train loss at 70M; reduces activation kurtosis). BUT bundled with muP, NorMuon, value residuals, LN-scaling — so its LR effect is entangled with muP. Uses kurtosis, not an LR sweep.
• Relevance: kills the "isolation at small scale is un-done" claim. Our cleaner angle: isolate + sweep, muP-off. Cite as the paper we're verifying/disentangling.

6. OLMo 2 Team (2025) — "2 OLMo 2 Furious" (modern adoption)

• arXiv:2501.00656 — https://arxiv.org/abs/2501.00656 — RMSNorm on Q,K improves loss + stability. Also used by ViT-22B, Gemma 2/3, Qwen3, Chameleon.

7. Small-model ablations that DISPUTE the benefit (the tension — but see verdict)

• codelion, "The Optimal Architecture for Small LMs": https://huggingface.co/blog/codelion/optimal-model-architecture (19 configs @ ~70M/1B tokens; "modern tricks no measurable benefit at 70M").
• "Architectural Trade-offs in SLMs Under Compute Constraints" arXiv:2512.20877 — https://arxiv.org/abs/2512.20877
• Relevance: the apparent contradiction with IMU-1 — but it's already reconciled by the LR-conditionality finding (QK-norm helps mainly at high LR / when instability would otherwise occur).

B. ADJACENT / CONTEXT

• 2410.16682 "Methods of improving LLM training stability" — https://arxiv.org/abs/2410.16682 — 830M; QK-norm + softmax cap → ~1.5× higher LR before divergence.
• QK-Normed MLA (2026) arXiv:2606.16310 — https://arxiv.org/abs/2606.16310 — head-to-head high-LR stress test: QK-norm keeps training while QK-Clip NaNs.
• SimpleGPT (2602.01212) & "Controlling changes to attention logits" (2511.21377) — 181M–1.3B QK-norm LR ablations (GA-QK-Norm / QK-Norm / baseline).
• Yang et al. (2022) μP / μTransfer arXiv:2203.03466 — LR transfer across width; the confound we hold OFF.
• "Enhanced QKNorm with the Lp norm" arXiv:2602.05006 — a QK-norm variant (possible extension).
• Peri-LN (2502.02732) / HybridNorm (2503.04598) — normalization-placement neighbours.

C. NON-ACADEMIC MATERIAL

Directly on-topic (assess overlap carefully)

• aflah02 — "On Stability with QK Norm during Pretraining with NoPE": https://aflah02.substack.com/p/on-stability-with-qk-norm-during — isolates QK-norm, 3-point LR sweep (4e-5/4e-4/4e-3), Llama-3.2-1B, FineWeb-Edu 100B. Closest blog; 1B & coarse & NoPE-specific, so our <100M dense sweep still differs.
• Ishan Mukherjee — "QK norm is probably a free lunch": https://ishanjmukherjee.github.io/qk-norm — advocacy/survey (adoption across Llama4/Gemma3/Qwen3/DeepSeek-V3; Marin-32B loss-spike anecdote). Not empirical.

Explainers

• Sebastian Raschka — "The Big LLM Architecture Comparison" (§2.2 QK-Norm): https://magazine.sebastianraschka.com/p/the-big-llm-architecture-comparison ← best plain-English intro
• Sebastian Raschka — "LLM Architecture Gallery: QK-Norm": https://sebastianraschka.com/llm-architecture-gallery/qk-norm/
• emergentmind.com: /topics/query-key-normalization-qk-norm , /topics/attention-entropy-collapse

Code (experimental substrate)

• karpathy/nanoGPT: https://github.com/karpathy/nanoGPT ← RECOMMENDED clean base (muP-off, AdamW)
• karpathy/nanochat + Issue #362 (QK-norm vs logit softcap debate): https://github.com/karpathy/nanochat/issues/362
• KellerJordan/modded-nanogpt: https://github.com/kellerjordan/modded-nanogpt

Video / courses (confirm URLs before citing)

• Karpathy — "Let's reproduce GPT-2 (124M)", "Let's build GPT: from scratch" (YouTube).
• Stanford CS336 "Language Modeling from Scratch".

D. SYNTHESIS — ESTABLISHED vs WHAT'S LEFT

Established (do NOT re-argue as if new): QK-norm bounds attention-logit growth → prevents entropy collapse → widens the stable-LR band, especially at high LR; cheap; standard in frontier models. This reproduces at small scale (Wortsman; IMU-1 @70M; multiple 181M–1.3B ablations; a 1B blog). The low-LR-no-benefit / high-LR-necessary reconciliation is already the accepted mechanism.

What's genuinely left (thin): a dense LR-sensitivity curve for isolated, muP-free QK-norm in the 10–100M band, released as a laptop-reproducible artifact; and (sharper) charting the LR threshold at which QK-norm flips from slightly-harmful to necessary vs. model depth.

Cleanest defensible claim: "We reproduce, with a dense LR sweep at 10–100M on commodity hardware, the max-stable-LR / LR-sensitivity benefit of isolated QK-norm, and locate the LR×depth boundary at which it becomes necessary." Framed as verification + boundary-charting, not discovery.

E. QUICK REFERENCE LIST (for notebook Section 2 — verify all before final use)

• Henry, Dahl, Bhoopchand, Bica (2020). Query-Key Normalization for Transformers. EMNLP Findings. arXiv:2010.04245
• Dehghani et al. (2023). Scaling ViTs to 22B Parameters. arXiv:2302.05442
• Zhai et al. (2023). Preventing Attention Entropy Collapse. ICML. arXiv:2303.06296
• Wortsman et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. ICLR 2024. arXiv:2309.14322
• Grigorev et al. (2026). IMU-1: Sample-Efficient Pre-training of Small LMs. arXiv:2602.02522
• OLMo 2 Team (2025). 2 OLMo 2 Furious. arXiv:2501.00656
• (2024). Methods of improving LLM training stability. arXiv:2410.16682
• Yang et al. (2022). Tensor Programs V (μP / μTransfer). arXiv:2203.03466
• (2025). Architectural Trade-offs in SLMs Under Compute Constraints. arXiv:2512.20877

F. SURROUNDING LITERATURE — added 18/08/2026

[TO VERIFY] = surfaced via web search this round; confirm each against Scopus / the primary PDF before use. arXiv IDs below were taken from live result links (reliable); author names & venues come from search summaries (treat as unconfirmed). Papers without an arXiv ID need the ID found too.

F.1 Why attention logits explode — massive activations & attention sinks (mechanism depth)

• Sun et al. (2024), "Massive Activations in Large Language Models," ICML 2024. arXiv:2402.17762 — a few tokens develop extreme-magnitude activations → dominant key vectors → "attention sink"; the logit-blow-up mechanism QK-norm targets. [TO VERIFY]
• "When Attention Sink Emerges in Language Models: An Empirical View." arXiv:2410.10781 [TO VERIFY]
• "The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks" (2026). arXiv:2603.05498 — "write–sink–cancel" cycle. [TO VERIFY]
• "A Single Layer to Explain Them All: Understanding Massive Activations in LLMs" (2026). arXiv:2605.08504 [TO VERIFY]
• Relevance: mechanistic spine for Component 2 — links Zhai (entropy collapse) + IMU-1 (kurtosis); QK-norm = magnitude control.

F.2 Competing / complementary logit-control & stability techniques (supports the "bundled" gap)

• Attention logit soft-capping — Gemma 2 report: z ← c·tanh(z/c), c=50 (attn)/30 (final). The direct alternative to QK-norm (cf. nanochat issue #362). [TO VERIFY — find arXiv ID + exact citation]
• z-loss — PaLM (Chowdhery et al., 2022): auxiliary loss keeping softmax normalizer log Z ≈ 0; the fix Wortsman bundles with qk-layernorm. [TO VERIFY — find arXiv ID]
• NormFormer (Shleifer et al., 2021), extra normalization. arXiv:2110.09456 [TO VERIFY authors]
• Peri-LN (2025) arXiv:2502.02732; HybridNorm (2025) arXiv:2503.04598 — normalization-placement neighbours. [TO VERIFY]
• Relevance: QK-norm is one of several logit-control methods, almost always studied bundled → reinforces the isolated-comparison gap (Component 3).

F.3 Recent QK-norm-specific work (2025–2026) — it's a live topic

• QK-Normed MLA (2026). arXiv:2606.16310 — QK-norm in Multi-head Latent Attention; reaffirms QK-norm "permits significant increases in learning rate" (supports H2); proposes per-head LR scaling (QuacK) where QK-norm doesn't apply. [TO VERIFY]
• "Controlling changes to attention logits" (2025). arXiv:2511.21377 [TO VERIFY]
• "Enhanced QKNorm with the Lp norm" (2026). arXiv:2602.05006 — QK-norm variant (possible extension). [TO VERIFY]

F.4 Reproducible small-model pretraining & data (method + significance)

• TinyStories (Eldan & Li, 2023). arXiv:2305.07759 — sub-10M-param models produce coherent English; grounds both dataset choice and the tiny-model scale. [TO VERIFY]
• Cramming (Geiping & Goldstein, ICML 2023). arXiv:2212.14034 — near-BERT on a single GPU in one day; precedent for the accessible/laptop-reproducible framing. [TO VERIFY]
• BabyLM Challenge (Warstadt et al.). arXiv:2301.11796 (findings arXiv:2504.08165) — sample-efficient pretraining benchmark; relevant to the depth-vs-width data-constrained fallback. [TO VERIFY]
• SmolLM2 (Allal et al., 2025) — over-training small models; context for why small models matter and are trained long. [TO VERIFY — find citation/arXiv ID]