Optimizers & Learning Rates

Plain SGD follows the current batch gradient; momentum accumulates a velocity so persistent directions are amplified and contradictory batch noise is damped. Adam keeps moving estimates of both gradient direction and scale, adapting the step for each parameter, while AdamW applies weight decay separately so its regularising effect remains predictable. A schedule changes the global learning rate over time: warmup prevents unstable early updates, and decay encourages later updates to refine rather than overshoot. Optimiser state is part of a resumable training run and must be checkpointed with the model.

Optimiser and schedule choices can decide whether the same architecture trains at all. They are hyperparameters, so compare them on validation data rather than tuning on the test set. Start with a sensible, documented baseline such as AdamW plus warmup and decay, then change one factor at a time while recording the exact effective batch size and update count.

Momentum keeps one extra tensor per parameter — a velocity — and updates it before stepping: v ← βv + g, then θ ← θ − lr·v, with β typically 0.9. Because the velocity is a decaying sum of past gradients, a direction the last several batches agreed on builds up to roughly 1/(1 − β) times a single gradient, while a direction where consecutive gradients point opposite ways cancels itself out to almost nothing.

It fixes the geometry that plain descent handles worst: a long narrow valley, where the gradient points mostly across the valley rather than along it. Plain SGD zigzags between the walls and creeps forward; momentum cancels the crosswise component and accumulates the forward one, often converging several times faster for the same learning rate. momentum=0.9 is close enough to optimal in most settings that it is a default rather than something to tune.

Adam keeps two running averages per parameter: m, a momentum-like average of the gradient, and v, an average of the squared gradient. The update divides one by the square root of the other — θ ← θ − lr · m̂ / (√v̂ + ε) — so each parameter's step is normalised by how large its own gradients have recently been. Both averages are bias-corrected early on, because they start at zero. Defaults: betas=(0.9, 0.999), eps=1e-8, lr=1e-3.

It makes training far less sensitive to badly scaled parameters, which is why it is the default for almost everything today. A parameter whose gradients are consistently tiny — an embedding row for a rare token, a weight deep in an unnormalised stack — gets a step of comparable size to one with large gradients, so no part of the model is left behind by a single global learning rate. The cost is memory: two extra tensors the size of the model, so an Adam run needs roughly four times the parameter bytes before any activation is stored.

Weight decay is meant to pull every weight gently toward zero, independent of its gradient. Plain Adam implements it by adding wd · θ to the gradient — and then divides the whole thing by √v̂, so a parameter with large gradients has its decay divided away while one with tiny gradients gets decayed hard. The intended uniform regularisation becomes an arbitrary per-parameter one. AdamW fixes it by applying the decay directly to the weights, outside the adaptive scaling: θ ← θ − lr·m̂/(√v̂ + ε) − lr·wd·θ.

It is a rare case where a one-line change to an optimiser measurably improved results across the board, and it is why AdamW — not Adam — is the default in every modern training recipe. It also means the weight_decay value does not transfer between the two: the same number means something different in each, so a recipe tuned for Adam(weight_decay=0.01) is not a recipe for AdamW(weight_decay=0.01). Conventionally, biases and normalisation parameters are excluded from decay entirely.

On a typical run the three behave predictably. Plain SGD descends steadily but slowly, and its progress depends heavily on getting the one learning rate right. Momentum reaches a similar loss in noticeably fewer steps by damping the zigzag. Adam drops fastest in the first few hundred steps because every parameter gets a usefully sized step immediately, then often flattens to a final loss close to — occasionally slightly worse than — a well-tuned SGD-with-momentum run.

It sets expectations for what an optimiser swap can and cannot buy. Adam buys speed to a good-enough answer and robustness to your choice of learning rate, which is why it dominates research and prototyping. Well-tuned SGD with momentum still wins the last fraction of a percent on some vision benchmarks, which is why production image models sometimes still use it. Neither buys you a better model than your data and architecture allow, and a run that fails with all three has a problem elsewhere.

A schedule makes the learning rate a function of the step count. Step decay multiplies it by a factor (often 0.1) at fixed milestones. Cosine annealing sweeps it smoothly from the initial value down to near zero over the planned length of the run. Warmup does the opposite at the start: it ramps from near zero up to the target rate over the first few hundred or few thousand steps, then hands over to the decay. warmup + cosine is the standard recipe for Transformer training.

Early and late training want different step sizes. At the start the weights are random, the gradients are large and poorly aligned, and Adam's variance estimates are still unreliable — a full-size step then can wreck the run, which is what warmup prevents. Late in training the model is near a minimum and a large step only bounces around it, so decaying the rate is what converts 'roughly right' into 'converged'. A fixed rate has to compromise between the two and does worse than either.

A larger batch gives a less noisy gradient estimate, which makes a larger step safe. The two common heuristics are linear scaling — multiply the learning rate by the same factor as the batch size — and square-root scaling, which keeps the *noise* per step constant and is often the better fit for adaptive optimisers. Both hold over roughly an order of magnitude of batch size and then break down, at which point warmup becomes necessary to survive the early steps at the higher rate.

It is why 'I increased the batch size to use the GPU better and training got worse' is such a common report. At a fixed learning rate, doubling the batch halves the number of updates per epoch while leaving the step size alone, so the model simply travels half as far per epoch. Scaling the rate restores the distance travelled; warmup makes the larger initial steps survivable.

lr=1e-3 is Adam's default because it is a reasonable starting point for a small, normalised network — not because it is correct in general. Transformers at scale typically train at 1e-4 to 3e-4 with warmup; fine-tuning a pretrained model usually wants 1e-5 to 5e-5, because 1e-3 will destroy the pretrained features in the first few hundred steps; SGD on a convolutional network wants something in the 0.01 to 0.1 range, two orders of magnitude away from Adam's default. The usable rate depends on the optimiser, the depth, the normalisation, and the batch size.

Because the default is *almost* right for small models, it survives into larger ones and then fails in ways that look like other problems: a fine-tune that forgets everything it knew, a Transformer whose loss spikes in the first thousand steps, a deep network that plateaus high. A short learning-rate sweep — three runs, one order of magnitude apart, a few hundred steps each — is the cheapest and highest-value experiment available, and it should precede any architectural change.