Neural network basics

Optimisers: SGD, momentum and Adam

The same slope, a different stride: how training turns a gradient into a step.

Intermediate · 6 steps

Step 1 of 6· A toy two-weight valley, L = ½(x² + 100y²); the SGD, momentum and Adam update rules are real and run live

use the steps or arrow keys

Loading visual…

in one minute

Every training step starts from the same raw material, the gradient, but optimisers turn it into a step in different ways. Plain stochastic gradient descent (SGD) multiplies the gradient by one learning rate. In a long, narrow valley that rate must be small enough for the steep walls, so progress along the gentle floor is painfully slow. Momentum carries most of each step into the next, so directions that stay consistent build up speed. Adam goes further: it tracks how large each weight's gradients usually are and scales every weight's step to match. Widely used language models, including Llama, are trained with AdamW, a variant of Adam.

Why it matters for your product

The optimiser decides how quickly a training budget turns into a better model, and how much GPU memory that takes. Adam's two extra numbers per weight, on top of the weights and their gradients, are a large part of why full fine-tuning needs far more memory than running a model. For most teams the practical advice is to keep the optimiser the base model was trained with, tune the learning rate first, and consider LoRA, which trains a small fraction of the weights and so needs far less optimiser state.

For engineersShow the maths

v ← β·v + g, θ ← θ − η·v

Momentum: the velocity v keeps a fraction β of itself and adds the new gradient g; the weights θ move along v.

Worked example: With β = 0.9 and a gradient that stays at 1, v grows 1, 1.9, 2.71, … towards 1 ÷ (1 − 0.9) = 10, so a steady slope ends up moving ten times faster.

m ← β₁·m + (1 − β₁)·g, v ← β₂·v + (1 − β₂)·g², θ ← θ − α·m̂ / (√v̂ + ε)

Adam averages the gradient (m) and the squared gradient (v), corrects both for starting at zero (m̂ and v̂), then steps by their ratio.

Worked example: On the first step m̂ = g and v̂ = g², so each weight moves by α·g / |g|, a step of exactly α whatever the gradient's size (ignoring ε): 0.001 with the paper's defaults.

η < 2 / λₘₐₓ

On a bowl-shaped loss, plain gradient descent is stable only if the learning rate is below 2 divided by the largest curvature.

Worked example: Across this valley the curvature is 100, so η must stay under 0.02; along the floor, where the curvature is 1, such a step closes under 2% of the remaining distance.

where it stops working

A race on a toy valley does not transfer directly: real losses have millions of directions, noisy mini-batch gradients and shifting curvature, and well-tuned SGD with momentum still matches Adam on some tasks. Adam's per-weight scaling costs memory, and with plain L2 regularisation it tends to generalise worse, which is what AdamW fixes.

Key terms