The training loop: loss, gradients, learning rate

Lesson 03 fitted logistic regression with a loop that computed a gradient and stepped against it. That loop is not specific to logistic regression. With a different model and a different loss it trains a convolutional network, a transformer or a diffusion model, and the only things that change are the size of the parameter vector and who computes the gradient. This lesson isolates the loop, works its arithmetic on a two-parameter model, and then runs it with three learning rates so you can see the three ways a training run behaves. By the end you will be able to look at a loss curve and say what is wrong with it.

The loop in four steps

Every training run repeats the same four steps until the loss stops falling or the budget runs out.

One iteration of the training loop. The forward pass computes predictions, the loss compares them with labels, the backward pass computes the gradient of the loss with respect to every parameter, and the update moves each parameter a small step against its gradient.

The forward pass runs the model on a batch of inputs with the current parameters and produces predictions. The loss compares predictions with labels and reduces the comparison to one number. The backward pass computes how much that number would change if each parameter were nudged, which is the gradient. The update moves each parameter by a small multiple of its gradient, in the direction that lowers the loss. That multiple is the learning rate, and it is the single most important number you will set.

The gradient by hand

Take a model with two parameters, a slope w and an offset b, predicting w * x + b, with the mean squared error as the loss. Two examples: x equal to 1 and 2, labels 1 and 2. Start at w equal to 0 and b equal to 0.

The forward pass predicts 0 for both. The errors (prediction minus label) are minus 1 and minus 2. The loss is the mean of 1 and 4, which is 2.5.

For the backward pass, write the loss out: the mean over examples of (w * x + b minus y) squared. Differentiate with respect to w: the chain rule gives 2 times the error times x, averaged. Here that is the mean of (2 times minus 1 times 1) and (2 times minus 2 times 2), which is the mean of minus 2 and minus 8, so minus 5. With respect to b the inner derivative is 1, so the gradient is 2 times the error averaged: the mean of minus 2 and minus 4, which is minus 3.

A negative gradient means that increasing the parameter decreases the loss. So the update adds. With a learning rate of 0.1, w becomes 0 plus 0.5 and b becomes 0 plus 0.3. Run the forward pass again: predictions 0.8 and 1.3, errors minus 0.2 and minus 0.7, loss the mean of 0.04 and 0.49, which is 0.265. One step cut the loss by a factor of nine. The next step would cut it less, because the gradient is now smaller, and so on down to the bottom of the bowl near w equal to 1 and b equal to 0.

“Against the gradient” is the whole idea. The gradient points uphill, towards higher loss, because that is what a derivative measures. Subtracting it walks downhill. The learning rate decides how big each stride is.

Checking a gradient with finite differences

Hand derivations have sign errors. The check is to nudge one parameter by a tiny h, measure how much the loss changed, and divide by h. It is slow, since it needs one forward pass per parameter, but it is the ground truth, and the habit of using it on a new loss or a new model pays for itself. The block below computes the analytic gradient and the finite difference side by side.

import numpy as np

# One weight, one bias, a straight-line model, mean squared error.
rng = np.random.default_rng(0)
x = rng.uniform(-1, 1, size=100)
y = 0.8 * x - 0.3 + rng.normal(0, 0.1, size=100)

def loss_and_grads(w, b):
    pred = w * x + b
    err = pred - y
    loss = np.mean(err ** 2)
    dw = np.mean(2 * err * x)      # d loss / d w
    db = np.mean(2 * err)          # d loss / d b
    return loss, dw, db

w, b = 0.0, 0.0
print("start           w b:", w, b, " loss:", round(loss_and_grads(w, b)[0], 4))
loss, dw, db = loss_and_grads(w, b)
print("gradient at start  dw:", round(dw, 4), " db:", round(db, 4))

# Nudge each parameter by a tiny amount and watch the loss move the way the gradient says.
h = 1e-4
print("finite difference dw:", round((loss_and_grads(w + h, b)[0] - loss) / h, 4))
print("finite difference db:", round((loss_and_grads(w, b + h)[0] - loss) / h, 4))

