Warmup and Cosine Decay
The warmup-then-cosine schedule ramps the learning rate up briefly, then lowers it along a cosine curve to a small floor, a widely used default.
A two-phase schedule
Warmup with cosine decay is one of the most common learning-rate schedules for training deep networks. It has two phases. In the first, the learning rate increases linearly from a small value to its peak over a set number of warmup steps. In the second, the rate follows the shape of half a cosine wave, descending smoothly from the peak toward a small final value over the remaining steps. The combination gives a gentle start, a strong middle, and a soft landing.
The cosine curve
During decay, the rate at step t is set to a floor plus half the difference between peak and floor, times one plus the cosine of pi times the fraction of decay steps completed. Early in decay the cosine is flat near its top, so the rate stays high and the model keeps making rapid progress. Near the end the curve flattens again at the bottom, so the rate lingers at small values, letting the model settle into a minimum without overshooting.
import math
def lr_at(step, warm, total, peak, floor=0.0):
if step < warm:
return peak * step / warm
p = (step - warm) / (total - warm)
return floor + 0.5*(peak-floor)*(1 + math.cos(math.pi*p))
- Warmup avoids early instability while optimizer statistics settle
- The cosine spends time at both high and low rates without sharp jumps
- A nonzero floor prevents the rate from reaching exactly zero too early
- Widely used for both pretraining and long training runs
Variations
Cosine schedules can include restarts, where the rate jumps back to a high value and decays again, forming warm restarts that sometimes improve results by re-exploring the loss surface. For continued or resumed training, practitioners sometimes extend the schedule rather than restarting cold. The choice of peak rate and warmup length still dominates outcomes, so those are tuned first; the cosine shape is a robust default once they are set. It pairs naturally with AdamW.