Backpropagation and the Computational Graph

Quick Context

Gradient descent needs a gradient: for each weight in the network, how much the loss would change if that weight moved a little. Backpropagation is how that count is obtained, for millions of weights, at the cost of roughly one extra forward pass.

It is not a learning rule and not an optimiser. It computes derivatives, and nothing else. What to do with them is the optimiser's business.

The mechanism is the chain rule of calculus, applied systematically. A neural network is a composition of functions: the input is transformed by layer 1, the result by layer 2. And so on until the loss function gives a single count. The chain rule says the derivative of the loss with respect to a weight deep inside the network is the product of the local derivatives along the path from that weight to the loss.

The forward pass computes the output layer by layer, caching the intermediate vals. The backward pass starts at the loss, computes the local gradient at each layer. And multiplies it by the gradient flowing back from the layer above. By the time the backward pass reaches the first layer, each weight in the network has its gradient.

The cost is roughly two forward passes: one to compute the outputs and one to compute the gradients. This is remarkably cheap — without backpropagation, each weight would need its own forward pass to estimate its gradient by finite gaps. And a network with a million weights would need a million forward passes per training step. Backpropagation reduced that to two, which is what made deep learning real.

In modern frameworks like PyTorch and TensorFlow, backpropagation is handled automatically. You set the forward computation, call loss.backward(). And the framework traverses the computation graph in reverse, accumulating gradients. Grasp what it is doing matters when gradients vanish, explode. Or behave unexpectedly — all of which are backpropagation problems, not optimiser problems.

All is a graph

A network is a long expression. And any expression can be drawn as a graph of small steps: multiply, add, apply a function. The graph above is a single neuron with two inputs. It is now enough to show each mechanic that matters.

m₁ = w₁x₁   m₂. = w₂x₂   s = m₁+m₂.   z = s+b   a = σ(z)   L = (a − y)²

The forward pass fills in a value at each node. Crucially, those vals are kept — the backward pass needs them. It is just why training a network takes far more memory than running one.

Each step knows one thing

Each node needs just its local derivative: how its output responds to its own inputs. It has no idea what the rest of the network looks like. And does not need to.

• Add passes the gradient through unchanged to both inputs. This is why ∂L/∂b equals ∂L/∂z just, each time.
• Multiply hands each input the gradient scaled by the other input. So ∂L/∂w₁. = ∂L/∂m₁ · x₁ — a weight's gradient is proportional to the input it was multiplying. It is why an input of zero gives a gradient of zero and why unscaled features distort learning.
• The sigmoid adds a(1 − a). It peaks at 0.25 and falls to nearly nothing when the neuron saturates.

The chain rule then says: the gradient arriving at a node, times the node's local derivative, is the gradient leaving it. Backpropagation is that single sentence applied repeatedly, in reverse topological order.

Giving blame to each weight

Training needs one count per weight: if I nudge this weight slightly, how much does the loss change? With ten million weights, computing that apart ten million times is impossible.

Backpropagation gets all of them in a single backward pass, at roughly the cost of one forward pass. That efficiency is the reason deep learning is real at all.

The mechanism is the chain rule. A network is nested functions. So the effect of an early weight on the loss is a product of factors, one per layer in between:

∂L/∂w⁽¹⁾ = ∂L/∂a⁽³⁾ × ∂a⁽³⁾/∂a⁽²⁾ × ∂a⁽²⁾/∂a⁽¹⁾ × ∂a⁽¹⁾/∂w⁽¹⁾

Computed naively, weights in the same layer would each redo most of that product. Backpropagation works from the loss backwards, carrying the accumulated product with it. And hands each layer just the factor it needs. Nothing is recomputed.

One neuron, worked through with counts

A single neuron: z = wx + b, a = sigmoid(z), loss L = (a − y)².

Take x = 2, w = 0.5, b = 0, y = 1.

Forward: z = 1.0, a = sigmoid(1.0) = 0.731, L = (0.731 − 1)² = 0.072.

Backward, three factors:

• ∂L/∂a = 2(a − y) = 2(−0.269) = −0.538
• ∂a/∂z = a(1 − a) = 0.731 × 0.269 = 0.197
• ∂z/∂w = x = 2

