The Chain Rule for Neural Networks

Consider a differentiable function $f: \mathbb{R}^n \to \mathbb{R}^m$. Its derivative at a point $\mathbf{x}$ is not a single number but an $m \times n$ matrix, the Jacobian, whose $(i,j)$ entry is $\partial f_i / \partial x_j$. This matrix is the best linear approximation to $f$ near $\mathbf{x}$: for a small perturbation $\Delta \mathbf{x}$,

Now suppose we compose two such functions, $h = g \circ f$, with $f: \mathbb{R}^n \to \mathbb{R}^m$ and $g: \mathbb{R}^m \to \mathbb{R}^p$. The multivariable chain rule states that the Jacobian of the composition is the matrix product of the Jacobians:

The order matters: $J_g$ is a $p \times m$ matrix and $J_f$ is $m \times n$, so the product is well-defined and gives a $p \times n$ matrix, exactly the shape a Jacobian of $h: \mathbb{R}^n \to \mathbb{R}^p$ should have. This is the exact generalization of the familiar scalar identity $\frac{dh}{dx} = \frac{dg}{df}\cdot\frac{df}{dx}$; the only change is that scalar multiplication becomes matrix multiplication, and the order of the factors is now fixed by matrix-shape compatibility rather than being commutative. A neural network is a long composition $L = \ell \circ f_K \circ f_{K-1} \circ \cdots \circ f_1$, where $f_k$ is the function computed by layer $k$ (an affine map followed by a nonlinearity, say) and $\ell$ is the loss function producing a single scalar. Applying the chain rule repeatedly,

Because $L$ is scalar-valued, $J_L$ is a row vector — this is exactly the gradient $\nabla_{\mathbf{x}} L$ (as a row) that we want for every parameter in the network. The key structural fact, used throughout this article, is that this long chain of matrix products can be evaluated in two very different orders — left-to-right ("forward mode") or right-to-left ("reverse mode") — and for networks with millions of parameters feeding into one scalar loss, right-to-left is dramatically cheaper. That asymmetry is the entire computational reason backpropagation is organized the way it is.

Setting up the composition. Model a network's forward computation as a sequence of intermediate variables:

Each $f_k$ takes the previous activation $\mathbf{z}_{k-1}$ and layer parameters $\boldsymbol{\theta}_k$ (weights, biases) and produces the next activation. We want $\partial L / \partial \boldsymbol{\theta}_k$ for every $k$, and often $\partial L/\partial \mathbf{x}$ as well.

Reverse accumulation. Define the upstream gradient at each stage as the row vector

By the chain rule applied one layer at a time,

where $J_{f_k}$ is the Jacobian of layer $k$'s output with respect to its input, evaluated at the activations recorded during the forward pass. Similarly, the gradient with respect to that layer's own parameters is

using the Jacobian of $f_k$ with respect to $\boldsymbol{\theta}_k$ instead. This recursive relation is initialized at $\bar{\mathbf{z}}_K = \partial L/\partial \mathbf{z}_K = J_\ell$ and run from $k=K$ down to $k=1$ — hence "reverse mode," "backward pass," or, in the neural network context, backpropagation. Crucially, layer $k$ never needs to construct or store the full Jacobian $J_{f_k}$ as a matrix. It only needs to compute the vector–Jacobian product $\bar{\mathbf{z}}_k J_{f_k}$ directly, which for common layers (linear, elementwise nonlinearity, convolution) has a closed form that costs about the same as the forward computation of that layer itself. ```

# Reverse-mode chain rule (as in backpropagation)

# Forward pass: compute and CACHE all intermediate activations z[0] = x for k = 1 to K: z[k] = f_k(z[k-1], theta[k]) # cache z[k] and any values needed

# to build the local Jacobian later L = loss(z[K])

# Backward pass: propagate the gradient from the scalar loss backward zbar[K] = d(loss)/d(z[K]) # local Jacobian of the loss, evaluated at z[K] for k = K downto 1:

# vector-Jacobian product: local Jacobian w.r.t. INPUT of layer k zbar[k-1] = zbar[k] @ J_input(f_k, z[k-1], theta[k])

# vector-Jacobian product: local Jacobian w.r.t. PARAMETERS of layer k grad_theta[k] = zbar[k] @ J_param(f_k, z[k-1], theta[k]) return grad_theta[1..K], zbar[0] # zbar[0] = dL/dx, if needed ``` Two design choices in this pseudocode deserve emphasis. First, the forward pass caches every intermediate $\mathbf{z}_k$: the backward pass needs these values to evaluate each local Jacobian, since $J_{f_k}$ is generally evaluated at the point $\mathbf{z}_{k-1}$ that was seen during the forward pass, not at some fixed reference point. Second, `J_input` and `J_param` are never materialized as dense matrices in practice — they are applied as functions (vector-Jacobian products), because for a linear layer with a $1000\times1000$ weight matrix, the Jacobian with respect to the input is that weight matrix, and for a layer with $10^6$ parameters, materializing the parameter-Jacobian explicitly would be enormous and wasteful.

