Backpropagation as a Computation Graph: A Two-Layer MLP by Hand

Backpropagation is often introduced as the mysterious engine powering deep learning, but at its core, it is simply reverse-mode automatic differentiation applied to a computation graph. It is not a neural-network-specific trick; rather, it is a disciplined, highly optimized way to move local gradients backward through program operations using the chain rule of calculus.

This article uses one hidden layer and two parameterised affine layers: x @ W1+b1 → ReLU → h @ W2+b2 → mean cross-entropy. Samples are rows of a batch. We derive the gradients, run a complete NumPy example, and check every parameter with finite differences. Random inputs and labels are numerical fixtures, not evidence of classification ability.

Choose the matching experiment. The run instructions above for src/mlp_backprop.py belong to the older series package and its single-sample column-vector example. For this article’s 32-sample batch and 1029-parameter checks, use the Neural Gradient Lab results, download and commands below. The fixtures differ, so their numbers are not interchangeable.

1. The Computation Graph Unveiled

To compute derivatives systematically, we decompose the complex neural network into a Directed Acyclic Graph (DAG) of primitive operations. Each node represents a simple mathematical operation (like matrix multiplication or ReLU), and edges represent the flow of tensors. Crucially, each node must be capable of doing two things: calculating its forward output, and calculating its local Vector-Jacobian Product (VJP) during the backward pass.

graph TD
    x["Input x"] --> z1["z1 = x @ W1 + b1"]
    W1["Weights W1"] --> z1
    b1["Bias b1"] --> z1
    z1 --> h["h = ReLU(z1)"]
    h --> logits["logits = h @ W2 + b2"]
    W2["Weights W2"] --> logits
    b2["Bias b2"] --> logits
    logits --> p["p = Softmax(logits)"]
    p --> L["Loss = CrossEntropy(p, target)"]
    target["Target y"] --> L

    classDef fwd fill:#e1f5fe,stroke:#039be5,stroke-width:2px;
    classDef param fill:#fce4ec,stroke:#d81b60,stroke-width:2px;
    class z1,h,logits,p,L fwd;
    class W1,b1,W2,b2 param;

The backward pass needs access to values such as x, z1 and h. They can be stored, or some activations can be recomputed during backward. Training also usually needs gradients and optimizer state; whether capacity, bandwidth or arithmetic limits performance requires profiling the particular workload.

2. The Softmax Cross-Entropy Shortcut

In theory, you could calculate the Jacobian of the Cross-Entropy loss with respect to the Softmax probabilities, and then multiply that by the Jacobian of the Softmax with respect to the logits. In practice, doing this explicitly is a recipe for numerical disaster and wasted compute.

For one sample with an integer class label, no class weights and no label smoothing, the logits gradient simplifies as follows. A mean loss over N samples divides each row by N; a summed loss does not. Weighting or ignored targets require rechecking the reduction rather than reusing this formula blindly.

dL/dlogits = p - one_hot(target)

With zero-based class indices, target 1 and probabilities [0.1,0.7,0.2] give single-sample gradient [0.1,-0.3,0.2]. PyTorch CrossEntropyLoss takes unnormalised logits, not a preceding softmax output; probabilities here are an intermediate in the derivation.

3. Mathematical Derivations of the MLP

Once we have the gradient of the loss with respect to the logits (dlogits), we propagate it backward using the chain rule. Notice that we never instantiate full Jacobian matrices; instead, we compute Vector-Jacobian Products efficiently using matrix transposes.

# Layer 2 gradients
dW2 = h^T dlogits
db2 = sum(dlogits, axis=0)
dh  = dlogits W2^T

# Layer 1 gradients
dz1 = dh * ReLU'(z1)  # Element-wise multiplication
dW1 = x^T dz1
db1 = sum(dz1, axis=0)

Each gradient must match its parameter shape exactly; bias gradients reduce over the batch dimension. Norms can flag unusual scales, but a single norm neither proves correctness nor rules out vanishing/exploding gradients throughout training. The fixed-input run and numerical checks below provide more specific evidence.

4. A Runnable NumPy Implementation

