Gradient descent

Gradient descent is the method that trains most AI models. The model makes a prediction, measures how far off it is, and adjusts each of its internal settings a little in the direction that makes the error smaller, over and over. Picture walking down a hill in fog by always stepping where the ground slopes down most. It also runs when a company fine-tunes a model on its own examples.

How does it work?

Take a model with one setting and plot its error against that setting. The slope of the curve shows which way is downhill, so you move the setting a small step that way and look at the slope again. The textbook Deep Learning by Goodfellow, Bengio and Courville traces the technique to Cauchy in 1847. A real model has many settings, called weights. The gradient holds one slope per weight and points uphill, so each step goes the other way: weight = weight - learning_rate * gradient. The learning rate sets the size of the step.

The error is measured by a loss function, and backpropagation computes the gradient. For the digit-recognition network in 3Blue1Brown's lesson, the gradient is a list of 13,002 nudges, one per weight and bias. For Meta's Llama 3 405B it has 405 billion entries, one per parameter (Llama 3 paper). Training starts from random weights and repeats the step until the loss stops falling, which Google's Machine Learning Crash Course calls convergence.

Batch, stochastic or mini-batch?

The variants differ in how many examples each step looks at. Chapter 8 of the same textbook calls methods that use the whole training set per step batch gradient descent, and methods that use one example at a time stochastic or online. Deep learning mostly sits in between, with a mini-batch of typically 32 to 256 examples (batch size), and this is now simply called stochastic gradient descent (SGD). According to the book, nearly all of deep learning is powered by it.

Other optimisers change the step itself. Momentum adds a decaying average of earlier gradients, which helps in a long, narrow valley where plain gradient descent zigzags. Adam (Kingma and Ba, 2014) gives each weight its own step size, based on running averages of its gradients and their squares. AdamW (Loshchilov and Hutter, 2017) changes how Adam applies weight decay, which shrinks the weights a little at every step. Meta trained Llama 2 with AdamW (Touvron et al., 2023).

Where does it go wrong?

• The learning rate is off. Too large, and gradient descent can increase the error instead of reducing it; too small, and training is slower and can get stuck at a high error (Deep Learning, chapter 11).
• It lands in a valley that is not the lowest. Where you end up depends on where you start, as 3Blue1Brown shows. Linear models are safe from this: their loss surface is convex, one single bowl, so gradient descent finds the lowest point. For large neural networks, the textbook writes that experts now suspect most local minima have a low loss, so finding the very lowest point matters little.
• It slows down on flat ground. In many high-dimensional problems, saddle points, where the surface curves up in one direction and down in another, are far more common than local minima. Dauphin et al. (2014) argue that the plateaus around saddle points are the bigger obstacle, because they slow learning and look like a minimum.

Example

Google's crash course predicts a car's fuel efficiency in miles per gallon from its weight, with seven cars and the mean squared error as the loss. The model starts with its weight and bias at 0, and the loss is 303.71. The slopes are -119.7 for the weight and -34.3 for the bias. With a learning rate of 0.01, the weight becomes 1.20 and the bias 0.34, and the loss drops to 170.84. By the sixth iteration it is 42.17. This loop reproduces those numbers:

1
x = [3.5, 3.69, 3.44, 3.43, 4.34, 4.42, 2.37]  # car weight in 1,000 pounds
2
y = [18, 15, 18, 16, 15, 14, 24]                # miles per gallon
3
w, b, lr = 0.0, 0.0, 0.01                       # model weight, bias, learning rate
4
for step in range(6):
5
    errors = [(w * xi + b) - yi for xi, yi in zip(x, y)]
6
    loss = sum(e * e for e in errors) / len(x)  # mean squared error
7
    grad_w = sum(2 * e * xi for e, xi in zip(errors, x)) / len(x)
8
    grad_b = sum(2 * e for e in errors) / len(x)
9
    print(step + 1, round(w, 2), round(b, 2), round(loss, 2))
10
    w, b = w - lr * grad_w, b - lr * grad_b     # the gradient descent step