Gradient Descent: 3D Loss Surface & Optimisers

A loss function is a landscape

Every training run of a neural network is, geometrically, a ball rolling downhill on a landscape defined by the loss function — height is loss, position is the model's parameters, and training is the search for the lowest point reachable from wherever you started. For a network with millions of parameters that landscape lives in a space you cannot draw, but its qualitative failure modes — narrow curved valleys, flat plateaus, fields of small local bumps — are exactly the same ones visible on classic two-parameter test functions, which is why optimisation researchers still lean on them to build intuition.

live demo · an optimiser's trajectory across a loss surface● LIVE

Three landscapes, three failure modes

The Rosenbrock function — f(x,y) = (a-x)² + b(y-x²)² — carves out a long, thin, curved parabolic valley: the gradient across the valley is steep while the gradient along it toward the minimum is nearly flat, an extreme case of what optimisation theory calls ill-conditioning. The Rastrigin function adds a cosine ripple on top of a bowl, producing thousands of shallow local minima around one global one — a stress test for whether an optimiser gets trapped. The Himmelblau function has four separate global minima of equal height, useful for showing that where you start decides which minimum you actually reach — gradient descent has no way to know a better one exists elsewhere.

Plain SGD: correct direction, wrong step

Vanilla gradient descent takes the simplest possible action: step directly opposite the gradient, scaled by a fixed learning rate.

theta = theta - lr * grad(theta)

On the Rosenbrock valley this fails visibly: because the gradient points mostly across the narrow direction rather than along the shallow one, SGD bounces wall-to-wall in a slow zig-zag, wasting almost all of its motion on a direction that does not matter and making painfully slow progress toward the actual minimum. Set the learning rate too high anywhere on this surface and the steps overshoot and diverge; set it too low and progress along the shallow valley floor is glacial. A single global learning rate simply cannot be right for both directions at once.

Momentum: remembering where you were heading

Momentum keeps a running, exponentially-decayed average of past gradients and steps in that direction instead of the raw instantaneous gradient:

v = beta * v + grad(theta)      // beta ~ 0.9, a running velocity
theta = theta - lr * v

Because the zig-zagging components across the Rosenbrock valley point in alternating directions, they partially cancel in the running average, while the small, consistent component along the valley floor accumulates and reinforces itself — the ball behaves more like a physical object with inertia than a memoryless point that reacts only to the local slope. The net effect is dramatically faster progress along shallow, consistent directions and damped oscillation across steep, alternating ones.

RMSprop: a learning rate per direction

RMSprop attacks the ill-conditioning problem from the opposite angle: instead of smoothing the direction, it tracks a running average of the squared gradient in each parameter direction and divides the step by its square root.

s = decay * s + (1 - decay) * grad(theta)^2
theta = theta - lr * grad(theta) / (sqrt(s) + eps)

Directions where the gradient is consistently large (the steep walls of the Rosenbrock valley) get their effective step size shrunk automatically; directions where the gradient stays small (the shallow floor) get their step size enlarged. Each parameter effectively earns its own adaptive learning rate, which is exactly what the ill-conditioned valley needs and what a single global rate cannot provide.

Adam: momentum and RMSprop combined

Adam (Kingma and Ba, 2014) tracks both a first moment (Momentum's running gradient average) and a second moment (RMSprop's running squared-gradient average), applies a bias correction to account for both starting at zero, and combines them into the update. It converges quickly on most practical loss surfaces with comparatively little learning-rate tuning, which is why it became the default optimiser for training neural networks — though a well-tuned SGD with momentum is still competitive, and sometimes generalises better, on some vision benchmarks. On the Rosenbrock demo, watch Adam settle into the valley floor almost immediately, while plain SGD is still zig-zagging across the walls.

Frequently asked questions

Why does plain SGD zig-zag on the Rosenbrock function?

The Rosenbrock valley is long, narrow and curved, so its gradient points almost across the valley (steep direction) rather than along it (shallow direction, toward the minimum). Plain gradient descent always steps exactly opposite the gradient, so it bounces between the steep walls in a zig-zag while making very slow net progress along the valley floor toward the actual optimum.

What is the difference between Momentum and RMSprop?

Momentum accumulates a running average of the raw gradient (first moment) and keeps moving in a consistent direction, which damps zig-zag and accelerates through shallow, consistent slopes. RMSprop instead accumulates a running average of the squared gradient (second moment) and divides the step by its square root, which shrinks steps in directions of large, noisy gradients and enlarges them in flat directions. They fix different problems and Adam combines both.

Is Adam always better than plain SGD for training neural networks?

Not always. Adam typically converges faster and needs less learning-rate tuning, which is why it is the default starting point. But a well-tuned SGD with momentum often generalises slightly better on some vision tasks, and several published results show Adam settling into a solution with higher test error unless carefully regularised, which is why practitioners still benchmark both rather than assuming Adam wins by default.

Try it live

Everything above runs in your browser — open Gradient Descent 3D and change the parameters while it is running. Nothing is installed, nothing is uploaded, the whole model lives in one tab.

▶ Open Gradient Descent 3D simulation

Gradient Descent in 3D: Reading a Loss Surface Like a Landscape