How models learn: the training loop

Chapter 1 fitted a line in one step with a formula. Most models have no such formula. They are trained by a loop that is almost embarrassingly simple, and the same loop trains a line, a spam filter and a language model with hundreds of billions of parameters. Learn it once.

The four steps#

1. Predict. Run the examples through the model with its current parameters.
2. Measure. A loss function turns the gap between predictions and true answers into one number: how wrong the model is. Lower is better.
3. Find the direction. The gradient tells you, for every parameter, whether increasing it makes the loss go up or down, and how fast (chapter 23).
4. Update. Move every parameter a small step against the gradient, so the loss goes down. The step size is the learning rate.

Then repeat, thousands or billions of times. This is gradient descent. The picture to hold is a ball rolling downhill on a landscape whose height is the loss.

The loop in code#

For a line y = w x + b with the mean squared error as the loss, the gradients have a simple closed form: the gradient of the loss with respect to w is the average of 2 × error × x, and with respect to b the average of 2 × error. The script generates points around the line y = 3x + 2, starts with w = b = 0 and runs the loop.

import numpy as np

rng = np.random.default_rng(0)
x = rng.uniform(-3, 3, 200)
y = 3.0 * x + 2.0 + rng.normal(0, 0.8, 200)          # the truth is w = 3, b = 2, plus noise

w, b = 0.0, 0.0                                       # start knowing nothing
learning_rate = 0.05

for step in range(101):
    prediction = w * x + b                            # 1. predict
    error = prediction - y
    loss = np.mean(error ** 2)                        # 2. measure
    grad_w = np.mean(2 * error * x)                   # 3. which way is downhill?
    grad_b = np.mean(2 * error)
    w -= learning_rate * grad_w                       # 4. step against the gradient
    b -= learning_rate * grad_b
    if step in (0, 1, 2, 5, 10, 25, 50, 100):
        print(f"step {step:3d}: loss {loss:8.3f}   w = {w:5.2f}   b = {b:5.2f}")

step   0: loss   36.411   w =  1.03   b =  0.26
step   1: loss   17.437   w =  1.72   b =  0.48
step   2: loss    8.867   w =  2.17   b =  0.65
step   5: loss    2.154   w =  2.78   b =  1.03
step  10: loss    1.008   w =  2.98   b =  1.41
step  25: loss    0.682   w =  2.98   b =  1.82
step  50: loss    0.667   w =  2.97   b =  1.93
step 100: loss    0.667   w =  2.97   b =  1.94

The loss collapses from 36 to about 1 in ten steps and the parameters settle near the true 3 and 2 (w = 2.97, b = 1.94: not exact, because of the noise in a finite sample). The loss bottoms out around 0.67, close to the variance of the noise we added (0.8 squared is 0.64), which is the best any line can do on this data. A real network does precisely this, only with millions or billions of parameters, and a method called backpropagation to compute all the gradients efficiently (chapter 64).

The learning rate decides everything#

The step size is the most important setting in the loop. Too small and training crawls. Too big and every step overshoots the valley, and the loss grows instead of shrinking. Try it yourself: in the explorable, start on either side, pick a learning rate, and press Step.

From a line to everything else#

Nothing in the loop depends on the model being a line. Replace the line by a neural network and the squared error by another loss, compute the gradients with backpropagation, and you have deep learning. Replace the data by the whole internet and the model by a transformer, and you have a language model. The rest of the course is about better models, better data, better losses, better ways of computing the gradient and better ways of knowing whether the result is any good.