Backpropagation and the Training Loop

Backpropagation is a bookkeeping algorithm for the chain rule. Nothing statistical, nothing approximate: given a composition of differentiable functions, it computes the exact derivative of the loss with respect to every parameter in one backward sweep. Work the smallest net that still has two layers, and the pattern generalizes to any depth.

x = 1.0 (input),   t = 0.0 (target)
layer 1:  z1 = w1·x + b1,   a1 = σ(z1)
layer 2:  z2 = w2·a1 + b2,  ŷ = σ(z2)
loss:     L = (ŷ − t)^2
start:    w1 = 0.6, b1 = 0.1, w2 = −0.4, b2 = 0.2

σ is the sigmoid, σ(z) = 1 / (1 + e^(−z)), and it has the famous derivative σ′(z) = σ(z)·(1 − σ(z)) — the slope comes from the output itself, which keeps the arithmetic short.

Forward pass first, because every gradient below names one of these values:

z1 = 0.6·1.0 + 0.1        = 0.7
a1 = σ(0.7)               = 0.6682
z2 = −0.4·0.6682 + 0.2    = −0.0673
ŷ  = σ(−0.0673)           = 0.4832
L  = (0.4832 − 0)^2       = 0.2335

Now backward, chain rule one link at a time. Each gradient is the upstream gradient multiplied by the local derivative:

dL/dŷ       = 2(ŷ − t)          = 0.9664
dŷ/dz2      = σ′(z2) = ŷ(1−ŷ)   = 0.2497
δ2 = dL/dz2 = 0.9664 · 0.2497   = 0.2413
dL/dw2      = δ2 · a1           = 0.1612
dL/db2      = δ2                = 0.2413
dL/da1      = δ2 · w2           = −0.0965
δ1 = dL/dz1 = −0.0965 · σ′(0.7) = −0.0214
dL/dw1      = δ1 · x            = −0.0214
dL/db1      = δ1                = −0.0214

Two structural facts to take away. Every layer’s gradient follows the same recipe: δ at that layer = upstream gradient × local derivative, and a weight’s gradient is that δ times the weight’s input. And look at the magnitudes: |δ1| = 0.0214 is roughly eleven times smaller than |δ2| = 0.2413, because backprop multiplied one more sigmoid slope (0.2217 at z1) on the way down. Two layers cost a factor of eleven. Chapter three is about what fifty layers cost.

One bookkeeping rule completes the method: when a tensor feeds several downstream paths, the chain rule sums the contributions — dL/dz = (dL/dy1)·(dy1/dz) + (dL/dy2)·(dy2/dz). Autograd does that addition for you, and in a hand derivation it is the step people forget. Every gradient above is a product along one path; deeper nets are the same products summed over many paths.

PyTorch computes exactly these numbers from the recorded computation graph:

import torch

w1 = torch.tensor(0.6, requires_grad=True)
b1 = torch.tensor(0.1, requires_grad=True)
w2 = torch.tensor(-0.4, requires_grad=True)
b2 = torch.tensor(0.2, requires_grad=True)

x, t = 1.0, 0.0
a1 = torch.sigmoid(w1 * x + b1)
y = torch.sigmoid(w2 * a1 + b2)
loss = (y - t) ** 2

loss.backward()
print(w1.grad, w2.grad, b1.grad, b2.grad)
# tensor(-0.0214) tensor(0.1612) tensor(-0.0214) tensor(0.2413)

The forward pass logged every operation on an autograd graph; backward() walks that graph in reverse, multiplying local derivatives, and deposits each result in the corresponding tensor’s .grad buffer. This is reverse-mode autodiff, and its cost property is why deep learning is affordable at all: one forward pass plus one backward pass yields the gradient of one scalar loss with respect to millions of parameters — a small constant multiple of the forward cost, not one derivative per parameter. Ask a symbolic system for dL/dw1 of a 100-million-parameter model and it computes dL/dw100,000,000 too, for the same price.

Every PyTorch training loop, from a course homework to a fine-tune of a large model, is the same four lines inside the same two loops:

model = MLP(n_in=20, n_hidden=64, n_out=3)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.CrossEntropyLoss()

for epoch in range(30):                 # one epoch = one full pass over the data
    for xb, yb in train_dl:             # one batch, e.g. xb: (32, 20), yb: (32,)
        logits = model(xb)              # forward: compute predictions, record the graph
        loss = loss_fn(logits, yb)      # scalar objective — the pairing from Output Layers and Loss Pairings
        optimizer.zero_grad()           # clear last batch gradients
        loss.backward()                 # fill every parameter .grad
        optimizer.step()                # apply gradients to the weights

The order is not stylistic. zero_grad must precede backward because .grad is an accumulating buffer: each backward() does grad += contribution, it does not overwrite. That design is deliberate — it enables gradient accumulation, where you split a batch too large for memory into micro-batches and step once at the end:

for i, (xb, yb) in enumerate(train_dl):
    loss = loss_fn(model(xb), yb) / 4   # scale each micro-batch before accumulating
    loss.backward()
    if (i + 1) % 4 == 0:               # effective batch = 4 × micro-batch
        optimizer.step()
        optimizer.zero_grad()

The same accumulation explains the classic beginner bug. Forget zero_grad() in the plain loop and every update uses the running sum of all gradients since the process started: early batches keep influencing step 500, the sums grow without bound, and the loss looks reasonable for a few dozen batches before rocketing to NaN. Newer PyTorch defaults zero_grad(set_to_none=True), which sets buffers to None instead of zeros — same training semantics, less memory, and a layer that never received a gradient now fails loudly instead of silently skipping updates.

What step() does is per-parameter arithmetic. SGD applies w ← w − lr·g to every weight, plus momentum bookkeeping if configured. Using the numbers from the first chapter: w1 = 0.6 with gradient −0.0214 and lr = 0.1 becomes 0.6 − 0.1·(−0.0214) = 0.6021. Swap SGD for Adam and the call site is identical — only the update rule behind step() changes, from a fixed learning rate to per-parameter adaptive scaling from running moment estimates.

Epochs and batches set the pace of that arithmetic. A 10,000-sample dataset at batch_size 32 delivers ceil(10000/32) = 313 updates per epoch; 30 epochs is 9,390 optimizer steps. Batch size trades noise against throughput: smaller batches give noisier gradient estimates, more steps per epoch, and a mild regularizing effect; larger batches give smoother estimates, fewer steps, and more memory per update. When comparing two runs, compare optimizer steps and wall-clock, not epochs — “it trained for 30 epochs” means nothing without the batch size attached. Log the running loss with loss.item() rather than the tensor itself, and evaluate a held-out loader under torch.no_grad() so the curve you tune against is out-of-sample.

Two more ordering traps worth memorizing: calling optimizer.step() before loss.backward() applies gradients from the previous batch (and errors outright once gradients are None), and calling backward() twice without zeroing accumulates two batches’ gradients — a bug in a plain loop, a feature in the accumulation pattern above. Each iteration’s loss comes from the pairings chosen in Output Layers and Loss Pairings; a mismatched pairing still produces gradients, they just descend the wrong objective. The scaffolding around this loop — datasets, modules, checkpoints — is Building a Network in PyTorch.

The scenario interviewers love to hand you: a 6-layer MLP with sigmoid hidden units trains for 30 epochs, the loss drops from 0.91 to 0.51 in the first few, then crawls. You print the weight updates and find layer 6 moving by 1e-2 per step while layer 1 moves by 1e-6. Raising the learning rate destabilizes layer 6 without waking layer 1. This is not a tuning problem; it is a gradient-flow problem, and the first chapter already contained its cause.

Backprop multiplies one local derivative per layer on the way down. The sigmoid derivative never exceeds 0.25, and tanh’s never exceeds 1, so a stack of squashing activations attenuates the signal multiplicatively:

gradient at layer k = dL/d(output) × Π (local derivatives from layer L down to k)

10 sigmoid layers:  ≤ 0.25^10 ≈ 1e-6 of the loss signal reaches layer 1
20 sigmoid layers:  ≤ 0.25^20 ≈ 1e-12 — the early layer is effectively frozen

Nothing malfunctioned. You asked for a derivative that is a product of twenty numbers smaller than one, and got one. The mirror failure is exploding gradients: a recurrent network unrolled over 50 time steps multiplies by a weight matrix with effective gain greater than 1 at each step, the norms double until they exceed 1e20, and the next optimizer step writes NaN into every parameter. An oversized learning rate produces the same runaway in plain feed-forward nets.

Diagnose with numbers, not vibes. Log the gradient norm per parameter group every few hundred steps:

for name, p in model.named_parameters():
    if p.grad is not None:
        print(f'{name:24s} grad_norm={p.grad.norm():.3e}')

Healthy training shows gradient norms within roughly one order of magnitude across layers. Vanishing shows head layers two or three orders of magnitude below the tail — the exact 11× gap from chapter one, compounded. Exploding shows any norm past 1e3, or NaN appearing after it once overflowed.

Each remedy fixes a specific mechanism, so name the mechanism when you name the fix:

Note what clipping is and is not: it stops a bad step from destroying the weights, it does not fix the layer or the initialization that produced the explosion. Print the norms again after clipping; if the head layers still sit at 1e-7, the architecture needs the fix, not the loop.

This is the question behind most “my network does not learn” reports, so practice the full answer out loud: print per-layer gradient norms, compare them, attribute the spread to multiplicative chain-rule attenuation (sigmoid stacks) or gain above one (recurrent unrolls, oversized lr), then prescribe from the table — activations and residual paths for vanishing, clipping and learning-rate surgery for exploding. The same norms live inside the loop you built in Building a Network in PyTorch, and the loss being differentiated came from the pairings in Output Layers and Loss Pairings.