Gradient descent, step by step

Every model you have heard of — linear regression, ResNet, GPT — was trained by the same procedure: guess, measure how wrong you are, and nudge the guess in the direction that reduces the wrongness. Repeat a few million times. That is gradient descent, and this lesson is about understanding it well enough to debug it.

The valley metaphor, made precise

Imagine you are standing somewhere on a hillside in thick fog. You want to reach the lowest point in the valley, but you can only see the ground immediately at your feet. The strategy that works: feel which way the ground slopes downward, take a step that way, and repeat.

Make the metaphor literal:

• Your position is the current set of model parameters — the weights and biases.
• Your altitude is the loss: a single number saying how badly the model is doing on the training data right now.
• The slope under your feet is the gradient: the derivative of the loss with respect to each parameter.
• Your step size is the learning rate.

The update rule is one line, and it is the same line in every deep learning framework ever written:

Read it aloud: the new weight is the old weight, minus the learning rate times the gradient of the loss with respect to that weight. The minus sign is the whole idea — the gradient points uphill, so we move against it.

Working one out by hand

Take the simplest possible model: ŷ = w·x, one parameter, no bias. Use squared error as the loss over n examples:

Differentiate with respect to w, using the chain rule — the outer function is "square", the inner is w·xᵢ − yᵢ:

Now put numbers in. Suppose our data is a single point, x = 2, y = 6, so the true answer is w = 3. Start at w = 0 with a learning rate of η = 0.1:

• Step 1: prediction = 0. Error = 0 − 6 = −6. Gradient = 2 × (−6) × 2 = −24. New w = 0 − 0.1 × (−24) = 2.4.
• Step 2: prediction = 4.8. Error = −1.2. Gradient = 2 × (−1.2) × 2 = −4.8. New w = 2.4 + 0.48 = 2.88.
• Step 3: prediction = 5.76. Error = −0.24. Gradient = −0.96. New w = 2.88 + 0.096 = 2.976.

2.4, 2.88, 2.976, 2.995… converging on 3. Notice the steps shrinking automatically: as the error falls, the gradient falls, so the steps fall too. Nobody had to schedule that; it is a property of the maths.

Lab: watch it converge, or not

The simulator below fits a line to real-looking data using gradient descent on two parameters, w and b. The left panel shows the data and the current fit; the right panel shows the loss surface with the path descent has taken. Change the learning rate and press run.

The three failure modes, and how to spot them

Too small a learning rate

The loss falls, but so slowly it looks flat. You cannot tell whether the model is learning badly or the rate is just tiny. Symptom: loss decreasing in the fourth decimal place after hundreds of steps. Fix: multiply η by 3 until progress becomes visible.

Slightly too large

Each step overshoots the minimum and lands on the other side of the valley, so the path zig-zags across it while creeping downhill. Symptom: loss curve that bounces up and down while trending down. Fix: halve η, or add momentum, which averages recent gradients and cancels the zig-zag.

Far too large

Each overshoot is bigger than the last. The loss climbs, then becomes inf, then NaN, and every weight in the model is poisoned. Symptom: NaN loss within the first few hundred steps. Fix: reduce η by an order of magnitude, and add gradient clipping.

Batch, stochastic, and mini-batch

The formula above sums over every training example before taking one step. With a million examples that is a million computations per step, which is unusable. Three variants exist:

• Batch gradient descent. Use all the data for each step. Accurate gradient, very slow, and it needs the whole dataset in memory.
• Stochastic gradient descent (SGD). Use one example per step. Fast and noisy — the noise actually helps, because it can knock the model out of a poor local minimum.
• Mini-batch. Use 32, 64 or 256 examples per step. Almost all real training does this. It gets a reasonably accurate gradient while keeping the GPU fully occupied.

When you read "batch size 64" in a paper, this is what it refers to. One pass through the full dataset is called an epoch.

Beyond plain descent

Vanilla gradient descent has two weaknesses: it treats every parameter with the same step size, and it has no memory. Two fixes stack on top of it:

Momentum keeps a running average of past gradients and moves along that instead. A ball rolling downhill builds speed in consistent directions and damps out the side-to-side zig-zag. Typically β = 0.9.

Adaptive rates (RMSProp, Adam) track how large each parameter's gradients have recently been, and give parameters with small gradients larger steps. Adam combines momentum with adaptive rates, and is the default optimiser almost everywhere for good reason.

The code

Full linear regression by gradient descent, twenty-five lines. Runnable in the playground.

import numpy as np

# synthetic data: y = 2.5x + 1.0 + noise
rng = np.random.default_rng(0)
X = rng.uniform(-3, 3, 200)
y = 2.5 * X + 1.0 + rng.normal(0, 0.8, 200)

w, b = 0.0, 0.0
lr = 0.05

for step in range(201):
    y_hat = w * X + b
    error = y_hat - y

    loss  = np.mean(error ** 2)          # mean squared error
    dw    = 2 * np.mean(error * X)       # dL/dw
    db    = 2 * np.mean(error)           # dL/db

    w -= lr * dw                         # the update rule
    b -= lr * db

    if step % 50 == 0:
        print(f"step {step:3d}  loss={loss:6.3f}  w={w:.3f}  b={b:.3f}")

# step   0  loss=22.871  w=0.843  b=0.152
# step  50  loss= 0.667  w=2.467  b=0.933
# step 100  loss= 0.667  w=2.466  b=0.936
# step 200  loss= 0.667  w=2.466  b=0.936
# true values were w=2.5, b=1.0 — the gap is the noise we added

Change lr to 0.5 and run it again. Then to 0.9. Watching your own code produce nan teaches the lesson better than any paragraph can.

Check yourself

1. If the gradient at your current position is +8 and η = 0.01, in which direction and how far does w move? (Down by 0.08 — against the gradient.)
2. Why does the loss stop changing once you reach a minimum, even though you keep running updates? (The gradient is zero there, so the update term is zero.)
3. Your loss goes 4.2 → 3.9 → 5.1 → 2.8 → 6.0. What is going on and what would you change? (Oscillation from too high a learning rate — halve it.)
4. Why does mini-batch descent produce a noisier loss curve than batch descent, and why is that not necessarily bad? (Each batch estimates the true gradient imperfectly; the noise aids escape from sharp poor minima.)