Optimizer race

Training a model means finding parameters $x$ that make a loss $f(x)$ small. An optimizer does this step by step: it looks at the gradient $g = \nabla f(x)$, the direction in which $f$ grows fastest, and moves the parameters a little against it. How far it moves is set by the learning rate $\eta$ (the step size).

Real models have millions of parameters, but the main effects are already visible with two. Here $x = (x_1, x_2)$, and the loss is drawn as a contour map: each line joins points with the same loss, like the height lines of a hiking map, and darker regions are lower. Seven optimizers start at the same point and race to the minimum:

• GD: plain gradient descent;
• momentum and Nesterov: gradient descent with a memory of the previous steps;
• AdaGrad, RMSProp, Adam and AdamW: adaptive methods, which give each coordinate its own step size.

You will see where each one is fast, where it is slow, and where it fails. The update rules and the math behind the pictures are in section 3.

1. The race #

Interactive figure: loading…

The six surfaces:

Surface	$f(x_1, x_2)$	What it tests
Rosenbrock	$(1 - x_1)^2 + 100\,(x_2 - x_1^2)^2$	a long, curved, flat valley; minimum $(1, 1)$
Beale	$\sum_{i=1}^{3} (c_i - x_1 + x_1 x_2^{\,i})^2$, $c = (1.5, 2.25, 2.625)$	flat plateaus next to steep walls; minimum $(3, 0.5)$
Saddle	$\tfrac12 x_1^2 - \tfrac12 x_2^2 + \tfrac14 x_2^4$	a saddle point at $(0, 0)$; minima $(0, \pm 1)$
Ill-conditioned	$\tfrac12\big(\tfrac1\kappa u^2 + w^2\big)$, valley turned by the angle $\theta$	the condition number $\kappa$; axis-aligned or tilted valley
Noisy	$\tfrac12 x^\top A x$, with gradient noise of size $\sigma$	stochastic gradients, as with minibatches
Himmelblau	$(x_1^2 + x_2 - 11)^2 + (x_1 + x_2^2 - 7)^2$	four global minima; which one do you reach?

2. Try this #

1. Momentum in a curved valley. Use Rosenbrock with the default settings. GD, momentum and Nesterov all have the same learning rate $\eta = 10^{-3}$. Question: which of them reach the finish line in 1500 steps? What to see: GD does not; after 1500 steps it is still in the valley near $(0.60, 0.36)$ with $f - f^\star \approx 0.16$. Momentum finishes at $t = 836$ and Nesterov at $t = 873$. Now set $\beta = 0.5$: momentum does not finish either. With $\beta = 0.95$ it finishes at $t = 407$. Along the valley the gradient keeps its direction, so momentum makes the step about $1/(1 - \beta)$ times longer. Finally set AdamW's weight decay to $\lambda = 0.3$: AdamW now stops near $(0.48, 0.23)$, because the decay pulls it toward $0$.
2. Too large a step. Choose Ill-conditioned ($\kappa = 25$, $\theta = 0°$). The side panel shows the largest stable learning rate of GD, momentum and Nesterov. Question: move GD's $\eta$ to $1.5$, $2$, $3$. What happens? What to see: at $1.5$ GD zigzags across the valley but still finishes ($t = 59$). At exactly $\eta = 2$ it jumps between two points forever. At $\eta = 3$ the board shows diverged after 8 steps, and the curve in the chart goes up. Momentum diverges above $3.8$ and Nesterov above $1.36$ (with $\beta = 0.9$). Adam does not blow up even with $\eta = 5$ (it finishes at $t = 82$): its steps are normalized. On Rosenbrock, the curvature near the start is about $1500$, so GD with $\eta = 0.007$ diverges after 4 steps.
3. Adam rescales the coordinates. Stay on Ill-conditioned and set $\kappa = 100$, $\theta = 0°$. Question: compare GD and Adam, then turn the valley to $\theta = 45°$. Who suffers from the turn? What to see: with the axis-aligned valley Adam finishes at $t = 68$ and AdaGrad at $t = 26$, while GD needs $290$ steps. Adam's number does not change when you raise $\kappa$. At $45°$ Adam and AdaGrad no longer finish in 400 steps. GD and momentum do not slow down: they even finish earlier ($t = 161$ and $t = 85$), only because the start point now has a smaller part along the flat direction. Their speed does not depend on the angle. Adam scales each coordinate; that fixes a valley along an axis, not a tilted one.
4. Escaping a saddle point. Choose Saddle. The start is $(-2, 0.01)$, very close to the line $x_2 = 0$ that leads into the saddle. Question: which optimizers leave the saddle first? What to see: the adaptive methods turn toward a minimum at once ($|x_2| > 0.5$ after 3 steps for AdaGrad and 11 steps for RMSProp and Adam). Momentum and Nesterov need about 40 steps. GD needs 83 steps and spends 33 of them near the saddle, where the gradient is tiny. Now drag the start point onto the axis $x_2 = 0$ (it snaps to steps of $0.01$). All optimizers run into the saddle and stop there with $f - f^\star = 0.25$: the gradient has no $x_2$ part, so nothing pushes them off the line.
5. Noise and the step size. Choose Noisy. Each optimizer sees the true gradient plus random noise. Question: set GD's $\eta$ to $0.2$, then to $0.02$. How do the cloud of points at the end and the mean column change? What to see: with $\eta = 0.05$ the mean loss of the last 100 steps is about $0.02$; with $\eta = 0.2$ it is about $0.1$; with $\eta = 0.02$ it drops to about $0.01$, but GD needs more steps to arrive. With a constant learning rate, the optimizer does not converge; it moves around in a noise ball whose size grows with $\eta$. AdaGrad's steps shrink over time, so it usually ends lowest. Press New noise to see another random run, and set $\sigma = 0$ to remove the noise.
6. Many minima. Choose Himmelblau. Question: from the default start $(-1, -0.6)$, which minimum does each optimizer reach? What to see: GD, momentum and Nesterov go down-left to star 3; AdaGrad, RMSProp, Adam and AdamW go up to star 2. The first steps of the adaptive methods follow the sign of the gradient, not the gradient itself, so they take another route. Move the start to $(-1.2, -0.4)$: GD still ends at star 3, but momentum and Nesterov now reach star 2. In a nonconvex problem the solution you get depends on the start and on the optimizer.

