Learning rate schedules: warmup and cosine decay

Summary

A learning-rate schedule is a function $\eta (t)$ that varies the optimizer’s step size over training. The dominant 2026 default for LLMs is linear warmup for a few hundred to a few thousand steps, followed by cosine decay down to ~10% of the peak LR.

Why warmup

Early training can combine poorly estimated optimizer moments, high curvature near initialization, and rapidly changing activations. Jumping immediately to the target LR can make those transient updates unstable. Warmup limits the initial step size while the model and optimizer state enter a regime that can support the peak LR.

Typical warmup: $W=1000$ to $4000$ steps for pretraining; ~100 steps is enough for fine-tuning.

Why cosine

After warmup, decay from the peak LR toward a lower floor over the remaining horizon. Cosine decay,

$$\eta (t)={\eta }_{\min }+\frac{1}{2}({\eta }_{\max }-{\eta }_{\min })\left(1+\cos \left(\pi \cdot \frac{t-W}{T-W}\right)\right),$$

starts and ends with a shallow slope while changing fastest near the middle of the decay interval (Loshchilov & Hutter, 2017).

Common practice: cosine to ${\eta }_{\min }=0.1\cdot {\eta }_{\max }$ over the full training horizon $T$.

Learning objective

Where are the three defining checkpoints of warmup plus cosine decay?

How to set the peak LR

For Adam/AdamW on transformer training, the default starting point is ${\eta }_{\max }=3\times 10^{-4}$ (Karpathy’s “magic constant”) for moderate batch sizes. Larger batches scale up roughly linearly until the LR-batch tradeoff breaks ($\sim 10^{-3}$ for very large batch).

For SFT or task-specific fine-tuning of a pretrained model: 10–100× lower than pretraining ($1\times 10^{-5}$ to $5\times 10^{-5}$).

A LR range test (Smith, 2017): sweep $\eta$ from $10^{-7}$ to $10^{-1}$ over a few hundred steps, plot loss vs. LR. Pick the point at the steepest descent (typically ~10× below where the loss diverges).

Other schedules

• Constant: useful for online learning / RL where the data distribution shifts.
• Inverse square root (original transformer paper): $\eta (t)\propto 1/\sqrt{t}$ after warmup. Largely superseded by cosine.
• One-cycle (Smith, 2018): warmup to a high peak, then decay aggressively. Used in some vision training; uncommon in LLMs.
• WSD (Warmup-Stable-Decay): warmup, hold constant for most of training, decay sharply at the end. Used in some recent LLM training (Hu et al., 2024) for easier checkpoint resumption.

Common pitfalls

• Skipping warmup. Adam without warmup on a transformer routinely diverges.
• Decaying too fast. Aggressive decay limits how far the model can move; cosine-to-10% is a reliable default.
• Forgetting LR scales with batch. Doubling the batch usually requires roughly doubling the LR.
• Resuming a cosine schedule from a checkpoint. If the cosine is parameterized over total steps, resuming with a different total breaks the schedule. Save schedule state with the checkpoint.