Vol. XI · Optimisation at Scale Plate 3 of 10AdamW, beta2, epsilon, second moment, update size, instability, large-scale trainingLE

AdamW at Scale.

Adam scales every step by a running average of past gradient sizes — so a gradient far larger than that average produces a far larger step.

AdamW (plate III.4) keeps two running averages per weight: m, of the gradient, with β₁ ≈ 0.9; and v, of its square, with β₂. The update is m/(√v + ε). Because it divides by the typical gradient size, Adam takes steps of roughly the learning rate in every direction, whatever the scale of the gradients — which is what makes it robust. Two consequences appear only at scale.

First, β₂. With β₂ = 0.999 the running average v remembers about a thousand steps. If a weight’s gradient has been tiny for a long time and then suddenly becomes large, v still reflects the tiny past and the step is much larger than the learning rate — a known contributor to loss spikes. Large runs therefore use β₂ = 0.95 (GPT-3, Llama 2). Second, ε: if gradients become comparable to ε, as they can in large models late in training, the update shrinks below the learning rate and Adam behaves more like plain SGD. Llama 2 used ε = 10⁻⁵, GPT-3 10⁻⁸.

Below: left, the size of Adam’s update around one sudden large gradient after a long quiet stretch, for three values of β₂. Right, the steady-state update against the gradient’s scale, for the chosen ε.

What you will learn
Workshop
Left: one weight’s gradient is 10⁻³ for 300 steps, then 1 for one step, then 10⁻³ again; update magnitude |m̂/(√v̂ + ε)| (in units of the learning rate) for β₂ = 0.95, 0.99 and 0.999, ε = 10⁻⁸. Right: steady-state update for a constant gradient g, which is g/(g + ε), against g.

Normalised steps and their failure modes

Before the spike, every value of β₂ gives an update of 1.00 learning rates: Adam has normalised the tiny gradient to a full-sized step. When a gradient a thousand times larger arrives, β₂ = 0.95 lets v catch up quickly and the step is 0.45 learning rates; with β₂ = 0.999 the long memory of small gradients makes it 1.63, 3.6 times larger, and the effect decays slowly. In a network where many weights change regime at once — a new feature switching on, an unusual batch — such oversized steps can push the model into an unstable region, one proposed mechanism for the spikes of plate XI.5. GPT-3 and Llama 2 both used β₂ = 0.95.

The right panel shows ε’s role. The update is close to a full step for gradients well above ε and shrinks for gradients near it. At the default 10⁻⁸ a gradient of 10⁻⁶ still gets 0.990 of a step; at Llama 2’s 10⁻⁵ it gets 0.091. A larger ε damps updates for weights whose gradients are tiny — a deliberate stabiliser in some recipes — but if it is larger than typical gradients, Adam quietly turns into a smaller-step SGD. Both settings look like details; at scale they are part of what keeps a run stable.

The “W” in AdamW is decoupled weight decay: shrinking each weight by η λ w directly, rather than adding λ w to the gradient where Adam’s normalisation would rescale it. Most large runs use λ = 0.1.

Terms

Second-moment estimate v
A running average of squared gradients; its square root is the typical gradient size.
β₂
How long v remembers: about 1/(1 − β₂) steps.
ε
A small constant added to √v; gradients near it get smaller steps.
Decoupled weight decay
Shrinking weights directly, outside Adam’s normalisation.
Adam normalises steps by recent gradient size; a long memory (high β₂) turns a sudden large gradient into an oversized step, and a large ε quietly shrinks small ones.

Sources

  1. Loshchilov and Hutter, “Decoupled Weight Decay Regularization”, ICLR 2019.
  2. Brown et al., NeurIPS 2020, Appendix B; Touvron et al., “Llama 2”, arXiv:2307.09288, 2023, section 2.2.
  3. Molybog et al., “A Theory on Adam Instability in Large-Scale Machine Learning”, arXiv:2304.09871, 2023.