Learning Rate Scheduling

Instead of one fixed learning rate for the whole run, shrink it on a schedule so early steps move fast and later steps settle in carefully.

Over 1000 steps from 0.001, cosine glides through 0.000854, 0.000500 and 0.000146 while step decay jumps 0.001 to 0.0005 to 0.00025 to 0.000125

Why Does This Exist?

Train a 12-layer transformer from scratch on a few million sentence pairs with one fixed learning rate and neither end of training is happy. Size the rate for the start, when random weights sit far from anything useful, and that same rate is still firing a hundred thousand steps later, when the model is close to a good solution and each step should be a small refinement. The oversized steps keep kicking the weights past the minimum, and the loss curve jitters instead of settling, sometimes for the rest of the run.

Size the rate for the end instead and the first several thousand steps, the ones with the most room to improve, crawl. The step size that is right at step 100 is wrong at step 100,000, because how far you are from done changes everything. The update rule itself lives in stochastic gradient descent, and this page covers the separate decision of how the global rate multiplying every update moves over time.

Think of It Like This

Parking a car across an empty lot

You drive at normal speed across the open lot to cover distance fast, slow as the parking spot comes into view, and creep the final foot with small wheel adjustments until the car is centered.

Holding full speed into the final foot clips the curb. Creeping the whole way across the lot takes ten times longer than needed. The right speed falls as the distance shrinks, smoothly rather than in sudden lurches.

The lot is the loss surface, your speed is the learning rate, and the parking spot is the minimum. Step decay is lifting off the gas in jerks. Cosine is easing off it. Where it stops: a car can see the spot, but training never knows how far it is from the minimum, so the schedule is a guess about distance, not a measurement of it.

How It Actually Works

Call the learning rate at step $t$ by $\eta_t$ , meaning the single global multiplier on every parameter update that step. Call the total planned steps $T$ , the peak rate $\eta_{\max}$ , and the floor $\eta_{\min}$ .

Step decay: flat stretches with sudden drops

Step decay multiplies the rate by a fixed factor, often 0.1 or 0.5, every fixed number of steps. With drops every 250 steps and factor 0.5 from 0.001, the visited rates are 0.001, then 0.0005, then 0.00025, then 0.000125. It is the oldest schedule: simple, predictable, and visibly jerky, because each drop halves every subsequent step at one boundary.

Cosine annealing: one smooth glide to zero

Cosine annealing follows half a cosine curve from peak to floor over the whole run:

\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\!\left(\pi \, \tfrac{t}{T}\right)\right)

Near $t = 0$ the curve is nearly flat, so the rate barely moves at first. It falls fastest mid-run and flattens again into the floor, easing to rest instead of snapping there. In a 1000-step run from 0.001 to 0, the rate reads 0.001000 at step 0, 0.000854 at step 250, 0.000500 at step 500, 0.000146 at step 750, and 0.000000 at step 999. At step 250 cosine has already eased down while step decay has jumped to its next plateau at 0.0005, which is the practical difference between the shapes, not just smoothness.

Warmup-stable-decay: the shape big runs actually use

Cosine needs $T$ fixed before training starts, which hurts runs measured in weeks: extend the run and the curve already sits near zero, so extra steps barely move the weights. Warmup-stable-decay splits the run into three phases instead, a short warmup ramping up to the peak, a long stable phase holding it flat, and a decay phase easing down at the end. The plateau can be extended or cut short without reshaping anything else. Adam and AdamW already scale each parameter from its own gradient history, and the schedule stacks on top as the shared global pace. The full walkthrough of the ramp-then-glide picture lives in the scheduling visual.

Code

Cosine annealing and step decay evaluated at the same five points in a 1000-step run from 0.001.

1import math2
3
4def cosine_lr(step: int, total_steps: int, lr_max: float = 1e-3) -> float:5    progress = step / total_steps6    return 0.5 * lr_max * (1 + math.cos(math.pi * progress))7
8
9def step_decay_lr(step: int, drop_every: int = 250, factor: float = 0.5,10                  lr_max: float = 1e-3) -> float:11    return lr_max * (factor ** (step // drop_every))12
13
14for step in [0, 250, 500, 750, 999]:15    cos_r = cosine_lr(step, 1000)16    step_r = step_decay_lr(step)17    print(f"step {step:4d} cosine {cos_r:.6f} step-decay {step_r:.6f}")18# -> step    0 cosine 0.001000 step-decay 0.00100019# -> step  250 cosine 0.000854 step-decay 0.00050020# -> step  500 cosine 0.000500 step-decay 0.00025021# -> step  750 cosine 0.000146 step-decay 0.00012522# -> step  999 cosine 0.000000 step-decay 0.000125

The printout matches the hand computation: cosine already easing at step 250 while step decay sits on its next plateau, and cosine touches zero at the end while step decay holds 0.000125.

Watch Out For

Extending a cosine run past its planned horizon

Cosine needs the total step count fixed in its formula before training starts. New data arrives or the loss has not plateaued, the run gets extended, and the curve already reached zero partway through the extension, so every extra step trains at a rate near zero and barely moves the weights.

Symptom: loss flatlines exactly when the extension begins despite healthy gradients. Fix: restart with a fresh schedule over the new horizon, or use warmup-stable-decay from the start so the plateau absorbs the extra steps.

Reading loss jumps at drop points as an architecture bug

A step-decay run shows a brief transient in the loss right at each drop boundary, and it is tempting to blame the model or the data.

Symptom: a kink in the loss curve landing exactly on a schedule boundary. Fix: check the schedule first. A 2x or 10x change in every step size genuinely disturbs the curve for a while, and that transient is the schedule working as designed.

The Quick Version

A fixed learning rate is a compromise that fits neither the fast early steps nor the careful late ones.
Step decay holds the rate flat, then drops it by a fixed factor at set intervals, visiting 0.001, 0.0005, 0.00025 and 0.000125 in a 1000-step run with drops every 250.
Cosine annealing glides 0.001 down through 0.000854, 0.000500 and 0.000146 to zero over the whole planned run.
Warmup-stable-decay splits big runs into ramp-up, a long extendable plateau, and ramp-down, because cosine cannot survive a changed horizon.
The schedule sets the shared global pace over time, while Adam scales each parameter separately, and the two stack.