Loading models...

3D view unavailable

This simulation needs WebGL2. Try another browser.

Neural Network Loss Landscape

A real two-weight neural network trained four ways at once. Four marbles - SGD, Momentum, RMSprop, Adam - roll down the same loss surface from the same start. Click the map to move the start, and find the symmetric saddle, the plateau and the step size that blows up.

-SGD: loss | state

-Momentum: loss | state

-RMSprop: loss | state

-Adam: loss | state

-training step

-step-size limit at a minimum

Independent checks

Each row compares the backpropagation code with a separate calculation: central finite differences, the exact symmetry and zero of the loss, a hand-computed first Adam step, equality with plain SGD, second differences for the curvature, and the 2/lambda stability limit of gradient descent.

The science

The network. A tiny neural network with two hidden tanh units: y(x) = tanh(w1 x + b) + tanh(w2 x + b), with the bias b = 0.4 and the output weights fixed. Only the two weights w1, w2 are trained, so the whole loss function fits on a surface. The data are 48 points y = tanh(2.0 x + b) + tanh(-1.5 x + b) from a "teacher" network, loss L = mean (y_pred - y)^2, which is exactly zero at the teacher weights.

Backpropagation. With e = y_pred - y and the tanh derivative 1 - tanh^2, dL/dw_k = mean 2 e (1 - tanh^2(w_k x + b)) x. The marbles move by this exact gradient (checked against a finite difference below). The height drawn is 2.4 ln(1 + L/0.15) / ln(1 + 4/0.15), a logarithmic compression so that both the deep valleys and the high plateaus are visible.

Optimisers. SGD: w -= lr g. Momentum: v = beta v + g, w -= lr v. RMSprop: s = 0.9 s + 0.1 g^2, w -= lr g / (sqrt(s) + 1e-8). Adam: m and s are running averages of g and g^2 (0.9, 0.999), bias-corrected, w -= lr m / (sqrt(s) + 1e-8). With a minibatch smaller than 48 the gradient is computed on a random subset, which makes the path noisy.

What to look for. The loss surface is symmetric: swapping w1 and w2 gives the same network, so there are two equal minima (white dots, the flags) and a ridge between them along w1 = w2. Starting exactly on the diagonal, the gradient has no component across it and every optimiser slides to the saddle at about (0, 0) and stays there. Starting at large positive weights, tanh saturates, the gradient vanishes and learning stalls on a plateau. Near a minimum a gradient-descent step is stable only for lr < 2/lambda_max, where lambda_max is the largest curvature (Hessian eigenvalue); the experiment measures that limit.

Simplifications. Two trainable weights out of a real network's millions: real loss surfaces are in very high dimensions and this slice is a teaching model, not a picture of one. No weight decay, batch normalisation or learning-rate schedule; the minibatch is sampled without replacement from the same seeded generator for all four optimisers. Adam and RMSprop use their usual constants; the rates are yours to change.