3. The math behind the race #

All seven optimizers use the gradient $g_t = \nabla f(x_{t-1})$ at step $t = 1, 2, \dots$ (on the noisy surface, $g_t$ is the noisy gradient). The playground uses exactly the update rules of PyTorch's torch.optim. Squares, square roots and divisions of vectors act per coordinate.

Gradient descent and the step size #

Gradient descent (GD) moves against the gradient:

The clearest case is a quadratic $f(x) = \tfrac12 x^\top A x$ with a symmetric positive definite $A$. Its gradient is $Ax$. Write $x$ in the eigenvectors of $A$, the Hessian of $f$ (the matrix of second derivatives). Then each coordinate $e_i$ follows its own rule:

where $\lambda_i$ is the eigenvalue, the curvature in that direction. Call the largest eigenvalue $L$ and the smallest $\mu$. Every factor $|1 - \eta\lambda_i|$ must be below 1, so GD converges only if

For a larger $\eta$ the steep direction overshoots more at every step and the iterates grow: this is divergence (experiment 2). But the flat direction shrinks only by the factor $1 - \eta\mu$, which is close to 1. With the best fixed step $\eta = 2/(\mu + L)$ the error shrinks by $(\kappa - 1)/(\kappa + 1)$ per step, where $\kappa = L/\mu$ is the condition number. So the number of steps grows in proportion to $\kappa$. A long, narrow valley is a valley with a large $\kappa$. In the Ill-conditioned surface $L = 1$ and $\mu = 1/\kappa$.

Momentum and Nesterov: memory along the valley #