Multiply: ∂L/∂w = −0.538 × 0.197 × 2 = −0.212

Negative, so raising w lowers the loss. With a learning rate of 0.1 the update is w ← 0.5 + 0.0212 = 0.521 — a small step in the right way. It is just what gradient descent is.

Notice a(1 − a): it peaks at 0.25 and ways zero when the neuron saturates near 0 or 1. A confidently wrong sigmoid neuron therefore learns very slowly. It is precisely the argument for cross-entropy loss, whose derivative cancels that term.

Why the forward pass must be remembered

The backward pass needs the activations computed on the way forward — a(1 − a) needs a. So each intermediate result is stored in the forward pass and kept until the backward pass consumes it.

That storage is activation memory. And it is what often limits batch size, not the weights. A ResNet-50's weights are about 100MB. Its activations for a batch of 64 pics run to a few gigabytes.

Two standard replies. Gradient checkpointing stores just some activations and recomputes the rest in the backward pass — trading roughly 30% more compute for a large memory saving. And mixed precision stores activations in 16-bit, halving the memory and speeding up the arithmetic.

This also explains why torch.no_grad() matters at inference: without gradients, nothing needs to be stored. And memory use drops sharply.

Exploration guide

1. Start from the loss. The page opens with the forward pass now done: vals left to right, ending on a single count. Press Reset and Step if you want to watch that part happen too — but it holds no surprises, it is just the network making a prediction.
2. Keep stepping. The gradient starts at the loss as 1, then flows right to left. Each node prints the multiplication it done. So you can read the chain rule rather than take it on faith.
3. Check the two weight gradients. They differ just by their input: ∂L/∂w₁. Carries x₁ and ∂L/∂w₂ carries x₂. Set Input x1 to 0 and its weight's gradient goes to just zero — that weight cannot learn from this case at all.
4. Check the answer. The gradient-check panel recomputes ∂L/∂w₁. Numerically, by nudging w₁. And dividing the change in loss by the change in weight. It agrees with the chain rule to a few decimal places. It is the test to reach for whenever a hand-written backward pass looks suspicious.
5. Saturate the neuron. Push Weight w1 and Bias b to their maximums so that z is large. The sigmoid flattens, a(1 − a) collapses towards zero. And each gradient behind it shrinks with it. That is the vanishing gradient, visible in one count.
6. Close the gap. Drag Target y towards the prediction a. The loss collapses and each gradient in the graph shrinks with it, since the whole chain is scaled by 2(a − y) — when there is nothing left to fix, nothing moves.

Why it is fast

The obvious way to get gradients is to nudge each weight and re-run the network, just as the gradient check does. That costs one forward pass per weight. On a net with a million weights it is hopeless.

Backpropagation gets all of them in one backward pass, since the gradient arriving at a node is shared by all feeding it. Work is reused rather than repeated — the same insight that makes dynamic programming fast. This is reverse-mode automatic differentiation. And it is what each deep learning framework builds underneath loss.backward().

What trips people up

• Forgetting to zero the gradients. Frameworks accumulate into .grad by design. So that a batch can be split across a few backward passes. Omit the reset and you are descending on the sum of each batch so far.
• Saturation. Sigmoid and tanh have tiny derivatives at their extremes. And a product of tiny counts vanishes. This is the reason ReLU and its relatives took over. And the reason weight initialisation is treated so carefully.
• Exploding gradients. The mirror pic: repeated multiplication by counts above one, common in deep or recurrent stacks. Gradient clipping exists for this.
• Detaching by accident. Turn a tensor to a plain count mid-network and the graph is cut there. All behind it silently gets no gradient at all.
• Trusting the derivation. If you write a backward pass by hand, check it numerically. It takes five lines and it finds sign errors at once.

Where that leaves you

Backpropagation is the chain rule applied to a computational graph in reverse order: each step adds just its own local derivative. And the gradient arriving from downstream is multiplied by it on the way past. An add passes the gradient through untouched, a multiply scales it by the other input. And a saturating activation shrinks it — which is where vanishing gradients come from. Since each node reuses the gradient now computed for the nodes it feeds, one backward pass yields each weight's gradient at roughly the cost of one forward pass. And that efficiency is the reason training deep networks is possible at all. The forward vals have to be kept in memory for it to work, and a numeric gradient check is the fastest way to prove your backward pass is right.

