The chain rule: products of Jacobians and computational graphs

Before this lesson: lesson 01 (the 1-D chain rule and linear approximation), lesson 02 (partial derivatives, the gradient), lesson 03 (the derivative as a linear map, the Jacobian and its shape), and linear algebra lesson 01 (the product $AB$ is the map "first $B$, then $A$"). Time: about 70 minutes, plus the notebook.

---

Why this matters #

Backpropagation is the chain rule, applied in a clever order. Every deep-learning framework (PyTorch, JAX, TensorFlow) stores your model as a computational graph and runs the chain rule on it. When you understand the chain rule as a product of Jacobians, you can:

• derive the gradient of any new layer, loss or custom operation by hand (and check it);
• see why a residual connection $y = x + F(x)$ helps the gradient reach early layers: the two paths add, and the identity path passes the gradient unchanged;
• see why a weight used in many places (an RNN cell, tied input and output embeddings) gets the sum of the gradients of all its uses;
• understand the adjoint models of science (weather data assimilation, seismic imaging): they are the same backward sweep, run through a simulator.

1. Predict first #

2. The picture #

One variable first: rates multiply. Take $y = (3x+1)^2$. Write it as two steps: $u = 3x+1$ and $y = u^2$. At $x = 1$: $u = 4$, $du/dx = 3$, and $dy/du = 2u = 8$. A small step $\Delta x$ gives $\Delta u \approx 3\,\Delta x$, and this step in $u$ gives $\Delta y \approx 8\,\Delta u$. So

Check: $y = 9x^2 + 6x + 1$ gives $dy/dx = 18x + 6 = 24$ at $x = 1$.

Interactive figure: loading…

Many variables: linear maps compose. Near a point $x$, a differentiable map $f$ acts on a small step $h$ like its Jacobian (lesson 03): $f(x + h) \approx f(x) + Df(x)\,h$. Near $u = f(x)$, the map $g$ acts like $Dg(u)$. So the step $h$ goes through two matrices:

Composing two linear maps multiplies their matrices (linear algebra lesson 01). The map applied first stands on the right. The figure shows this for two maps of the plane.

Interactive figure: loading…

A program is a graph. A formula like $f(x_1, x_2) = \ln x_1 + x_1x_2 - \sin x_2$ is computed in small steps. Give each step a name: $v_1 = \ln x_1$, $v_2 = x_1x_2$, $v_3 = \sin x_2$, $v_4 = v_1 + v_2$, $v_5 = v_4 - v_3$. Draw one node per step and an arrow from each value to the steps that use it. This is the computational graph of $f$ (the example is from Baydin et al. 2018).

To get the derivatives, sweep the graph twice. The forward run computes all values from the inputs to the output. The backward run starts at the output with the number 1 and walks back. At each node it multiplies by the local derivative on the edge (the derivative of one small step). Where a value feeds two nodes, the two contributions add.

Interactive figure: loading…

3. Definitions and main results #

We use the notation of lesson 03. Vectors are columns. For $f:\R^n\to\R^m$, the Jacobian $Df(x)\in\R^{m\times n}$ has entries $[Df(x)]_{ij} = \partial f_i/\partial x_j$ (outputs × inputs). $f$ is differentiable at $x$ if

In words: the error $r(h)$ of the linear approximation is small even compared with the step $h$. We write this as $r(h) = o(\|h\|)$.

Theorem (chain rule) #

Let $f:\R^n\to\R^m$ be differentiable at $x$, and $g:\R^m\to\R^p$ differentiable at $u = f(x)$. Then $g\circ f$ is differentiable at $x$, and

Why it is true. Write $A = Df(x)$ and $B = Dg(u)$. A step $h$ moves $f$ by $k = Ah + r_f(h)$. This step $k$ moves $g$ by $Bk + r_g(k) = BAh + \big(B\,r_f(h) + r_g(k)\big)$. The main part is $BAh$, a linear function of $h$. The rest is a sum of two errors, and both are $o(\|h\|)$: $B\,r_f(h)$ because $B$ is a fixed matrix, and $r_g(k)$ because $k$ is at most a constant times $\|h\|$. So $BA$ is the derivative. The full proof makes "the rest is small" precise.

Full proof (six steps)

Let $\|\cdot\|$ be the Euclidean norm, and for a matrix the induced operator norm $\|A\| = \max_{h\ne0}\|Ah\|/\|h\|$ (linear algebra lesson 11), so that $\|Ah\| \le \|A\|\,\|h\|$.

