Momentum and adaptive methods

The gradient-descent lesson ended with a diagnosis: it pays for ill-conditioning linearly — $\kappa$ ten times worse, ten times the iterations — and wastes most of its motion zig-zagging across ravines. This lesson is the treatment arc: first give the iterate memory (momentum), then give each coordinate its own learning rate (adaptive methods), then combine both and fix the resulting interaction with weight decay (AdamW). By the end, the optimizer line in every modern training script should read as a sequence of deliberate design decisions rather than folklore.

Momentum: the heavy ball

Polyak’s 1964 idea: keep an exponentially-decaying memory of past gradients — a velocity — and move along it:

Two complementary readings:

As physics. These are the discretized equations of a ball with mass rolling on the loss surface with friction $1-\beta$. Along the ravine floor, where the gradient points the same way step after step, velocity accumulates: the geometric series gives a terminal speed of $\eta \lVert\nabla f\rVert / (1-\beta)$ — an effective learning rate of $\eta/(1-\beta)$, so $\beta = 0.9$ is a 10× amplifier in consistent directions ($0.99$ is 100×, which is why cranking $\beta$ without dropping $\eta$ diverges). Across the ravine, where the gradient flips sign each step, successive contributions cancel inside the velocity — the zig-zag is averaged away. One mechanism, both benefits.

As mathematics. On the quadratic, the two-step recurrence has spectral radius minimized at $\beta^\star = \big(\tfrac{\sqrt{\kappa}-1}{\sqrt{\kappa}+1}\big)^2$ with optimal $\eta$, giving

The condition number enters through its square root: at $\kappa = 100$, GD needs ~100 iterations per digit, momentum ~10. This is not a constant-factor tweak; it is a different complexity class, and (Nesterov proved) the best possible one for first-order methods.

Nesterov’s variant evaluates the gradient at the lookahead point $x_t + \beta v_t$ rather than at $x_t$ — the ball measures the slope where it is about to be, and brakes before overshooting instead of after. Same cost, provably optimal $O(1/t^2)$ rate on smooth convex problems, and a slightly better-damped trajectory you can see in the race below.

Adaptive methods: a learning rate per coordinate

Momentum still applies one global $\eta$ to every parameter. But real models mix parameters with wildly different gradient scales — embeddings of rare tokens receive sparse, tiny gradients while bias terms receive dense, large ones. The adaptive family normalizes each coordinate by a running estimate of its own gradient magnitude.

RMSProp (Hinton, famously published as a lecture slide) keeps an exponential moving average of squared gradients and divides by its root:

Coordinates with persistently large gradients get their steps shrunk; quiet coordinates get boosted. On a quadratic this is implicit preconditioning — it equalizes the effective curvature across axes, attacking $\kappa$ directly rather than tolerating it as momentum does.

Adam = momentum on the first moment + RMSProp on the second + a correction for the zero-initialization of both EMAs:

The bias correction is not a detail: without it, $m_0 = s_0 = 0$ makes the first steps systematically tiny (an EMA warm-starting from zero underestimates by the factor $1-\beta^t$), and with $\beta_2 = 0.999$ that shrinkage lingers for a thousand steps. Defaults $\beta_1{=}0.9$, $\beta_2{=}0.999$, $\varepsilon{=}10^{-8}$ have survived a decade of scrutiny largely intact. A useful invariance to internalize: because $g$ appears homogeneously in numerator and denominator, Adam’s step size is insensitive to the overall scale of the gradient — rescale the loss by 100× and the update barely changes, which is much of why Adam is so forgiving to tune.

AdamW: the weight-decay fix

For SGD, adding L2 regularization $\tfrac{\lambda}{2}\lVert x \rVert^2$ to the loss and “weight decay” $x \leftarrow (1 - \eta\lambda)x$ are the same update. For Adam they are not: an L2 term enters through $g_t$ and gets divided by $\sqrt{\hat s_t}$ like everything else, so heavily-updated parameters — precisely the ones you most want to regularize — receive the weakest effective decay. Loshchilov & Hutter’s fix is simply to take decay out of the gradient path:

The empirical gap this closes is why AdamW — not Adam — is the default in every serious transformer codebase, and why setting weight_decay in Adam-with-L2 mode silently does less than its name promises.

Choosing, in practice

Starting points that respect the theory: SGD+momentum $\eta \sim 10^{-1}$–$10^{-2}$ (effective rate is $\eta/(1{-}\beta)$ — remember the 10×), AdamW $\eta \sim 10^{-3}$–$3{\cdot}10^{-4}$, $\lambda \sim 10^{-2}$. Every one of these is a starting point for a sweep, not a conclusion — the SGD lesson’s noise-ball logic still sets the endgame, and schedules from that lesson compose with everything here.

Exercises