start           w b: 0.0 0.0  loss: 0.294
gradient at start  dw: -0.5386  db: 0.4613
finite difference dw: -0.5386
finite difference db: 0.4614

The two agree to the fourth decimal, and the last-digit difference on db is the finite difference’s own error from using a step of 0.0001 rather than an infinitely small one. The signs say that w should go up (towards 0.8) and b should go down (towards minus 0.3), which is correct.

The loop as code, with three learning rates

Now the loop proper. train starts both parameters at zero, runs twenty steps, and records the loss before each update. It is run three times with learning rates of 0.05, 0.5 and 1.1.

import numpy as np

rng = np.random.default_rng(0)
x = rng.uniform(-1, 1, size=100)
y = 0.8 * x - 0.3 + rng.normal(0, 0.1, size=100)

def loss_and_grads(w, b):
    err = w * x + b - y
    return np.mean(err ** 2), np.mean(2 * err * x), np.mean(2 * err)

def train(learning_rate, steps=20):
    w, b = 0.0, 0.0
    history = []
    for step in range(1, steps + 1):
        loss, dw, db = loss_and_grads(w, b)
        w -= learning_rate * dw
        b -= learning_rate * db
        history.append(round(float(loss), 4))
    return w, b, history

for lr in (0.05, 0.5, 1.1):
    w, b, history = train(lr)
    print(f"lr {lr:<5} final w {w:7.3f} b {b:7.3f}")
    print("   losses:", history)

lr 0.05  final w   0.404 b  -0.226
   losses: [0.294, 0.2695, 0.2476, 0.2279, 0.2101, 0.1941, 0.1795, 0.1662, 0.1541, 0.1431, 0.1329, 0.1237, 0.1151, 0.1073, 0.1001, 0.0934, 0.0873, 0.0816, 0.0763, 0.0715]
lr 0.5   final w   0.793 b  -0.307
   losses: [0.294, 0.111, 0.0507, 0.0262, 0.0162, 0.0122, 0.0105, 0.0099, 0.0096, 0.0095, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094, 0.0094]
lr 1.1   final w   2.596 b  11.598
   losses: [0.294, 0.0725, 0.0899, 0.131, 0.194, 0.2896, 0.4348, 0.6553, 0.99, 1.4983, 2.2699, 3.4414, 5.2201, 7.9205, 12.0205, 18.2453, 27.6961, 42.0448, 63.8299, 96.9052]

The chart plots those three histories on a log scale, which is the right scale for a loss because a drop from 0.3 to 0.03 matters as much as a drop from 0.03 to 0.003.

Loss per step for the three learning rates, from the printed histories above. Log scale on the loss axis.

Too slow

At 0.05 the loss falls steadily and is still falling at step 20, from 0.294 to 0.0715. The parameters have reached 0.404 and minus 0.226 on the way to 0.8 and minus 0.3. Nothing is wrong, and another hundred steps would get there. But every step costs a forward and backward pass over the data, and on a real model that is the difference between an hour and a day. A loss curve that is a smooth, shallow slope is the signature of a learning rate that could be larger.

About right

At 0.5 the loss drops to 0.0094 by step 11 and stays there. That number is the noise floor: the labels were generated with noise of standard deviation 0.1, so the mean squared error cannot go below about 0.01 for any w and b. The fitted 0.793 and minus 0.307 are as close to 0.8 and minus 0.3 as a hundred noisy points allow. A curve that drops fast and then flattens is what a healthy run looks like, and when it flattens you can stop.

Diverging

At 1.1 the first step is the best of the three, dropping the loss to 0.0725. Then it climbs, slowly at first and then by a factor of about 1.5 per step, reaching 97 by step 20 with b at 11.6 and heading nowhere useful. The step overshot the bottom of the bowl and landed on the far side, higher than where it started, and each subsequent step overshoots further.

