Deriving the Gradients

Assumes you know from ML-101

You have completed ML-101 lesson 14 — Neural Networks & Backprop, where the backprop algorithm was presented as a four-step recipe: forward pass, compute loss, backward pass, update weights. You have also completed ML-201 lesson 2 (The Math You Actually Need), where the matrix-as-function picture and the chain rule in matrix form were established. This lesson does not re-teach what a forward pass is or what a weight update looks like. It derives the backward pass from the chain rule of calculus, applied to the computational graph of a two-layer network. Every gradient in this lesson is one matrix product along one path in the graph, and every shape rule from lesson 2 is enforced by construction.

Learning Objectives

1. Write the computational graph of a two-layer fully-connected network as a sequence of differentiable maps, and apply the matrix chain rule at every edge to obtain the gradient of the loss with respect to every parameter.
2. State the Jacobian shape rules from lesson 2 explicitly for each parameter (input weights, hidden weights, biases, input), and predict the shape of every gradient tensor before computing it.
3. Derive the gradient of the cross-entropy loss through a softmax layer, and show that the simplified expression $\partial L / \partial z = \hat{p} - y$ is a special case of the general chain rule, not a hand-waved shortcut.
4. Implement a finite-difference gradient checker for an arbitrary computation graph, and use it to debug a backprop implementation by comparing numerical and analytical gradients to within $\sim 10^{-5}$.
5. Diagnose the failure mode of a backprop implementation that returns NaN, by identifying which gradient is large enough to overflow float32 and which shape rule has been violated.

The computational graph, written down

A computational graph is a directed acyclic graph whose nodes are intermediate values and whose edges are elementary operations. Every parameterised model that is assembled from differentiable building blocks has a computational graph, and every gradient in the model is the matrix chain rule applied to that graph. There is no other source of gradients; backpropagation is the chain rule, not an analogy.

For a two-layer fully-connected network with one hidden layer of width $h$, the graph has nodes for the input $x \in \mathbb{R}^m$, the first linear pre-activation $z^{(1)} = W^{(1)} x + b^{(1)} \in \mathbb{R}^h$, the hidden activation $a^{(1)} = \sigma(z^{(1)}) \in \mathbb{R}^h$, the second linear pre-activation $z^{(2)} = W^{(2)} a^{(1)} + b^{(2)} \in \mathbb{R}^{n}$, and the prediction $\hat{y} = \mathrm{softmax}(z^{(2)}) \in \mathbb{R}^{n}$. The loss $L(y, \hat{y})$ is computed from the prediction and the label. The graph is

There are two parameter nodes, $W^{(1)} \in \mathbb{R}^{h \times m}$ and $W^{(2)} \in \mathbb{R}^{n \times h}$, and two bias nodes, $b^{(1)} \in \mathbb{R}^h$ and $b^{(2)} \in \mathbb{R}^n$. The chain rule, applied along the path from $L$ back to each parameter, gives the gradient with respect to that parameter. There is no other computation in backpropagation.

The deeper point is that the graph is not a list of operations in code; it is a mathematical object — the composition $L \circ \mathrm{softmax} \circ W^{(2)} \cdot \sigma \circ W^{(1)} \cdot \mathrm{affine}(x)$, with the affine maps absorbing the biases. Every backprop derivation in ML-201 starts from this graph and applies the chain rule. The shape of every intermediate gradient tensor is determined by the chain rule and the shape of every intermediate forward tensor, with no further freedom.

Forward pass as composition of maps

The forward pass is the composition

and the parameters of the composition are the four tensors $\Theta = \{W^{(1)}, b^{(1)}, W^{(2)}, b^{(2)}\}$. The composition is differentiable at every intermediate value provided the activation $\sigma$ is differentiable at the evaluation point. ReLU evaluated exactly at zero is a non-differentiable point, and the chain rule has nothing to multiply there; this is the single assumption whose violation breaks the entire derivation in this lesson.

import numpy as np

def sigmoid(z):
    # Stable sigmoid; exp(z) overflows for large positive z in float32.
    out = np.empty_like(z, dtype=float)
    pos = z >= 0
    out[pos] = 1.0 / (1.0 + np.exp(-z[pos]))
    ex = np.exp(z[~pos])
    out[~pos] = ex / (1.0 + ex)
    return out

def softmax(z):
    z = z - z.max(axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)

def forward(x, W1, b1, W2, b2):
    """Two-layer network forward pass.

    Shapes:
      x  : (m,)
      W1 : (h, m), b1 : (h,)
      W2 : (n, h), b2 : (n,)
    Returns dict of intermediates so the backward pass can reuse them.
    """
    z1 = W1 @ x + b1                # (h,)
    a1 = sigmoid(z1)                # (h,)
    z2 = W2 @ a1 + b2               # (n,)
    p  = softmax(z2)                # (n,)
    return {"z1": z1, "a1": a1, "z2": z2, "p": p}

