Backpropagation, in depth: reverse-mode autodiff from scratch, a hand-traced example, the MLP backward pass and gradient checking

Why reverse mode

Training needs dL/dtheta for every parameter theta. The naive way is finite differences: nudge one parameter, rerun the forward pass, and measure the change. With P parameters that is P (or 2P for central differences) forward passes per step. A model with 100 million parameters would need 100 million forward passes for one gradient. Forward-mode differentiation has the same problem: it pushes one directional derivative through the graph per pass.

Reverse mode turns this around. One forward pass records every intermediate value. One backward pass then computes the adjoint dL/dv of every intermediate v, starting from dL/dL = 1. Because the loss is a single scalar, one backward pass yields the derivative with respect to every input at once. The cost of the backward pass is a small constant multiple of the forward pass, typically about two: for a matrix multiply Y = XW, the backward needs two matrix multiplies of the same size, one for dX and one for dW. The price is memory, because the intermediates must be kept until the backward pass uses them.

Computation graphs and accumulated adjoints

Write the forward pass as a sequence of primitive operations, each producing a value from earlier values. That sequence is a directed acyclic graph. For each primitive you know its local derivative: for m = x * w, dm/dx = w and dm/dw = x; for z = m + b, both partials are 1; for p = sigmoid(z), dp/dz = p(1 - p).

The chain rule says that if v feeds into several later nodes u1, u2 and so on, then dL/dv = sum over consumers u of dL/du * du/dv. Two rules follow. First, a node's adjoint is only complete once every consumer has contributed, so nodes must be processed in reverse topological order. Second, contributions are summed, which is why every autodiff engine accumulates gradients with += rather than assigning them. This is also why PyTorch requires you to zero gradients between steps: the accumulation that is correct within one graph becomes a bug across two.

A scalar autodiff engine

Here is a complete reverse-mode engine for scalars. Each operation creates a node that remembers its parents and a closure that pushes its adjoint to them. The backward method orders the graph topologically with a depth-first search and runs the closures in reverse.

import math

class Value:
    def __init__(self, data, parents=(), backward=lambda: None):
        self.data, self.grad = data, 0.0
        self._parents, self._backward = parents, backward

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other))
        def backward():
            self.grad += out.grad                 # d(a+b)/da = 1
            other.grad += out.grad
        out._backward = backward
        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other))
        def backward():
            self.grad += other.data * out.grad    # d(ab)/da = b
            other.grad += self.data * out.grad
        out._backward = backward
        return out

    def sigmoid(self):
        s = 1.0 / (1.0 + math.exp(-self.data))
        out = Value(s, (self,))
        def backward():
            self.grad += s * (1.0 - s) * out.grad
        out._backward = backward
        return out

    def log(self):
        out = Value(math.log(self.data), (self,))
        def backward():
            self.grad += out.grad / self.data
        out._backward = backward
        return out

    def backward(self):
        order, seen = [], set()
        def visit(v):                             # post-order DFS = topological order
            if id(v) not in seen:
                seen.add(id(v))
                for p in v._parents:
                    visit(p)
                order.append(v)
        visit(self)
        self.grad = 1.0
        for v in reversed(order):
            v._backward()

Note what the closures capture: the forward values (self.data, s). That is the memory cost of backprop in miniature. Real frameworks do the same with tensors, saving activations for the backward pass, and the recursive DFS here would be replaced by an iterative one to avoid recursion limits on deep graphs.

Worked example: one logistic neuron

Take one logistic neuron with input x = 2, weight w = -0.5, bias b = 0.3 and true label y = 1, so the cross-entropy loss is -log p. Forward: m = x * w = -1.0, z = m + b = -0.7, p = sigmoid(-0.7) = 0.3318, L = -log(0.3318) = 1.1032.

Backward, starting from dL/dL = 1: dL/dp = -1/p = -3.0138. The sigmoid rule multiplies by p(1 - p) = 0.2217, giving dL/dz = -0.6682, which equals p - y exactly, the well-known logistic regression gradient. The sum copies it: dL/db = dL/dm = -0.6682. The product sends each input the other's value: dL/dw = x * dL/dm = -1.3364 and dL/dx = w * dL/dm = 0.3341. Running the engine on the same expression prints exactly these numbers. A gradient-descent step would now increase w and b, raising p towards the label.

A second test checks accumulation. For f = a * a + a at a = 3, the true derivative is 2a + 1 = 7. Node a feeds the product twice and the sum once; the engine returns 7.0 because each contribution is added. Replace += with = in __mul__ and you get the wrong answer silently, which is the most common bug in hand-written autodiff.

From scalars to matrices: the MLP backward pass

Scalar graphs are too slow for real networks; frameworks use tensor operations with hand-written backward rules. Two rules carry most of the weight. The shape rule: the gradient with respect to a tensor has that tensor's shape, which tells you where transposes go. For Z = XW + b with X of shape (B, d) and W of shape (d, h): dW = X.T @ dZ, dX = dZ @ W.T and db = dZ.sum(axis=0) because b was broadcast over the batch. And the fused softmax with cross-entropy has the simple gradient (P - onehot(Y)) / B on the logits, which is both cheaper and more numerically stable than differentiating softmax and log separately.