The threshold is not arbitrary. For a squared-error bowl, the gradient along a direction grows in proportion to the distance from the bottom, with a constant that is the bowl’s curvature in that direction. A step of size lr times that gradient lands at (1 minus lr times curvature) times the original distance. If that factor is between minus 1 and 1 the distance shrinks. If it is below minus 1 the distance grows, oscillating in sign. Here the curvature along b is 2 (the second derivative of the mean of err squared with respect to b), so the factor is 1 minus 2 times 1.1, which is minus 1.2, and the error in b grows by 20 % per step while flipping sign. Any learning rate below 1.0 would converge along b. Along w the curvature is 2 times the mean of x squared, about 0.67, so w is stable for learning rates below 3 and was never the problem.

A loss that goes up and keeps going up, or becomes nan, is a learning rate that is too large. The fix is to divide it by ten and try again. Nothing else in the model needs to change.

Batches, steps and epochs

The loop above used all 100 examples for every gradient. That is full-batch gradient descent, and it is exact but expensive when there are millions of examples. The usual alternative computes the gradient on a batch of, say, 32 or 256 examples at a time, which is a noisy estimate of the true gradient but far cheaper, so you can take many more steps in the same time. That is stochastic gradient descent, and the noise turns out to help more than it hurts, by shaking the parameters out of narrow pits in the loss surface.

The vocabulary follows from this. A step (or iteration) is one pass through the four boxes on one batch. An epoch is one pass through the whole training set, so with 10,000 examples and batches of 100 an epoch is 100 steps. Training runs are described in epochs, learning rates are often adjusted per epoch, and the loss curves in papers are usually plotted per step or per epoch. Lesson 07 writes a loop with explicit batches.

What changes when the model gets big

The loop stays the same. What changes is how the gradient is computed and how the step is taken.

For a two-parameter line, the gradient fits on one line of code. For a network with ten million parameters there is no hand derivation, so frameworks compute it automatically: you write the forward pass, and reverse-mode automatic differentiation (backpropagation) produces the gradient of the loss with respect to every parameter at about twice the cost of the forward pass. The PyTorch training loop reference shows the same four steps as loss.backward() and optimizer.step().

The update rule also gets refinements. Momentum keeps a running average of recent gradients so that consistent directions accelerate and oscillating ones cancel. Adam, the optimiser used by most deep learning work since 2015, additionally scales each parameter’s step by the recent size of its gradient, so parameters with tiny gradients still move. Learning-rate schedules start high and decay, or warm up from zero, to get the fast early progress of a large rate and the precision of a small one. All of these are extras on the loop above, and all of them still have a learning rate that can be too large or too small. The deep learning hub covers the optimiser family.

One more difference matters in practice. The squared-error bowl on this page has one bottom. The loss surface of a neural network has many valleys, ridges and flat plateaus, and gradient descent finds a valley, not necessarily the lowest one. That it finds good valleys at all, reliably, from random starting points, is one of the empirical facts the field is built on and one of the least well understood.

Exercise

Change the learning rate to 0.95 and run twenty steps. The loss should oscillate but still fall. Explain why using the stability factor from the diverging section.
Replace the full-batch gradient with a gradient computed on a random batch of 10 examples per step (rng.choice(100, size=10, replace=False)), run 200 steps at learning rate 0.5 and plot or print the loss on all 100 examples every 20 steps. Describe the difference from the full-batch curve.
Add a third parameter c so the model is w * x + c * x ** 2 + b, derive its gradient, add it to loss_and_grads and confirm with a finite difference. Then train and report what value c settles near, given that the data was generated with no squared term.

Summary

Training is a loop of four steps: forward pass, loss, backward pass (gradient), update. Every model from a line to a language model runs it.
The gradient points uphill in the loss, so the update subtracts it. A finite-difference check catches sign errors in a hand-derived gradient.
The learning rate scales each step. Too small is slow but safe. Too large overshoots, and for a bowl the step diverges once the learning rate exceeds 2 divided by the curvature.
Read the curve: shallow and still falling means raise the rate, fast drop then flat means stop when flat, rising or nan means lower the rate by ten.
Batches make the gradient cheap and noisy, an epoch is one pass through the data, and frameworks add automatic differentiation and optimisers such as Adam without changing the loop.