1. By the definition of the derivative, $f(x+h) = u + Ah + r_f(h)$ with $\|r_f(h)\|/\|h\| \to 0$, and $g(u+k) = g(u) + Bk + r_g(k)$ with $\|r_g(k)\|/\|k\| \to 0$.
2. Set $\rho(k) = \|r_g(k)\|/\|k\|$ for $k \ne 0$ and $\rho(0) = 0$. Then $\|r_g(k)\| = \rho(k)\,\|k\|$ for every $k$ (also for $k = 0$, because $r_g(0) = 0$), and $\rho(k) \to 0$ as $k \to 0$. This trick avoids dividing by $\|k\|$, which can be 0 even when $h \ne 0$.
3. Put $k(h) = f(x+h) - u = Ah + r_f(h)$. Then $g(f(x+h)) = g(u) + BAh + R(h)$ with $R(h) = B\,r_f(h) + r_g(k(h))$.
4. First part: $\|B\,r_f(h)\| \le \|B\|\,\|r_f(h)\|$, so $\|B\,r_f(h)\|/\|h\| \to 0$.
5. Second part: for small $h$, $\|r_f(h)\| \le \|h\|$, so $\|k(h)\| \le (\|A\| + 1)\|h\|$. Hence $k(h) \to 0$ and $\|r_g(k(h))\|/\|h\| = \rho(k(h))\,\|k(h)\|/\|h\| \le \rho(k(h))\,(\|A\| + 1) \to 0$.
6. So $R(h) = o(\|h\|)$. By the definition (and the uniqueness) of the derivative, $D(g\circ f)(x) = BA$. $\blacksquare$

Shape check. Always check $(p\times m)(m\times n)$: the inner sizes must agree, and the result has the outer sizes. If the shapes do not fit, the order is wrong. If the shapes fit by accident (square matrices), the order can still be wrong: matrix products do not commute.

Longer chains. For $F = f_L\circ\cdots\circ f_1$ with $x_0 = x$, $x_k = f_k(x_{k-1})$ and $J_k = Df_k(x_{k-1})$, induction on $L$ gives $DF(x) = J_L\,J_{L-1}\cdots J_1$. Each $J_k$ is evaluated at the input of its own step.

The entry form #

Read the product $Dg\,Df$ entry by entry. Entry $(i, j)$ is row $i$ of $Dg$ times column $j$ of $Df$:

In words: sum over every intermediate variable $u_k$ through which $x_j$ can change $g_i$. Each term is the product of the rates along one path $x_j \to u_k \to g_i$.

The case most people learn first. Let $f(x, y)$ be a function of two variables, and let $x(t)$, $y(t)$ be functions of time. Then $t \mapsto (x(t), y(t))$ has the $2\times1$ Jacobian $(x'(t), y'(t))^\top$, and $f$ has the $1\times 2$ Jacobian $(\partial_x f, \partial_y f)$. Their product is

Example. $f(x, y) = x^2y$ with $x(t) = \cos t$, $y(t) = \sin t$, at $t = 0$: $x = 1$, $y = 0$, $x' = -\sin 0 = 0$, $y' = \cos 0 = 1$. So $\frac{d}{dt}f = 2xy\cdot 0 + x^2\cdot 1 = 1$. Check: $f(x(t), y(t)) = \cos^2 t\,\sin t$ has derivative $-2\cos t\sin^2 t + \cos^3 t$, which is 1 at $t = 0$.

A gradient through a map is a vector-Jacobian product #

In ML the last function is almost always a scalar loss $\ell:\R^m\to\R$. Its Jacobian is the row $\nabla\ell^\top$ (lesson 02). The chain rule gives the $1\times n$ row $\nabla\ell(f(x))^\top\,Df(x)$. Transpose it to get a column:

This is a vector-Jacobian product (VJP): a vector times the Jacobian, written as $Df^\top u$ (or $u^\top Df$ as a row). We write the gradient of the loss with respect to any intermediate value $u$ as $\bar u = \partial\ell/\partial u$. It has the same shape as $u$ and is called the adjoint of $u$. The formula above says:

The adjoint of the output, multiplied by the transposed local Jacobian, gives the adjoint of the input. Three VJPs appear in every network.

