Research map / Architectures and efficiency / RNNs and normalization
Small-scale proxies for large-scale Transformer training instabilities
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Reproduces Transformer training instabilities (attention logit growth, output logit divergence) in small models at high learning rates and studies interventions (qk-layernorm, z-loss, warm-up, weight decay, muParam) via a learning-rate-sensitivity metric.
Problem
Large Transformer training instabilities are costly to reproduce and study; the paper asks whether they can be reproduced and predicted at small scale.
Approach
Measure LR sensitivity (expected deviation from optimal final loss across three orders of magnitude of learning rate) as a summary of LR-vs-loss curves across model scales, use it to reproduce known instabilities in small models, and extrapolate scaling trends of activation/gradient norms to predict instabilities.
Key results
- Qk-layernorm allows training a 1.2B parameter model at learning rate 0.3 and reduces LR sensitivity, though LR sensitivity still increases with scale.
- A quadratic fit of max attention logit vs. parameter count predicted that a 4.8B parameter model would diverge at LR 1e-2 without qk-layernorm; the model diverged as predicted.
- Gradient RMS approaches the default AdamW epsilon (1e-8) at 4.8B parameters and LR 0.3; decreasing epsilon to 1e-15 improves loss while increasing it to 1e-6 causes divergence.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add qk-layernorm (ablation) | small Transformers · LR sweep 3e-4 to 3e-1 | LR sensitivity / training stability | Helped |
| add z-loss regularization (coefficient 1e-4) (ablation) | small Transformers with output logit divergence · LR vs loss curves | instability from output logit divergence | Helped |
| use independent weight decay (decoupled from LR) (ablation) | small Transformers · LR sweep | LR sensitivity | Helped |
| longer warm-up (ablation) | small Transformers with qk-layernorm · LR sweep, model sizes | LR sensitivity | Helped |
| scale depth instead of width (ablation) | small Transformers · LR sensitivity scaling | LR sensitivity | Hurt |
| muParam (simple) (ablation) | Transformers with qk-layernorm · LR sweep | LR sensitivity and loss | No clear effect |
| decrease AdamW epsilon from 1e-8 to 1e-15 (ablation) | 4.8B Transformer at LR 0.3 · training loss | loss | Helped |
| increase AdamW epsilon from 1e-8 to 1e-6 (ablation) | 4.8B Transformer at LR 0.3 · training loss | divergence | Hurt |
| increase weight decay from 1e-4 (ablation) | small Transformers · LR sweep | optimal LR location | No clear effect |
| change attention scaling from 1/sqrt(d_h) to 1/d_h (muParam full change ii) (ablation) | small Transformers · LR sweep | loss | Mixed |
| increase batch size from 256 to 512 or 1024 (ablation) | small Transformers · LR sensitivity | LR sensitivity | No clear effect |
| change number of training steps from 1e5 to 5e4 or 2e5 (ablation) | small Transformers · LR sweep | LR sensitivity | No clear effect |
Surprising findings
- Muparam stabilizes the optimal learning rate but does not reduce LR sensitivity or improve loss, and does not remove the need for qk-layernorm.
- Scaling depth increases LR sensitivity faster than scaling width, though depth scaling gives lower loss at the largest scale tested.
Experimental setup
Datasets: C4
Metrics: validation loss LR sensitivity max attention logit gradient RMS update RMS
Baselines: baseline Transformer without qk-layernorm muParam
Builds on: qk-layernorm z-loss muParam AdamW GPT-2
Closest papers by meaning
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.88)
- Taming Transformer Without Using Learning Rate Warmup (ICLR 2025 · similarity 0.88)
- On the Surprising Effectiveness of Large Learning Rates under Standard Width Scaling (NeurIPS 2025 · similarity 0.87)
- From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers (NeurIPS 2025 · similarity 0.87)
- Infinite Limits of Multi-head Transformer Dynamics (NeurIPS 2024 · similarity 0.87)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.