What goes wrong, and why

Since the chain rule multiplies one factor per layer, deep networks multiply many counts together — and that is numerically unstable in both ways.

Vanishing gradients. Sigmoid's derivative peaks at 0.25. Ten layers gives 0.25¹⁰ ≈ 0.000001, so early layers barely move. The fixes all target the size of those factors: ReLU (derivative just 1 when active), residual connections (an identity path whose derivative is 1). And normalisation layers (which keep activations where the factors stay near 1).

Exploding gradients. If the factors average above 1, the product grows without bound and weights turn NaN. Gradient clipping caps the gradient's norm right and is standard in recurrent nets and transformers.

Dead ReLUs. A unit stuck on the negative side has a derivative of just 0, gets no gradient. And never recovers. Leaky ReLU, careful initialisation and a less learning rate all reduce it.

Diagnosing all three is the same action: log the gradient norm per layer. A healthy network shows close magnitudes throughout. A vanishing one shows orders of magnitude of decay from output to input.

Automatic differentiation in practice

You will never build backpropagation by hand. Frameworks record each step into a graph as the forward pass runs, then walk it backwards applying each step's known derivative rule.

loss = criterion(model(x), y)

optimizer.zero_grad()    # clear the previous step's gradients
loss.backward()          # backpropagation: fills .grad on every parameter
optimizer.step()         # apply the update

The zero_grad() line is not optional. PyTorch accumulates gradients into .grad rather than replacing them. So forgetting it means each step uses the sum of all last gradients — training appears to work and converges to nonsense. (The accumulation behaviour is deliberate: it is how you simulate a large batch on a small GPU, by calling backward() a few times before step().)

To verify a hand-written gradient, check it with a numerical one: (f(x+h) − f(x−h)) / 2h with h around 1e-5. If they disagree, the analytic version is wrong.

Each gradient, by hand, checked numerically

Backpropagation is the chain rule applied backwards through a network. Here is a two-layer net where each partial derivative is computed twice -- analytically and numerically -- so you can see they agree.

import numpy as np

rng = np.random.default_rng(0)
X = rng.normal(size=(6, 3))
y = rng.integers(0, 2, (6, 1)).astype(float)

W1 = rng.normal(0, 0.5, (3, 4)); b1 = np.zeros((1, 4))
W2 = rng.normal(0, 0.5, (4, 1)); b2 = np.zeros((1, 1))

def sigmoid(z):
    return 1 / (1 + np.exp(-z))

def forward(W1, b1, W2, b2):
    z1 = X @ W1 + b1
    a1 = np.tanh(z1)
    z2 = a1 @ W2 + b2
    a2 = sigmoid(z2)
    loss = -np.mean(y * np.log(a2) + (1 - y) * np.log(1 - a2))
    return z1, a1, z2, a2, loss

z1, a1, z2, a2, loss = forward(W1, b1, W2, b2)
print("forward pass on %d rows: loss %.6f" % (len(X), loss))
print("  x %s -> z1 %s -> a1 %s -> z2 %s -> a2 %s"
      % (X.shape, z1.shape, a1.shape, z2.shape, a2.shape))
print()

Building on the same setup, the next block shows the left cases.

import numpy as np

print("now walk it backwards. each line is one application of the chain rule:")
n = len(X)
# dL/dz2 for sigmoid + binary cross-entropy collapses to (a2 - y)/n
dz2 = (a2 - y) / n
print("  dL/dz2 = (a2 - y)/n           shape %s" % (dz2.shape,))
dW2 = a1.T @ dz2
db2 = dz2.sum(axis=0, keepdims=True)
print("  dL/dW2 = a1.T @ dL/dz2        shape %s" % (dW2.shape,))
print("  dL/db2 = sum(dL/dz2)          shape %s" % (db2.shape,))
da1 = dz2 @ W2.T
print("  dL/da1 = dL/dz2 @ W2.T        shape %s" % (da1.shape,))
dz1 = da1 * (1 - np.tanh(z1) ** 2)
print("  dL/dz1 = dL/da1 * tanh'(z1)   shape %s" % (dz1.shape,))
dW1 = X.T @ dz1
db1 = dz1.sum(axis=0, keepdims=True)
print("  dL/dW1 = X.T @ dL/dz1         shape %s" % (dW1.shape,))
print("  dL/db1 = sum(dL/dz1)          shape %s" % (db1.shape,))
print()
print("every gradient has the same shape as the thing it updates. that is not")
print("a coincidence -- it is what makes 'W -= lr * dW' type-check.")
print()

