6.5 Optimizers and learning-rate schedules

What this is and why it exists

The optimiser and its schedule are the difference between a model that trains and one that sits there. They are also where the single most sensitive setting in deep learning lives: the learning rate, which cannot be copied between models, between batch sizes, or between datasets, and which a short experiment will find for you in minutes. This topic is what each optimiser is doing, what a schedule is for, and the one-line safeguard that stops a bad batch destroying a long run.

The vocabulary

• Step — one update of every parameter, from one batch.
• Learning rate — how far a step moves, per unit of gradient.
• Momentum — carrying part of the previous step into this one.
• Adaptive method — an optimiser that gives each parameter its own effective step size.
• Weight decay — pulling every parameter slightly towards zero at each step.
• Warmup — starting with a small learning rate and raising it over the first steps.
• Decay schedule — lowering the learning rate over the course of training.
• Gradient clipping — capping the size of an update before it is applied.

The mental model

Plain gradient descent takes a step downhill and nothing else, and its weakness is a surface that is much steeper in one direction than another: it bounces across the steep direction while creeping along the shallow one. Momentum fixes this by keeping a running average of recent steps and moving along that instead. Consistent directions accumulate, oscillating ones cancel, and progress along the shallow direction accelerates. The Nesterov variant computes the gradient at the point the momentum was already going to carry you to, so it responds to the slope ahead rather than the slope behind — a small change that makes the correction a little more timely.

Adaptive methods add a second idea: per-parameter step sizes. They keep a running estimate of how large each parameter's gradients have recently been, and divide by it, so parameters with consistently small gradients take proportionally larger steps and vice versa. Adam combines this with momentum and is the practical default: it works acceptably out of the box on almost anything, which is worth a great deal when you have one experiment to run.

AdamW is the correction you should actually use, and the reason is worth stating because it is small to say and consequential. Weight decay is meant to pull parameters towards zero by a fixed proportion. Implemented as an addition to the gradient, it then passes through the adaptive division, so parameters with large gradient histories get less decay than intended and the regularisation is no longer what you asked for. PyTorch describes AdamW as implementing the algorithm "where weight decay does not accumulate in the momentum nor variance" — the decay is applied to the parameters directly, outside the adaptive machinery, which is what "decoupled" means. Its documented defaults are a learning rate of 0.001 and a weight decay of 0.01. If you are using Adam with weight decay, you almost certainly want AdamW instead.

Schedules frequently matter more than the choice of optimiser. Three pieces, each solving something specific.

Warmup exists because the first steps of training are the most dangerous. The adaptive estimates are based on almost no history and are unreliable, the weights are random, and gradients are large; a full-size step here can push the model somewhere it never recovers from. Raising the learning rate from near zero over the first few hundred or few thousand steps costs almost nothing and prevents that. It is close to mandatory for transformers and for any large batch.

Decay exists because the step size that makes rapid early progress is too large to settle with. Cosine decay lowers the rate smoothly along a curve to near zero by the end of training, which lands the model gently in a minimum rather than bouncing around it. Step decay — cutting the rate by a factor at fixed points — is the older approach and still works. One-cycle combines both into a single shape: warm up to a peak well above your usual rate, then decay all the way down, on the argument that the large middle phase explores and the long decay refines. It trains many models faster than a constant rate, and it needs the peak chosen with care.

Gradient clipping is one line and it belongs in every sequence model. A single unlucky batch can produce an enormous gradient that moves every parameter far out of its useful range, and the run never recovers. Clipping caps the update instead. PyTorch's function is documented as clipping "the gradient norm of an iterable of parameters", where "the norm is computed over the norms of the individual gradients of all parameters, as if the norms of the individual gradients were concatenated into a single vector" — so the direction of the update is preserved and only its length is limited.

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
optimizer.zero_grad()

Now the setting everything above depends on. The learning rate interacts with the batch size, the architecture, the initialisation and the normalisation, so a value copied from another project is a guess wearing borrowed authority. The range test takes minutes: start at a very small rate, increase it smoothly over a few hundred batches, and plot the loss against the rate. There is a region where the loss falls steeply, and it ends where the loss turns upward; the useful value is inside that falling region, comfortably before the turn. Run this instead of arguing about it. And when you change the batch size, re-run it — the two are coupled, and the common rule of scaling the rate with the batch size is a starting point rather than an answer.

Finally, a diagnostic that saves hours. A loss that is flat from the first step usually means the rate is far too small, or nothing is connected. A loss that becomes meaningless within a few steps means it is far too large. And a loss that falls and then plateaus higher than it should is where a schedule earns its place — the rate that got you there is now too large to go further.

What you should now be able to explain or do

Say what momentum does and what problem it solves. Explain what makes an optimiser adaptive. State the difference between Adam and AdamW in a sentence and say which to reach for. Give the reason warmup exists and when it is close to mandatory. Describe cosine decay and one-cycle, and what each is for. Add gradient clipping correctly and say what it preserves. Run a learning-rate range test and read the result. Diagnose a flat, a divergent and a plateaued loss curve.

Check yourself