Notice that the shapes are pinned: $x \in \mathbb{R}^m$, $W^{(1)} \in \mathbb{R}^{h \times m}$, $b^{(1)} \in \mathbb{R}^h$, $W^{(2)} \in \mathbb{R}^{n \times h}$, $b^{(2)} \in \mathbb{R}^n$. Every shape mismatch is a bug; the 201-level reading is that the shape is constrained by the chain rule, not a matter of convention.

Backward pass as chain rule on the graph

The backward pass is the chain rule applied to the graph in reverse topological order. Start from $\partial L / \partial \hat{y}$, compute $\partial L / \partial z^{(2)}$ by the softmax chain rule, then $\partial L / \partial W^{(2)}$ and $\partial L / \partial b^{(2)}$ by the affine map, then $\partial L / \partial a^{(1)}$ by the matrix chain rule, then $\partial L / \partial z^{(1)}$ by the activation chain rule, then $\partial L / \partial W^{(1)}$ and $\partial L / \partial b^{(1)}$.

For a generic affine map $z = W a + b$ with $a \in \mathbb{R}^{h}$, $W \in \mathbb{R}^{n \times h}$, and upstream gradient $\partial L / \partial z \in \mathbb{R}^n$, the chain rule gives

These three identities are the only backprop step that a linear layer contributes. The shape of $\partial L / \partial W$ is identical to the shape of $W$, which is the rule from lesson 2 that every backprop implementation honours whether it knows it or not.

For the activation $\sigma$, with $a = \sigma(z)$ and $\partial L / \partial a \in \mathbb{R}^h$, the chain rule gives

where $\odot$ is Hadamard (elementwise) product. The Jacobian of $\sigma$ is diagonal, with $\sigma'(z)$ on the diagonal, and the chain rule collapses the matrix product into a Hadamard product. For sigmoid, $\sigma'(z) = \sigma(z) \odot (1 - \sigma(z))$, which simplifies the multiplication further. For ReLU, $\sigma'(z) = \mathbb{1}[z > 0]$; at $z = 0$ the derivative does not exist and the sub-gradient is taken as either $0$ or $1$ by convention.

def cross_entropy(p, y):
    # y is one-hot; p is softmax output. The clip keeps log finite.
    p = np.clip(p, 1e-12, 1.0)
    return float(-(y * np.log(p)).sum())

def backward(cache, y, W2):
    """Reverse-mode autodiff through the two-layer network.

    Every intermediate gradient has its shape pinned by the chain rule;
    if a shape is wrong, the derivation is wrong.
    """
    p   = cache["p"]
    z1  = cache["z1"]
    a1  = cache["a1"]

    # dL/dp : gradient of cross-entropy w.r.t. softmax output.
    p_clip = np.clip(p, 1e-12, 1.0)
    dL_dp = -(y / p_clip)                       # shape (n,)

    # Softmax Jacobian: J_softmax(z) = diag(p) - p p^T. dL/dz = (p - y).
    dL_dz2 = (p - y)                            # shape (n,) — the famous shortcut
    dL_dW2 = dL_dz2[:, None] @ a1[None, :]      # (n, 1) @ (1, h) -> (n, h)
    dL_db2 = dL_dz2                             # (n,)
    dL_da1 = W2.T @ dL_dz2                      # (h,)

    # Sigmoid derivative; the chain rule through sigma is a Hadamard product.
    sigp  = a1 * (1.0 - a1)                     # (h,)
    dL_dz1 = dL_da1 * sigp                      # (h,)
    dL_dW1 = dL_dz1[:, None] @ cache.get("x",
                  np.zeros(W2.shape[1]))[None]  # stands in for x; see below
    dL_db1 = dL_dz1                             # (h,)

    return {"dL_dW2": dL_dW2, "dL_db2": dL_db2,
            "dL_dW1": dL_dW1, "dL_db1": dL_db1,
            "dL_dz2": dL_dz2}

The softmax + cross-entropy shortcut, derived

The simplification $dL / dz^{(2)} = \hat{p} - y$ is the single most backprop'd identity in deep learning. It is usually presented as "well-known" or asserted as a hand-wave. It is in fact a special case of the chain rule applied to the composition $\mathrm{softmax} \circ \mathrm{cross\_entropy}$.

For $L = -\sum_k y_k \log p_k$ and $p_k = \exp(z_k) / \sum_j \exp(z_j)$, the Jacobian of the softmax is

which is the $K \times K$ matrix $J_{\mathrm{softmax}}(z) = \mathrm{diag}(p) - p p^\top$. The chain rule gives

where the second equality uses $\sum_k y_k = 1$ for a one-hot label. The shortcut $dL / dz^{(2)} = \hat{p} - y$ is the chain rule on the softmax–cross-entropy pair, with the particular Jacobian of the softmax collapsing the sum to two terms. It is not a magic property of the loss; it is a property of the algebra $\sum_k y_k = 1$.

