Vol. XI · Optimisation at Scale Plate 1 of 10learning-rate schedule, warmup, cosine decay, warmup–stable–decay, WSD, noise floor, annealingLC

Learning-Rate Schedules.

Go fast while far from the answer, slow down near it — and the loss drops the moment you do.

Almost every large run uses the same shape: a linear warmup over the first few hundred million tokens, so that early gradients — large and unreliable — do not throw the weights far, then a long decay. GPT-3 decayed with a cosine to 10% of the peak; many runs since decay further. The reason for the decay is noise. Each step uses a gradient estimated from one batch, and at a constant rate the weights cannot settle closer to the minimum than that noise allows; lowering the rate shrinks the noise floor, and the loss drops.

If the decay produces most of the final gain, it can be postponed. Warmup–stable–decay (WSD) holds the rate constant for most of training and decays it only over the last part, typically 10–20%. A run can then be extended by continuing from a stable-phase checkpoint and decaying later, which a cosine schedule, tied to a fixed length, does not allow. The workshop runs the four common schedules on the noisy quadratic model, a standard stand-in for the loss surface of a network trained with noisy gradients.

Below: left, the four schedules over 3 000 steps. Right, the loss each one reaches on the noisy quadratic model, on a log scale.

What you will learn
Workshop
Momentum SGD on the noisy quadratic model: 20 directions with curvatures from 0.001 to 1 and gradient noise proportional to the square root of the curvature, batch 4, 3 000 steps from the same start and the same random noise. Left: learning rate against step. Right: loss against step for each schedule.

The drop when the rate comes down

At a constant rate the loss falls quickly and then stalls at 5.68 × 10⁻²: the weights keep moving with every noisy batch and cannot settle below that floor. Every decaying schedule ends far lower — cosine at 4.17 × 10⁻³, linear to zero at 3.93 × 10⁻³, WSD at 8.06 × 10⁻³. Warmup–stable–decay shows the mechanism most plainly: its loss tracks the constant schedule exactly until the decay begins — 1.38 × 10⁻¹ just before it — and then drops by an order of magnitude within the last fifth of training.

Real runs behave the same way, which is why loss curves of WSD runs show a sharp fall at the end and why decaying “to 10%” or “to zero” is a choice that matters. The practical attraction of WSD is flexibility: the stable phase can be checkpointed and extended, and several decays can branch from one run to compare data mixes for the final phase (plate XII.2), while a cosine schedule must know its length in advance. GPT-3’s recipe — warmup over 375 million tokens, cosine decay to 10% of the peak — remains a common default.

Warmup has a second job in Adam-family optimisers: their estimate of the gradient’s second moment starts from zero and is unreliable for the first steps, so large early updates are damped until it settles.

Terms

Warmup
Raising the learning rate from near zero over the first steps.
Decay
Lowering the learning rate towards the end of training.
Noise floor
The loss a constant learning rate cannot get below, set by gradient noise.
Warmup–stable–decay (WSD)
A schedule held constant for most of training and decayed only at the end.
Decaying the learning rate lowers the noise floor; WSD postpones the decay to the end, which makes runs extendable.

Sources

  1. Brown et al., “Language Models are Few-Shot Learners”, NeurIPS 2020, Appendix B.
  2. Hu et al., “MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies”, arXiv:2404.06395, 2024 (warmup–stable–decay).
  3. Zhang et al., “Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model”, NeurIPS 2019.