There is a line the learning rate cannot cross

The learning rate is usually taught like this: too small is slow, too large diverges. True, and useless, because it never says how large is large.

With two parameters the boundary can be computed exactly. Compute it once and the same reason explains why the learning rate cannot be raised on a model with a hundred million parameters.

Building a loss surface

Four points, one straight line. The parameters are the slope w and the intercept b.

import numpy as np

x = np.array([1., 2., 3., 4.])
y = np.array([2., 4., 5., 8.])

def loss(w, b):
    return np.mean((w * x + b - y) ** 2)

def grad(w, b):
    e = w * x + b - y
    return 2 * np.mean(e * x), 2 * np.mean(e)

The answer is known in advance: least squares gives w = 1.9, b = 0, with loss 0.175. The data carries noise, so the loss does not reach zero.

Two parameters make the loss a surface. This one is exactly quadratic, so the contours are ellipses - not round ones, but stretched more than sevenfold in one direction. That elongation decides everything in the next two sections.

One step

Gradient descent is one line. Step against the gradient, by the learning rate.

w = b = 0.0
lr = 0.05
for step in range(5):
    gw, gb = grad(w, b)
    w, b = w - lr * gw, b - lr * gb
    print(f"step {step+1}: w={w:.4f} b={b:.4f} loss={loss(w, b):.4f}")

step 1: w=1.4250 b=0.4750 loss=0.9647
step 2: w=1.6625 b=0.5463 loss=0.2478
step 3: w=1.7041 b=0.5510 loss=0.2267
step 4: w=1.7133 b=0.5449 loss=0.2247
step 5: w=1.7171 b=0.5371 loss=0.2232

The first step dwarfs the rest: the loss falls from 27.25 to 0.96. After that it looks stopped - five steps move it from 0.2478 to 0.2232, about one percent.

It has not stopped. Two hundred steps from the start reach w=1.8904, b=0.0284, loss 0.1751, which is essentially the minimum. Across the valley it is fast, along the valley it is slow. The contours being long ellipses rather than circles shows up directly in the walk.

The boundary is a calculation

Raise the learning rate.

lr=0.11   loss after 200 steps 0.175
lr=0.119  loss after 200 steps 0.343
lr=0.121  loss after 200 steps 100803
lr=0.13   loss after 200 steps 7.3e+28

Something snaps between 0.119 and 0.121. That something is the curvature of the loss.

For a quadratic loss the curvature is one matrix.

H = 2 * np.array([[np.mean(x*x), np.mean(x)],
                  [np.mean(x),   1.0        ]])
print(H)                        # [[15.  5.] [ 5.  2.]]
print(np.linalg.eigvalsh(H))    # [ 0.2994 16.7006]
print(2 / 16.7006)              # 0.11976

Two eigenvalues. The larger, 16.70, is the curvature of the steepest direction; the smaller, 0.30, the gentle one along the valley. On this surface gradient descent converges under exactly one condition.

lr < 2 / largest eigenvalue = 2 / 16.7006 = 0.1198

That is why 0.119 survived and 0.121 exploded. Not a feel for it - a line, and it holds to three decimal places.

Why 2/lambda is clear if you isolate one direction. If the loss along it is lambda/2 * d^2, the gradient is lambda*d, and one step leaves the distance at d(1 - lr*lambda). Shrinking needs |1 - lr*lambda| < 1, that is lr < 2/lambda. Above that, every step lands further out on the other side. Trace lr=0.13 and that is exactly the shape: (0, 0) jumps to (3.7, 1.24), returns to (-0.62, -0.26), then goes out to (4.46, 1.44). The straight climb of that curve in the figure is this ricochet.

Why this is the same story at scale

With a hundred million parameters the Hessian cannot be written down - it would need the square of that many entries. But a Hessian-vector product costs two backward passes, and power iteration on top of it recovers the largest eigenvalue. You cannot see all of it; you can see the one that sets stability. The structure is unchanged.

• Stability is set by one direction, the steepest. However gentle the rest, if that one diverges everything goes with it
• Slowness comes from the other end, the gentlest direction. The learning rate has to respect the steepest, so the gentlest crawls at that rate
• The ratio of the eigenvalues, 16.7 / 0.30 = 56, is this problem’s condition number. The larger it is, the longer the ellipse and the more laborious gradient descent becomes

This is also where normalising the input becomes precise. It is often explained as “shrink the inputs and you can raise the learning rate”, which is half right. Measured on this data:

x as is             eigenvalues [0.299, 16.70]   condition 55.8   lr ceiling 0.120
x times 0.1         eigenvalues [0.024,  2.13]   condition 90.4   lr ceiling 0.940
x minus its mean    eigenvalues [2.00,   2.50]   condition  1.2   lr ceiling 0.800
centred and scaled  eigenvalues [2.00,   2.00]   condition  1.0   lr ceiling 1.000

Shrinking alone lifts the ceiling eightfold and makes the conditioning worse (55.8 to 90.4). The surface got longer, so whatever the higher ceiling buys is handed straight back.

What actually works is subtracting the mean. Centring x sends mean(x) to zero, which kills the off-diagonal term of the Hessian. The entanglement between w and b comes apart and the condition number falls from 55.8 to 1.2. Standardise as well and it is 1.0, a perfect circle. The value of normalising is in squaring up the axes, not in the scale.

Momentum and Adam attack the same problem differently: they change the step rather than the landscape. Momentum accumulates steps along directions that keep their sign, so it travels further along the gentle axis; Adam divides each axis by its recent gradient magnitude, evening out the imbalance between axes. Adam does not remove the conditioning, though - it corrects per axis, so a landscape tilted off the coordinate axes keeps its elongation.

So

• A loss surface can be drawn, and with two parameters it really can be
• The learning-rate ceiling is 2 / lambda_max, and the experiment agrees to three decimal places
• A good share of training that will not converge is not the algorithm but the elongation of the surface

The next part is the gradient itself. Above, grad was written by hand; with layers stacked that stops being possible. It follows backpropagation through a three-node graph by hand, then checks the result against numerical differences.