Learning Rate Schedules Visualization | Warmup, Step and Cosine Decay

The idea: η changes with the step

The learning rate η sets how far each SGD step moves: w ← w − η·g. A learning rate schedule makes η a function of the step number, η(t), chosen before training starts. The optimizer just reads η(t) at each step. Almost every modern training run uses one: a short warmup that ramps η up from near zero, then a decay that brings it down towards zero by the end. The page trains the same model three times side by side, with the same mini-batches, so the only difference between the lanes is the schedule.

Six small plots of the learning rate against the step from 0 to 200. Constant is a flat line. Step decay is flat, then drops to a tenth at step 100 and again at 150. Cosine falls smoothly from the peak to zero along half a cosine wave. Warmup plus cosine ramps up linearly first, then follows the cosine. Warmup plus constant ramps up, then stays flat. Warmup-stable-decay ramps up, stays flat until step 160, then falls in a straight line to zero — The six schedules on the page: decay brings η down at the end, warmup brings it up at the start.

The model on the page

Sixteen parameters start at w_i = 3 on the loss L(w) = Σ (a_i·w_i²/2 + b·w_i⁴/4), with curvatures a_i from 1 down to 0.16 and b = 0.5. Each step uses a mini-batch gradient: the true gradient plus noise of size σ = 0.4, drawn from a fixed seed so every run is the same. Training lasts T = 200 steps. Two things make the schedule matter, as in real networks: the gradient is noisy, and the surface is sharp at the start (the w⁴ term) and flat near the minimum.

Why decay: the noise floor

Far from the minimum the true gradient is large and the noise hardly matters. Near the minimum the true gradient is tiny, so each step is mostly noise: η·σ·ξ. With a constant η the parameters keep being kicked around the minimum and the loss stops falling at a noise floor that grows with η (about η·σ²·D/4 for D parameters). The only way down is to make the kicks smaller, which means shrinking η. That is decay.

Two copies of the same valley. On the left, with constant eta, the points jump in from the side and then keep bouncing back and forth at the same height above the minimum. On the right, with decayed eta, the jumps shrink and the points settle in a tight cluster at the minimum — With constant η the noise keeps the loss at a floor; shrinking η shrinks the kicks and lets the loss go lower.

At peak η = 0.3 the constant lane ends at loss 0.137, step decay (×0.1 at steps 100 and 150) at 0.055 and cosine at 0.034 (Demo 1). Each cut of step decay shows as a drop to a lower, smoother floor. Cosine decay, η·½(1 + cos(πt/T)), has no corners: it stays near the peak for a while, then falls gently to zero.

Decay is not free. While η is small, the flat directions (small a_i) move slowly. If η is already small, its floor is already low, and decaying it early leaves those directions unfinished. At η = 0.05 the constant lane wins: 0.024 against 0.054 for step decay and 0.060 for cosine (Demo 4).

Why warmup: the sharp start

Along a direction with curvature λ, one step multiplies the distance to the minimum by 1 − η·λ. If η·λ > 2 the step overshoots by more than it started with. At the start here the sharpest curvature is a + 3b·w² = 14.5, so a peak η of 0.5 gives η·λ = 7.3. The first step throws w far out, where the w⁴ wall is even steeper, and the loss explodes to 4.7·10⁵ at step 2 (Demo 2). The lane cards show λmax and η·λmax, in red while it is above 2.

Two copies of a valley with a flat bottom and steep walls. On the left, a step at the peak learning rate from the steep wall jumps to the other wall, higher up, then even higher and off the top. On the right, small warmup steps walk down the wall to the flat bottom, where larger steps are safe — Warmup keeps η small while the surface is sharp, so the peak η arrives where it is flat.

Warmup starts at η/W and ramps up linearly over W steps. The first steps are short, the parameters walk off the steep wall, and by the time η reaches its peak the curvature is small. With 5 or 20 warmup steps the same peak η = 0.5 trains fine and ends at 0.039 and 0.042. In real networks the reasons are similar: the loss is sharp early in training, and Adam's variance estimate is still poor in its first steps. Warmup is only needed when the peak η is large. At η = 0.15 all three lanes train, ending at 0.032, 0.033 and 0.033 (Demo 3); warmup just costs a few slow steps.

The common schedules

Schedule	η(t)	Where you meet it
Constant	η	quick experiments; a baseline
Step decay	×0.1 at fixed milestones	ResNet on ImageNet (PyTorch `MultiStepLR`)
Cosine	η·½(1 + cos(πt/T))	PyTorch `CosineAnnealingLR`; SGDR adds restarts
Warmup + cosine	linear ramp over W steps, then cosine	most transformer and LLM training (`get_cosine_schedule_with_warmup`)
Warmup + inverse square root	ramp, then η·√(W/t)	the original Transformer paper
Warmup–stable–decay (WSD)	ramp, flat, short decay at the end	recent LLM pre-training

Warmup–stable–decay

Cosine needs the run length T in advance: stop early and η is still high, so the loss is still on its floor. WSD keeps η at the peak and decays only in a short cool-down at the end (here the last 40 steps). The loss stays on the high-η floor during the stable phase, then drops during the cool-down. You can save a checkpoint from the stable phase and cool down a copy whenever you want a finished model, or keep training. On this page WSD ends at 0.089: far below warmup + constant (0.251), above warmup + cosine (0.042), which spends more steps at a small η (Demo 5).

What to try

Step twice at peak η 0.5 with the warmup comparison and watch η·λmax in the cards.
Train to End with each comparison at η 0.3, then at 0.05 and 0.5.
Watch the loss curve of step decay at steps 100 and 150, and of WSD at step 160.

See also SGD vs Momentum vs Adam for what an optimizer does with one gradient, and Gradient Descent for one learning rate that is too small or too large.

What the page leaves out

The model is a sum of sixteen one-dimensional problems with plain SGD; real networks have millions of coupled parameters, and most schedules there drive Adam or AdamW. The page shows only training loss, not how well the model generalizes, which is where decay often helps most. It leaves out cyclical and one-cycle schedules, cosine with restarts, ReduceLROnPlateau, schedules for batch size and weight decay, and how the best peak η depends on model width.