Why reverse mode is asymptotically the right choice. Suppose the network has $K$ layers and, across all layers, $P$ scalar parameters in total, with $n$ the input dimension. There are two natural ways to accumulate the chain of Jacobian products $J_\ell J_{f_K} \cdots J_{f_1}$:

• Forward-mode accumulation multiplies the chain left-to-right starting from the input side, propagating a Jacobian (a full matrix, not just a vector) forward through the network for each input direction or parameter direction you care about. Computing $\partial L/\partial \theta_k$ this way for one parameter costs roughly the cost of one forward pass; doing this independently for every one of the $P$ parameters costs $O(P \cdot C)$, where $C$ is the cost of one forward pass.
• Reverse-mode accumulation multiplies the chain right-to-left, starting from the scalar loss. Because $L$ is scalar, every intermediate quantity $\bar{\mathbf{z}}_k$ is a vector, not a matrix (a $1\times m$ row rather than an $m\times m$ or larger object). Each step is one vector–Jacobian product, costing about the same as the corresponding forward step. A single backward pass over all $K$ layers therefore costs $O(C)$ total and yields the gradient with respect to every parameter simultaneously.

This is the Baur–Strassen-type observation underlying all reverse-mode automatic differentiation: computing the gradient of a scalar function of $n$ (or $P$) variables costs a small constant multiple of the cost of evaluating the function itself, independent of $n$ or $P$. For neural networks with $P$ in the millions or billions, the difference between $O(P \cdot C)$ and $O(C)$ is the difference between infeasible and routine.

