Optimizers: SGD, Momentum, Adam

Backprop produced the gradients — but how exactly do you take the step? The choice of optimizer decides whether a model trains in hours or weeks, and whether it converges at all. We'll walk through the evolution: SGD → Momentum → RMSProp → Adam → AdamW.

Intuition

Imagine a ball rolling down the landscape of the loss function. Stochastic gradient descent (SGD) is a massless ball: each step it moves strictly against the current gradient, computed on a random mini-batch:

This ball has two problems. The first is ravines: if the surface is elongated (steep in one direction, flat in another), the ball jumps across the ravine and barely crawls along it. The second is noise: a mini-batch gradient is only a noisy estimate of the true one, and the trajectory jerks around.

How it works

Momentum gives the ball mass: instead of the raw gradient we accumulate velocity — an exponential moving average of gradients with coefficient β ≈ 0.9:

Oscillations across the ravine have opposite signs and cancel each other, while the contribution along the ravine accumulates — the ball accelerates in the right direction. β = 0.9 is roughly an average over the last ~10 gradients.

RMSProp attacks a different problem: different parameters have different gradient scales. It keeps a moving average of the squares of the gradients and divides the step by its square root: parameters with persistently large gradients take modest steps, those with tiny ones take bolder ones. That is a per-parameter adaptive learning rate.

Adam (adaptive moment estimation) = momentum + RMSProp. It tracks two moments:

where m̂ = m/(1−β₁^t) and v̂ = v/(1−β₂^t) are the bias correction: m and v start at zeros and are underestimated in the first steps; the correction compensates for that. Typical values: β₁ = 0.9, β₂ = 0.999.

AdamW is Adam with fixed weight decay: the penalty for large weights is applied as a separate term θ ← θ − η·λ·θ, not mixed into the gradient (as in L2 regularization), where the adaptive normalization would distort it. Today AdamW is the default for training transformers.

What to choose in practice

For transformers and LLMs the answer is almost non-negotiable: AdamW with warmup and a cosine schedule. SGD with momentum is still alive in computer vision (ResNets often generalize a bit better with it), but needs more careful tuning. Typical starting values worth remembering: η ≈ 3·10⁻⁴ for fine-tuning small models, β₁ = 0.9, β₂ = 0.999, weight decay ≈ 0.01–0.1 — and almost nobody touches them except the learning rate.

On the learning rate schedule, two tricks are enough to know. Warmup — for the first hundreds/thousands of steps η grows linearly from zero: early in training Adam's moment estimates are still noisy and the gradients are large, so a full step can blow the model up. Cosine schedule — after warmup η smoothly decays by a cosine nearly to zero: large steps early, careful polishing at the end. Check in the interactive: too large an η makes Momentum oscillate across the whole ravine, and too small a one leaves SGD crawling along the bottom almost forever.

All interview questions in this section: Neural Networks — interview questions