UgraByte

3D Loss Landscape — Watch SGD, Momentum and Adam Descend

Rotate a real loss surface and race SGD, momentum and Adam down it — then raise the learning rate until they break.

Loading the interactive version…

What this actually is

Training a model means minimising a loss function, and a loss function is a landscape. Every setting of the parameters is a location, the loss there is the height, and training is a search for a low point. With two parameters that landscape is a surface you can look at — which is why this page can exist at all.

Gradient descent is the rule for walking it: find which way is downhill, step, repeat. Everything that makes training hard is a property of the shape. A long narrow valley makes the steepest direction almost perpendicular to the one you need. A saddle flattens the gradient to nearly nothing without being a minimum. Ripples create dips lower than their surroundings but nowhere near the lowest point.

What you’re watching

The surface is the loss, with height and colour both showing it — blue low, orange high. Drag to rotate; the shape is the point, so look from a few angles first. The panel beside it is the same surface from above, banded like a contour map, and tapping it moves the start.

The coloured trails are three optimisers descending from that start. SGD steps straight downhill. Momentum accumulates velocity, so it carries through flat ground and cuts corners. Adam gives each parameter its own step size. Watch where they separate rather than where they finish.

Then push the learning rate up. On the ravine and the bowl there is a point where each step overshoots the valley and lands higher up the far wall than it started, the loss climbs, and the run leaves the plot. That is divergence — the most common reason a real run suddenly produces NaN.

How to drive it

  1. Drag the surface to rotate it and get a feel for the shape before running anything.

  2. Press Run and watch where the three paths separate. Compare the “steps to 1%” column rather than just the final loss.

  3. Raise the learning rate until a run diverges, then lower it until it does not — the boundary between the two is the single most important hyperparameter in training.

  4. Tap the contour map to start somewhere else, and try the saddle and the rippled surfaces to see what defeats plain gradient descent.

How this relates to real systems

Two honest simplifications. Real landscapes have millions of dimensions, which changes what matters: up there local minima are rare and saddles are everywhere, so the saddle surface is the representative one and the rippled surface the least. Published loss-landscape figures are themselves 2D slices through that space — pictures of this kind, not photographs.

These surfaces are also fixed and exactly differentiable, where a real one is re-estimated from a different mini-batch every step. That is the stochastic in SGD: the path jitters and can rattle out of a shallow dip by itself. The optimisers are otherwise identical to PyTorch's, and the learning-rate failure you can trigger here is the one that breaks real runs.

Frequently asked questions

What exactly is a loss landscape?
A surface where every horizontal position is one setting of the model's parameters and the height is how wrong the model is at that setting. Training walks downhill on it. Here there are two parameters so the landscape is an ordinary surface; a real network has one dimension per parameter, often hundreds of millions, and the same maths applies to a shape nobody can picture.
Why does a high learning rate make the loss go up?
Because the gradient only tells you which way is downhill, not how far downhill continues. In a narrow valley a step larger than the valley is wide crosses the bottom and lands part way up the opposite wall — higher than where it began. The next step does the same in reverse, larger. The loss grows geometrically and overflows to infinity or NaN within a few dozen steps.
Is momentum always better than plain SGD?
It is better on every surface here, and that generalises reasonably well, but it is not free. Momentum can overshoot a narrow minimum it would otherwise have settled into, and it adds a second hyperparameter to get wrong. What it reliably fixes is the two problems visible here: zigzagging in an ill-conditioned valley, and crawling across the near-flat region around a saddle.
Why do saddle points matter more than local minima?
Because of how dimensionality works. A point is a local minimum only if the surface curves upward along every single axis; if even one axis curves down, it is a saddle. With millions of parameters, the odds of every axis curving the same way are vanishingly small, so almost every point where the gradient vanishes is a saddle rather than a trap. That is why modern optimisers are designed to keep moving through flat regions.
Is this a real 3D renderer?
It is real 3D projection, drawn on a plain 2D canvas with no 3D library. The surface is a height field, so sorting its quads back to front gives correct occlusion directly — which is why the descent path disappears behind a ridge instead of being painted over it. The projection is orthographic rather than perspective on purpose: parallel projection keeps equal losses at equal screen heights, and comparing heights is the entire point.