The example fixes seed=42, a 32×10 input, 64 hidden units and 5 output classes, with NumPy’s generated float64 data. It computes mean cross-entropy, four parameter gradients and the loss after one update. This code matches mlp_example.py in the download. The loss uses log-sum-exp rather than taking a logarithm of an underflowed probability. This is an educational implementation with explicit input constraints, not a general training library.

"""One-hidden-layer ReLU MLP, row-batch convention, mean unweighted CE."""

import numpy as np


def fixture(seed=42):
    rng = np.random.default_rng(seed)
    x = rng.standard_normal((32, 10))
    params = {
        "W1": rng.standard_normal((10, 64)) * 0.1,
        "b1": np.zeros((1, 64)),
        "W2": rng.standard_normal((64, 5)) * 0.1,
        "b2": np.zeros((1, 5)),
    }
    targets = rng.integers(0, 5, size=32)
    return x, targets, params


def cross_entropy(logits, targets):
    if logits.ndim != 2 or min(logits.shape) < 1:
        raise ValueError("Expected nonempty (batch, classes) logits")
    if targets.shape != (logits.shape[0],) or targets.dtype.kind not in "iu":
        raise ValueError("Expected one integer class index per row")
    if not np.isfinite(logits).all() or np.any(targets < 0) or np.any(targets >= logits.shape[1]):
        raise ValueError("Nonfinite logits or out-of-range target")
    shifted = logits - logits.max(axis=1, keepdims=True)
    exps = np.exp(shifted)
    normalizer = exps.sum(axis=1, keepdims=True)
    probs = exps / normalizer
    # Compute the loss from logits, not log(probs), which can underflow to log(0).
    loss = np.mean(np.log(normalizer[:, 0]) - shifted[np.arange(len(targets)), targets])
    dlogits = probs.copy()
    dlogits[np.arange(len(targets)), targets] -= 1
    dlogits /= len(targets)
    return float(loss), probs, dlogits


def loss_and_grads(x, targets, params):
    W1, b1, W2, b2 = (params[k] for k in ("W1", "b1", "W2", "b2"))
    if x.ndim != 2 or W1.ndim != 2 or W2.ndim != 2:
        raise ValueError("Expected matrices for x, W1, W2")
    if x.shape[1] != W1.shape[0] or W1.shape[1] != W2.shape[0]:
        raise ValueError("Incompatible layer shapes")
    if b1.shape != (1, W1.shape[1]) or b2.shape != (1, W2.shape[1]):
        raise ValueError("Biases must have shape (1, units)")
    if not all(np.isfinite(a).all() for a in (x, W1, b1, W2, b2)):
        raise ValueError("Nonfinite input or parameter")
    z1 = x @ W1 + b1
    h = np.maximum(z1, 0)
    logits = h @ W2 + b2
    loss, probs, dlogits = cross_entropy(logits, targets)
    dh = dlogits @ W2.T
    dz1 = dh * (z1 > 0)  # Choose derivative 0 at the ReLU kink.
    grads = {
        "W1": x.T @ dz1,
        "b1": dz1.sum(axis=0, keepdims=True),
        "W2": h.T @ dlogits,
        "b2": dlogits.sum(axis=0, keepdims=True),
    }
    for key in params:
        if grads[key].shape != params[key].shape:
            raise AssertionError("Gradient shape mismatch: " + key)
    return loss, grads, {"z1": z1, "h": h, "logits": logits, "probs": probs}


def step(params, grads, learning_rate):
    return {key: value - learning_rate * grads[key] for key, value in params.items()}


if __name__ == "__main__":
    x, targets, params = fixture()
    loss, grads, cache = loss_and_grads(x, targets, params)
    next_loss = loss_and_grads(x, targets, step(params, grads, 0.1))[0]
    print(f"seed=42 dtype={x.dtype} batch=32, features=10, hidden=64, classes=5")
    print(f"mean CE: {loss:.12f} -> {next_loss:.12f}")
    for key, grad in grads.items():
        print(f"{key}: shape={grad.shape}, norm={np.linalg.norm(grad):.12f}")

