Gradient descent: how the network learns

Training a model means choosing parameters that minimize a loss function. Gradient descent does this in steps: it computes the derivative of the loss with respect to the parameter and moves a little in the opposite direction, w = w - rate * gradient. A neural network does the same with millions of weights.

One weight, step by step

The model is y = w·x and the loss is the mean squared error, MSE = (1/n)·Σ(w·x − y)². The derivative with respect to w is 2/n · Σ x·(w·x − y). The data follow y = 2x, so the right weight is 2.

xs = [1, 2, 3]
ys = [2, 4, 6]

def gradient(w):
    n = len(xs)
    return 2 / n * sum(x * (w * x - y) for x, y in zip(xs, ys))

w = 0.0
rate = 0.1
for step in range(1, 6):
    w = w - rate * gradient(w)
    print(step, round(w, 3))

1 1.867
2 1.991
3 1.999
4 2.0
5 2.0

The step shrinks on its own: near the minimum the gradient is small, so w approaches 2 without any manual adjustment.

Pitfall

A learning rate that is too large makes the step overshoot the minimum, and the error grows on every iteration instead of shrinking. In this problem, the gradient is 9.33·(w − 2). Each step multiplies the distance to the minimum by 1 − 9.33·rate, and the method only converges when that factor has absolute value less than 1, that is, a rate below about 0.21.

xs = [1, 2, 3]
ys = [2, 4, 6]

def gradient(w):
    return 2 / len(xs) * sum(x * (w * x - y) for x, y in zip(xs, ys))

for rate in (0.1, 0.3):
    w = 0.0
    for _ in range(5):
        w = w - rate * gradient(w)
    print(rate, round(w, 2))

0.1 2.0
0.3 39.79

With 0.3, after 5 steps w is at 39.79, far from 2, and it keeps moving away. If the loss goes up or becomes inf/nan, reduce the rate. Another caution: with features on very different scales, the rate that is safe for one is too large for another, so standardize the features before training.