Why do LLMs use learning-rate warmup, and how do cosine and warmup–stable–decay schedules compare?

30-second answer

Early in training, Adam's second-moment estimates rest on few samples and activations haven't settled, so full-size steps can be destructive; a linear warmup over the first ~1–2k steps keeps updates small until statistics are reliable (post-norm models need much more). Cosine decay to ~10% of peak must be sized to the run length. Warmup–stable–decay holds the peak rate for most of training then decays quickly at the end; it matches cosine quality while letting you branch decays from any checkpoint, extend training and run scaling studies cheaply.

Try answering out loud before reading on. The walkthrough below is paced for a 40-minute interview round.

0–5 min Clarify and frame

Framing: "I'll explain why warmup is needed, compare cosine with warmup–stable–decay, then cover choosing the peak learning rate, its interaction with batch size, and how schedules affect scaling studies."

5–15 min Core explanation

Warmup. Linearly increase the learning rate from ~0 to its peak over the first ~0.5–2k steps.

Adam's v̂ is estimated from very few samples early on; underestimates make m̂/√v̂ large and steps destructive.
Randomly initialized networks have large, poorly scaled gradients at first; residual-stream statistics settle over the first steps.
Post-norm transformers are especially sensitive and need long warmup; pre-norm models need little, but still benefit.

Cosine decay. η_t = η_min + ½(η_max − η_min)(1 + cos(π · t / T)) after warmup, usually to ~10% of the peak. Much of the final loss improvement arrives during the decay. The total length T must match the planned training steps.

Warmup–stable–decay (WSD). Warm up, hold η_max for most of training, then decay over the last ~10–20% (linear, cosine or 1 − √ shapes).

15–25 min Deep dive

Scroll sideways to see the whole diagram.

Learning-rate schedules. Both schedules warm up linearly over 2k steps to a peak of 3e-4. Cosine decays smoothly to 3e-5 at 100k steps and must know the run length in advance. Warmup–stable–decay holds the peak until 85k steps, then decays; a decay can be branched from any stable checkpoint.

Why decay helps. At a constant learning rate, SGD-type optimizers bounce around in a noisy region whose size scales with the learning rate. Decaying reduces that noise floor, letting the model settle into a lower-loss region. That's why WSD's loss drops sharply during the final decay.

WSD's practical advantages:

Continue training without knowing the final length: branch a decay from any stable checkpoint.
Cheap scaling studies: one long stable run with decays branched at several points replaces many separate cosine runs of different lengths.
Data annealing: pair the decay phase with high-quality data.

Choosing the peak LR. Sweep on small models; the optimum shrinks as width grows unless you use μP, which makes it transfer across widths. Large models typically use smaller peaks (e.g. ~1.5e-4 for very large dense models vs ~3e-4 at a few billion parameters).

Batch size coupling. Larger batches reduce gradient noise and typically allow larger learning rates, roughly with the square root of the batch ratio for Adam, up to the critical batch size, beyond which more batch gives little.

25–32 min Trade-offs

Schedule	Pros	Cons
Cosine	Strong, well understood	Must fix run length in advance; no clean continuation
WSD	Continuation, cheap scaling studies, annealing	Needs a decay phase to realize gains; shape matters
Constant	Simple	Leaves loss on the table
Inverse square root	Classic for Transformers	Usually worse than cosine/WSD for LLMs

Mismatched cosine length is a classic pitfall: comparing a 10k-step cosine run with the 10k-step checkpoint of a 100k-step cosine run unfairly favours the former. This mistake contributed to early scaling-law estimates that favoured larger models.

32–37 min Production example

A team needed to extend a model's pre-training by 40% after seeing strong evaluations. Their run used cosine decay, already near its minimum LR, so continuing meant re-warming the learning rate, which caused a loss bump and wasted tokens. The next model used WSD: they extended the stable phase and decayed at the end, and also used branched decays to estimate how much extra data would help.

37–40 min Likely follow-ups

What's the role of the final LR value? Too high leaves noise; decaying to zero can over-fit to the last data; ~10% of peak is common for cosine, lower for WSD's final decay.
Do you re-warm for continued pre-training? Usually yes, briefly, with a lower peak; watch for forgetting.
How do you detect a too-high LR? Rising clip rate, growing update-to-weight ratios, loss spikes.

What interviewers listen for

Mechanistic reasons for warmup, not just "it helps".
Understanding that decay realizes gains, and WSD's operational benefits.
Coupling between LR, batch size and width.

Red flags

Treating the schedule as unimportant.
Comparing runs with mismatched schedule lengths.