Gradient descent

Rolling downhill

A model learns by making a number called its loss smaller. The loss measures how wrong the model's answers are, and it depends on the model's settings. Picture it as a landscape: the settings are where you stand, the loss is the height, and learning means walking down to the lowest point.

Start with the simplest landscape there is, the curve $f(x) = x^2$, whose lowest point is at $x = 0$. Press Step a few times, then try the three presets.

The small learning rate creeps down, the medium one overshoots the bottom and zig-zags in, and the large one overshoots further every time until the ball flies off. By the end of this lesson you'll be able to predict all three.

Which way is downhill?

The derivative $f'(x)$ is the slope of the curve at $x$. If the slope is positive, the curve goes up to the right, so downhill is to the left; if it's negative, downhill is to the right. Either way, downhill is the direction of $-f'(x)$.

Each step moves against the slope, and moves further where the curve is steeper.

The theorem doesn't say how small "small enough" is: that depends on how sharply the curve bends. The next step works it out exactly for $x^2$.

The learning rate

For $f(x) = x^2$ one step is $x_{k+1} = x_k - \eta \cdot 2x_k = (1 - 2\eta)\, x_k$. Every step multiplies $x$ by the same number, so after $k$ steps

The sign of $r$ decides how the ball moves:

The presets in the first step are $\eta = 0.1$ (multiplies by $0.8$), $\eta = 0.9$ (by $-0.8$) and $\eta = 1.05$ (by $-1.1$). Go back and try $\eta = 0.5$ and $\eta = 1$ on the slider.

The gradient

A real loss depends on many settings at once. For a function of two variables, the slope depends on which way you walk.

For example, $f(x, y) = x^2 + 4y^2$ has $\nabla f = (2x, \, 8y)$, so at the point $(2, 1)$ the gradient is $(4, 8)$.

The directions where the height doesn't change run along the contour lines of $f$, the curves of equal height. So the gradient always crosses contour lines at right angles.

Gradient descent in two dimensions

With many variables, the rule is the same with the gradient in place of the slope.

Here it is on a long, narrow valley, $f(x, y) = 0.1x^2 + 2y^2$. The rings are contour lines, the arrow is the gradient, and the dashed ring shows where the next step lands. Try the presets.

Why does the medium preset zig-zag? The gradient is $\nabla f = (0.2x, \, 4y)$, so a step treats each coordinate on its own:

Across the valley, $y$ behaves like the steep curve from the last step: it zig-zags once $\eta > 0.25$ and diverges once $\eta > 0.5$. Along the valley, $x$ shrinks by a factor of at least $0.9$ for any $\eta$ up to $0.5$, so it crawls. The steepest direction limits the learning rate, and then the gentle direction takes forever. That's why machine learning rescales its inputs so that no direction is much steeper than the others.

Learning a line

Now a real model. Four data points $(x_i, y_i)$ are $(0, 1)$, $(1, 3)$, $(2, 2)$ and $(3, 4)$. The model predicts $\hat{y} = wx + b$, and learning means finding the slope $w$ and intercept $b$ that fit best.

To walk downhill we need its gradient. Write $r_i = w x_i + b - y_i$ for the error at point $i$. By the chain rule, the derivative of $r_i^2$ is $2 r_i$ times the derivative of $r_i$, and $r_i$ changes by $x_i$ per unit of $w$ and by $1$ per unit of $b$. So

Here are the same four points and the same starting line. Press Step once to check the numbers above, then Play. The dashed lines are the errors $r_i$.

In code

Here is the whole training loop in NumPy: compute the errors, the loss and the gradient, then step. Press Run:

import numpy as np

x = np.array([0.0, 1.0, 2.0, 3.0])
y = np.array([1.0, 3.0, 2.0, 4.0])
w, b, eta = 0.0, 0.0, 0.1
losses = []

for step in range(200):
    r = w * x + b - y                 # the errors
    losses.append(np.mean(r ** 2))    # the loss L(w, b)
    dw = 2 * np.mean(r * x)           # dL/dw
    db = 2 * np.mean(r)               # dL/db
    w, b = w - eta * dw, b - eta * db

loss = np.mean((w * x + b - y) ** 2)
print(f"after 1 step: loss {losses[1]:.2f}")
print(f"after 200 steps: w = {w:.3f}, b = {b:.3f}, loss = {loss:.3f}")
show(losses, title="Loss after each step")

It prints a loss of $0.81$ after one step, as we found by hand, and $w = 0.800$, $b = 1.300$ with loss $0.450$ after 200 steps. The chart shows the loss falling fast at first, then levelling off. Try eta = 0.3: the steps overshoot and the loss explodes.

For a line, there's also an exact formula for the best fit (least squares, in Linear Algebra for ML). NumPy's np.polyfit uses it, and gradient descent found the same line:

import numpy as np

x = np.array([0.0, 1.0, 2.0, 3.0])
y = np.array([1.0, 3.0, 2.0, 4.0])
print(np.polyfit(x, y, 1))  # [slope, intercept]

Where it's used in ML

Want the full story, with convergence proofs, Newton's method and Adam? It's all in Calculus & Optimization for ML.