Backpropagation

Backpropagation is the calculation that lets a neural network, the type of AI model behind chatbots and photo recognition, learn from its mistakes. When the model gets a training example wrong, backpropagation works backwards through it and finds out how much each internal setting contributed to the error, so each setting can be nudged the right way. It runs whenever a model is trained, including when a company fine-tunes one.

How does it work?

Training repeats four steps, for each example or each small batch of examples:

1. Forward pass: the input runs through the network layer by layer and produces a prediction.
2. Loss: a loss function scores how far the prediction is from the right answer.
3. Backward pass: backpropagation goes from the output back to the input and computes the gradient, which says for every weight how much the loss changes when that weight changes a little.
4. Update: an optimiser such as gradient descent moves each weight a small step against its gradient, scaled by the learning rate.

The backward pass rests on the chain rule from calculus. A network is a chain of functions, one per layer, so the effect of an early weight on the loss is a product of the effects along the way. Backpropagation stores the intermediate results and reuses them, which keeps the calculation cheap. The textbook Deep Learning by Goodfellow, Bengio and Courville notes that, strictly, the term covers only this gradient calculation, and that another algorithm, such as stochastic gradient descent, does the update.

In PyTorch, the autograd engine records every operation of the forward pass as a graph. Calling backward() on the loss walks that graph back and stores a gradient in the .grad attribute of every parameter:

1
prediction = model(data)            # forward pass
2
loss = loss_fn(prediction, labels)  # how far off is it?
3
optimizer.zero_grad()               # gradients add up, so clear the previous ones
4
loss.backward()                     # backpropagation: a gradient for every weight
5
optimizer.step()                    # gradient descent: adjust the weights

Where does it go wrong?

Backpropagation makes a training step more expensive than a prediction. For a transformer, the backward pass takes about twice the computing of the forward pass, so a training step costs about three times as much (Kaplan et al., 2020). It also needs the intermediate results of the forward pass, and those take memory. Gradient checkpointing keeps only some of them and recomputes the rest: Chen et al. (2016) cut the memory for a 1,000-layer network from 48 GB to 7 GB at the cost of one extra forward pass.

Because the gradient of an early layer is a product of many factors, it can shrink towards zero or grow out of control: the vanishing and exploding gradient problems. Pascanu et al. (2013) analysed them for recurrent networks and proposed clipping gradients that grow too large. According to He et al. (2015), careful initialisation and normalisation layers had largely solved the problem for networks with tens of layers, and their residual connections made networks of up to 152 layers easier to train.

Where does it come from?

According to the same textbook, the chain rule dates from the 17th century, and efficient ways to apply it appeared in control engineering in the 1960s and 1970s. Paul Werbos proposed using them for neural networks in 1981. David Rumelhart, Geoffrey Hinton and Ronald Williams were among those who rediscovered the method in the mid-1980s. Their 1986 paper in Nature showed that hidden units, the parts of a network between input and output, learn by themselves to represent important features of the task.

Example

In a 3Blue1Brown lesson, a small network learns to recognise handwritten digits. It has 784 inputs, one per pixel, two hidden layers of 16 neurons and 10 outputs: 13,002 weights and biases in total. Before training, its outputs for an image of a 2 are effectively random. Backpropagation works out how to nudge each of the 13,002 numbers so that the output for 2 goes up and the other nine go down, with the largest nudges where an output is furthest off. Gradient descent applies the average of these nudges over many examples.