5. Diagnose Numerical Failures

These observations are supported by this page’s executable tests and public documentation, not personal credentials or an unreported CUDA performance study.

• Bias broadcasting. Here b1 has shape (1,64), so db1 is dz1.sum(axis=0, keepdims=True). A parameter stored as (64,) can instead use a one-dimensional gradient. Match the parameter; keepdims=True is not mandatory in every implementation. The test rejects a mistakenly batch-shaped (32,64) bias.
• Finite differences. Compute (L(p+eps)-L(p-eps))/(2eps), changing and restoring one coordinate at a time. Error depends on step size, precision and crossing ReLU’s zero point. An error above 1e-4 does not by itself prove a bug. PyTorch gradcheck documentation explicitly warns about precision and nondifferentiable points.
• Stable loss. For logits [1000,-1000,0] and target 1, even max-shifted softmax underflows the target probability to zero. Taking its negative logarithm gives Inf; the logit-based implementation returns loss 2000 and gradient [1,-1,0]. This does not protect every intermediate operation against arbitrarily extreme finite inputs.
• Activation recomputation. Gradient checkpointing saves stored activations by recomputing them. Whether the memory/time trade-off helps needs actual measurements; it does not establish bandwidth as the bottleneck of all training.

6. Visualizing the Flow

When watching the animation, do not only watch arrow direction. Observe which forward values each node must store and reuse during the backward pass. The next article studies how these gradients move parameters and why optimizers take different paths to the minima.

7. Backpropagation Verification Matrix

When reproducing this article, split the forward and backward pass into auditable stages. The table below turns “is the formula correct?” into visible evidence, so a decreasing loss is not mistaken for proof that every gradient is correct.

8. Executed Per-Parameter Checks

This run uses Python 3.13.9, NumPy 2.3.5, CPU, float64 and seed=42. Initial mean CE is 1.647801573097. Using one gradient snapshot and learning rate 0.1 gives 1.628354091165 after updating. These random labels have no train/test split: this is not an accuracy or generalisation experiment.

All 1029 coordinates are checked, including biases. Each must satisfy abs(a-n) ≤ 1e-7 + 1e-5 × max(abs(a),abs(n)), with analytical gradient a and numerical gradient n. This tolerance belongs to this double-precision fixture, not every model.

The first two rows cross ReLU kinks and cannot serve as smooth local derivative checks. The last two have no crossings and pass all 1029/1029 coordinates. Display values are rounded; full precision and every coordinate remain in the downloadable CSV and JSON files.

The smallest |z1| is about 1.34039e-5, but a W1 coordinate probe moves z1 by x_j × eps, not necessarily eps, so even 1e-5 can cross a kink. Smaller steps increase sensitivity to subtraction rounding: the 1e-7 error here exceeds the 1e-6 error. Exactly at z=0, ReLU’s central difference is 0.5 while this implementation chooses subgradient 0; a fixed threshold cannot settle correctness at that nondifferentiable point.

To check that the audit detects actual mistakes, it deliberately reverses all gradient signs and separately omits division by batch size; both mutations are rejected. Repeating the same batch leaves the correct mean loss and gradients unchanged. Six invalid label, numerical or shape inputs are also rejected. Passing fixtures are concrete evidence, not a formal proof for arbitrary networks.

Download and reproduce

Download the backpropagation gradient-check lab. The reference/mlp-arrays.npz file contains inputs, targets, parameters, forward intermediates and analytical gradients. Four CSVs list all numerical probes; reference/audit.json records summaries and source hashes. The older series package retains its single-sample column-vector example; its numbers should not be confused with this 32-sample batch.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt
python mlp_example.py
python audit_gradients.py --output run-results

The audit ends with NEURAL_GRADIENT_AUDIT_PASSED. It needs no network, dataset or GPU during execution; dependency installation may contact a package index. BLAS, CPU or NumPy versions can change low-order digits, so compare with tolerances rather than demanding cross-machine byte-identical JSON. Framework autograd and CUDA parity were not run.

For a smaller hand-computed exercise, start with the five-parameter update and initialisation counterexamples.