NLP Foundations
Warmup and Why It Helps
The first few hundred to few thousand steps of a large training run are its most fragile, and a short linear ramp on the learning rate is the cheapest insurance against wrecking them.
intermediate · 7 min read
Every major LLM pretraining recipe starts with the learning rate near zero and ramps it linearly up to its peak over the first fraction of training, typically well under 5% of total steps. Skip this and, at small scale, training often still works. At the scale of a real pretraining run it frequently does not: the first few hundred steps are the likeliest place for a run to destabilise before it has even found its footing, and warmup is the single cheapest defence against that.
What is actually happening early on
At initialisation, a network's weights are set by a scheme designed to keep activation variance roughly stable at the start, not to guarantee it stays stable as weights move. The very first gradient steps, taken at full learning rate with a randomly initialised network, can be unusually large and poorly aligned with any sensible descent direction, because the loss landscape near a random point looks nothing like the landscape a partially trained network will later see. A large early step can push the network into a bad region it then has to spend a long time escaping, if it escapes at all.
The Adam-specific version of the problem
Adam (and AdamW, see optimisers-adam-adamw-lion) makes this worse by construction. Its per-parameter step size divides by sqrt(v_hat), a running estimate of squared gradient magnitude. In the first handful of steps, v_hat is built from very few samples and its estimate is statistically noisy: a parameter that happens to get one unusually small early gradient can end up with an artificially small v_hat, which then inflates its next update far beyond what the raw gradient alone would justify. Liu et al.'s RAdam paper makes this precise: the variance of Adam's adaptive learning rate is provably large in the first steps and only settles down after enough samples accumulate. Warmup is the blunt-instrument fix, it keeps the actual learning rate small during exactly the window where the adaptive rate is least trustworthy, buying Adam's moment estimates time to become reliable before they are allowed to drive large updates. RAdam is the same diagnosis with a more surgical fix, rectifying the adaptive rate directly instead of just capping the global learning rate from outside.
Warmup and normalisation placement
Warmup's necessity also depends on where the normalisation layers sit. The original post-norm transformer needed careful, often lengthy warmup to train at all; without it the un-normalised residual stream could grow large enough early on to destabilise the whole stack. Pre-norm architectures (see layernorm-residual-connections if available) are noticeably more forgiving, because the residual highway stays close to its input scale regardless of what the sublayers are doing, which is part of why nearly every model since GPT-2 uses pre-norm. Forgiving is not the same as unnecessary, though: essentially every large modern pretraining run still uses warmup, just typically a shorter one than post-norm models required.
How long is enough
There is no universal number, but the practical range across published recipes is narrow: on the order of a few hundred to a few thousand steps, or roughly the first 0.1% to a few percent of total training steps, depending on model and batch size. Too short and the fragile-early-steps problem is not actually solved; too long and the run spends real compute training at a learning rate well below its useful peak. Linear warmup is by far the most common shape; alternatives (square-root, exponential) show up occasionally but rarely with a decisive advantage.
When it falls down
- Warmup reduces spike risk, it does not eliminate it. Large runs still see loss spikes well after warmup ends, driven by causes warmup does not touch (see loss-spikes-and-divergence).
- It is not free. Every step spent below peak learning rate is a step of comparatively slow progress; an oversized warmup window is measurable wasted compute at scale.
- The right length is not obviously predictable in advance. Teams typically inherit a warmup length from a smaller prior run rather than deriving it from first principles for the specific model size and batch size at hand.
- Warmup does not fix a badly chosen peak learning rate. It only smooths the approach to that peak; if the peak itself is too high, warmup delays the resulting instability rather than preventing it.
Further reading
- Liu et al., 2019, On the Variance of the Adaptive Learning Rate and Beyond, arXiv:1908.03265 - the RAdam paper and the variance argument for warmup.
- Kingma and Ba, 2014, Adam: A Method for Stochastic Optimization, arXiv:1412.6980 - the bias-correction machinery warmup is compensating for.
- Xiong et al., 2020, On Layer Normalization in the Transformer Architecture, arXiv:2002.04745 - why post-norm needed warmup more acutely than pre-norm.