James Phoenix
Training

The optimiser race

SGD, momentum and Adam released together on an ill-conditioned valley. Drag the condition number and watch the exact learning rate at which SGD must explode, while Adam barely notices.

Condition number 12:1
Learning rate 0.080
Run
minimumSGDMomentumAdam
loss 14.186
SGD
loss 14.186
Momentum
loss 14.186
Adam
0.167ok
SGD stability limit

All three see the same gradients and the same learning rate. SGD zigzags across the narrow axis while crawling along the flat one; momentum damps the zigzag by averaging it out; Adam divides each coordinate by its own RMS gradient, which flattens the valley from its point of view.

Real loss surfaces are badly scaled: some directions are steep, others nearly flat, often by factors of thousands. This bowl compresses that fact into one slider. The condition number is the ratio between the steepest and flattest curvature, and every pathology on this page comes from it.

Watch SGD first. The learning rate that is safe for the steep axis is timid for the flat one, so the trajectory zigzags violently across the valley while creeping along it. The stability boundary is not vague: on this bowl SGD converges if and only if the learning rate is below 2 divided by the condition number, and the readout tracks that limit live. Nudge the rate across it and the zigzag flips from shrinking to growing mid-flight.

Momentum helps for an averaging reason. The across-valley component of the gradient alternates sign once the rate is large enough to overshoot, so a running average cancels it, while the along-valley component always points the same way, so the average reinforces it. Same information, better spent.

Adam does something more radical: it divides each coordinate's step by the root-mean-square of that coordinate's own gradient history. Steep directions get shrunk, flat ones get boosted, and the valley becomes approximately round from the optimiser's point of view. That is why the green trail heads almost straight for the minimum, and why its behaviour barely changes as you drag the condition number from 1 to 40.

The honest caveat: this bowl is axis-aligned, and Adam's normalisation is per-coordinate, so this is its best case. Rotate the valley 45 degrees and Adam's advantage shrinks, because the ill-conditioning no longer lines up with the coordinates it normalises. Real transformers sit somewhere between, which is why Adam variants dominate but do not make conditioning irrelevant.

The maths

The bowl and its condition number
L(x,y)  =  12(x2+κ y2),κ  =  λmaxλminL(x, y) \;=\; \tfrac{1}{2}\big(x^2 + \kappa\, y^2\big), \qquad \kappa \;=\; \frac{\lambda_{max}}{\lambda_{min}}

Kappa is the slider. The contours are exact ellipses squashed by root kappa, which is why the drawing narrows as you drag it.

SGD stability, exactly
yt+1  =  (1−ηκ) yt⇒converges  ⟺  η<2κy_{t+1} \;=\; (1 - \eta\kappa)\,y_t \quad \Rightarrow \quad \text{converges} \iff \eta < \tfrac{2}{\kappa}

Along the steep axis each step multiplies the coordinate by (1 - eta kappa). Past 2/kappa that factor is below -1: the sign alternates and the magnitude grows, which is the widening zigzag you can trigger with the slider.

Adam, per coordinate
Δθi  =  −η m^iv^i+ε,v^i≈E[gi2]\Delta\theta_i \;=\; -\eta\, \frac{\hat{m}_i}{\sqrt{\hat{v}_i} + \varepsilon}, \qquad \hat{v}_i \approx \mathbb{E}[g_i^2]

Each coordinate is divided by its own RMS gradient, so on a coordinate with gradients kappa times larger it takes a step roughly kappa times smaller than SGD would, cancelling the conditioning. The step size becomes about eta everywhere, which is why Adam tolerates rates that destroy SGD here.

Related terms

More visualisations

Building with language models?

These explainers come out of the work. If you want the same thinking applied to your own system, that is what I do.

See how I can help