Space complexity. The tradeoff for this time saving is memory: reverse mode requires caching every intermediate activation $\mathbf{z}_k$ (and any auxiliary values a layer's Jacobian needs) computed during the forward pass, because the backward pass needs them to evaluate local Jacobians at the correct point. Memory cost is $O(K \cdot \bar{m})$ where $\bar{m}$ is average activation width — linear in depth. This is precisely why very deep or memory-constrained networks use gradient checkpointing (recomputing some activations during the backward pass instead of storing all of them), trading extra compute for lower memory.

Correctness argument. The correctness of the whole scheme rests on two facts, both provable from the definition of the derivative as a local linear approximation: 1. Chain rule for two functions. If $f$ is differentiable at $\mathbf{x}$ and $g$ is differentiable at $f(\mathbf{x})$, then $g\circ f$ is differentiable at $\mathbf{x}$ with $J_{g\circ f}(\mathbf{x}) = J_g(f(\mathbf{x}))J_f(\mathbf{x})$. This follows because the composition of two best-linear-approximations is itself a best linear approximation to the composed function: writing $f(\mathbf{x}+\Delta) = f(\mathbf{x}) + J_f\Delta + o(\|\Delta\|)$ and substituting into the analogous expansion of $g$ shows the first-order term is exactly $J_g J_f \Delta$, with all higher-order terms absorbed into the $o(\|\Delta\|)$ remainder. 2. Induction over the composition length. Applying fact (1) repeatedly to the $K{+}1$-fold composition $\ell\circ f_K\circ\cdots\circ f_1$ gives $J_L = J_\ell J_{f_K}\cdots J_{f_1}$ by induction on $K$. The reverse recursion $\bar{\mathbf{z}}_{k-1} = \bar{\mathbf{z}}_k J_{f_k}$ is exactly this product evaluated right-to-left, one factor at a time, which is legal because matrix multiplication is associative — the value of the product does not depend on the order in which we group the multiplications, only the cost of computing it does. So reverse-mode backpropagation is not an approximation or a heuristic: it computes the exact analytic gradient (up to floating-point rounding), for the same reason $2\times 3\times 4$ can be computed as $(2\times3)\times4$ or $2\times(3\times4)$ and get the same answer either way, just possibly at different cost.

Honest caveats. The chain rule requires each $f_k$ to be differentiable at the relevant point; nonsmooth functions like ReLU or max-pooling have a measure-zero set of non-differentiable points, handled in practice by choosing a subgradient (e.g., defining ReLU's derivative to be $0$ at $x=0$) — this is technically not a chain rule but a chain rule for subgradients, and it works empirically but is not covered by the classical smooth chain-rule proof. Also, reverse mode is efficient specifically because the output is low-dimensional (scalar loss); if instead one needed the full Jacobian of a high-dimensional output with respect to a low-dimensional input, forward mode can be cheaper — the choice of mode is a genuine engineering tradeoff, not a universal rule.

Take a tiny two-layer network with scalar output, so every matrix is small enough to write out fully. Let the input be $\mathbf{x} = [1,\,2]^\top$, and define:

with concrete numbers

Forward pass.

Backward pass. Work from $L$ back to $\mathbf{x}$, one vector–Jacobian product at a time. Step 1 — loss w.r.t. $y$: $\bar{y} = \dfrac{\partial L}{\partial y} = 2(y-t) = 2(1-5) = -8.$ Step 2 — through $y = \mathbf{w}_2^\top \mathbf{a}_1$. The Jacobian of $y$ w.r.t. $\mathbf{a}_1$ is the row vector $\mathbf{w}_2^\top = [1,\,1]$, and w.r.t. $\mathbf{w}_2$ it is $\mathbf{a}_1^\top = [0,\,1]$:

Step 3 — through ReLU. Its Jacobian is diagonal with entries $\mathbb{1}[z_{1,i}>0]$: since $z_1 = [-1, 1]^\top$, the diagonal is $\mathrm{diag}(0, 1)$. Elementwise multiply:

Step 4 — through $\mathbf{z}_1 = W_1\mathbf{x}+\mathbf{b}_1$. The Jacobian w.r.t. $\mathbf{x}$ is $W_1$, w.r.t. $\mathbf{b}_1$ is the identity, and w.r.t. each row of $W_1$ pairs with $\mathbf{x}^\top$:

The table below traces every quantity, forward then backward, so the two passes can be checked against each other:

One can sanity-check Step 3 with a finite-difference estimate: perturbing $x_1$ by $\epsilon=10^{-4}$ and recomputing $L$ numerically gives $\Delta L/\epsilon \approx -16.0$, matching $\partial L/\partial x_1 = -16$ computed above — confirming the chain of local Jacobian products reproduces the true derivative, not merely a plausible-looking approximation.

!A mathematician's contribution to differential calculus underlies the chain rule used throughout deep learning

Every deep learning framework — PyTorch's `autograd`, TensorFlow's `GradientTape`, JAX's `grad` — is a generic engine for exactly the reverse-mode chain-rule recursion described above, applied to an arbitrary computational graph rather than a fixed layer stack. Training image classifiers, language models, diffusion models, and reinforcement learning policies all reduce, at the lowest level, to constructing a forward computational graph and running this backward recursion to get gradients for gradient descent. Compiler-level optimizations in these frameworks (operator fusion, memory planning, gradient checkpointing) are all optimizations of this chain-rule computation, not alternatives to it. Understanding the chain rule precisely is also what lets practitioners diagnose vanishing or exploding gradients — a direct consequence of repeatedly multiplying many Jacobians whose eigenvalues are consistently below or above one — and motivates architectural fixes like residual connections and normalization layers, which are explicitly designed to keep the product of local Jacobians well-conditioned across many layers.

Ask the tutor: Can you walk me through why reverse-mode accumulation is cheaper than forward-mode when there are many parameters but only one output?