Raghuram Shankar
← Language Modelling

Part 5 · The Maths III: Training at Scale

AdamW, and why the W matters

Adam keeps running estimates of the gradient’s first and second moments and steps by their ratio:

mt = β₁mt−1 + (1−β₁)gt     vt = β₂vt−1 + (1−β₂)gt²
θt = θt−1 − η m̂t√v̂t + ε − ηλθt−1

gt the gradient · m, v running estimates of its mean and squared magnitude · β₁, β₂ their decay rates (0.9, 0.95–0.999) · m̂, v̂ the same after bias correction · η the learning rate · λ the weight decay · ε a small constant (~10−8) keeping the denominator finite

The ratio makes the step scale-free: multiply every gradient by a million and the first step is unchanged, because the numerator and denominator scale together. That is why Adam needs no per-layer learning-rate tuning.

Why the last term is separate. Classic L2 regularization adds λθ to the gradient, so it passes through the 1/√v̂ normalizer and gets divided by the gradient magnitude — parameters with large gradients end up barely regularized. AdamW applies the decay directly to the weights instead, decoupled from the adaptive scaling. That one change is the difference between Adam and the optimizer every large model is actually trained with.

The schedules themselves, exactly as implemented. Warmup exists because the second-moment estimate v is meaningless for the first few steps — a full-size step on a near-zero denominator is how runs diverge in the first hundred iterations. WSD holds a plateau and decays late, which lets you branch a run at any point rather than committing to a horizon in advance.