The shape of the loss surface

You know that training follows the gradient downhill. This section is about the surface it is walking on, because the picture most people carry from two dimensions is misleading in millions: nearly every flat point is a saddle rather than a local minimum, which makes getting permanently stuck far less of a worry than it sounds. Once you can see the landscape, the optimiser settings stop being magic numbers. Momentum, learning rate and the schedules around them are each an answer to a specific feature of that surface, and knowing which feature is causing trouble is how you rescue a run that will not converge.

Picture a landscape#

The loss is a surface over millions of dimensions, one per parameter. Training walks downhill. Intuition from two dimensions is partly wrong here, and the most useful correction is that in high dimensions most flat points are saddle points rather than local minima.

Local minimum   downhill in no direction      rare in high dimensions
Saddle point    downhill in some directions   common, and they slow training
Plateau         nearly flat in all directions small gradients, slow progress

The reason saddles dominate is a counting argument worth carrying around. At a flat point, each direction is either curving up or curving down. A local minimum requires every one of them to curve up. With two parameters that is a coin landing heads twice; with a million parameters it is a coin landing heads a million times. Somewhere in all those directions there is almost always one that still goes down, and a point with at least one downhill direction is a saddle, not a minimum.

The smallest saddle is the surface x squared minus y squared. It curves up along x and down along y, so the origin is flat but not a minimum. Start near it, slightly off the ridge, and watch the two coordinates behave completely differently: x is pulled straight to zero, while y, which started at 0.01, takes a long time to make anything of the downhill direction it is on.

import numpy as np

def saddle_run(momentum, lr=0.1, steps=200):
    p, v = np.array([1.0, 0.01]), np.zeros(2)
    for i in range(1, steps + 1):
        g = np.array([2 * p[0], -2 * p[1]])      # gradient of x^2 - y^2
        v = momentum * v + g
        p = p - lr * v
        if abs(p[1]) > 1.0:
            return i, p.round(4)
    return None, p.round(4)

print(saddle_run(0.0))     # (26, array([0.003 , 1.1448]))
print(saddle_run(0.9))     # (13, array([0.4992, 1.2975]))

Momentum#

Plain gradient descent zigzags across narrow valleys and stalls on flat regions. Momentum keeps a running average of past gradients, so consistent directions build up speed and oscillating ones cancel out. It carries the optimiser through saddle points and plateaus that would otherwise trap it.

Build the valley and count the steps. A loss of one half of x squared plus twenty y squared is a long trough: gently sloped along x, steeply walled along y. The ratio of those two curvatures, 20 to 1, is called the condition number, and it is the single number that decides how badly plain gradient descent will behave. The step size is limited by the steep direction, at 2 divided by 20, which is 0.1 — above that it diverges — but progress has to be made in the gentle direction, where a step of that size is tiny.

def valley(lr, momentum, steps=2000):
    p, v = np.array([1.0, 1.0]), np.zeros(2)
    for i in range(1, steps + 1):
        g = np.array([1.0 * p[0], 20.0 * p[1]])         # gradient of 0.5(x^2 + 20y^2)
        v = momentum * v + g
        p = p - lr * v
        if 0.5 * (p[0] ** 2 + 20 * p[1] ** 2) < 1e-4:
            return i
    return None

print(valley(0.02, 0.0), valley(0.02, 0.9))      # 211 83
print(valley(0.09, 0.0), valley(0.09, 0.9))      # 46 98

That second line is the part people are not told. Momentum is not free speed; it multiplies the effective step size by roughly 1 divided by 1 minus the momentum coefficient, which at 0.9 is ten. Turning momentum on without lowering the learning rate can take a run that was stable and make it oscillate, and the symptom looks identical to a learning rate that was simply set too high.

Learning rate is step size#

On a curved surface, too large a step overshoots the valley and bounces up the far side; too small and progress is glacial. The right step depends on local curvature, which varies across the surface and across parameters. Adaptive optimisers like Adam exist to set the step per parameter.

The condition number is why the per-parameter part matters. In the valley above, one learning rate has to serve a direction that is twenty times steeper than the other, so whatever you choose is wrong for one of them. Real networks have condition numbers far worse than 20, and the parameters are not conveniently aligned with the axes. Adam's per-parameter scaling is an attempt to give the gentle directions larger steps and the steep ones smaller, without ever computing the curvature it would need to do this properly.

The size of that expense is worth knowing rather than assuming. The curvature of a loss with n parameters is an n by n matrix, so for a model with a billion parameters it has a billion squared entries — an amount of memory that does not exist. Newton's method would need to solve a system with that matrix at every step. Every practical second-order method is therefore an approximation that never forms the matrix, and the cheapest of those approximations, a diagonal estimate, is essentially what Adam maintains.

What to do when a run stalls#

symptom                                  most likely cause         first thing to try
loss flat from step 0                    bug, not landscape        overfit one batch
loss falls then sits still for a while   saddle or plateau          wait; check gradient norms
loss oscillates, amplitude steady        learning rate near limit   halve it
loss oscillates, amplitude growing       learning rate over limit   halve it, add warmup
loss falls smoothly but very slowly      rate far below limit       raise it, or add momentum
loss fell, then jumped and stayed high   one bad batch or spike     clip gradients

What is worth memorising from this section is short. Flat does not mean finished. The stable learning rate is set by the steepest curvature in the problem, which is why a rate that worked on one model fails on another. Momentum lengthens the effective step by a factor of about ten at the usual setting. And nothing here requires you to compute a second derivative by hand, ever.