Sandbox

Pick a loss surface, choose whether to compare batch size or optimizers, and watch how the pinned dimension changes the descent.

Optimizers
0.01
60 iter/s

Click anywhere on the contour map to move the start.

Convexity

A function is convex if the line segment between any two points on its graph never dips below the graph itself: every local dip is the same dip as the global minimum.

$$f(t x + (1-t) y) \le t f(x) + (1-t) f(y), \quad t \in [0,1]$$

For a twice-differentiable function like the surfaces here, this is equivalent to the Hessian (the matrix of second derivatives) being positive semi-definite everywhere: curving upward, never downward, in every direction, at every point. When a surface is convex, any local minimum gradient descent finds is automatically the global minimum, and where you start only changes how long the descent takes, never which minimum it ends at.

The elongated bowl and quadratic bowl are both convex: simple paraboloids, curving upward in every direction everywhere, with a single minimum at the origin regardless of start point. Rosenbrock is not convex, its curved valley bends the wrong way along parts of its length, but it is still unimodal: only one minimum exists, at $(1, 1)$, so a non-convex surface does not by itself mean multiple competing minima. The saddle surface is both non-convex and multimodal: the origin is a saddle point, curving upward along one axis and downward along the other, with two separate, equally deep minima sitting on opposite sides of it. Try setting the start point above the origin versus below it on the saddle surface: the same optimizer settles into two different minima depending only on which side of the saddle it started, not on which optimizer is used.

The algorithms

SGD

At each step, SGD moves the parameters directly opposite the gradient, scaled by the learning rate $\alpha$.

$$\theta_{t+1} = \theta_t - \alpha \nabla L(\theta_t)$$

When a surface curves much more steeply in one direction than another, a single learning rate cannot serve both directions well: small enough to avoid diverging on the steep direction means far too small for efficient progress on the shallow one. SGD has no way to tell the two apart, so on any such surface it bounces across the steep axis while crawling along the shallow one. This same update rule is what every batch-size line in the sandbox uses when SGD is the pinned optimizer; only how noisy the gradient it receives changes.

Momentum

Momentum accumulates a velocity vector $v$ that carries the optimizer across flat regions and dampens oscillations across curved ones, fixing the bouncing SGD shows above. $\beta$ (typically 0.9) controls how much of the previous velocity survives each step.

$$v_{t+1} = \beta v_t - \alpha \nabla L(\theta_t)$$
$$\theta_{t+1} = \theta_t + v_{t+1}$$

On the elongated bowl, the accumulated velocity along the long axis lets momentum accelerate where SGD is slow. The overshoot into the steep walls eventually dampens, but the path oscillates more dramatically before settling.

AdaGrad

Momentum still applies the same effective step size to every parameter, so a steeply curved axis and a shallow one are still stuck sharing one learning rate. AdaGrad addresses this by giving each parameter its own rate, shrinking it in proportion to the cumulative sum of squared gradients that parameter has seen so far.

$$G_t = G_{t-1} + (\nabla L(\theta_t))^2$$
$$\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{G_t} + \epsilon} \nabla L(\theta_t)$$

Parameters with consistently large gradients get their rate shrunk quickly; parameters with small gradients keep a larger effective rate. Because $G_t$ only ever grows, the effective rate keeps shrinking for the rest of the run, and on a long descent it can decay so far that progress effectively stalls.

RMSProp

RMSProp keeps AdaGrad's idea of a per-parameter rate but fixes the stall: instead of an ever-growing sum of squared gradients, it keeps an exponential moving average, so old gradients are gradually forgotten.

$$E[g^2]_t = \beta E[g^2]_{t-1} + (1-\beta)(\nabla L(\theta_t))^2$$
$$\theta_{t+1} = \theta_t - \frac{\alpha}{\sqrt{E[g^2]_t} + \epsilon} \nabla L(\theta_t)$$

Because the average can rise again if gradients grow again, RMSProp's effective rate does not monotonically decay the way AdaGrad's does, so it keeps making progress over long runs.

Adam

Adam combines momentum's directional smoothing with RMSProp's per-parameter scaling: a running estimate of both the first moment (mean gradient $m$) and the second moment (uncentered variance $v$) of the gradient. The parameter update is scaled by the per-dimension gradient variance, so large gradients get smaller effective steps and small gradients get larger ones.

$$m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla L$$
$$v_t = \beta_2 v_{t-1} + (1-\beta_2)(\nabla L)^2$$
$$\theta_{t+1} = \theta_t - \alpha \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$$

The $\hat{m}_t$ and $\hat{v}_t$ are bias-corrected estimates that account for the zero-initialization of the moment vectors at $t=0$. Adam navigates elongated landscapes efficiently because the per-dimension scaling automatically compensates for the difference in curvature between axes. Its effective learning rate for each parameter is approximately $\alpha$, regardless of gradient magnitude.

Batch size

Every update rule above needs a gradient to work with, and how that gradient is computed is a separate choice from which update rule consumes it. Full-batch gradient descent computes the exact gradient over the entire dataset every step: the direction is precise, but one step costs a full pass over the data. Stochastic gradient descent estimates the gradient from a single example: each step is cheap, but the estimate is noisy, so the path jitters. Mini-batch gradient descent averages over a small batch of examples, trading some of SGD's noise for some of full-batch's cost, and is what most training in practice actually uses.

The surfaces in this sandbox are analytic functions, not real datasets, so there is nothing to literally subsample. The batch-size comparison instead adds simulated Gaussian noise to the true analytic gradient, scaled down as the batch size grows ($\sigma \propto 1/\sqrt{n}$), to illustrate the effect a real batch size would have.