How training actually works

A fresh network has random weights and produces nonsense. This section is about the loop that changes that: run a batch forwards, measure how wrong the answer was, push that error backwards to work out each weight's share of the blame, and nudge every one of them. The optimiser and the learning rate decide how big those nudges are, and reading a loss curve is how you tell whether any of it is working.

The loop#

for epoch in range(epochs):
    for batch in loader:
        optimizer.zero_grad()          # clear last batch's gradients
        output = model(batch.x)        # forward
        loss = criterion(output, batch.y)
        loss.backward()                # backward: compute gradients
        optimizer.step()               # update the weights

Here is the whole thing as something you can run, on data small enough to finish instantly. Two hundred steps of full-batch gradient descent on a two-layer network fitting a linear function with noise added.

import torch
import torch.nn as nn

torch.manual_seed(0)
X = torch.randn(256, 10)
true_w = torch.randn(10, 1)
y = X @ true_w + 0.1 * torch.randn(256, 1)      # a linear rule plus noise

def train(lr, steps=200, zero_grad=True, seed=0):
    torch.manual_seed(seed)
    model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    losses = []
    for step in range(steps):
        if zero_grad:
            opt.zero_grad()
        loss = ((model(X) - y) ** 2).mean()
        loss.backward()
        opt.step()
        losses.append(loss.item())
    return losses

losses = train(0.03)
print([round(losses[i], 4) for i in (0, 5, 20, 50, 199)])
# [5.8346, 4.2446, 0.4697, 0.0706, 0.0402]

Watch the accumulation happen, because the description is vaguer than the numbers. Printing the norm of all gradients before each step, with zeroing removed, gives 3.44, 6.73, 9.81, 12.71, 15.44 — climbing by roughly the size of one step's gradient every step, because that is literally what is being added.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.SGD(model.parameters(), lr=0.03)

for step in range(1, 6):
    loss = ((model(X) - y) ** 2).mean()
    loss.backward()                                   # no zero_grad anywhere
    norm = torch.cat([p.grad.flatten() for p in model.parameters()]).norm()
    print(step, round(loss.item(), 4), round(norm.item(), 4))
    opt.step()

# 1 5.8346 3.4435
# 2 5.4877 6.7322
# 3 4.8473 9.8099
# 4 3.9718 12.7096
# 5 2.9005 15.4357

print([round(v, 4) for v in train(0.03, zero_grad=False)[::50]])
# [5.8346, 1.2349, 1.0015, 0.7041]   against [5.8346, 0.0706, 0.0523, 0.0447]

Backpropagation#

After the forward pass you have a loss. Backpropagation works backwards through the network applying the chain rule, computing for each parameter how much it contributed to that loss. The framework does this automatically by recording the operations performed.

The smallest network that shows the whole mechanism has one input, one hidden unit and one output, which is three weights and no matrices. Set the input to 2, the two weights to 0.5 and -1.5, the bias on the hidden unit to 0.1, and the target to 1. Forwards: the hidden pre-activation is 0.5 times 2 plus 0.1, which is 1.1; ReLU leaves it at 1.1; the output is -1.5 times 1.1, which is -1.65; the squared error against a target of 1 is 7.0225.

Backwards, one local derivative at a time. The loss changes with the output at a rate of 2 times the error, which is 2 times -2.65, or -5.3. The output weight's own gradient is that times the hidden value: -5.3 times 1.1 is -5.83. To reach the first weight the signal passes back through the output weight, giving -5.3 times -1.5, which is 7.95; through the ReLU, which multiplies by 1 because the unit is active; and then out to the weight by multiplying by the input, 2, giving 15.9. Every one of those numbers is a multiplication of things you already have, which is the entire reason backpropagation is cheap.

x, w1, b1, w2, target = 2.0, 0.5, 0.1, -1.5, 1.0

z = w1 * x + b1              # 1.1
h = max(0.0, z)              # 1.1   ReLU is active
out = w2 * h                 # -1.6500000000000001
loss = (out - target) ** 2   # 7.022500000000002

d_out = 2 * (out - target)           # -5.300000000000001
d_w2 = d_out * h                     # -5.830000000000001
d_h = d_out * w2                     #  7.950000000000001
d_z = d_h * (1.0 if z > 0 else 0.0)  #  7.950000000000001
d_w1 = d_z * x                       # 15.900000000000002
d_b1 = d_z * 1.0                     #  7.950000000000001

