Optimisers: SGD, Adam, AdamW and Why It Changed

An optimiser is a rule for turning a gradient into a step. Plain gradient descent uses the gradient itself; everything since adds memory of previous gradients. Three of those additions are worth deriving, because the third one changed the default training recipe for every transformer.

Momentum, and the factor it buys

Keep a running average of gradients and step along that instead:

v ← β·v + ∇L          (β typically 0.9)
w ← w − η·v

Suppose the gradient is the same value g for many steps. Then v becomes a geometric series: g(1 + β + β² + …) → g/(1 − β). At β = 0.9 that is 10g. So in a direction where the gradient is consistent, momentum builds up to a tenfold larger step; in a direction where the gradient alternates sign, the terms cancel and the step stays small.

That is precisely the fix for the conditioning problem in the stability bound: progress along the flat, consistent direction is amplified while the steep, oscillating direction is damped. It also means the effective learning rate is roughly η/(1 − β), which is why raising β without lowering η can push a stable run into divergence.

What Adam adds

Adam keeps two running averages — of the gradient and of the squared gradient — and divides one by the square root of the other:

m ← β₁·m + (1 − β₁)·g          first moment   (β₁ = 0.9)
v ← β₂·v + (1 − β₂)·g²         second moment  (β₂ = 0.999)

m̂ = m / (1 − β₁ᵗ)              bias correction
v̂ = v / (1 − β₂ᵗ)

w ← w − η · m̂ / (√v̂ + ε)

The bias correction is worth understanding rather than copying. At step 1, starting from m = 0, the update gives m = (1 − β₁)·g = 0.1g — ten times too small, purely because the average started at zero. Dividing by (1 − β₁ᵗ) = 0.1 restores it. The correction fades as β₁ᵗ shrinks, so it matters only in the first few dozen steps, and those are the steps where a wrong scale does the most damage.

The division is the real content. If a parameter’s gradients are consistently around g, then m̂ ≈ g, √v̂ ≈ |g|, and the step is approximately ±η — independent of the gradient’s magnitude. Adam takes steps of roughly fixed size in parameter space, which is why it works out of the box on losses where different layers have gradients differing by orders of magnitude, and why its typical learning rate (1e-3 to 1e-4) is so much smaller than SGD’s (1e-1): with SGD, η multiplies a gradient; with Adam, η more or less is the step.

The bug AdamW fixed

Now combine Adam with L2 regularisation the obvious way — add (λ/2)‖w‖² to the loss — and follow the λw term through the algorithm.

g = ∇L + λw                     the penalty enters the gradient

… so λw flows into m and into v …

w ← w − η · (m̂ + λw-ish) / (√v̂ + ε)

The decay term is divided by √v̂, the same as everything else.

Read what that means. A parameter with a long history of large gradients has a large v̂, so its decay is divided down. A parameter with small gradients gets a small v̂ and more decay. The strength of your regularisation has become a function of each parameter’s gradient history — which is not what anyone writing weight_decay=0.01 intends, and it means the parameters most in need of restraint are the ones restrained least.

Loshchilov and Hutter’s fix (ICLR 2019, “Decoupled Weight Decay Regularization”) is to take the decay out of the gradient entirely and apply it as its own term:

AdamW:   w ← w − η · m̂/(√v̂ + ε)  −  η·λ·w
                  ↑ adaptive data term      ↑ plain decay, undivided

Now every weight decays by the same factor per step, as the L2 derivation says it should, and the decay strength is decoupled from the gradient statistics. It also decouples λ from η in a useful way: with the coupled version, tuning the learning rate silently retuned the regularisation, so the two hyperparameters could not be searched independently.

Two footnotes worth carrying. First, the frameworks disagree on naming: passing weight_decay to a plain Adam implementation usually gives you the coupled L2 version, not decoupled decay, so check which class you are instantiating. Second, decay is normally not applied to biases and normalisation parameters — shrinking a LayerNorm gain toward zero is not regularisation, it is damage.

One more setting deserves naming, because it is the usual cause of a run that looks fine and then spikes. β₂ controls how long the second moment remembers: at the default 0.999 the effective window is roughly the last thousand steps, so a sudden burst of large gradients gets divided by a √v̂ estimated from a much calmer past, and the resulting step is enormous. Large transformer runs commonly lower β₂ to 0.95 for exactly this reason, and pair it with gradient clipping. The symptom to recognise is a flat loss curve, then a vertical spike, then either a slow recovery or NaN.

What the optimiser costs in memory

Adam stores m and v per parameter. Count the bytes for a naive fp32 setup — no sharding, no mixed precision, activations excluded:

weights   4 bytes
gradients 4 bytes
m         4 bytes
v         4 bytes
          ────────
          16 bytes per parameter

7B parameters × 16 = 112 GB     before a single activation is stored
7B parameters ×  2 =  14 GB     the same model served in bf16

An eightfold gap between serving a model and full-precision training of it, and two of those four lines are optimiser state. This is the arithmetic behind almost every efficiency technique in fine-tuning: LoRA trains a small adapter so that m and v exist for a tiny fraction of the parameters, QLoRA quantises the frozen backbone on top of that, and optimiser sharding across devices splits the state rather than replicating it. Adafactor (Shazeer and Stern, 2018) attacks the same cost by factorising the second moment, and 8-bit optimiser states (Dettmers et al., 2022) by quantising it.

What to actually use