Interactive Explainer

LR Schedule Visualizer

Four schedules, one canvas. Drag peak LR, warmup fraction, total steps — see exactly how the learning rate moves through training, and why Transformers cold-started at full LR diverge in their first 100 steps.

The same optimizer can succeed or fail depending on the learning rate schedule you wrap around it. Cold-starting Adam on a Transformer with lr = 3e-4 diverges within a few hundred steps. Warm it up linearly from 0 and the same architecture trains cleanly.

The playground

PyTorch for the current schedule

Why Transformers need warmup

Adam's second-moment estimate v̂_t is tiny and noisy at step 1. Dividing by √v̂_t wildly amplifies the first few steps. Meanwhile, random-init attention produces peaky distributions — large early gradients in a few heads. Combined, cold-starting at full LR produces huge unstable first steps, and the loss diverges.

Warmup linearly ramps LR from 0 to peak over 1–10% of training. Once v̂_t has a few steps to stabilize and attention has diffused, you can safely ride the peak LR.

The four-stage anatomy of a modern LR schedule

1. Warmup (0 → peak). Linear ramp for 1–10% of total steps. Stabilises Adam's moment estimates and gives initialisation time to settle.
2. Peak (constant or near-constant). The bulk of training. Some recipes (Llama, Chinchilla) hold constant; others (cosine, Noam) start decaying immediately after warmup.
3. Decay (peak → near-zero). Cosine, linear, or step. The exact shape matters less than the magnitude — what matters is reaching a sufficiently low final LR.
4. Tail (~0). Optional. Final low-LR steps for last-mile fine-tuning. Standard in LLM pretraining.

Schedule cheatsheet by use case

• Computer vision (ResNet, ViT) on ImageNet. Cosine decay over 100 epochs, no warmup needed for ResNet, ~5 epochs warmup for ViT. Peak LR ≈ 0.1 (SGD) or 1e-3 (AdamW).
• LLM pretraining. Linear warmup ~1% of total tokens, then cosine decay to 10% of peak. Peak LR ~3e-4 (small) down to ~1e-4 (70B+).
• LLM SFT / instruction tuning. Short linear warmup (~50 steps), cosine to zero. Peak LR ~2e-5 to ~5e-5. Much lower than pretraining to avoid catastrophic forgetting.
• LoRA fine-tuning. Constant or single cosine cycle. Peak LR ~1e-4. Adapter parameters tolerate higher rates than the base.
• Object detection / segmentation. Step schedule (multiply by 0.1 at 60%, 80% of training). The classical DETR / Mask R-CNN recipe.
• Diffusion training. Constant LR + EMA on weights. The EMA does most of the implicit "decay".

One-cycle vs cosine vs constant

Three regimes you'll see in production:

• Constant LR + decay-at-the-end. Used by Chinchilla, recent Llama runs. Empirically as good as cosine when training duration is well-tuned; simpler to reason about for very long runs.
• Cosine annealing (Loshchilov & Hutter, 2017). Smooth decay; the default in HuggingFace Trainer, PyTorch Lightning. Cosine with restarts ("SGDR") periodically resets to peak — was popular pre-2022 for vision, mostly abandoned for LLMs.
• 1-cycle (Leslie Smith, 2018). LR + momentum sweep — short warmup, long cosine decay, very low tail. Strong for short training runs (fewer than 1000 epochs). Mostly historical for LLMs but still useful for small-scale CNN training.

Reading list

• Loshchilov & Hutter (2017) — SGDR: Stochastic Gradient Descent with Warm Restarts.
• Smith (2018) — A Disciplined Approach to Neural Network Hyper-Parameters. The 1-cycle paper.
• Hoffer et al. (2017) — Train Longer, Generalize Better.
• Liu et al. (2020) — On the Variance of the Adaptive Learning Rate and Beyond (RAdam). Why warmup helps Adam.
• Hoffmann et al. (2022) — Training Compute-Optimal Large Language Models (Chinchilla). The constant-LR-with-decay recipe.