S098 · Backpropagation — Derived by Hand on a 2-Layer Net
The algorithm that made deep learning possible — the chain rule, systematically applied. Derive backpropagation for a 2-layer MLP by hand, implement it in ~50 lines of numpy that trains without any autograd, and understand why every gradient bug is either a shape bug or a sign bug.
🎯 Derive backprop for a 2-layer MLP + softmax + cross-entropy on paper, implement it in numpy, and reason about vanishing / exploding gradients from first principles.
Why this session exists
Backpropagation is not magic — it's the chain rule applied systematically, plus one clever observation about caching intermediate values. Every DL framework's autograd is a general-purpose backprop engine. If you never derive it by hand once, you'll spend your career treating gradient bugs like superstition ("try Adam? try a smaller learning rate?"). If you derive it once, you'll debug them in minutes. This session is that one derivation, done slowly and completely.
- Derive dL/dW and dL/db for every layer of a 2-hidden-layer MLP with ReLU + softmax + cross-entropy.
- Explain why softmax + cross-entropy gives the clean gradient (probs − y) — and why sigmoid + MSE does not.
- Implement backprop from scratch in numpy that matches PyTorch's autograd to 6 decimal places.
- Trace a vanishing or exploding gradient bug to its exact layer using shape and magnitude checks.
- Read any framework's backward pass and know what it's doing.
Prerequisites
- S096 · Perceptron & Activation Functions — one neuron, forward direction.
- S097 · MLP forward pass — you already wrote the forward pass; today is the backward.
- Comfort with the multivariable chain rule and matrix calculus (Jacobians, matmul derivatives).
(a) Intuition · 5 min
A car rolls off the assembly line with a paint defect. Whose fault? The paint booth, obviously. But the paint booth technician says "the primer was uneven." The primer station says "the surface prep was rough." The surface prep station says "the raw metal was bent." Blame propagates backward, one station at a time, and each station adjusts its process a little.
That's it. That's backprop. Loss at the output is the defect; you propagate blame backward one layer at a time, and each layer nudges its weights in the direction that reduces its share of the blame.
Mathematically: given loss L at the output, we want dL/dW for every weight matrix. The chain rule gives us that, IF we walk backward from the output and multiply the local derivative at each step. That's backprop.
The clever bit is caching: the forward pass saves every intermediate activation. The backward pass reuses them (a1, a2, z1, z2) without recomputing. This is why backprop is O(forward-pass-cost) instead of O(square of that).
The one sentence to memorise
- Start at the loss L. Compute dL/d(logits). For softmax + cross-entropy this is (probs − y_onehot) / batch_size.
- For each layer, going backward: dW = a_prev.T @ dz, db = dz.sum(axis=0), da_prev = dz @ W.T.
- Then propagate through the activation: dz_prev = da_prev * φ'(z_prev). For ReLU, φ' is 1 where z > 0, else 0.
- Repeat until you hit the input. Every weight matrix now has a gradient.
- Update: W -= lr × dW (or Adam / SGD / whatever). That's one training step.
A quick history
- 1960sAutomatic differentiation · Wengert, LinnainmaaThe core algorithm exists in the numerical analysis community for decades before ML picks it up.
- 1974Werbos · PhD thesisApplies backprop to neural networks. Nobody reads it for ten years.
- 1986Rumelhart, Hinton, WilliamsRediscover and popularise backprop for multi-layer perceptrons. Deep learning becomes feasible in principle.
- 2006Hinton · deep belief netsLayer-wise pre-training makes deep nets trainable in practice, before ReLU and better init made it easy.
- 2015Autograd libraries matureAutograd, then PyTorch (2016), then JAX (2018) — reverse-mode autodiff as a first-class primitive.
(b) Visual walkthrough · 15 min
The whole chain rule as a graph
Each backward edge is one line of code. The whole algorithm is on this diagram.
The derivation, step by step
z1 = x @ W1 + b1; a1 = ReLU(z1); z2 = a1 @ W2 + b2; a2 = ReLU(z2); logits = a2 @ W3 + b3; probs = softmax(logits).
L = -(1/n) · Σ y_onehot · log(probs). Its gradient w.r.t. logits is dz3 = (probs − y_onehot) / n. This one line is the magic — softmax + CE gives the clean gradient.
dW3 = a2.T @ dz3 (shape (H2, K)). db3 = dz3.sum(axis=0). da2 = dz3 @ W3.T.
dz2 = da2 * (z2 > 0). If z2 was negative, gradient is blocked (this is what ‘dead ReLU’ means algebraically).
dW2 = a1.T @ dz2 (shape (H1, H2)). db2 = dz2.sum(axis=0). da1 = dz2 @ W2.T.
dz1 = da1 * (z1 > 0).
dW1 = x.T @ dz1 (shape (D, H1)). db1 = dz1.sum(axis=0). Done — every weight now has a gradient.
Why (softmax + cross-entropy) has a clean gradient
dL/dz = (probs − y). Clean, non-vanishing.
- Log inside CE cancels the exponential inside softmax
- Result: dL/dz is just (predicted − true) — never zero unless perfect
- Standard pairing for multi-class classification
- Same for sigmoid + binary-cross-entropy → dL/dz = (σ(z) − y)
dL/dz = 2(σ(z)−y)·σ(z)(1−σ(z)). Vanishing.
- Extra factor σ(z)(1−σ(z)) is at most 0.25
- Gradient vanishes at both extremes even when very wrong
- Never use for classification
- Historical mistake in many old textbooks
Bad idea, don't do it
- No log-exp cancellation
- Multi-modal loss surface
- Never appears in modern practice
- Included here only so you recognise the antipattern
Correct for regression
- No activation on output layer
- dL/dz = 2(pred − y) / n — trivial
- The regression counterpart to softmax+CE
- Simple, unbeatable for continuous targets
Where gradients vanish or explode
Diagnosing gradient pathologies
"Backpropagation is the learning algorithm for neural networks — it's how the network learns from its errors and updates the weights."
Backpropagation computes gradients. That is its entire job. It is reverse-mode automatic differentiation applied to a computation graph, and it updates nothing. The learning algorithm is the optimiser — SGD, Adam — which consumes those gradients and decides the step. Backprop is exact and deterministic; the optimiser is where every hyperparameter and every training heuristic lives. Conflating them is why "backprop got stuck" gets said about problems that are actually optimiser or loss-surface problems.
The myth is sticky because the two are always used together and are usually taught in one breath: "compute the error, propagate it backward, update the weights". The third clause is a different algorithm and it gets absorbed into the phrase. Frameworks reinforce this — loss.backward() and optimizer.step() are adjacent lines that always appear together, so they feel like one operation. They are not, and the separation is the reason you can change optimisers without touching the model.
Call backward twice with no optimiser and watch gradients accumulate — proving backward only writes .grad and changes no parameter:
import torch
w = torch.tensor([2.0], requires_grad=True)
before = w.item()
loss = (w * 3).sum()
loss.backward()
print(w.grad, w.item() == before) # tensor([3.]) True -- w UNCHANGED
loss = (w * 3).sum()
loss.backward()
print(w.grad) # tensor([6.]) -- ACCUMULATED
# this is why zero_grad() exists, and why forgetting it
# silently sums gradients across batchesWhy go backward? You could compute derivatives forward, propagating the derivative of every parameter alongside the forward pass. Both directions give identical, exact answers. The choice of direction is a complexity decision, and it is decisive.
- 1The chain rule for a composition f(g(h(x))) is a product of Jacobians. Matrix products are associative, so you may bracket them in any order and the result is identical.forced by · associativity of matrix multiplication is what makes both modes valid
- 2But the cost depends entirely on the bracketing. Multiplying a (1×n) row vector by an (n×n) matrix costs O(n²); multiplying two (n×n) matrices costs O(n³). Same answer, wildly different bills.forced by · the cost of a chain of matrix products is dominated by the largest intermediate object you materialise
- 3Forward mode brackets from the input side. It propagates the derivative with respect to one input, so a full gradient requires one pass per parameter. With P parameters that is P forward passes.forced by · each forward pass carries the derivative along a single input direction
- 4Reverse mode brackets from the output side. Because the loss is a scalar, the leftmost object is a 1×n row vector, and every intermediate stays a vector rather than a matrix. One backward pass yields the derivative with respect to all parameters.forced by · starting from a scalar output keeps every partial product one-dimensional on the left
- 5Neural networks have millions to billions of parameters and exactly one scalar loss. So forward mode costs O(P) passes and reverse mode costs O(1) — a factor of millions. Deep learning at scale is only possible in the reverse direction.forced by · the asymmetry between many inputs and one output is precisely what reverse mode exploits
Therefore backprop goes backward because the loss is a scalar and the parameters are many. Reverse mode is optimal for many-inputs-to-one-output; forward mode is optimal for the reverse shape.
And note what this predicts: reverse mode must pay for its speed in memory, because computing the backward pass requires the forward activations, so every intermediate must be retained until it is used. That is exactly why activation memory scales with depth × batch size and why it, not parameter count, is usually what makes training OOM. It is also why gradient checkpointing works — discard activations and recompute them during the backward pass, trading compute for memory, a tradeoff the derivation says must exist.
The forward pass builds a graph: every operation is a node that remembers its inputs and knows its own local derivative. The loss sits at the far end as a single number.
The backward pass starts there with a gradient of 1 and walks backward. At each node it applies one rule: multiply the incoming gradient by this node's local derivative, then pass it to the node's inputs. Where a value fed several nodes, its gradients sum. That is the whole algorithm — a local rule applied repeatedly. No node knows anything about the network; each only knows its own operation.
- Every node needs only its local derivative. This locality is why frameworks can autodiff arbitrary code — you never write a global derivative.
- Branches sum on the way back (a value used twice receives two gradients); sums split evenly; products swap (∂(ab)/∂a = b).
- Backward requires the forward activations, so memory grows with depth × batch. Activation memory usually causes OOM long before parameter memory does.
- Gradients accumulate into
.gradrather than overwrite. That is a feature — it is how gradient accumulation simulates large batches — and it is whyzero_grad()is mandatory.
Fire this the moment you see: a missing zero_grad() · loss that decreases then goes NaN · gradients vanishing in early layers of a deep stack · an in-place operation that breaks autograd · a detached tensor silently stopping the gradient flow · OOM that scales with batch size rather than model size · a custom layer whose gradient you need to check numerically.
You are hitting GPU memory limits during training. Do you shrink the batch, use gradient checkpointing, or drop to mixed precision?
Enable mixed precision first, because it is nearly free and improves both axes at once. Then reduce batch size with gradient accumulation, which preserves optimisation behaviour at the cost of time. Save gradient checkpointing for when the model genuinely will not fit, since it buys memory with a real and permanent compute tax on every step.
The framing that generalises: activation memory is a direct consequence of reverse-mode autodiff needing forward values to compute backward ones. Every technique above is a different point on that same trade — recompute instead of store, store fewer bits, or store fewer examples. Knowing why the memory exists tells you immediately which lever applies, instead of trying them at random when the OOM appears at 2am.
(c) Hands-on · 25 min
Implement backprop from scratch on a 2-hidden-layer MLP + softmax + CE, then gradient-check against a numerical derivative to prove the derivation. Save as backprop_lab.py, uv run backprop_lab.py.
"""backprop_lab.py — backpropagation from scratch, with a gradient check.
Trains a 2-hidden-layer MLP on a small classification problem using our own
backward pass, then runs a numerical gradient check to prove the derivation.
"""
from __future__ import annotations
import numpy as np
from sklearn.datasets import make_classification
RNG = np.random.default_rng(0)
# --- activations ---
def relu(z): return np.maximum(0, z)
def relu_grad(z): return (z > 0).astype(z.dtype)
def softmax(z):
z = z - z.max(axis=1, keepdims=True)
e = np.exp(z)
return e / e.sum(axis=1, keepdims=True)
# --- model ---
def init_params(sizes: list[int]) -> dict:
p = {}
for i in range(len(sizes) - 1):
fan_in, fan_out = sizes[i], sizes[i + 1]
p[f"W{i+1}"] = RNG.normal(0, np.sqrt(2.0 / fan_in), (fan_in, fan_out))
p[f"b{i+1}"] = np.zeros(fan_out)
return p
def forward(X: np.ndarray, p: dict) -> tuple[np.ndarray, dict]:
z1 = X @ p["W1"] + p["b1"]; a1 = relu(z1)
z2 = a1 @ p["W2"] + p["b2"]; a2 = relu(z2)
z3 = a2 @ p["W3"] + p["b3"]; probs = softmax(z3)
return probs, {"X": X, "z1": z1, "a1": a1, "z2": z2, "a2": a2, "z3": z3, "probs": probs}
def loss_ce(probs: np.ndarray, y_onehot: np.ndarray) -> float:
return -float(np.mean(np.sum(y_onehot * np.log(probs + 1e-12), axis=1)))
def backward(y_onehot: np.ndarray, cache: dict, p: dict) -> dict:
n = y_onehot.shape[0]
# Output layer
dz3 = (cache["probs"] - y_onehot) / n
dW3 = cache["a2"].T @ dz3
db3 = dz3.sum(axis=0)
da2 = dz3 @ p["W3"].T
# Hidden layer 2
dz2 = da2 * relu_grad(cache["z2"])
dW2 = cache["a1"].T @ dz2
db2 = dz2.sum(axis=0)
da1 = dz2 @ p["W2"].T
# Hidden layer 1
dz1 = da1 * relu_grad(cache["z1"])
dW1 = cache["X"].T @ dz1
db1 = dz1.sum(axis=0)
return {"W1": dW1, "b1": db1, "W2": dW2, "b2": db2, "W3": dW3, "b3": db3}
def sgd_step(p: dict, grads: dict, lr: float) -> None:
for k in p:
p[k] -= lr * grads[k]
# --- gradient check ---
def numeric_grad(name: str, idx: tuple, p: dict, X, y_oh, eps: float = 1e-5) -> float:
p[name][idx] += eps
plus = loss_ce(forward(X, p)[0], y_oh)
p[name][idx] -= 2 * eps
minus = loss_ce(forward(X, p)[0], y_oh)
p[name][idx] += eps
return (plus - minus) / (2 * eps)
def grad_check(p, X, y_oh, n_checks: int = 5) -> None:
_, cache = forward(X, p)
analytic = backward(y_oh, cache, p)
print("\nGradient check (analytic vs numeric):")
for name in ["W1", "W2", "W3"]:
for _ in range(n_checks):
idx = tuple(RNG.integers(0, s) for s in p[name].shape)
g_num = numeric_grad(name, idx, p, X, y_oh)
g_an = analytic[name][idx]
rel = abs(g_num - g_an) / (abs(g_num) + abs(g_an) + 1e-12)
print(f" {name}{idx}: analytic={g_an:+.6f} numeric={g_num:+.6f} rel_err={rel:.2e}")
# --- data + training ---
def one_hot(y: np.ndarray, k: int) -> np.ndarray:
out = np.zeros((y.shape[0], k), dtype=np.float32)
out[np.arange(y.shape[0]), y] = 1.0
return out
def train(p: dict, X, y, *, epochs=40, batch=64, lr=0.1) -> None:
y_oh = one_hot(y, 3)
n = X.shape[0]
for ep in range(1, epochs + 1):
idx = RNG.permutation(n)
for start in range(0, n, batch):
b_idx = idx[start:start + batch]
probs, cache = forward(X[b_idx], p)
grads = backward(y_oh[b_idx], cache, p)
sgd_step(p, grads, lr)
probs, _ = forward(X, p)
loss = loss_ce(probs, y_oh)
acc = float((probs.argmax(axis=1) == y).mean())
if ep % 5 == 0 or ep == 1:
print(f" epoch {ep:>2}: loss={loss:.4f} acc={acc:.4f}")
if __name__ == "__main__":
X, y = make_classification(n_samples=2000, n_features=10, n_informative=6,
n_classes=3, n_clusters_per_class=2, random_state=0)
p = init_params([10, 32, 16, 3])
# 1. Prove derivation is correct
grad_check(p, X[:32], one_hot(y[:32], 3), n_checks=3)
# 2. Train
print("\nTraining:")
train(p, X, y, epochs=40, batch=64, lr=0.1)Anatomy of the script
Anatomy of the script
Introduce a subtle bug: in backward, change dW3 = cache["a2"].T @ dz3 to dW3 = cache["a1"].T @ dz3. The shapes might even still work if H1 == H2 (they don't here, so you'll get an immediate shape error). If they do match, the training loss will still drop somewhat because SGD is forgiving — but the gradient check will report relative errors of ~0.5, telling you the analytic gradient is nowhere near the numeric one. THIS is why gradient check exists — bugs that look like ‘it kind of trains’ are the worst kind.
(d) Production reality · 15 min
An engineer implements a custom loss function with a hand-written gradient (for perf reasons). Model trains fine but converges to a worse loss than the same loss with autograd. Weeks of debugging reveal an off-by-a-factor-of-2 error in the manual gradient.
Two structural changes that are now standard: (1) residual connections (He et al. 2015) — the output of each block is added to its input, so gradients have a direct path to earlier layers; (2) LayerNorm at each block — keeps activations at a healthy scale.
Every transformer since — BERT, GPT, T5, LLaMA — uses both.
Two protections: (1) validate input for NaN/inf before every forward pass; (2) after every backward pass, check torch.isfinite(grad).all() for each parameter — if any is not finite, skip the update. Also: keep the last N stable checkpoints and auto-rollback on production NaN.
Where this shows up in the rest of the plan
(e) Recall + stretch · 10 min
Explain-out-loud test
If you can't teach these three to a friend without notes, redo the session:
- Write dW2 and da1 for a 2-hidden-layer MLP with ReLU and softmax + CE.
- Why does softmax + cross-entropy have a clean gradient?
- What is a gradient check and why do you always do one?
What comes next
Hub: The 6-Month Learning Plan
Part of a 130-session evergreen learning series. Session structure: intuition → visual → hands-on → production war stories → recall. Duration: 90 minutes.