Optimisers from gradient descent to AdamW

Almost every large model in this course, and most robot policies, is trained with AdamW. This short video walks the family tree that leads there: plain gradient descent, momentum, RMSProp, Adam, and the one-line fix that turns Adam into AdamW. This guide explains each step from the original sources and gives you a 2D "optimiser race" to play with.

Video by Jia-Bin Huang · Watch on YouTube

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Watch

Topic guide

The description for this video has no chapter timestamps. It names five topics, listed here in the order it gives them, so this table has no time links. The whole video is just over seven minutes, so watching straight through is easy.

2. The ideas, explained

Gradient descent and the narrow-valley problem

Given a loss L(θ) and its gradient g = ∇L(θ), gradient descent repeats one update. In deep learning the gradient comes from a mini-batch rather than the whole dataset, which makes it stochastic gradient descent (SGD).

θ ← θ − η · g          η = learning rate

on a quadratic with curvature h along some direction:
  error_next = (1 − η·h) · error
  stable only if η < 2 / h_max

The last line is the whole problem. The steepest direction (largest curvature h_max) caps the learning rate, and with that small η the flattest direction (smallest curvature h_min) creeps along at a rate set by η·h_min. The ratio h_max/h_min, the condition number, says how badly a single learning rate fits. Neural-network losses are badly conditioned, so the two fixes below attack exactly this: momentum by accumulating speed along the flat direction, adaptive methods by giving each parameter its own step size.

Momentum

Heavy-ball momentum goes back to Polyak (1964). Sutskever, Martens, Dahl and Hinton (2013) showed that momentum with a well-chosen schedule matters a lot for training deep networks. The common form keeps a velocity:

v ← μ · v + g          μ ≈ 0.9
θ ← θ − η · v

Along a direction where the gradient keeps the same sign, v grows towards g/(1 − μ), so the effective step is up to 10× larger at μ = 0.9. Along a direction where the gradient flips sign each step (the valley walls), the terms partly cancel. Nesterov's variant evaluates the gradient at the look-ahead point θ − ημv and is often a little more stable.

RMSProp

RMSProp was never published as a paper. Geoffrey Hinton introduced it in lecture 6 of his 2012 Coursera course, Neural Networks for Machine Learning, and the slides (the video's first reference) are the usual citation. The idea is to divide each parameter's gradient by a running estimate of its typical size:

s ← ρ · s + (1 − ρ) · g²      (elementwise; the slides suggest ρ = 0.9)
θ ← θ − η · g / (√s + ε)

A parameter with consistently large gradients gets its step shrunk; one with tiny gradients gets it enlarged. In the narrow valley, both directions end up moving at a similar speed of roughly η per step, regardless of curvature. RMSProp descends from Rprop, which used only the sign of the gradient; dividing by the root mean square is a mini-batch-friendly way to get sign-like steps. AdaGrad (Duchi, Hazan and Singer, 2011) had divided by the root of the sum of all past squared gradients, which shrinks the step forever; RMSProp's moving average forgets old gradients and keeps going.

Adam

Adam (Kingma and Ba, ICLR 2015) keeps both a momentum-like first moment and an RMSProp-like second moment, and corrects both for starting at zero:

m ← β₁ · m + (1 − β₁) · g          first moment  (β₁ = 0.9)
v ← β₂ · v + (1 − β₂) · g²         second moment (β₂ = 0.999)
m̂ = m / (1 − β₁ᵗ),  v̂ = v / (1 − β₂ᵗ)    bias correction at step t
θ ← θ − α · m̂ / (√v̂ + ε)          α = 0.001, ε = 1e-8 (paper defaults)

Why bias correction: m and v start at zero, so early on they underestimate the true averages. At step 1, m = 0.1·g, and dividing by 1 − 0.9 = 0.1 recovers g. Without the correction for v, which uses β₂ = 0.999 and so warms up over about a thousand steps, the first steps would divide by a tiny √v and be far too large. The paper also notes that the step size is roughly bounded by α, which makes α easy to reason about: it is about how far each parameter can move per step.

Adam has a cost that matters at scale: two extra numbers per parameter, usually stored in fp32. With mixed precision that's 8 of the roughly 16 bytes per parameter that ZeRO shards across GPUs. Module 08 · Stage 06 does the accounting and has the ZeRO memory calculator.

AdamW: decoupled weight decay

Loshchilov and Hutter's "Decoupled Weight Decay Regularization" (ICLR 2019) starts from an observation. For plain SGD, adding an L2 penalty (λ/2)‖θ‖² to the loss and shrinking the weights directly by a factor each step are the same thing, up to rescaling λ by the learning rate. For adaptive methods like Adam they are not, and most libraries at the time implemented "weight decay" in Adam as the L2 version.

Adam + L2 (what libraries did):
  g ← ∇L(θ) + λ·θ                    decay enters the gradient …
  θ ← θ − α · m̂ / (√v̂ + ε)           … and is then divided by √v̂

AdamW (decoupled):
  g ← ∇L(θ)
  θ ← θ − α · ( m̂ / (√v̂ + ε) + λ·θ )   decay applied directly, not rescaled

With L2, the decay term passes through the adaptive scaling. Parameters with a history of large gradients have a large √v̂, so they are decayed less, which is the opposite of what a regulariser should do uniformly. AdamW applies the same relative shrink to every weight. The paper reports two effects: the best weight decay becomes much less dependent on the learning rate, so the two can be tuned separately, and Adam's generalisation improves enough to compete with SGD with momentum on image classification, where it had usually lost before. (In the paper the whole update is also multiplied by a schedule multiplier η_t; PyTorch's torch.optim.AdamW applies θ ← θ − lr·λ·θ before the Adam step.)

In LLM practice, AdamW with β₂ around 0.95, a warmup followed by cosine or linear decay, gradient clipping, and weight decay around 0.1 has been the standard recipe for years. Weight decay is usually not applied to biases, normalisation gains or sometimes embeddings. The next video, on Muon, is about the first optimiser to seriously challenge this default for the hidden-layer matrices.

3. Where the course goes deeper

Why it matters for robotics

Self-check

References

From the video's description:

• [RMSProp] Neural Networks for Machine Learning, lecture 6 slides (Hinton, with Srivastava and Swersky, 2012)
• [Adam] Adam: A Method for Stochastic Optimization (Kingma and Ba, ICLR 2015)
• [AdamW] Decoupled Weight Decay Regularization (Loshchilov and Hutter, ICLR 2019); also on arXiv

Also cited on this page:

• Polyak, B. T. (1964). Some methods of speeding up the convergence of iteration methods. USSR Computational Mathematics and Mathematical Physics 4(5).
• On the importance of initialization and momentum in deep learning (Sutskever, Martens, Dahl and Hinton, ICML 2013)
• Adaptive Subgradient Methods for Online Learning and Stochastic Optimization (AdaGrad) (Duchi, Hazan and Singer, JMLR 2011)