The gradient says which way is downhill; the optimiser decides how far, and in which direction exactly. SGD, momentum, Adam's two moments and bias correction, AdamW's decoupled weight decay, run live on a narrow valley and measured on the real model.
This deck follows chapter 5, optimisers. Measured numbers: results.md §5 (each optimiser checked against torch.optim over 25 steps, §3). The valley below is computed live in this page.
The valley \(f(x,y) = \tfrac12\big((x-2)^2 + 50\,(y-1)^2\big)\) is fifty times steeper across (\(y\)) than along (\(x\)); its minimum is at (2, 1). Real loss surfaces are full of such directions. All four start at (−2, 2):
Stochastic gradient descent steps by the learning rate \(\eta\) times the gradient. In the valley, \(\eta\) must be small enough not to overshoot across the steep walls: SGD is stable only for \(\eta < 2/50 = 0.04\). It is then far too small for the gentle slope along the floor, where the error shrinks only by a factor \(1-\eta\) per step: SGD zig-zags across and crawls along.
The largest stable learning rate is set by the steepest direction (the largest curvature); progress is set by the shallowest. When curvatures differ by orders of magnitude, as they do in a transformer, one global step size cannot suit them all.
Momentum keeps a running sum of past gradients and steps along it. Across the valley, successive gradients point in opposite directions and cancel; along it, they agree and add up, so the optimiser builds speed where it is useful (\(\mu = 0.9\) is the usual default).
One number per weight, carried from step to step: the first piece of optimiser state (Adam keeps two).
Methods that compress what workers exchange, such as DeMo, work on this buffer rather than on one step’s gradient (deck 12).
Adam (Kingma & Ba 2014) keeps, for every weight, a running mean \(m\) of the gradient (the first moment) and \(v\) of its square (the second moment: how large that weight’s gradients usually are). Stepping by \(m/\sqrt v\), each coordinate moves by roughly the learning rate whatever the scale of its gradient: in the valley, \(v\) for \(y\) is far larger, so the steep and the gentle direction get similar step sizes.
Bias correction divides by \(1-\beta^t\) because \(m\) and \(v\) start at zero and would be too small for the first steps. Near a minimum a constant-learning-rate Adam keeps taking steps of about \(\eta\) and jitters: one reason runs decay the learning rate at the end (deck 05). Defaults: \(\beta_1 = 0.9\), \(\beta_2 = 0.999\) (0.99 in the site’s default run), \(\epsilon = 10^{-8}\).
\(g_t \leftarrow g_t + \lambda\,\theta_{t-1}\) first. The decay then goes through the \(1/\sqrt v\) rescaling, so weights with large gradients are barely decayed at all.
\(\theta_{t-1} \leftarrow (1 - \eta\lambda)\,\theta_{t-1}\) first, then the Adam step: the same decay for every weight (Loshchilov & Hutter 2017).
In the valley the decay pulls AdamW’s end point towards the origin, away from the minimum of the loss alone: weight decay is a preference for small weights, traded against the loss. There \(\lambda = 0.5\) is set very high so that you can see it; the site’s default run uses 0.01, applied only to the weight matrices, not to LayerNorm’s \(\gamma, \beta\) or the biases. AdamW is the usual choice for large language models: Llama 3 405B was pretrained with it (Grattafiori et al. 2024, §3.4.1).
The Decoder’s model, each optimiser at three constant learning rates, 1,500 steps of 16 windows, clipping at 1.0, three data seeds each: final validation loss, mean ± standard deviation (results.md §5).
| Optimiser | Learning rate | Final validation loss | Best for it |
|---|---|---|---|
| sgd | 0.1 | 2.897 ± 0.005 | |
| sgd | 0.3 | 2.374 ± 0.014 | |
| sgd | 1 | 2.241 ± 0.034 | best |
| momentum | 0.03 | 2.342 ± 0.021 | |
| momentum | 0.1 | 2.213 ± 0.008 | best |
| momentum | 0.3 | 2.253 ± 0.055 | |
| adam | 0.003 | 2.310 ± 0.031 | |
| adam | 0.01 | 2.185 ± 0.041 | best |
| adam | 0.03 | 2.301 ± 0.026 | |
| adamw | 0.003 | 2.311 ± 0.038 | |
| adamw | 0.01 | 2.230 ± 0.041 | best |
| adamw | 0.03 | 2.357 ± 0.029 |
At their best rates the four end close together, differences of about one standard deviation: at this scale the data say no more than that all four work.
SGD at 0.1 instead of 1.0 ends at 2.897, still on the plateau of deck 01. Choosing the optimiser moved the result by about one standard deviation; choosing its learning rate moved it by far more. Comparing optimisers fairly means tuning each one’s learning rate and repeating over seeds (deck 05). Single runs are not enough: at the higher rates training is chaotic, and recomputing the same run on another computer moved its final loss by about 0.1.
SGD keeps nothing, momentum one number per weight, Adam and AdamW two. In FP32 that is 8 bytes per parameter for Adam’s moments, the biggest single item of the 16 bytes per parameter of mixed-precision training (deck 09), and what ZeRO shards first (deck 10).
Each of the site’s optimisers matches torch.optim over 25 steps on the model: largest difference 1.1e-16 for AdamW, 0 for SGD (results.md §3).
Animated: chapter 5, optimisers · optimisation theory: Mathematics for Machine Learning · next: deck 05, learning rate and batch size