Learning rate with warmup and cosine decay
Medium12 min12 hidden tests
A team trains a small language model with a learning rate schedule. During warmup, the rate climbs in a straight line from 0 at step 0 to peak_lr at step warmup_steps. It then follows half a cosine wave down to min_lr at step total_steps, and stays at min_lr after that. Write lr_at(step, peak_lr, warmup_steps, total_steps, min_lr) that returns the rate at one step.
Drag the step along the schedule
warmup, then a half cosine to min_lrDecay: (550 − 100) / (1000 − 100) = 0.5 of the way, so lr = 0.02 + (0.2 − 0.02) × ½(1 + cos(π × 0.5)) = 0.11.
Step 550, learning rate 0.11.
Examples
Halfway through warmup, half the peak
lr_at(50, 0.2, 100, 1000, 0.02)
returns 0.1
Halfway through the decay, halfway between the peak and the floor
lr_at(550, 0.2, 100, 1000, 0.02)
returns 0.11
Past the end, the rate stays at the floor
lr_at(1200, 0.2, 100, 1000, 0.02)
returns 0.02
Requirements
stepis an integer of 0 or more,warmup_stepsis 0 or more and belowtotal_steps, and0 <= min_lr <= peak_lr.- For
step < warmup_steps, the rate grows linearly withstep: 0 at step 0 andpeak_lrat stepwarmup_steps. Withwarmup_steps = 0there is no warmup. - From
warmup_stepstototal_steps, the rate is a half cosine frompeak_lrdown tomin_lr. Its position is the share of the decay steps done so far, from 0 atwarmup_stepsto 1 attotal_steps. - At and after
total_stepsthe rate ismin_lr. Return a float.
The idea behind it
A learning rate schedule changes the step size during training. Two pieces are standard for transformers. Warmup starts the rate near 0 and raises it linearly over the first steps. Early on, the weights are random and Adam's running estimates of gradient size are still rough, so a full-size step can throw training off course. Cosine decay then lowers the rate smoothly: lr = min_lr + 0.5 * (peak_lr - min_lr) * (1 + cos(pi * progress)), where progress goes from 0 to 1 across the decay. The cosine stays near the peak for a while, falls fastest in the middle, and flattens out near the floor, so the last steps make small, careful updates. A floor above 0, often a tenth of the peak, keeps those steps from vanishing. Libraries such as Hugging Face Transformers ship schedules of this shape, and the edges are where hand-written versions break.