Beyond SGD — Adam, Second-Order Methods, and Constrained Optimization

What exactly is momentum accumulating? What does each of Adam's four lines do? What did AdamW fix? And why does nobody train an LLM with second-order methods that are supposedly faster? Metaphor, equations, live figures, code, and production practice — no prerequisites assumed.

Adam: A Method for Stochastic Optimization

A metaphor: descending a narrow canyon by feel

You are in fog, finding the valley floor using nothing but the slope under your feet — the picture from Calculus for AI. Let's push it one step further.

This time the terrain is a narrow canyon: steep walls on both sides, nearly flat along the canyon floor. You dutifully walk in whatever direction is steepest downhill. But that direction is not along the floor — it is straight across the canyon. One step lands you on the opposite wall, where the steepest direction again points across. You bounce from wall to wall and creep only inches toward the exit.

This is not a fanciful picture; it is everyday life on a loss surface. If curvature differs by a factor of 100 between directions, your step size has to accommodate the steepest one, so the flat direction advances a hundred times slower. SGD is slow not because the gradient is wrong, but because "steepest" and "where you want to go" are different directions. There are four ways out: accumulate (momentum), give each direction its own step size (Adam), measure the curvature itself (second-order methods), and fence off where you're allowed to move (constrained optimization).

Recap: what SGD actually does

In plain words: from the current parameters $\theta_t$, move a distance set by the learning rate $\eta$ in the direction opposite to $g_t$, the "gets worse" direction measured on the data you happen to be holding. That is all. $B_t$ is the minibatch you drew this step — a bundle of tens to hundreds of examples. Put in words: measure the slope under your feet using a few hundred examples, take one fixed-length step the other way, and repeat forever.

Two problems live inside that line. One is the direction problem from the canyon. The other is noise: $g_t$ is not the true gradient over the full dataset but an estimate from a few hundred samples, so standing in the same spot on different batches gives you different arrows. Conveniently, one tool relieves both at once.

Momentum: keep it as velocity instead of discarding it

SGD has no memory of previous steps. Give it velocity.

Read aloud: velocity $v_t$ is "the previous velocity decayed by a factor $\beta$" plus "this step's gradient," and the parameters move by that velocity. $\beta$ is a decay rate between 0 and 1; 0.9 is the conventional choice. Put in words: rather than measuring a fresh step from scratch each time, you carry over nine tenths of the push you already had and add today's slope on top of it.

What 0.9 means falls out of $1/(1-\beta) = 10$: you are roughly summing the last ten steps' gradients. That sum kills both problems at once.

• Direction: the across-the-canyon component flips sign between the left and right wall, so the terms cancel. The along-the-floor component keeps the same sign every step, so it accumulates. The zigzag vanishes and the direction you actually want grows.
• Noise: random batch-to-batch jitter also washes out under averaging — ten steps buys you roughly a factor of $\sqrt{10}$ in noise resistance.

If SGD is a hiker who stops and re-measures the ground at every step, momentum is a heavy ball rolling down the slope.

The side effect is overshoot. Velocity survives arrival at the bottom, so the ball sails past it. Raising $\beta$ buys acceleration at the cost of stopping power — that tug-of-war sits behind everything that follows.

A step size per direction: AdaGrad to RMSProp

Momentum fixed the direction, but the step size is still shared by every parameter. So give each parameter its own. The raw material is how large that direction's gradient has been in the past: shrink the step where big gradients arrive often, stretch it where they rarely do. That is AdaGrad's idea. The catch is that AdaGrad sums all past squared gradients, so the denominator grows monotonically and the step size dies toward zero. RMSProp replaces "all" with "recent."

In plain words: keep a running average $v_t$ of recent squared gradients for each parameter, and divide the learning rate by its square root. Here $g_t^2$ is elementwise squaring and $\varepsilon$ is a small constant (around $10^{-8}$) that prevents division by zero. Since $\sqrt{v_t}$ is "the typical magnitude of the gradient in that direction," dividing by it means rescaling every gradient to roughly unit size before you step. Think of it in words: the directions that always shout get their volume turned down and the ones that rarely speak get turned up, so every direction reaches you at the same loudness before you decide where to walk.

That gives us two ingredients. Momentum answers which way to go; RMSProp answers how far to go. They track different statistics, so they can coexist without interfering.