print("the check that proves it. nudge each parameter and measure the loss:")
def numeric(param, i, j, h=1e-6):
    up = [W1.copy(), b1.copy(), W2.copy(), b2.copy()]
    dn = [W1.copy(), b1.copy(), W2.copy(), b2.copy()]
    up[param][i, j] += h
    dn[param][i, j] -= h
    return (forward(*up)[4] - forward(*dn)[4]) / (2 * h)

The next block covers the left steps.

import numpy as np

print("%-10s %6s %16s %16s %12s" % ("param", "index", "analytic", "numeric", "diff"))
checks = [("W1", 0, dW1, (0, 0)), ("W1", 0, dW1, (2, 3)), ("b1", 1, db1, (0, 1)),
          ("W2", 2, dW2, (1, 0)), ("W2", 2, dW2, (3, 0)), ("b2", 3, db2, (0, 0))]
worst = 0.0
for name, idx, grad, (i, j) in checks:
    a = grad[i, j]
    num = numeric(idx, i, j)
    worst = max(worst, abs(a - num))
    print("%-10s %6s %16.10f %16.10f %12.2e"
          % (name, "[%d,%d]" % (i, j), a, num, abs(a - num)))
print("  worst disagreement anywhere: %.2e" % worst)
print("  that is floating-point noise, not error. the derivation is right.")
print()

print("and it trains. 400 steps of plain gradient descent:")
lr = 0.5
for step in range(401):
    z1, a1, z2, a2, loss = forward(W1, b1, W2, b2)
    if step % 100 == 0:
        acc = ((a2 > 0.5).astype(float) == y).mean()
        print("  step %3d  loss %.6f  accuracy %.4f" % (step, loss, acc))
    dz2 = (a2 - y) / n
    dW2_ = a1.T @ dz2; db2_ = dz2.sum(axis=0, keepdims=True)
    dz1 = (dz2 @ W2.T) * (1 - np.tanh(z1) ** 2)
    dW1_ = X.T @ dz1; db1_ = dz1.sum(axis=0, keepdims=True)
    W2 -= lr * dW2_; b2 -= lr * db2_
    W1 -= lr * dW1_; b1 -= lr * db1_
print()
print("that is the entire algorithm. a forward pass that caches its")
print("intermediates, a backward pass that reuses them, and a subtraction.")
print("PyTorch's autograd is this, recorded automatically as you compute.")

Questions people ask

Is backpropagation the same as gradient descent? No. Backpropagation computes the gradients. Gradient descent uses them to update the weights. Two split steps.

Why do I need to call zero_grad()? Since gradients accumulate by default. Omitting it is one of the most common PyTorch bugs.

Does it work through any step? Any differentiable one. Non-differentiable steps (argmax, sampling, hard thresholds) block the gradient and need workarounds such as straight-through estimators or the reparameterisation trick.

How does it handle branching designs? Where a value feeds a few paths, the gradients from all of them are summed. That is the multivariable chain rule.

Why is training about three times a forward pass? Roughly: one forward, one backward of close cost, plus the weight update.

Is it biologically plausible? Not really — real neurons have no known mechanism for sending precise error signals backwards. It is an engineering algorithm, not a net of the brain.

Recap in one screen

• Backpropagation computes each weight's gradient in one backward pass with the chain rule.
• It reuses shared factors, which is why training costs about the same as a few forward passes rather than one pass per param.
• Activations from the forward pass must be stored. And that memory often limits batch size.
• Multiplying many factors causes vanishing and exploding gradients. ReLU, residuals, normalisation and clipping are the replies.
• Frameworks do it automatically — but call zero_grad(), since gradients accumulate.