Research map / Optimization / Implicit bias and generalization
Why Warmup the Learning Rate? Underlying Mechanisms and Improvements
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.
Systematic study showing the main benefit of learning rate warmup is letting networks tolerate larger target learning rates by reducing sharpness (or preconditioned sharpness for Adam), and proposes loss-catapult-based eta_init selection and GI-Adam initialization.
Problem
Learning rate warmup is standard in deep learning, but it is unclear why it helps and to what extent it is necessary, especially for SGD and adaptive optimizers like Adam.
Approach
Warmup mainly allows training at larger target learning rates by annealing (pre-conditioned) sharpness through early catapults or natural sharpness reduction; the regime depends on whether initialization is in progressive sharpening or sharpness reduction. Based on this, the paper selects eta_init via a forward-pass catapult line search and proposes GI-Adam, which initializes the second moment wit
Key results
- For SGD, longer warmup extends the maximum trainable target learning rate, with the benefit most pronounced for large (SP) initializations and MSE loss.
- For Adam, pre-conditioned sharpness starts around 10^5 to 10^6 at initialization; longer warmup reduces catapult severity and avoids training failures.
- GI-Adam reduces initial pre-conditioned sharpness by about a factor of sqrt(1-beta_2) and consistently improves over standard Adam, pushing the training failure boundary to higher target learning rates.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase warmup duration | SGD · WRNs and Transformers across the reported datasets | maximum trainable target learning rate | Helped |
| increase warmup duration at a fixed target learning rate | SGD · reported training settings | test accuracy | No clear effect |
| increase warmup duration | Adam · reported training settings | training failures | Helped |
| initialize the second moment with the squared initial gradient (GI-Adam) | Adam · reported datasets | performance | Helped |
| initialize the second moment with the squared initial gradient (GI-Adam) | Adam · reported training settings | initial pre-conditioned sharpness: around two orders of magnitude | Helped |
| initialize the second moment randomly with the same norm as the initial gradients (ablation) | Adam · reported training settings | performance | Helped |
| remove first-moment bias correction from GI-Adam (Flat-Adam) (ablation) | GI-Adam · reported training settings | initial pre-conditioned sharpness decrease | Hurt |
| set the initial learning rate to the estimated critical learning rate instead of zero | linear warmup · reported training settings | steps to reach the target learning rate: T_{\text{reach}} = T_{\text{wrm}} (1 - \nicefrac{\eta_{c}}{\eta_{\text{trgt}}} ) | Helped |
| use flat μP initialization | SGD · reported training settings | need for warmup to achieve optimal performance | Helped |
| use a longer warmup | SGD with large initialization · reported training settings | performance | Helped |
Surprising findings
- Adam models always exhibit training failures rather than divergence, often with layers or residual blocks outputting zero.
- Adam's test accuracy degrades well before the failure boundary, unlike SGD where optimal learning rates lie near divergence.
Limitations
- Experiments were on relatively small-scale datasets and models, so generalization to larger settings is untested.
- Dependence on Adam hyperparameters beta1, beta2, and epsilon was not explored.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 Tiny-ImageNet WikiText-2 WikiText-103
Benchmarks: CIFAR-10 CIFAR-100 Tiny-ImageNet WikiText-103
Metrics: test accuracy test loss sharpness pre-conditioned sharpness training divergence/failure boundary
Baselines: Adam RAdam SGD with constant learning rate linear warmup
Builds on: Lewkowycz et al. catapult dynamics Adam Maximal Update Parameterization RAdam
Closest papers by meaning
- Analyzing & Reducing the Need for Learning Rate Warmup in GPT Training (NeurIPS 2024 · similarity 0.91)
- When Will Gradient Regularization Be Harmful? (ICML 2024 · similarity 0.85)
- Where Does Warm-Up Come From? Adaptive Scheduling for Norm-Constrained Optimizers (NeurIPS 2026 · similarity 0.85)
- Surge Phenomenon in Optimal Learning Rate and Batch Size Scaling (NeurIPS 2024 · similarity 0.84)
- Adam-mini: Use Fewer Learning Rates To Gain More (ICLR 2025 · similarity 0.82)
- Deep Linear Network Training Dynamics from Random Initialization: Data, Width, Depth, and Hyperparameter Transfer (ICML 2025 · similarity 0.82)
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.