Backpropagation by hand

Introduction

In the previous posts, we covered why taking the derivative is linked to optimisation of a machine learning model and how gradient descent is performed with a simple model. This post focuses on the core algorithm used for training neural networks, including large language models: backpropagation.

It is simple to take the derivative of the loss with respect to the weight, when a model consists of a single parameter e.g., \( \hat{y} = wx \). What if the model has multiple layers of weights and they are dependent on each other? \[ h = w_1 x \] \[ \hat{y} = w_2 h \] With this model, we need to compute an internal representation \( h \), and then produce the model prediction \( \hat{y} \) with multiplication of the hidden \( h \) by the second weight \( w_2 \). Visualisation of this model is provided in the figure below.

How do we compute the derivative of the loss with respect to \( w_1 \), which sits in the first layer of the model? The answer is applying the chain rule of the derivative all the way from the loss to the first weight. That is the high level description of backpropagation.

Backpropagation

Suppose we have a model with an input layer, a hidden layer and an output layer, the model can be expressed as: \[ h^{(i)} = w_1 x^{(i)} \] \[ \hat{y}^{(i)} = w_2 h^{(i)} \] where \(x^{(i)}\) is the i-th input value, \(h^{(i)}\) is the hidden representation for the i-th input and \(\hat{y}^{(i)}\) is the model prediction for the input [1]. 

We then define the Mean Squared Error as a loss function: \[ L = \frac{1}{N} \sum^N_{i=1} (\hat{y}^{(i)} - y^{(i)})^2 \] The relationships of the two weights in the model are: \( w_2 \rightarrow \hat{y} \rightarrow L \) and \( w_1 \rightarrow h \rightarrow \hat{y} \rightarrow L \). Let's first work on the partial derivative of the loss with respect to \( w_2 \).

Partial derivative of the loss w.r.t. \(w_2\)

For easier application of the chain rule, define a variable \( u^{(i)} \) for the difference between the predicted value and the target value for the i-th input \[ u^{(i)} = \hat{y}^{(i)} - y^{(i)} \] The partial derivative of the loss for a single input value w.r.t. \( w_2 \) is as follows: \[ \frac{\partial L^{(i)}}{\partial w_2} = \frac{\partial L^{(i)}}{\partial u^{(i)}} \frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}} \frac{\partial \hat{y}^{(i)}} {\partial w_2} \] Then, we resolve each partial derivative term. \[ \frac{\partial L^{(i)}}{\partial u^{(i)}} = 2u^{(i)} = 2(\hat{y}^{(i)} - y^{(i)}) \] \[ \frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}} = 1 \] \[ \frac{\partial \hat{y}^{(i)}} {\partial w_2} = h^{(i)} \] Multiplying the terms altogether gives us the partial derivative of the loss for an i-th input value with respect to \(w_2\): \[ \frac{\partial L^{(i)}}{\partial w_2} = 2h^{(i)}(\hat{y}^{(i)} - y^{(i)}) \] Finally, averaging over all \(N\) input gives the partial derivative of the total loss: \[ \frac{\partial L}{\partial w_2} = \frac{2}{N} \sum^N_{i=1} h^{(i)}(\hat{y}^{(i)} - y^{(i)}) \]

Partial derivative of the loss w.r.t. \(w_1\)

The partial derivative of the loss w.r.t. \( w_1 \) is computed in a similar manner with one extra connection between the weight and the hidden variable \( w_1 \rightarrow h \). \[ \frac{\partial L^{(i)}}{\partial w_1} = \frac{\partial L^{(i)}}{\partial u^{(i)}} \frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}} \frac{\partial \hat{y}^{(i)}} {\partial h^{(i)}} \frac{\partial h^{(i)}}{\partial w_1} \] We have already solved \(\frac{\partial L^{(i)}}{\partial u^{(i)}}\) and \(\frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}}\) in the previous section. We only need to work on the other two terms: \[ \frac{\partial \hat{y}^{(i)}}{\partial h^{(i)}} = w_2 \] \[ \frac{\partial h^{(i)}}{\partial w_1} = x^{(i)} \] Multiplying all of the terms together gives us: \[ \frac{\partial L^{(i)}}{\partial w_1} =  2w_2x^{(i)}(\hat{y}^{(i)} - y^{(i)}) \] and the partial derivative for the total loss is: \[ \frac{\partial L}{\partial w_1} = \frac{2w_2}{N} \sum^N_{i=1} x^{(i)}(\hat{y}^{(i)} - y^{(i)}) \]

Weight updates

