8.4 KiB
8.4 KiB
marp, theme, paginate, size, header, style
| marp | theme | paginate | size | header | style |
|---|---|---|---|---|---|
| true | default | true | 16:9 | CS6513 · Research Proposal Pitch | section { font-size: 26px; } h1 { color: #1a3b5d; } h2 { color: #1a3b5d; } strong { color: #b0451f; } section.title h1 { font-size: 46px; } small { color: #666; } |
When Does a Small Model Need QK-Norm?
A reproducible, isolated learning-rate characterization of Query–Key Normalization at <100M parameters
Cormac Shannon · ID 23366303 MSc Immersive Software Engineering · CS6513 Research Proposal Pitch
Provisional working title — draft pitch
1 · The Problem & Topic
- Modern LLMs stack many components — RoPE, RMSNorm, GQA, SwiGLU, QK-norm — mostly validated at large scale. QK-norm normalises the query & key vectors so attention logits can't explode.
- What's already known: QK-norm improves training stability, especially at high learning rates — shown by dense LR-sweeps but with QK-norm bundled with other fixes (Wortsman, ICLR 2024), and as the top single architectural win in 70M ablations (IMU-1, 2026 — by cutting extreme activation outliers, not an LR sweep).
- What's still fuzzy — my focus: the qualitative rule is known; a clean, dense measurement is not. Small-scale evidence is patchy and inconsistently measured (larger-scale proxies; looser signals — extreme activation outliers, a single "≈1.5× LR" figure, a coarse 3-point blog sweep at 1B). Un-charted: the dense LR-sensitivity curve for isolated QK-norm at 10–100M, and how the LR at which it becomes necessary moves with model depth.
2 · Significance — the "So What?"
- The field is going small — on-device/edge, cost, energy, reproducibility all push toward sub-100M LMs. Practitioners need concrete rules for that regime.
- A usable rule of thumb: if I can map the learning-rate (× depth) threshold where QK-norm flips from harmless-to-drop to necessary, a small-model builder gets a direct answer — "below this LR you can skip it; above it you can't." The qualitative flip is known; the precise threshold — and how it moves with depth — isn't charted at this scale.
- Evidence quality: current small-scale results use looser metrics; a dense LR-sensitivity sweep (à la Wortsman) puts the small-model claim on firmer footing.
- Reproducibility: a laptop-runnable, open artifact adds trustworthy, re-runnable evidence to a literature dominated by expensive, hard-to-reproduce large-scale results.
3 · Conceptual Framework — Questions & Objectives
Primary RQ. At <100M parameters, does isolated QK-norm's stability benefit reproduce under a dense learning-rate sweep, and at what learning rate does it become necessary?
Sub-RQ. How does that necessity threshold shift with model depth?
Hypotheses
- H1 — QK-norm's effect is LR-conditional: negligible/slightly negative at low LR, decisive at high LR.
- H2 — QK-norm raises the max stable LR and lowers LR-sensitivity at <100M.
- H3 — the LR at which QK-norm becomes necessary falls as depth increases (deeper ⇒ needs it sooner).
Objectives: isolate a muP-free QK-norm toggle → dense LR sweep × a few depths → quantify max-stable-LR, LR-sensitivity, activation outliers → chart the LR × depth necessity boundary → release the artifact.
4 · Methodological Approach
Quantitative, experimental — a controlled ablation (per Creswell, Ch. 1).
| Element | Choice |
|---|---|
| Independent variable | QK-norm {on, off} |
| Moderators | Learning rate (dense, log-spaced, safe → divergent) · depth (2–3 values) |
| Controls | identical width/data/tokens/optimizer (AdamW) · μP explicitly OFF · fixed seeds |
| Model | plain nanoGPT decoder, ~10–15M params |
| Data | TinyStories / FineWeb-Edu slice (fixed token budget) |
| Measures | max stable LR · LR-sensitivity (Wortsman) · steps-to-threshold · perplexity · attention entropy · extreme activation outliers |
Analysis: loss-vs-LR curves per depth (≥2 seeds, mean ± spread); locate the divergence boundary and the QK-norm necessity threshold. Base = plain nanoGPT, μP off so QK-norm is the only variable.
5 · Feasibility & Timeline
- Runs on my hardware: M4 Pro (48 GB), PyTorch MPS. ~10–15M model on ~100–150M tokens ≈ ~1–1.5 h/run. Core matrix: QK{2} × LR{6} × seed{2} ≈ 24 runs → 2–3 overnight batches.
- Stretch (the depth boundary): add 2 more depths near the threshold (
+18 runs) — or a rented RTX 4090 ($0.40/hr) clears it in a few dollars. - Not GPT-2 scale — an ablation needs relative effects, so tiny models are correct, not a compromise.
| Week | Milestone |
|---|---|
| 5 | Fork nanoGPT · isolated μP-off QK-norm toggle · reproduce baseline · pilot runs |
| 6 | Core dense LR sweep (≥2 seeds) · log curves |
| 7 | Analysis + (stretch) depth-threshold runs · figures |
| 8 | Write Literature Review + Introduction · finalise notebook |
6 · Ethics & Data Access
- No human subjects, no personal or sensitive data → no ethical-approval barriers.
- Public, openly-licensed datasets (TinyStories, FineWeb-Edu) — licences to be confirmed and respected; no data-access constraints.
- Low environmental footprint — a few GPU-hours total; energy-efficiency is itself a study motivation.
- Reproducibility & integrity: code, configs, seeds, and logs released openly; AI-tool use declared in the research notebook per module policy.
7 · Summary
- Not "does QK-norm help" (known) — but when a small model needs it, charted cleanly.
- A dense, isolated, μP-free LR sweep at <100M → the max-stable-LR / LR-sensitivity benefit and the LR × depth necessity boundary, reproducible on a laptop.
- An honest verification-plus-boundary-charting study: firms up patchy small-scale evidence and hands practitioners a usable rule.
Thank you — questions welcome
Key references: Henry et al. 2020 · Dehghani et al. 2023 (ViT-22B) · Wortsman et al. (ICLR 2024) · IMU-1 (2026) · OLMo 2 (2025) — venues/details to verify against primary sources