Gradient Descent Lab: Interactive AI Lab

About

Training a neural network means finding the parameters that make a loss function small. Real networks have millions or billions of parameters, so nobody can look at the whole landscape. What an optimizer can do is measure the slope where it stands (the gradient) and take a step downhill. Repeat that a few hundred thousand times and you have trained GPT, a ResNet, or AlphaFold.

This lab shrinks the problem to two parameters, x and y, so you can see the landscape. Colour is height: dark purple is low loss, yellow is high. The thin white lines are contours, like a hiking map. Yellow crosses mark the true minima.

Five optimizers start from the same point and race. Each one only ever sees the gradient at its current position, exactly as in real training. The paths they take, the zigzags, overshoots and stalls, are the same behaviours that decide whether a real model trains in hours or not at all.

Things to try

Ravine (2)Run with the defaults. SGD bounces between the walls while Momentum glides along the floor. Now halve the SGD learning rate: the bouncing stops but progress along the floor slows to a crawl.
Rosenbrock (3)SGD barely moves because the safe learning rate is tiny. Adam normalises each axis and follows the curved valley.
Saddle (4)Click a start exactly on the ridge (y = 0). Every optimizer slides into the saddle and stops, because the gradient never points off the ridge. Turn on minibatch noise and watch the optimizers fall off.
Himmelblau (5)Click different starting points. The same algorithm can end in a different minimum, just as networks trained from different random seeds do.
Bumpy bowl (6)Plain SGD, RMSProp and Adam get stuck in a local dip; the heavy-ball methods coast over the ripples.
Any surfaceDrag the learning rate up until something explodes. A cross in the optimizer’s colour marks a run that diverged.

How it works

Notation: θ is the parameter vector (here the point (x, y)), g = ∇f(θ) is the gradient at the current point, η is the learning rate. All operations on vectors are element-wise. These are the exact update rules this lab runs every step.

SGD

θ ← θ − η·g

Step against the gradient, scaled by the learning rate. With a full-batch gradient this is plain gradient descent; with minibatch noise it is stochastic gradient descent.

Momentum (Polyak heavy ball, 1964)

v ← μ·v − η·gθ ← θ + v

v is a velocity: an exponentially decaying sum of past gradients. Consistent directions add up, oscillating ones cancel. μ = 0.9 means roughly the last 10 steps matter.

Nesterov momentum

v ← μ·v − η·∇f(θ + μ·v)θ ← θ + v

Measure the gradient where the velocity is about to carry you, not where you are. That look-ahead corrects overshoot a step earlier. Form from Sutskever et al. (2013).

RMSProp (Hinton, 2012)

s ← ρ·s + (1 − ρ)·g²θ ← θ − η·g / (√s + ε)

s tracks the recent mean of squared gradients per parameter. Dividing by its square root makes steep directions take smaller steps and flat ones larger steps.

Adam (Kingma and Ba, 2014)

m ← β₁·m + (1 − β₁)·gv ← β₂·v + (1 − β₂)·g²m̂ = m / (1 − β₁ᵗ), v̂ = v / (1 − β₂ᵗ)θ ← θ − η·m̂ / (√v̂ + ε)

Momentum (m) and RMSProp scaling (v) together. Both averages start at zero, so early on they are too small; dividing by 1 − βᵗ corrects that bias. Defaults β₁ = 0.9, β₂ = 0.999, ε = 1e−8.

Minibatch noise

Real training estimates the gradient from a small random batch, so every step sees g plus noise. The noise toggle adds Gaussian noise with standard deviation σ times a typical gradient size for the surface. It is a simulation of that effect, not a real dataset.

Learning rates

The learning rate is the single most important hyperparameter in deep learning. Take the simplest possible loss, f(w) = w², whose gradient is 2w. One gradient step gives w ← w − η·2w = (1 − 2η)·w, so the distance to the minimum is multiplied by |1 − 2η| every step. Below η = 0.5 you approach from one side; between 0.5 and 1 you overshoot and oscillate but still converge; above 1 you diverge. These four panels are the real iterates, starting at w = 2.

In more dimensions each direction has its own curvature (the eigenvalues λ of the Hessian). Plain gradient descent is stable only if η < 2/λ_max, but progress along the flattest direction goes like 1 − ηλ_min. The ratio λ_max/λ_min, the condition number, is therefore what makes a problem hard. The ravine surface here has curvatures 0.1 and 5, a condition number of 50: any learning rate safe for the steep axis is painfully slow on the flat one. Momentum and adaptive methods exist largely to fix this.

Optimizers compared

Optimizer	Strengths	Weaknesses	Where you meet it
SGD	Simplest, least memory, well understood theory. With a good schedule it often generalises very well.	Zigzags in ravines, crawls on plateaus, one learning rate for every parameter.	Classic image models (ResNets on ImageNet were trained with SGD + momentum).
Momentum	Damps oscillation across a ravine and accelerates along it. Can coast over small bumps.	Overshoots and orbits around the minimum when μ is high.	Default partner of SGD in computer vision.
Nesterov	Same benefits as momentum, with earlier braking thanks to the look-ahead gradient.	Tighter stability limit on very steep directions: try the ravine with a high learning rate.	Common option in SGD optimisers (e.g. nesterov=True in PyTorch).
RMSProp	Per-parameter step sizes: badly scaled problems become roughly equally scaled.	Near a minimum the normalised step stays about η in size, so it jitters unless η decays.	Early deep reinforcement learning (e.g. DQN) and RNNs.
Adam	Robust default, little tuning, handles sparse and badly scaled gradients.	Extra memory (two numbers per parameter). Can settle in sharper minima than SGD.	The default for transformers; LLMs are trained with AdamW, Adam plus decoupled weight decay.

No optimizer wins everywhere, and a 2D picture can mislead: in a million dimensions, true local minima are rare and saddle points and long flat valleys are the common obstacles. Still, the qualitative behaviours you see here (oscillation across steep directions, acceleration from momentum, per-axis rescaling) are exactly what happens at scale.

Gradient Descent Lab