Stochastic optimization: SGD, momentum and Adam

No one computes the exact gradient of a language model's loss over its full training set: one gradient step would require a pass over trillions of tokens. Training instead uses cheap, noisy gradient estimates and adds two refinements, momentum and per-parameter step sizes. This lesson explains each piece, why it works, and the optimizer settings large models actually use.

Noisy gradients are enough

The training loss is an average over examples, $L(\theta) = \frac1N\sum_i \ell_i(\theta)$. Pick a random minibatch $B$ of $b$ examples and use

in place of $\nabla L$. This estimate is unbiased: its expectation is the true gradient. Its variance falls as $1/b$, so its standard deviation falls as $1/\sqrt b$. Quadrupling the batch only halves the noise, while it quadruples the compute per step.

That the method converges at all goes back to Robbins and Monro (1951). They studied finding the root of a function that can only be observed through noisy measurements. They showed that iterating with step sizes $\eta_k$ satisfying

converges, for example with $\eta_k \propto 1/k$.^[1] The first condition lets the iterates travel arbitrarily far; the second makes the accumulated noise finite. Stochastic gradient descent (SGD) is this scheme applied to $\nabla L = 0$. In practice, deep-learning schedules do not follow $1/k$. They use a roughly constant rate followed by a decay, but the logic is the same: large steps to make progress, small steps to average out the noise.

The noise is not only a cost. Keskar et al. found that very large batches tended to converge to sharper minima and to generalize worse than small batches, and attributed the difference to the noise in small-batch gradients.^[2] The interpretation is debated (see Generalization). Practitioners have also found ways to use large batches without that loss. Goyal et al. trained ResNet-50 on ImageNet with batches of 8,192 images without losing accuracy. They scaled the learning rate linearly with batch size and added a warmup phase.^[3]

Momentum: averaging out the zig-zag

Gradient descent in a long, narrow valley zig-zags across the steep walls and creeps along the floor. Polyak's heavy-ball method (1964) adds a velocity term:^[4]

The velocity is an exponentially weighted sum of past gradients. Components that keep the same sign, along the valley floor, add up, so with a constant gradient the effective step approaches $\eta/(1-\beta)$. That is 10 times larger for the common choice $\beta=0.9$. Components that flip sign, across the walls, largely cancel.

Worked example: the same valley, with momentum

Return to $L = \tfrac12(x^2 + 100y^2)$, with condition number $\kappa=100$. Plain gradient descent with its best step contracts the error by $(\kappa-1)/(\kappa+1) = 0.980$ per step, so a 100-fold reduction takes about 230 steps. On quadratics, heavy-ball momentum with tuned $\eta$ and $\beta$ achieves the rate $(\sqrt\kappa-1)/(\sqrt\kappa+1) = 9/11 \approx 0.818$, a standard result for quadratics. The same 100-fold reduction then takes

a tenfold speed-up. The dependence on conditioning has gone from $\kappa$ to $\sqrt\kappa$. With stochastic gradients the gain is smaller, because momentum also averages noise, but the qualitative benefit remains.

Adam: a step size per parameter

Parameters in a deep network see gradients on very different scales. Embedding rows for rare tokens are updated seldom, and some layers have far larger gradients than others. A single learning rate suits none of them well. Adam (Kingma and Ba, 2014) keeps two running averages per parameter, the mean and the uncentred variance of its gradient:^[5]

corrects them for their zero initialisation,

and updates

The paper's suggested defaults are $\beta_1=0.9$, $\beta_2=0.999$ and $\epsilon=10^{-8}$.^[5]

Why bias correction? At the first step, $m_1 = (1-\beta_1)g_1 = 0.1\,g_1$, which is far too small. Dividing by $1-\beta_1^1 = 0.1$ restores $\hat m_1 = g_1$. Likewise $\hat s_1 = g_1^2$, so the first update is $\eta\, g_1/|g_1| = \pm\eta$ per coordinate.

That last observation is the key property. The ratio $\hat m/\sqrt{\hat s}$ is roughly unit-scale whatever the raw size of the gradient, so each parameter moves by about $\eta$ per step, scaled down when its gradient is noisy (large $\hat s$ relative to $\hat m^2$). Rescaling the loss by 1,000 leaves Adam's updates almost unchanged. Adam behaves like momentum combined with a per-coordinate correction for conditioning.

AdamW: decoupling weight decay

Weight decay shrinks parameters toward zero every step. For plain SGD this is the same as adding an L2 penalty $\tfrac\lambda2\|\theta\|^2$ to the loss. Loshchilov and Hutter pointed out that for Adam the two are not equivalent.^[6] If the penalty's gradient $\lambda\theta$ is added to $g$, it gets divided by $\sqrt{\hat s}$ along with everything else. Parameters with large gradient variance are then barely regularised. AdamW applies the decay directly to the weights, outside the adaptive scaling:

They reported that decoupling made the best decay value largely independent of the learning rate. It also substantially improved Adam's generalization, enough to match SGD with momentum on image-classification benchmarks where Adam had lagged.^[6]

AdamW is widely used to train large transformers. In one documented example, Llama 2 was pretrained with AdamW using $\beta_1=0.9$, $\beta_2=0.95$, $\epsilon=10^{-5}$, weight decay 0.1, gradient clipping at 1.0, 2,000 warmup steps and cosine decay to 10% of the peak learning rate.^[7] The lower $\beta_2$ shortens the memory of the variance estimate from about 1,000 steps to about 20, so it responds faster to changes in gradient scale.

Weight decay matters beyond final accuracy. In this lab's small-network continual-learning experiments, loss of plasticity over long training streams behaved like an Adam-related weight-norm pathology. Plain weight decay fixed it better than several specialised methods at that scale. The authors note that other work finds L2 insufficient for deeper networks and longer streams.^[8] The full results are in Escaping the saturated game.

Key takeaways

• Minibatch gradients are unbiased and have noise that falls as $1/\sqrt{b}$. Robbins–Monro step-size conditions explain why noisy steps still converge.
• Momentum accumulates consistent gradient directions and cancels oscillating ones. On quadratics it improves the dependence on conditioning from $\kappa$ to $\sqrt\kappa$, from about 230 to about 23 steps in the worked example.
• Adam normalises each parameter's step by a running estimate of its gradient scale, so updates are roughly $\eta$ in size and insensitive to gradient magnitude. Bias correction fixes the early steps.
• With Adam, L2 regularisation and weight decay differ. AdamW decouples them and is the standard optimizer for large transformers.

Sources