We now have the partial derivatives of the loss with respect to \(w_1\) and \(w_2\): \[ \frac{\partial L}{\partial w_2} = \frac{2}{N} \sum^N_{i=1} h^{(i)}(\hat{y}^{(i)} - y^{(i)}) \] \[ \frac{\partial L}{\partial w_1} = \frac{2w_2}{N} \sum^N_{i=1} x^{(i)}(\hat{y}^{(i)} - y^{(i)}) \] The learning rate \( \tau \) controls how large a step to update the weights at one time. A larger value can speed up training, but it can cause instability of the process. The parameter updates can be expressed as: \[ w_2 \leftarrow w_2 - \tau \frac{\partial L}{\partial w_2} \] \[ w_1 \leftarrow w_1 - \tau \frac{\partial L}{\partial w_1} \] The next section shows that the same principle can be applied to a bigger model.

Backpropagation with More Weights

The figure below shows a model bigger than the previous one. Here, the input is a two dimensional vector \( \mathbf{x} = ( x_1, x_2 ) \).

The whole model can be expressed as below: \[ h_1 = w_1 x_1 + w_2 x_2 \] \[ h_2 = w_3 x_1 + w_4 x_2 \] \[ \hat{y} = w_5 h_1 + w_6 h_2 \] Let's use the Mean Squared Error and the intermediate variable \(u\) again. The partial derivative of the loss with respect to \( w_1 \) for a single train example can be computed by the same chain rule of differentiation:

\[ \frac{\partial L^{(i)}}{\partial w_1} = \frac{\partial L^{(i)}}{\partial u^{(i)}} \frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}} \frac{\partial \hat{y}^{(i)}}{\partial h_1^{(i)}} \frac{\partial h_1^{(i)}}{\partial w_1} \] \[ \frac{\partial L^{(i)}}{\partial u^{(i)}} = 2u^{(i)} = 2(\hat{y}^{(i)} - y^{(i)}) \] \[ \frac{\partial u^{(i)}}{\partial \hat{y}^{(i)}} = 1 \] \[ \frac{\partial \hat{y}^{(i)}}{\partial h_1^{(i)}} = w_5 \] \[ \frac{\partial h_1^{(i)}}{\partial w_1} = x^{(i)}_1 \] Multiplying all of the terms gives us: \[ \frac{\partial L^{(i)}}{\partial w_1} = 2w_5x^{(i)}_1(\hat{y}^{(i)} - y^{(i)}) \] The partial derivative of the total loss with respect to \(w_1\) is therefore: \[ \frac{\partial L}{\partial w_1} = \frac{2w_5}{N} \sum^N_{i=1} x^{(i)}_1(\hat{y}^{(i)} - y^{(i)}) \] So, even with more weights, the procedure to compute the partial derivative is still similar.

Backpropagation by Hand

Let's calculate one iteration of backpropagation by hand. Using the model with two weights from the first section, define initial weights and a pair of input and target values as follows: \( x = 2, y= 4, w_1 = 2, w_2 = 3, N = 1 \). Substituting these values in the model gives us the Mean Squared Error: \[ h = w_1x = 2 \times 2 = 4 \] \[ \hat{y} = w_2h = 3 \times 4 = 12 \] \[ L = \frac{1}{N} \sum_{i=1}^N (\hat{y}-y)^2 = 8^2 = 64 \] So far, so good. We already know from the first section how to compute the partial derivatives of the loss with respect to the two weights. Substituting the known values into the formula gives us the following partial derivatives: \[ \frac{\partial L}{\partial w_2} = \frac{2}{N} \sum^N_{i=1} h(\hat{y} - y) = 2 \times 4 (12 - 4) = 64 \] \[ \frac{\partial L}{\partial w_1} = \frac{2w_2}{N} \sum^N_{i=1} x(\hat{y} - y) = 2 \times 3 \times 2 (12-4) = 96 \] We can check if this is correct with PyTorch.

import torch

x = torch.tensor([2.0], requires_grad=True)
y = torch.tensor([4.0], requires_grad=False) # no need of gradient with respect to the target
w_1 = torch.tensor([2.0], requires_grad=True)
w_2 = torch.tensor([3.0], requires_grad=True)

h = w_1 * x
y_hat = w_2 * h

loss = (y_hat - y) ** 2
print(f"loss: {loss}")
loss.backward()

print(f"w_2.grad: {w_2.grad}")
print(f"w_1.grad: {w_1.grad}")

Output loss: tensor([64.], grad_fn=<PowBackward0>)
w_2.grad: tensor([64.])
w_1.grad: tensor([96.])

This confirms that our computation of the partial variables of the loss with respect to \( w_1 \) and \( w_2 \) was correct :)

Summary

Backpropagation is the algorithm that is still fundamental to train neural network models, including large language models. Backpropagation uses the chain rule of derivatives to propagate model prediction errors all the way from top to early layers for weight updates. In this post, we covered two examples of backpropagation, one with a small model with two weights and the other with a large model with six weights.

[1] In my previous posts, I used the subscript e.g., \( x_i \) to denote an i-th data example. The subscript \(x_i\) is reserverd for for i-th value of a vector e.g., \(x_i\) in this post.