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]`