import numpy as np

def forward_backward(params, X, Y):
    W1, b1, W2, b2 = params
    H_pre = X @ W1 + b1                       # (B, h)
    H = np.maximum(H_pre, 0.0)                # ReLU
    logits = H @ W2 + b2                      # (B, k)
    logits = logits - logits.max(axis=1, keepdims=True)   # stability shift
    P = np.exp(logits); P /= P.sum(axis=1, keepdims=True)
    B = X.shape[0]
    loss = -np.log(P[np.arange(B), Y]).mean()
    dlogits = P.copy(); dlogits[np.arange(B), Y] -= 1.0; dlogits /= B
    dW2 = H.T @ dlogits
    db2 = dlogits.sum(axis=0)
    dH = dlogits @ W2.T
    dH_pre = dH * (H_pre > 0)                 # ReLU passes gradient where input was positive
    dW1 = X.T @ dH_pre
    db1 = dH_pre.sum(axis=0)
    return loss, [dW1, db1, dW2, db2]

Trained with plain gradient descent (learning rate 1.0, 64 hidden units, He-scaled initialisation) on 300 points of a three-arm spiral, the loss fell from 1.152 at step 0 to 0.203 at step 100 and 0.022 at step 2,000, with training accuracy going from 0.517 to 1.0. That is training accuracy only; it shows the gradients are right, not that the model generalises. See Gradient Descent, in depth for what the optimiser does with these gradients.

Gradient checking

Never trust a hand-written backward pass until it agrees with finite differences. For each parameter entry, compute (L(theta + h) - L(theta - h)) / 2h with h around 1e-6 in float64, and compare to the analytic gradient using a relative error, |num - ana| / (|num| + |ana|). Values near 1e-7 or below are good; 1e-4 is suspicious; 1e-2 is a bug.

def grad_check(params, X, Y, h=1e-6):
    _, grads = forward_backward(params, X, Y)
    worst = 0.0
    for P_, G in zip(params, grads):
        it = np.nditer(P_, flags=["multi_index"])
        for _ in it:
            i = it.multi_index
            old = P_[i]
            P_[i] = old + h; lp, _ = forward_backward(params, X, Y)
            P_[i] = old - h; lm, _ = forward_backward(params, X, Y)
            P_[i] = old
            num = (lp - lm) / (2 * h)
            worst = max(worst, abs(num - G[i]) / max(1e-12, abs(num) + abs(G[i])))
    return worst

On a batch of 8 random examples the worst relative error was 3.9e-8. Dropping the division by B on dlogits produces gradients exactly 8 times too large, the batch size; a ratio equal to a known constant is a strong hint about which line is wrong. Check in float64 on small inputs, and expect occasional larger errors when an input sits right at a ReLU kink, where the function is not differentiable.

Failure modes

• Vanishing and exploding gradients. Deep products of local derivatives shrink or grow geometrically; sigmoid derivatives are at most 0.25. Use ReLU-family activations, careful initialisation, normalisation, residual connections, and gradient clipping.
• Forgetting to zero gradients. Accumulation across steps quietly multiplies the effective learning rate; useful on purpose for gradient accumulation, a bug otherwise.
• In-place edits of saved tensors. Overwriting an activation that the backward pass needs gives wrong gradients; PyTorch detects some cases with a version counter.
• Detached paths. Converting to NumPy, calling .item() or using a non-differentiable op (argmax, rounding) cuts the graph and leaves upstream gradients at zero.
• Dead ReLUs. Units whose input is always negative get zero gradient forever; watch the fraction of active units per layer.
• NaN and overflow. Unshifted softmax, log of zero or a too-high learning rate produce NaN that backprop spreads to every parameter in one step.

Operational guidance

In practice you will use a framework's autograd, and the skill that matters is debugging it. Log the global gradient norm and per-layer norms every step; a sudden spike precedes most divergences. Check that every parameter receives a non-zero gradient after the first backward pass; a frozen layer is usually a detached path. When you write a custom operation, implement its backward and gradient-check it in float64 before using it at lower precision. When memory is the limit, remember that backprop trades memory for compute: activation checkpointing recomputes parts of the forward pass during backward to store fewer activations.

What to do next

1. Type in the Value class, reproduce the worked example (dL/dw = -1.3364) and the a * a + a check (7.0).
2. Add tanh, exp and division to the engine, each with its own backward rule, and test each one against central differences.
3. Run grad_check on the NumPy MLP, then break one line on purpose and see how large the error becomes.
4. Implement the same network in PyTorch and compare its .grad tensors with yours on identical weights.
5. Log gradient norms per layer during a real training run and learn what healthy looks like for your model.
6. Read the transformer-specific backward rules and the memory trade-offs in the two transformer_math articles linked above.