Chapter 5 of 16

Optimisers

SGD, momentum, Adam and AdamW, on a valley and on the real model.

Slides for this chapter: Transformer Training 04, Optimisers

The gradient says which way is downhill; the optimiser decides how far to step, and in which direction exactly. Below, four optimisers, the same code that trains the model in chapter 1, start from the same point in a long, narrow valley.

Loading the animation…

Concept

SGD (stochastic gradient descent) steps straight downhill, by the learning rate times the gradient: the learning rate $\eta$ sets the size of every step, and is the most important number to tune. In a valley much steeper across than along, the learning rate must be small enough not to overshoot across the steep walls, and is then far too small for the gentle slope along the floor: SGD zig-zags and crawls. Real loss surfaces are full of such directions, with curvatures that differ by orders of magnitude.

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. The running sum is the momentum buffer: one number per weight, kept from step to step as part of the optimiser's state (Adam keeps two). Methods that compress what workers exchange, such as DeMo (chapter 13), work on this buffer rather than on a single step's gradient.

Adam (Kingma & Ba 2014) keeps two running averages for every single weight: $m$ , of the gradient (momentum), and $v$ , of its square (how large that weight's gradients usually are). It steps by $m/\sqrt v$ , so each coordinate moves by roughly the learning rate whatever the scale of its gradient: the steep direction and the gentle one get similar step sizes. The numbers under the picture are Adam's $m$ and $v$ for $x$ and $y$ ; $v$ for $y$ is far larger, because the valley is steeper in $y$ . The bias correction divides by $1 - \beta^t$ because $m$ and $v$ start at zero and would otherwise be too small for the first steps. Near the minimum, a constant-learning-rate Adam keeps taking steps of about the learning rate and jitters around it: this is one reason runs decay their learning rate at the end (chapter 6).

AdamW (Loshchilov & Hutter 2017) adds weight decay the direct way: each step first shrinks every weight by a factor $1-\eta\lambda$ . With plain Adam, adding the decay $\lambda\theta$ to the gradient (an L2 penalty) puts it through the $1/\sqrt v$ rescaling, so weights with large gradients are barely decayed at all; decoupling it makes the decay the same for every weight. 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. Here $\lambda$ is set very high so that you can see it; in the default run it is 0.01, and it is applied only to the weight matrices, not to LayerNorm's $\gamma, \beta$ or the biases.

On the real model. The chart below trains the Decoder's model with each optimiser at the best of three constant learning rates, 1,500 steps, three data seeds each (results.md §5). At their best rates the four end close together: Adam 2.185 ± 0.041, momentum 2.213 ± 0.008, AdamW 2.230 ± 0.041, SGD 2.241 ± 0.034 (mean ± standard deviation over the seeds). Those differences are about one standard deviation: at this scale the data say no more than that all four work. The learning rate matters far more: SGD at 0.1 instead of 1.0 ends at 2.897, still on the plateau. Single runs are not enough here: at the higher learning rates training is chaotic, and recomputing the same run on another computer moved its final loss by about 0.1. Comparing optimisers fairly means tuning each one's learning rate and repeating over seeds (chapter 6). For large language models AdamW is the usual choice: Llama 3 405B, for example, was pretrained with it (Grattafiori et al. 2024, §3.4.1).

SGD (lr 1): 2.241 ± 0.034 · Momentum (lr 0.1): 2.213 ± 0.008 · Adam (lr 0.01): 2.185 ± 0.041 · AdamW (lr 0.01): 2.230 ± 0.041.Mean over three data seeds (± their standard deviation) at each optimiser’s best of three constant learning rates, 1,500 steps; results.md §5.

Newer optimisers, for context rather than measured here: Lion keeps only the sign of a momentum (Chen et al. 2023); Shampoo (Gupta et al. 2018) and SOAP (Vyas et al. 2024) precondition each weight matrix with its gradient's row and column statistics instead of per-weight scales; Muon orthogonalises each matrix's momentum update (Liu et al. 2025 report it at scale).

Maths

For a weight $\theta$ with gradient $g_t$ at step $t$ (all operations element by element):

\begin{aligned} &\text{SGD:} && \theta_t = \theta_{t-1} - \eta\, g_t \\ &\text{momentum:} && b_t = \mu\, b_{t-1} + g_t,\ \ b_1 = g_1; \quad \theta_t = \theta_{t-1} - \eta\, b_t \\ &\text{Adam:} && m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t,\quad v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \\ & && \theta_t = \theta_{t-1} - \eta\, \hat m_t / (\sqrt{\hat v_t} + \epsilon),\quad \hat m_t = m_t/(1-\beta_1^t),\ \hat v_t = v_t/(1-\beta_2^t) \\ &\text{Adam + L2:} && g_t \leftarrow g_t + \lambda\, \theta_{t-1} \text{ first (coupled)} \\ &\text{AdamW:} && \theta_{t-1} \leftarrow (1 - \eta\lambda)\,\theta_{t-1} \text{ first (decoupled)} \end{aligned}

Defaults: $\mu = 0.9$ , $\beta_1 = 0.9$ , $\beta_2 = 0.999$ (0.99 in the default run), $\epsilon = 10^{-8}$ . The valley is $f(x, y) = \tfrac12\big((x-2)^2 + 50(y-1)^2\big)$ : SGD is stable only for $\eta < 2/50$ , and then its error along $x$ shrinks by a factor $1 - \eta$ per step. These are PyTorch's definitions: each optimiser is checked against torch.optim over 25 steps on the model (largest difference 1.1 × 10⁻¹⁶ for AdamW).

Code

The Adam/AdamW update for one row of a weight, src/lib/training/optimizers.ts:

for (let j = 0; j < row.length; j++) {
  let pj = row[j]!;
  let gj = g[j]!;
  if (o.kind === "adamw") pj = pj * (1 - lr * wd);
  else if (wd) gj = gj + wd * pj;
  m[j] = o.beta1 * m[j]! + (1 - o.beta1) * gj;
  v[j] = o.beta2 * v[j]! + (1 - o.beta2) * gj * gj;
  const mhat = m[j]! / bc1;
  const vhat = v[j]! / bc2;
  row[j] = pj - (lr * mhat) / (Math.sqrt(vhat) + o.eps);
}

The two moment buffers are why Adam needs two extra numbers per weight: with FP32 master weights and gradients, this is where the "16 bytes per parameter" of mixed-precision training comes from (chapter 10).

Comments

Be the first to leave a comment on this section.