Gradient descent

Gradient descent is a method that improves the parameters of a model in many small steps, where each step moves every parameter a little in the direction that makes the loss smaller¹. For example, the taxi model of Chapter 1 can start from $a = 0$ and $b = 0$ and get within 0.01 of the smallest loss after 144 steps.

Trying all values needs $11^k$ tries for $k$ parameters (see Section 1.2). The cost of one step of gradient descent grows only linearly with $k$, because the step calculates one derivative for each parameter. For this reason, gradient descent can train models with billions of parameters.

2.1One parameter

Suppose that the training data has only one ride of 1 km that cost 3 dollars, and that the model is $y = a \cdot x$. The loss is $L(a) = (a - 3)^2$, and it is smallest when $a = 3$.

The derivative $L'(a)$ tells how fast the loss changes when $a$ grows. In this case, $L'(a) = 2(a - 3)$. When the derivative is negative, a larger $a$ gives a smaller loss, and when it is positive, a smaller $a$ does. For example, $L'(0) = -6$, so the loss falls when $a$ grows from 0.

Each step of gradient descent changes the parameter as follows: $$a \leftarrow a - \eta \cdot L'(a)$$ where $\eta$ is the learning rate. As an example, consider the steps when $\eta = 0.25$ and the training starts from $a = 0$:

In this case, each step halves the distance between $a$ and 3, because $a - 0.25 \cdot 2(a - 3) = a - (a - 3)/2$. The following picture shows the loss and the first three steps, where the number next to each point is its step:

The steps become shorter near the minimum, because the derivative becomes smaller.

2.2The learning rate

The learning rate sets the length of the steps. With $\eta = 0.25$, the parameter gets within 0.03 of 3 after 7 steps. With $\eta = 0.05$, each step removes only 10% of the distance, and the same accuracy needs 44 steps.

However, a learning rate that is too large makes the training fail. With $\eta = 1$, the step is $a \leftarrow 6 - a$, so $a$ jumps between 0 and 6 and never gets nearer to 3. With $\eta = 1.25$, the values of $a$ are 0, 7.5, −3.75 and 13.125, and the distance to 3 grows by half at every step.

In practice, we can choose the learning rate by trial on short training runs, and the learning rate is often made smaller as the training goes on.

2.3Many parameters

The gradient of the loss is the vector of its derivatives with respect to each parameter. It points in the direction where the loss grows fastest, so gradient descent moves the parameters against it. For the taxi model $y = a \cdot x + b$ with $n$ rides, the derivatives are as follows:

$$\frac{\partial L}{\partial a} = \frac{2}{n} \sum_{i=1}^{n} (a x_i + b - y_i)\, x_i \qquad \frac{\partial L}{\partial b} = \frac{2}{n} \sum_{i=1}^{n} (a x_i + b - y_i)$$

For example, at $a = 0$ and $b = 0$, the errors of the three rides are −5, −6 and −9. In this case, $\partial L / \partial a = (2/3)(-5 \cdot 1 - 6 \cdot 2 - 9 \cdot 3) = -88/3$ and $\partial L / \partial b = (2/3)(-5 - 6 - 9) = -40/3$. The following table shows the training with $\eta = 0.05$:

The loss falls fast in the first steps and slowly after that. Note that $a$ first grows past its best value 2 and then comes back, while $b$ grows slowly to 8/3. With $\eta = 0.2$, the same training diverges.

Each step calculates the error of each ride and one derivative for each parameter, so a step takes $O(nk)$ time for $n$ examples and $k$ parameters. For the taxi model, we can write the derivatives by hand. For a neural network, Chapter 3 shows how backpropagation calculates them.

Gradient descent stops moving at a point where the gradient is 0. The loss of the taxi model has only one such point, but the loss of a neural network has many, and the result depends on the start.

2.4Batches

When the training data has millions of examples, a step over all of them is slow. Stochastic gradient descent² calculates each step from a random batch of $m$ examples instead. For example, the following table shows the gradient at $a = 0$ and $b = 0$ for batches of one ride:

Each batch gives a different gradient, but the average of the three batch gradients is exactly the gradient of all rides. Thus, the steps of stochastic gradient descent are noisy but correct on average, and each step costs only $m/n$ of a full step. In practice, language models are trained on batches of millions of tokens.

Most large models are trained with Adam³, a variant of stochastic gradient descent. Adam keeps two running averages for each parameter, one of its gradient and one of its squared gradient, and it scales each step by the recent size of the gradient. These averages and the gradient are the extra numbers per parameter that Section 11.3 counts.