Momentum (Polyak's heavy-ball method) adds up the past gradients in a buffer $b$:

This is PyTorch's SGD(momentum=β). If the gradient stays the same, $b_t$ grows to $g/(1 - \beta)$: the effective step is $\eta/(1 - \beta)$, ten times larger for $\beta = 0.9$. In a valley the gradient across the valley changes sign at every step, so these parts cancel in $b$. The gradient along the valley keeps its sign, so it adds up. The method speeds up exactly where GD crawls (experiment 1). With the best $\eta$ and $\beta$ on a quadratic, the error shrinks by about $(\sqrt\kappa - 1)/(\sqrt\kappa + 1)$ per step: about $\sqrt\kappa$ steps per digit of accuracy instead of $\kappa$.

Nesterov's method takes the gradient at a look-ahead point, where momentum is about to carry the iterate: $v_t = \beta v_{t-1} - \eta \nabla f(x_{t-1} + \beta v_{t-1})$, $x_t = x_{t-1} + v_t$. PyTorch (nesterov=True) writes the same method in the variables of the look-ahead point:

On a quadratic with largest eigenvalue $L$, the three methods are stable for

For $\beta = 0.9$ this gives $3.8$ and about $1.36$: the limits in experiment 2.

Adaptive methods: one step size per coordinate #

The adaptive methods divide each coordinate of the step by a running size of its own gradients.

AdaGrad sums all squared gradients:

RMSProp uses an exponential moving average instead, so old gradients are forgotten:

Adam also keeps a moving average of the gradients (momentum) and corrects both averages for their start at zero (bias correction):

AdamW adds decoupled weight decay: $x_t = x_{t-1} - \eta\big(\hat m_t/(\sqrt{\hat v_t} + \varepsilon) + \lambda x_{t-1}\big)$. The term $\lambda x$ pulls every coordinate toward 0. In a neural network this is a regularizer; on a 2D loss it moves the end point away from the minimum of $f$ (experiment 1). The playground uses $\varepsilon = 10^{-8}$, and all states start at zero.

Why Adam rescales the coordinates. At the first step $\hat m_1 = g_1$ and $\hat v_1 = g_1^2$, so the step is $\eta\, g_1/(|g_1| + \varepsilon) \approx \eta \operatorname{sign}(g_1)$: every coordinate moves by about $\eta$, whether its gradient is $1000$ or $0.001$. Later, as long as a coordinate's gradient keeps a similar size, $\hat m_t/\sqrt{\hat v_t}$ stays of order one. So Adam acts like gradient descent with a diagonal preconditioner: a different step size for each coordinate. If the valley lies along an axis, this removes the bad scaling, and $\kappa$ no longer matters. If the valley is tilted, a steep and a flat direction share each coordinate, and no per-coordinate scaling can separate them (experiment 3).

What goes wrong with a too-large step #

• GD and momentum diverge. On a quadratic the steep direction is multiplied by $|1 - \eta L| > 1$ at every step, so the error grows exponentially. On a non-quadratic surface the limit $2/L$ uses the local curvature: on Rosenbrock the largest Hessian eigenvalue at the start $(-1.2, 1)$ is about $1500$, so GD needs $\eta$ below about $0.0013$ there.
• Adam and RMSProp do not blow up, but they do not settle. Their step in each coordinate is at most a few times $\eta$ (with the default $\beta$ and $\rho$), however large the gradient is. With a large $\eta$ they jump over the minimum and wander. With a constant $\eta$, RMSProp ends up stepping back and forth within about $\eta$ of the minimum (look at its f − f* column).
• Noise needs a small or a decaying step. With stochastic gradients $\tilde g = \nabla f + \sigma\xi$ and a constant $\eta$, the iterates never stop moving. For GD on a quadratic in $d$ dimensions, the average loss gap settles at

which grows in proportion to $\eta$ (experiment 5). This is why training uses learning-rate schedules that lower $\eta$ at the end, and why AdaGrad, whose step shrinks like $1/\sqrt t$, ends low.

Where the noise-ball formula comes from

In the eigenvector coordinates of $A$, one GD step with noise is $e \leftarrow (1 - \eta\lambda)\,e - \eta\sigma\xi$ with $\xi \sim \mathcal N(0, 1)$, independent at every step. If the variance $V$ of $e$ has settled, it is the same before and after the step: $V = (1 - \eta\lambda)^2 V + \eta^2\sigma^2$. So $V = \eta\sigma^2 / \big(\lambda(2 - \eta\lambda)\big)$, and this direction adds $\tfrac12 \lambda V = \eta\sigma^2 / \big(2(2 - \eta\lambda)\big)$ to the mean loss gap. In the playground ($\lambda = 0.25$ and $1$, $\sigma = 1$, $\eta = 0.05$) the formula gives $0.0254$; long simulations give $0.025$ as well.

4. Learn more #

These lessons explain the ideas of this page in depth, with proofs, worked examples and notebooks:

• Optimization: loss surfaces and gradient descent, convexity, smoothness and the condition number, convergence rates of gradient descent, momentum and Nesterov acceleration, stochastic gradient descent and the noise ball, AdaGrad, RMSProp, Adam and AdamW and learning-rate schedules, warmup and clipping.
• Calculus: the gradient and the Hessian and Taylor's theorem.
• Linear algebra: eigenvalues and the spectral theorem, which give the eigen-directions of a quadratic.
• Deep learning: the training loop and mini-batches, weight decay as a regularizer and training dynamics: warmup, clipping and loss spikes.