Gradient descent: learning by small steps

In the last lesson a formula handed us the best line in one go. But a neural network with millions of weights has no such formula. How can a computer find good values for millions of knobs when nobody can tell it where the best setting is? The answer is almost embarrassingly simple: it feels which way is downhill and takes a small step. Then it does it again.

Try this

• Choose a small learning rate and press Step a few times. Watch the line move a little at each step and the loss curve slowly go down.
• Press Run and let it go. Does the line end up close to the least squares line from the previous lesson?
• Start over with the largest learning rate and press Step again. Watch the loss curve: instead of falling, it climbs, and the line swings wildly. That is divergence.
• Try the learning rates in between and find the largest one that still makes the loss fall steadily.

Now that you've tried it, let's see what was actually going on.

The idea in plain words

Imagine standing on a hillside in thick fog. You want to reach the valley, but you can see only a metre around you. A sensible strategy: feel the ground with your feet, find the direction in which it slopes down most steeply, take a step that way, and repeat. You never see the whole landscape, yet step by step you get lower.

In machine learning the "landscape" is the error of the model. Every possible pair of slope and intercept is a position on the map, and the height at that position is the mean squared error from lesson 3. The valley floor is the best line. Gradient descent is the foggy walk: start somewhere, measure which way the error decreases, move the parameters a little in that direction, and repeat until the error stops falling.

Why it works

The direction of steepest ascent of a function is given by its gradient — a list of derivatives, one for each parameter, saying how much the error would grow if that parameter were nudged up. Go the opposite way and the error goes down. For each parameter w one step is:

w ← w − η · ∂L/∂w

Here L is the loss (the error we minimise), ∂L/∂w is how strongly the loss reacts to w, and η (the Greek letter eta) is the learning rate, which sets the size of the step. Where the slope is steep, steps are big; near the bottom the slope flattens and the steps shrink by themselves.

Everything depends on the learning rate. If it is too small, each step is tiny and the walk takes forever. If it is just right, the error drops quickly and settles near the minimum. If it is too large, each step overshoots the valley and lands higher on the opposite slope; the next step overshoots even more, and the error grows without limit. This is called divergence, and it is the most common reason why a training run "explodes".

The beauty of the method is that it needs only one thing: the ability to compute the gradient at the current position. It does not need a formula for the answer, and it works the same for two parameters as for two billion. That is why it is the engine behind almost every neural network, including the language models of lesson 12. In practice the gradient is usually estimated on a small random batch of examples rather than on all of them, which makes each step cheaper and noisier but works remarkably well.

Common mistake

When the loss suddenly shoots up or turns into "NaN" (not a number), beginners often suspect a bug in their code. Much more often the learning rate is simply too large. The cure is to lower it, not to rewrite the model. The opposite mistake is less visible: a learning rate so small that the loss is still falling when training stops, so the model looks worse than it could be. A related trap is features on very different scales — say, age in years next to income in euros — which turns the valley into a long, narrow canyon where no single learning rate suits every direction. Scaling the features first makes the walk much easier.

We now have both halves of the learning recipe: a score and a way to improve it. In the next lesson we use them for a new kind of question — not "how much?" but "which one?".

import numpy as np
from sklearn.datasets import load_diabetes

X, y = load_diabetes(return_X_y=True, as_frame=True, scaled=False)
bmi = X["bmi"].to_numpy()
z = (bmi - bmi.mean()) / bmi.std()  # standardised BMI, as in the lab
mse = lambda w, b: np.mean((y - (b + w * z)) ** 2)
for lr in [0.01, 0.1, 0.8, 1.1]:
    w, b = 0.0, 0.0  # start at the origin
    for _ in range(5):  # five full-batch steps
        r = y - (b + w * z)
        w, b = w + lr * 2 * np.mean(r * z), b + lr * 2 * np.mean(r)
    print(f"lr {lr}: loss after 5 steps {mse(w, b):.0f}")
w_best, b_best = np.polyfit(z, y, 1)
print(f"minimum: {mse(w_best, b_best):.0f}")

lr 0.01: loss after 5 steps 24468
lr 0.1: loss after 5 steps 6595
lr 0.8: loss after 5 steps 4043
lr 1.1: loss after 5 steps 159823
minimum: 3890