Vol. III · Learning Machinery Plate 4 of 10SGD, momentum, Adam, AdamW, learning rate, weight decay, optimiserL

From SGD to AdamW.

The gradient says which way is downhill. How far to step, and how to smooth the steps, is the optimiser’s job.

Plain stochastic gradient descent moves every parameter a fixed multiple — the learning rate — against its gradient, estimated on a small batch of data. On a loss surface shaped like a long narrow valley, that is a poor strategy: a step size safe across the steep walls is painfully small along the gentle floor. Momentum keeps a running average of recent gradients, so that the zig-zag across the valley cancels and the steady push along it accumulates.

Adam (Kingma and Ba, 2014) adds a second running average, of squared gradients, and divides each parameter’s step by its square root: parameters with large, noisy gradients take smaller steps, parameters with small consistent ones take larger steps, and the step size becomes nearly independent of the gradient’s scale. AdamW (Loshchilov and Hutter, 2017) applies weight decay — shrinking each weight a little every step — separately from that division. AdamW is what almost every large language model is trained with.

Below: three optimisers on a long narrow valley (one direction 100 times steeper than the other), 100 steps each. Left: the final loss for every learning rate from 0.001 to 1 — where each one works and where it blows up. Right: the path of the chosen optimiser at the chosen learning rate.

What you will learn

Write the update rules of SGD, momentum and Adam.
Explain why a single learning rate struggles on an ill-conditioned surface.
Compare how wide a range of learning rates each optimiser tolerates.
Say what AdamW changes about weight decay.

Workshop

Left: log₁₀ of the loss after 100 steps against the learning rate, for SGD (blue), momentum (green) and Adam (red); a missing stretch means the run diverged. Right: the chosen optimiser’s path from the start (dot) towards the minimum at the origin; the faint line is the valley floor.

Step sizes that adapt

On this valley SGD is caught between its two directions. Its largest stable learning rate is set by the steep one — above 0.018 it diverges — and at such small rates it crawls along the gentle one: of the 25 learning rates tried, only 1 gets the loss below 0.01 within 100 steps. Momentum widens the usable range to 11, because the averaged steps cancel across the valley and build up along it. Adam, whose steps are scaled per parameter by the size of recent gradients, gets there with 6 of the 25 — on this smooth two-parameter bowl, well-tuned momentum is the better method — but Adam never diverges at all, even at a learning rate of 1: its failures are slow, not catastrophic. On a real network, where millions of parameters have different curvatures and noisy gradients and nobody can tune each one, that robustness and the per-parameter scaling are why Adam-family optimisers took over.

Two settings matter beyond the learning rate. β₂, the memory of the squared-gradient average, is usually 0.95 for language models rather than the default 0.999, because a long memory reacts too slowly when gradients suddenly grow; ε keeps the division safe (plate XI.3). And weight decay: in plain Adam an L2 penalty is added to the gradient and then divided by √v like everything else, so heavily-updated weights are barely decayed; AdamW subtracts ηλθ directly, so every weight shrinks by the same fraction. The price of Adam is memory: two extra numbers, m and v, for every parameter — a large part of the training memory budget of Vol. XIV.

Adam’s name stands for adaptive moment estimation. The paper by Diederik Kingma and Jimmy Ba was posted in December 2014 and became one of the most cited in computer science; AdamW followed in 2017, and the large language models of the 2020s were trained almost exclusively with it until new optimisers such as Muon began to appear (plate XI.9).

Terms

Learning rate: The multiplier on each step; too large diverges, too small crawls.
Momentum: A running average of gradients, which damps oscillation and accumulates consistent directions.
Adam: An optimiser that scales each parameter’s step by the root of its recent squared gradients.
Weight decay: Shrinking every weight by a small fraction each step; in AdamW it is applied separately from the gradient.

SGD steps against the gradient; momentum smooths the steps; Adam scales each parameter’s step to its own gradients, which makes it robust to the learning rate; AdamW decays weights separately.