Research map / Architectures and efficiency / RNNs and normalization
Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training
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.
Analyzes why learning rate warmup helps GPT-2 training with AdamW/Lion by measuring update size in l2-norm, angular update, and relative representation change, and proposes optimizers (LionA, LionAR, RRC-corrected, high-momentum variants) that reduce or eliminate the need for warmup.
Problem
Warmup is widely used when training neural networks, especially at large batch sizes, but its benefits are poorly understood; the paper asks why and by which criteria early updates are too large in GPT-2 training with adaptive optimizers.
Approach
Argues warmup keeps the effective update size limited, and measures update size via l2-norm, angular update, and relative representation change (RRC). Modifies Lion into LionA (fixed l2 update), LionAR (rotational, fixed angular update), adds an SNR-based RRC correction, and shows higher momentum and inverse bias correction can remove warmup need.
Key results
- AdamW without warmup shows a large permanent validation-loss gap versus short warmup on GPT-2 124M; run-to-run variance about 0.05 validation loss.
- LionAR with angular-update control reduces but does not eliminate the warmup benefit; LionAR with momentum 0.98, Nesterov and inverse bias correction shows no warmup performs best.
- On Llama2-124M, LionAR alone was already sufficient to eliminate the need for warmup.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a short learning-rate warmup | GPT2-124M trained with AdamW · OpenWebText | validation loss | Helped |
| Control the L2 update norm with LionA instead of AdamW | GPT2-124M training · OpenWebText | validation loss | Mixed |
| Control angular updates with LionAR | GPT2-124M training · OpenWebText | validation loss | Helped |
| Add an SNR-based relative representation change correction | LionAR · GPT2-124M training | validation loss | Mixed |
| Remove momentum (ablation) | LionAR · GPT2-124M training | validation loss | Hurt |
| Increase momentum to β=0.98, use Nesterov momentum, and apply inverse bias correction (ablation) | LionAR · GPT2-124M training | validation loss | Helped |
| Increase momentum to β=0.98, use Nesterov momentum, and apply inverse bias correction (ablation) | LionA · GPT2-124M training | validation loss | No clear effect |
| Use higher-momentum LionAR with Nesterov momentum and momentum correction (ablation) | GPT2-124M · SlimPajama | validation loss | Helped |
| Use LionAR to control angular updates (ablation) | Llama2-124M · OpenWebText | validation loss | Helped |
| Replace ReLU with Leaky ReLU (α=0.1) (ablation) | ResNet-20 · CIFAR-10 | test accuracy | Helped |
| Freeze biases at the onset of training (ablation) | ResNet-20 · CIFAR-10 | test accuracy | Helped |
| Use random gradient directions at the start of training instead of actual gradients (ablation) | ResNet-20 · CIFAR-10 | test accuracy | No clear effect |
Surprising findings
- Controlling the l2-norm of updates exactly (LionA) does not remove the warmup benefit, showing the update norm alone does not explain the need for warmup.
- The paper's RRC-based automatic warmup closes the gap but yields slightly worse overall performance, and the authors still recommend a short warmup in practice.
Limitations
- Main experiments use a single small GPT-2 124M network and OpenWebText, which may not generalize.
- Additional dataset and architecture experiments are limited in scale and hyperparameter range.
- The RRC correction is not tuned and is impractical; it distorts the decay phase and requires large learning rates.
Experimental setup
Datasets: OpenWebText SlimPajama CIFAR-10
Benchmarks: OpenWebText validation loss SlimPajama validation loss CIFAR-10 test accuracy
Metrics: final validation loss update l2-norm angular update size relative representation change fraction of dead ReLUs stable rank
Baselines: AdamW with warmup RAdamW Lion LionA SGD with momentum
Builds on: AdamW Lion Rotational Optimizers Weight Normalization NanoGPT
Closest papers by meaning
- Why Warmup the Learning Rate? Underlying Mechanisms and Improvements (NeurIPS 2024 · similarity 0.91)
- Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers (NeurIPS 2026 · similarity 0.87)
- When Will Gradient Regularization Be Harmful? (ICML 2024 · similarity 0.87)
- Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training (NeurIPS 2025 · similarity 0.85)
- Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation is Wasteful (NeurIPS 2025 · similarity 0.85)
- Where Do Large Learning Rates Lead Us? (NeurIPS 2024 · similarity 0.85)
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.