The deeper point is that the shortcut holds only because the label is one-hot. For a soft label — a probability vector with multiple non-zero entries — the result generalises to $dL / dz^{(2)} = \hat{p} - y$ but with $\sum_k y_k$ not equal to 1; the chain rule still gives the answer, but the intermediate algebra is longer. The one-hot assumption is the place where the result is simple, and the 201-level reading is that the assumption is not optional — it is a property of the label representation, not of the loss.

Shape rules for every parameter

The shapes of the gradients for the two-layer network are pinned by the chain rule and the forward shapes. The full table is:

Each row's "Gradient shape" column equals the parameter's shape, which is the general rule from lesson 2: the gradient of a scalar with respect to a tensor has the same shape as the tensor. The "Gradient expression" column is the chain rule applied along the unique path from the parameter to the loss. The five shape rules from lesson 2 are all the shape rules the backprop of a two-layer network needs.

A learner who can read this table and reproduce the expressions without memorising the expressions has the right mental model: each expression is the matrix chain rule along one path of the graph, with the shape check built in. A learner who has memorised the expressions but cannot reproduce them from the graph has memorised a recipe, not a derivation; their implementation breaks the moment the architecture changes.

The deeper point is that every backprop derivation in deep learning is the same derivation. The gradient through attention is the chain rule through a factorisation of the matrix product. The gradient through a convolution is the chain rule through an im2col reshape followed by a matrix product. The gradient through a recurrent cell is the chain rule through a graph with cycles, applied via the backpropagation-through-time unfolding. The "advanced" gradients are not advanced; they are the chain rule, with the graph's topology changing but the rule unchanged. This is the sense in which deriving the gradients is the title of the lesson: every gradient is a derivation, and every derivation is the chain rule.

Gradient checking, the right way

A gradient checker compares an analytical gradient from backprop to a numerical estimate from finite differences, and reports the relative error. The numerical estimate uses central differences:

with $h$ small enough that higher-order terms are negligible but large enough that roundoff does not dominate. A typical choice is $h = 10^{-5}$ for float64. The relative error between the analytical and numerical estimates should be below $\sim 10^{-5}$ for a clean implementation.

import numpy as np

def numerical_gradient(f, theta, h=1e-5):
    """Central-difference estimate of grad f at theta."""
    grad = np.zeros_like(theta)
    for i in np.ndindex(*theta.shape):
        e = np.zeros_like(theta); e[i] = h
        grad[i] = (f(theta + e) - f(theta - e)) / (2 * h)
    return grad

def relerr(a, b):
    """Relative error, robust to small denominators."""
    denom = np.maximum(np.abs(a) + np.abs(b), 1e-12)
    return float(np.abs(a - b).max() / denom.max())

The 201-level reading is that gradient checking is not a unit test for backprop; it is a diagnostic protocol. The procedure is:

1. Pick a small subset of parameters (a single weight matrix, not the whole network) and check that layer first. The error in any one layer propagates to every later layer's gradient, so isolating the layer closest to the loss is the most informative choice.
2. Compare on a single example, not a minibatch. The minibatch gradient is an average of per-example gradients, and per-example noise obscures the deterministic error.
3. Disable dropout, batch normalisation, and any stochastic regulariser before checking. The numerical estimate is for a deterministic forward; the analytical estimate with stochastic regularisation is for a stochastic forward; they are not comparable.
4. Use float64 throughout the check. Float32 roundoff in the central difference makes $h$ below $10^{-5}$ unreliable.

A common 101-level mistake is to compare the gradients on a minibatch with dropout active, see "non-zero relative error", and conclude that backprop is wrong. The 201-level reading is that the comparison protocol itself was wrong; a clean protocol on a single example with deterministic forward gives a relative error of $10^{-7}$ or smaller for a correct backprop implementation.

The deeper point is that gradient checking cannot certify correctness absolutely. The central-difference estimate is exact up to $O(h^2)$ relative to the true gradient; an analytical gradient that is wrong by a constant factor of $1 + 10^{-3}$ will pass a check with $h = 10^{-5}$ if the constant factor's derivative is also wrong by the same factor. Gradient checking catches the blatant errors — wrong sign, wrong shape, missing terms — and not the subtle ones. The practitioner who treats the gradient checker as a proof of correctness has misunderstood what finite differences measure.

Key Takeaways

• A computational graph is the composition of differentiable maps. Every gradient in a neural network is the matrix chain rule applied to a path in that graph.
• For a linear layer $z = W a + b$, the backprop identities $dL/dW = dL/dz \cdot a^\top$, $dL/db = dL/dz$, and $dL/da = W^\top \cdot dL/dz$ are the only backprop step that layer contributes; everything else is composition.
• The shortcut $dL/dz = \hat{p} - y$ is the chain rule on softmax and cross-entropy, with the one-hot label collapsing the algebra. It is a derivation, not a hand-wave.
• Gradient checking is a diagnostic protocol, not a unit test. Use central differences on a single example with deterministic forward in float64, and isolate one layer at a time.
• Every "advanced" gradient in deep learning is the chain rule on a graph with a different topology. The rule never changes; the graph does.