(a) Linear map, $z = Wx$ with $W\in\R^{m\times n}$. The Jacobian with respect to $x$ is $W$ itself, so $\bar x = W^\top \bar z$.

(b) Elementwise map, $a = \varphi(z)$ (for example $\tanh$ or ReLU applied to each entry). The Jacobian is $\operatorname{diag}(\varphi'(z))$ (lesson 03), a diagonal matrix, so $\bar z = \varphi'(z)\odot\bar a$. Here $\odot$ is the elementwise product. Never build the diagonal matrix in code: it has $m^2$ entries, and only $m$ of them are nonzero.

(c) The weights, $z = Wx$ as a function of $W$. Now $W$ is the variable and $x$ is fixed.

Proposition. If $z = Wx$ and $\bar z = \partial\ell/\partial z$, then $\bar W = \partial\ell/\partial W = \bar z\,x^\top$, an $m\times n$ matrix (the shape of $W$).

Proof.

1. $z_k = \sum_j W_{kj}x_j$, so $\partial z_k/\partial W_{ij} = x_j$ if $k = i$, and 0 otherwise.
2. By the entry form of the chain rule, $\partial\ell/\partial W_{ij} = \sum_k \bar z_k\,\partial z_k/\partial W_{ij} = \bar z_i\,x_j$.
3. The matrix with entries $\bar z_i x_j$ is the outer product $\bar z\,x^\top$. $\blacksquare$

An outer product of two vectors has rank at most 1. So for one sample, the gradient of a linear layer is a rank-1 matrix. Lesson 05 returns to this fact.

Fan-out: the contributions of all paths add #

A value fans out when it is used by more than one later step. For example, $x$ feeds two maps, $u = f(x)$ and $w = g(x)$, and the loss is $\ell = F(u, w)$.

Theorem (paths add). With $\bar u = \nabla_u F$ and $\bar w = \nabla_w F$,

Proof.

1. Stack the two maps into one: $s(x) = (f(x), g(x))$. Its Jacobian is the stacked block matrix $Ds = \begin{pmatrix} Df \\ Dg\end{pmatrix}$, because the rows of $Ds$ are the rows of $Df$ and then the rows of $Dg$.
2. $\ell = F(s(x))$, so the VJP rule gives $\bar x = Ds^\top\,\nabla F = \begin{pmatrix} Df^\top & Dg^\top\end{pmatrix}\begin{pmatrix}\bar u\\ \bar w\end{pmatrix}$.
3. Multiply the blocks: $\bar x = Df^\top\bar u + Dg^\top\bar w$. $\blacksquare$

So the gradient through each path is computed as if the other path did not exist, and then the results are added. The same is true for any number of paths.

Interactive figure: loading…

Three places where fan-out matters:

• Residual connections. $y = x + F(x)$. The identity path and the $F$ path add: $\bar x = \bar y + DF(x)^\top\bar y$. Even if $DF^\top\bar y$ is tiny, the gradient $\bar y$ still reaches $x$ through the identity path.
• Shared weights. A parameter used in several places (the same RNN weight at every time step, tied input and output embeddings) gets the sum of the gradients of all its uses.
• A variable used twice in a formula. In $y = x^2\sin x$, the product rule $2x\sin x + x^2\cos x$ is exactly "path through $x^2$" plus "path through $\sin x$".

Computational graphs #

Definitions.

• A primitive operation is a small function whose derivative we know in closed form: $+$, $-$, $\times$, $/$, $\exp$, $\ln$, $\sin$, $\tanh$, a matrix product, a sum over an axis.
• A program computes intermediate variables $v_i$, one per primitive operation. We number the $n$ inputs $v_{1-n}, \dots, v_0$ and the later variables $v_1, v_2, \dots$ in the order they are computed. This order is a topological order: every $v_i$ uses only variables with smaller index.
• The computational graph has one node per variable and an edge $j \to i$ when the operation that makes $v_i$ uses $v_j$ directly. Then $j$ is a parent of $i$, and $i$ is a child of $j$. The graph has no cycles: it is a directed acyclic graph (DAG).
• The local derivative on edge $j\to i$ is $\partial v_i/\partial v_j$: the derivative of the one operation that makes $v_i$, with respect to one of its arguments, the others fixed.

The example. For $f(x_1, x_2) = \ln x_1 + x_1x_2 - \sin x_2$:

variable	operation	parents	local derivatives
$v_{-1} = x_1$	input
$v_0 = x_2$	input
$v_1 = \ln v_{-1}$	$\ln$	$v_{-1}$	$\partial v_1/\partial v_{-1} = 1/v_{-1}$
$v_2 = v_{-1}\,v_0$	$\times$	$v_{-1}, v_0$	$\partial v_2/\partial v_{-1} = v_0$, $\ \partial v_2/\partial v_0 = v_{-1}$
$v_3 = \sin v_0$	$\sin$	$v_0$	$\partial v_3/\partial v_0 = \cos v_0$
$v_4 = v_1 + v_2$	$+$	$v_1, v_2$	$1$, $1$
$v_5 = v_4 - v_3$	$-$	$v_4, v_3$	$1$, $-1$

$x_1$ has two children ($v_1$ and $v_2$) and $x_2$ has two children ($v_2$ and $v_3$): both fan out.

Adjoints and tangents. Let $y$ be the output. For every node:

• the adjoint $\bar v_j = \partial y/\partial v_j$ says how much $y$ changes when $v_j$ changes and everything computed from $v_j$ is computed again;
• the tangent $\dot v_i = \frac{d}{dt}v_i(x + t\dot x)\big|_{t=0}$ says how fast $v_i$ changes when the inputs move in a chosen direction $\dot x$.

Theorem (the two sweeps).

1. Backward sweep. Set $\bar y = 1$. Then, for $j$ in reverse topological order,

At the end, $\bar v_{1-n}, \dots, \bar v_0$ are the partial derivatives $\partial y/\partial x_1, \dots, \partial y/\partial x_n$. 2. Forward sweep. Set the input tangents to $\dot x$. Then, for $i$ in topological order,

At the end, $\dot y = \nabla y^\top\dot x$, the directional derivative of $y$ along $\dot x$.

Proof of the backward sweep.

1. $v_j$ enters the rest of the program only through its children. Give each use its own copy: child $i$ uses a copy $c_i$, and all copies equal $v_j$. Then $y$ is a function $H$ of the copies (all other earlier variables stay fixed), and $y(v_j) = H(v_j, \dots, v_j)$.
2. By the fan-out rule (or the entry form with one input and several paths), $\partial y/\partial v_j = \sum_i \partial H/\partial c_i$.
3. The copy $c_i$ changes $y$ only through $v_i$. By the 1-D chain rule, $\partial H/\partial c_i = (\partial y/\partial v_i)\,(\partial v_i/\partial c_i) = \bar v_i\,\partial v_i/\partial v_j$.
4. Children come later in the topological order. In reverse order, all $\bar v_i$ of the children are known when we reach $j$. Induction from the output (where $\bar y = 1$) to the inputs gives the claim. $\blacksquare$

Proof of the forward sweep. Move the inputs along $x(t) = x + t\dot x$. Each $v_i$ is a function of its parents, so the entry form of the chain rule gives $\frac{d}{dt}v_i = \sum_j \frac{\partial v_i}{\partial v_j}\frac{d}{dt}v_j$. The parents come earlier, so their tangents are known. At the output, the chain rule gives $\dot y = Dy\,\dot x = \nabla y^\top\dot x$. $\blacksquare$

Graph and sweep are different things. The graph says what is computed. A sweep says in which order the local derivatives are combined. One backward sweep gives the derivatives of one output with respect to all inputs. One forward sweep gives the derivatives of all outputs along one input direction.

Why the order matters (a preview of lesson 05) #

For a chain $F = f_L\circ\cdots\circ f_1$, the Jacobian is $J_L\cdots J_1$. Matrix multiplication is associative, so every way of placing the brackets gives the same matrix. But the costs differ a lot:

• Forward: push an input direction $v$ through the chain, $J_L(\cdots(J_2(J_1 v)))$. Every step is a matrix-vector product. The result is $DF\,v$, a Jacobian-vector product (JVP).
• Backward: pull an output weight $u$ back, $(\cdots((u^\top J_L)J_{L-1})\cdots)J_1$. Every step is again a matrix-vector product. The result is $u^\top DF$, a VJP.

For a scalar loss, the last Jacobian $J_L$ is a single row. Starting there, every intermediate result stays a single row, and one backward sweep gives the whole gradient. This is why training uses the backward order. Lesson 05 turns the two sweeps into the forward and reverse modes of automatic differentiation and counts their costs. Linear algebra lesson 15 gives a short first view of JVPs and VJPs.

Answers to section 1. (1) $2\times 4$, and $Dg\,Df$ is the right product: $(2\times3)(3\times4)$. The other order $Df\,Dg$ would be $(3\times 4)(2\times 3)$, which is not defined. (2) Add. $x_1$ reaches $f$ along two paths, through $\ln x_1$ and through $x_1x_2$, and the two contributions add: $\partial f/\partial x_1 = 1/x_1 + x_2$. At $(2, 5)$ this is $0.5 + 5 = 5.5$.

4. Worked example #

(a) A chain rule with Jacobians. Let $f:\R^2\to\R^3$ and $g:\R^3\to\R$ be

1. Forward: $u = f(1, 2) = (2, 3, 4)$.
2. Jacobian of $f$ (rows = outputs): $Df(x) = \begin{pmatrix} x_2 & x_1 \\ 1 & 1 \\ 0 & 2x_2\end{pmatrix} = \begin{pmatrix} 2 & 1 \\ 1 & 1 \\ 0 & 4\end{pmatrix}$, a $3\times 2$ matrix.
3. Gradient of $g$: $\nabla g(u) = (u_2, u_1, 1) = (3, 2, 1)$.
4. VJP: $\nabla(g\circ f)(x) = Df^\top\nabla g = \begin{pmatrix} 2 & 1 & 0 \\ 1 & 1 & 4\end{pmatrix}\begin{pmatrix}3\\2\\1\end{pmatrix} = \begin{pmatrix} 6 + 2 + 0 \\ 3 + 2 + 4\end{pmatrix} = \begin{pmatrix}8\\9\end{pmatrix}$.
5. Check directly: $g(f(x)) = x_1x_2(x_1 + x_2) + x_2^2 = x_1^2x_2 + x_1x_2^2 + x_2^2$. Then $\partial_1 = 2x_1x_2 + x_2^2 = 4 + 4 = 8$ and $\partial_2 = x_1^2 + 2x_1x_2 + 2x_2 = 1 + 4 + 4 = 9$. ✓

(b) The graph of $f(x_1, x_2) = \ln x_1 + x_1x_2 - \sin x_2$ at $(2, 5)$.

Forward run (values).

step	value
$v_{-1} = x_1$	$2$
$v_0 = x_2$	$5$
$v_1 = \ln v_{-1}$	$\ln 2 = 0.693$
$v_2 = v_{-1}v_0$	$10$
$v_3 = \sin v_0$	$\sin 5 = -0.959$
$v_4 = v_1 + v_2$	$10.693$
$v_5 = v_4 - v_3$	$11.652 = f(2, 5)$

Backward run (adjoints), in reverse order. Each node sends $\bar v_i$ times the local derivative to each parent.

1. Seed: $\bar v_5 = 1$.
2. Node $v_5 = v_4 - v_3$: $\bar v_4 = \bar v_5\cdot 1 = 1$ and $\bar v_3 = \bar v_5\cdot(-1) = -1$.
3. Node $v_4 = v_1 + v_2$: $\bar v_1 = \bar v_4\cdot 1 = 1$ and $\bar v_2 = \bar v_4\cdot 1 = 1$.
4. Node $v_3 = \sin v_0$: $\bar v_0 = \bar v_3\cos v_0 = -\cos 5 = -0.284$.
5. Node $v_2 = v_{-1}v_0$: $\bar v_{-1} = \bar v_2\cdot v_0 = 5$, and add to $\bar v_0$: $\bar v_0 = -0.284 + \bar v_2\cdot v_{-1} = -0.284 + 2 = 1.716$.
6. Node $v_1 = \ln v_{-1}$: add to $\bar v_{-1}$: $\bar v_{-1} = 5 + \bar v_1/v_{-1} = 5 + 0.5 = 5.5$.

So $\partial f/\partial x_1 = 5.5$ and $\partial f/\partial x_2 = 1.716$. Check with the formulas: $\partial f/\partial x_1 = 1/x_1 + x_2 = 0.5 + 5 = 5.5$ and $\partial f/\partial x_2 = x_1 - \cos x_2 = 2 - 0.284 = 1.716$. ✓ The two "add" steps are the fan-outs of $x_2$ (step 5) and $x_1$ (step 6).

Forward sweep (tangents) with $\dot x = (1, 0)$. $\dot v_{-1} = 1$, $\dot v_0 = 0$; $\dot v_1 = \dot v_{-1}/v_{-1} = 0.5$; $\dot v_2 = \dot v_{-1}v_0 + v_{-1}\dot v_0 = 5$; $\dot v_3 = \cos(v_0)\,\dot v_0 = 0$; $\dot v_4 = 0.5 + 5 = 5.5$; $\dot v_5 = 5.5 - 0 = 5.5$. One forward sweep gives one number, $\partial f/\partial x_1 = 5.5$. For $\partial f/\partial x_2$ you need a second sweep with $\dot x = (0, 1)$. The backward sweep gave both at once.

(c) One neuron with a squared loss. $z = w^\top x + b$, $a = \sigma(z)$, $\ell = \tfrac12(a - t)^2$, with the sigmoid $\sigma(z) = 1/(1 + e^{-z})$ and $\sigma' = \sigma(1-\sigma)$. Take $x = (1, 2)$, $w = (0.5, -0.5)$, $b = 0.5$, $t = 1$.

1. Forward: $z = 0.5 - 1 + 0.5 = 0$, $a = \sigma(0) = 0.5$, $\ell = \tfrac12(0.5 - 1)^2 = 0.125$.
2. $\bar a = \partial\ell/\partial a = a - t = -0.5$.
3. Elementwise VJP: $\bar z = \sigma'(z)\,\bar a = 0.5\cdot 0.5\cdot(-0.5) = -0.125$.
4. Weights (the 1-row case of $\bar z\,x^\top$): $\bar w = \bar z\,x = (-0.125, -0.25)$, and $\bar b = \bar z = -0.125$.
5. Input: $\bar x = \bar z\,w = (-0.0625, 0.0625)$.

5. In code #

NumPy checks the three hand computations with central differences $\big(f(x + \eta e_i) - f(x - \eta e_i)\big)/(2\eta)$ (lesson 07 explains this check).

import numpy as np

def num_grad(f, x, eta=1e-6):
    """Central differences, one input at a time."""
    g = np.zeros_like(x)
    for i in range(x.size):
        e = np.zeros_like(x); e[i] = eta
        g[i] = (f(x + e) - f(x - e)) / (2 * eta)
    return g

# (a) gradient through a map is a VJP: Df^T grad g
x = np.array([1.0, 2.0])
f = lambda x: np.array([x[0] * x[1], x[0] + x[1], x[1] ** 2])
g = lambda u: u[0] * u[1] + u[2]
Df = np.array([[x[1], x[0]], [1.0, 1.0], [0.0, 2 * x[1]]])
u = f(x)
print(Df.T @ np.array([u[1], u[0], 1.0]))        # [8. 9.]
print(num_grad(lambda x: g(f(x)), x))            # [8. 9.] (up to ~1e-9)

# (b) backward sweep on the graph of ln x1 + x1 x2 - sin x2 at (2, 5)
x1, x2 = 2.0, 5.0
v1, v2, v3 = np.log(x1), x1 * x2, np.sin(x2)
v4 = v1 + v2; v5 = v4 - v3
b5 = 1.0
b4, b3 = b5 * 1.0, b5 * -1.0
b1, b2 = b4 * 1.0, b4 * 1.0
b0 = b3 * np.cos(x2)              # from v3
bm1 = b2 * x2                     # from v2
b0 += b2 * x1                     # from v2: fan-out of x2, add
bm1 += b1 / x1                    # from v1: fan-out of x1, add
print(round(v5, 3), round(bm1, 3), round(b0, 3))   # 11.652 5.5 1.716
F = lambda x: np.log(x[0]) + x[0] * x[1] - np.sin(x[1])
print(num_grad(F, np.array([2.0, 5.0])).round(3))  # [5.5   1.716]

# (c) one neuron, squared loss
xn, w, b, t = np.array([1.0, 2.0]), np.array([0.5, -0.5]), 0.5, 1.0
sig = lambda z: 1 / (1 + np.exp(-z))
z = w @ xn + b; a = sig(z)
zbar = (a - t) * a * (1 - a)
print(zbar * xn, zbar, zbar * w)                 # [-0.125 -0.25 ] -0.125 [-0.0625  0.0625]
loss_w = lambda w: 0.5 * (sig(w @ xn + b) - t) ** 2
print(num_grad(loss_w, w).round(6))              # [-0.125 -0.25 ]

6. Where it shows up #

• Backpropagation. A neural network is a long chain of layers. Its gradient is the product of the transposed layer Jacobians, applied to $\bar\ell = 1$ from the loss backward. Lesson 06 and deep learning lesson 03 derive the full backward pass of an MLP.
• Every autodiff framework. PyTorch, JAX and TensorFlow record the computational graph of your code. Each primitive operation comes with a VJP rule (its "backward" function), and loss.backward() runs the backward sweep of section 3.
• Residual networks (He et al., 2016). In $y = x + F(x)$ the gradient is $\bar y + DF^\top\bar y$. The identity path lets the gradient reach early layers of very deep networks. Transformers use a residual connection around every attention and MLP block.
• Shared parameters. Recurrent networks use the same weights at every time step; convolutions use the same kernel at every position; many language models tie the input embedding and the output projection. In each case the gradient is the sum over all uses.
• ML for science. In weather prediction, the backward sweep through a forecast model is the adjoint model; 4D-Var data assimilation gets the gradient of the misfit with respect to the initial state from one adjoint run. Seismic full-waveform inversion does the same through a wave simulation. See the ML for science track.

7. Common mistakes #

• Wrong order. $D(g\circ f) = Dg\,Df$, not $Df\,Dg$. The map applied first stands on the right. The shape check catches this unless the matrices are square.
• Wrong point. The outer Jacobian is evaluated at $f(x)$, not at $x$: $Dg(f(x))$. In a network, each layer's Jacobian uses that layer's own input, stored from the forward run.
• Missing transpose. The gradient through $z = Wx$ is $W^\top\bar z$, not $W\bar z$. The gradient for the weights is $\bar z\,x^\top$, with the shape of $W$. Check shapes after every step.
• Not adding at a fan-out. When a value is used twice (a residual, a shared weight, $x$ in $x^2\sin x$), its adjoint is the sum of the contributions. Overwriting instead of adding (= instead of +=) is a classic bug in hand-written backward passes.
• Mixing up the local and the total derivative. On the edge $x_1 \to v_2 = x_1x_2$, the local derivative is $x_2 = 5$. The total derivative $\partial f/\partial x_1 = 5.5$ also includes the path through $\ln x_1$.
• Building diagonal Jacobians. For an elementwise map, use $\varphi'(z)\odot\bar a$. A $\operatorname{diag}$ matrix of size $m\times m$ wastes memory and time.

8. Check yourself #

1. $f:\R^5\to\R^2$ and $g:\R^2\to\R^3$. Give the shapes of $Df(x)$, $Dg(f(x))$ and $D(g\circ f)(x)$.
   Answer

   $Df$ is $2\times 5$, $Dg$ is $3\times 2$, and $D(g\circ f) = Dg\,Df$ is $(3\times2)(2\times5) = 3\times 5$.
2. $f(x, y) = e^{x}y$, $x(t) = t^2$, $y(t) = 3t$. Compute $\frac{d}{dt}f(x(t), y(t))$ at $t = 1$ with the entry form, then check by substitution.
   Answer

   At $t = 1$: $x = 1$, $y = 3$, $x' = 2t = 2$, $y' = 3$. $\partial_x f = e^{x}y = 3e$ and $\partial_y f = e^{x} = e$. So $\frac{d}{dt}f = 3e\cdot 2 + e\cdot 3 = 9e \approx 24.46$. Check: $f = 3t\,e^{t^2}$, derivative $3e^{t^2} + 6t^2e^{t^2} = 9e$ at $t = 1$. ✓
3. $a = \tanh(z)$ elementwise, with $z\in\R^4$. Write $\bar z$ in terms of $\bar a$ and $a$. Why is it enough to store $a$ from the forward run?
   Answer

   $\bar z = (1 - a\odot a)\odot\bar a$, because $\tanh' = 1 - \tanh^2$. The derivative is a function of the output $a$, so $z$ is not needed.
4. A residual block computes $y = x + Wx$ with a square matrix $W$. Given $\bar y$, what is $\bar x$?
   Answer

   Two paths add: $\bar x = \bar y + W^\top\bar y = (I + W)^\top\bar y$. The identity path passes $\bar y$ unchanged.
5. Draw the computational graph of $f(x, y) = (x + y)\,e^{x}$ and run the backward sweep at $(1, 0)$.
   Answer

   $v_1 = x + y = 1$, $v_2 = e^{x} = e$, $v_3 = v_1v_2 = e$. Backward: $\bar v_3 = 1$; $\bar v_1 = v_2 = e$, $\bar v_2 = v_1 = 1$; $\bar y = \bar v_1\cdot 1 = e$; $\bar x = \bar v_1\cdot 1 + \bar v_2\cdot e^{x} = e + e = 2e \approx 5.437$ ($x$ fans out to $v_1$ and $v_2$). Check: $\partial_x f = e^{x} + (x + y)e^{x} = 2e$ and $\partial_y f = e^{x} = e$ at $(1, 0)$. ✓
6. For one sample, why does $\partial\ell/\partial W$ of a linear layer $z = Wx$ have rank at most 1?
   Answer

   It equals $\bar z\,x^\top$, an outer product. Every column $j$ is $x_j\,\bar z$, a multiple of the same vector $\bar z$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. State the chain rule for $g\circ f$ with Jacobians.

   Show answer

2. $f:\R^4\to\R^3$, $g:\R^3\to\R^2$. What is the shape of $D(g\circ f)$?

   Show answer

3. Write the entry form of the chain rule.

   Show answer

4. What is a vector-Jacobian product (VJP), and where does it appear?

   Show answer

5. For $z = Wx$, give $\bar x$ and $\bar W$ in terms of $\bar z$.

   Show answer

6. For an elementwise map $a = \varphi(z)$, what is $\bar z$?

   Show answer

7. Why do gradients add when a value is used twice?

   Show answer

8. What is the adjoint $\bar v_j$ of a node in a computational graph, and how is it computed?

   Show answer

9. For $f(x_1, x_2) = \ln x_1 + x_1x_2 - \sin x_2$, what is $\partial f/\partial x_1$ at $(2, 5)$, and which two paths give it?

   Show answer

10. What does one forward sweep give, and what does one backward sweep give?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Explain the chain rule for vector functions." — $D(g\circ f) = Dg(f(x))\,Df(x)$; show the shapes; the first map stands on the right because composition of linear maps is the matrix product.
• "What is the gradient of $\ell(Wx)$ with respect to $x$ and with respect to $W$?" — $W^\top\bar z$ and $\bar z\,x^\top$ with $\bar z = \nabla\ell(Wx)$. Mention the rank-1 structure of the weight gradient.
• "Why do residual connections help training?" — fan-out: $\bar x = \bar y + DF^\top\bar y$; the identity path carries the gradient back even when $DF$ is small.
• "Backprop through this small expression on the whiteboard." — draw the graph, compute the values, seed $\bar y = 1$, go backward, multiply by local derivatives, add at fan-outs, then check one partial derivative by hand or with finite differences.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. chain_jacobians: multiply a list of Jacobians in the right order, with shape checks.
2. Forward and backward sweeps by hand on $\ln x_1 + x_1x_2 - \sin x_2$, compared with central differences.
3. vjp_linear and vjp_elementwise: the two basic VJPs, without building any diagonal matrix.
4. Fan-out: the derivative of $g(x)\,h(x)$ by summing two paths, compared with the product rule.
5. A second graph: one neuron $\sigma(w^\top x + b)$ with a squared loss, backpropagated by hand.

11. Watch and read #

Backpropagation calculus | Deep Learning Chapter 43Blue1Brown · 10 min

The chain rule on a tiny network, with the "rates multiply" and "paths add" pictures of this lesson.

Multivariable chain ruleKhan Academy · 10 min

The entry form for f(x(t), y(t)), derived from small changes.

Lecture 4 | Introduction to Neural NetworksStanford University School of Engineering · CS231n (2017) · 74 min

Computational graphs and the backward pass with local gradients, on small examples like section 4.

• CS231n course notes, Backpropagation, intuitions: local gradients, the add, multiply and max gates, and gradients that add at a fan-out.
• Baydin, Pearlmutter, Radul and Siskind, Automatic differentiation in machine learning: a survey, JMLR 2018: §3 and Tables 2–3 use the same example graph as this lesson.
• Bright, Edelman and Johnson, Matrix Calculus (for Machine Learning and Beyond): §2.5 "The Chain Rule" (the derivative of a composition as a product of linear maps).
• Kevin Murphy, Probabilistic Machine Learning: An Introduction, §13.3 (backpropagation as the chain rule on a computational graph).