def loss_of(w):                       # numerical check on w1
    return (max(0.0, w * x + b1) * w2 - target) ** 2

print(d_w1, (loss_of(w1 + 1e-6) - loss_of(w1 - 1e-6)) / 2e-6)
# 15.900000000000002 15.900000001067838

Two details in that trace are worth keeping. The ReLU passes the gradient through untouched because the unit was active; had z been negative, that one multiplication by zero would have cut the signal to w1 and b1 dead, which is what a dead unit means. And the gradient for w1, 15.9, is three times the gradient for w2, purely because of the input value 2 and the weight -1.5 it passed through. Nothing is deciding which weights matter; the arithmetic is.

Optimisers#

• SGD: step against the gradient by the learning rate. Simple, and sensitive to the learning rate.
• SGD with momentum: accumulate a running direction, which smooths noisy gradients and moves faster along consistent slopes.
• Adam: adapts the step size per parameter. The default for most work, and usually the right first choice.
• AdamW: Adam with weight decay handled correctly. The standard for transformer training.

Learning rate#

Of all the hyperparameters, this one dominates. Too high and the loss oscillates or diverges. Too low and training crawls or settles somewhere poor. It is worth tuning before anything else.

Loss diverges or becomes NaN         learning rate far too high
Loss oscillates without improving    too high
Loss falls very slowly               too low
Loss falls then plateaus early       try a schedule

Those four rows are not a metaphor. Run the same function at four rates and every one of them appears. Note the order in the last row: the loss goes to infinity first and only then to NaN, because once a weight is infinite the next subtraction is infinity minus infinity, which is undefined. A NaN loss is therefore usually the second symptom, and the overflow that caused it happened a step earlier.

for lr in (0.001, 0.03, 0.3, 1.0):
    losses = train(lr)
    print(lr, [round(losses[i], 4) for i in (0, 5, 20, 50, 199)])

# 0.001  [5.8346, 5.7758, 5.6045, 5.2774, 3.704]    barely moving after 200 steps
# 0.03   [5.8346, 4.2446, 0.4697, 0.0706, 0.0402]   healthy
# 0.3    [5.8346, 17.385, 6.1761, 1.1906, 1.0414]   jumped to 17 first, settles 25x worse
# 1.0    [5.8346, inf,    nan,    nan,    nan]      gone

The 0.3 row deserves a second look. It did not diverge, and after 200 steps it reports a loss of 1.04, which someone reading a single number would accept. It is twenty-five times worse than the same model trained at 0.03. A too-high learning rate does not always announce itself; often it just costs you most of the model's quality quietly.

A schedule that warms up and then decays is standard for large models. Warmup avoids enormous early updates when the weights are random; decay lets it settle rather than bouncing around a minimum.

Reading a loss curve#

Training falls, validation falls, gap small     healthy
Training falls, validation rises               overfitting, stop earlier
Neither falls                                  learning rate, bug, or
                                               a model too small
Loss is NaN                                    exploding gradients or
                                               bad input data

What passing that check looks like is worth seeing, because the target is not a low loss but an absurdly low one. Eight examples, a model with several hundred parameters: it should reach a loss that is numerically zero, and if it stalls at 0.3 the problem is in your code rather than in your data.

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 1))
opt = torch.optim.Adam(model.parameters(), lr=0.01)
Xb, yb = X[:8], y[:8]                      # eight examples, nothing else

for step in range(1, 301):
    opt.zero_grad()
    loss = ((model(Xb) - yb) ** 2).mean()
    loss.backward()
    opt.step()
    if step in (1, 50, 100, 200, 300):
        print(step, round(loss.item(), 8))

# 1   12.29048252
# 50  0.06314296
# 100 0.00020716
# 200 1e-08
# 300 0.0

What to memorise, and what to look up#

Memorise the loop in the order zero, forward, loss, backward, step, because saying it out of order in an interview is noticed. Memorise that gradients accumulate and must be zeroed, that the learning rate is tuned before anything else, and the overfit-one-batch check. Everything else — the exact defaults of Adam, which scheduler a paper used, the parameters of a particular data loader — is documentation, and reaching for it in an interview is the correct behaviour rather than a gap.