Backpropagation via Jacobian-vector products

Setup: a layered network

$h_0=x$, and for $\ell=1,\dots,L$:

$$ z_\ell = W_\ell h_{\ell-1} + b_\ell, \qquad h_\ell = \sigma(z_\ell) \quad \text{(elementwise)}, $$

with scalar loss $\mathcal{L}(h_L)$ at the end. Two derivatives per layer (Gradients, Jacobians, and Hessians):

$$ \underbrace{\frac{dz_\ell}{dh_{\ell-1}}}_{n_\ell\times n_{\ell-1}} = W_\ell, \qquad \underbrace{\frac{dh_\ell}{dz_\ell}}_{n_\ell\times n_\ell} = \operatorname{diag}\big(\sigma'(z_\ell)\big). $$

Intuition: an elementwise nonlinearity's derivative is diagonal, since output $i$ only ever depends on input $i$; the linear map's derivative is just $W_\ell$ itself.

The chain, written once, then substituted

$h_{\ell-1}$ affects $\mathcal{L}$ only through $h_{\ell-1}\to z_\ell\to h_\ell\to\cdots\to\mathcal{L}$. Write that whole dependency as one chain-rule equation first, before substituting anything for any term in it:

$$ \underbrace{\frac{d\mathcal{L}}{dh_{\ell-1}}}_{1\times n_{\ell-1}} = \underbrace{\frac{d\mathcal{L}}{dh_\ell}}_{1\times n_\ell}\cdot\underbrace{\frac{dh_\ell}{dz_\ell}}_{n_\ell\times n_\ell}\cdot\underbrace{\frac{dz_\ell}{dh_{\ell-1}}}_{n_\ell\times n_{\ell-1}}. $$

Abbreviate $g_\ell:=\dfrac{d\mathcal{L}}{dh_\ell}$, so this reads $g_{\ell-1}=g_\ell\cdot\dfrac{dh_\ell}{dz_\ell}\cdot\dfrac{dz_\ell}{dh_{\ell-1}}$. Base case: $g_L=\dfrac{d\mathcal{L}}{dh_L}$, the loss's own gradient (e.g. $g_L=(h_L-y)^T$ for $\mathcal{L}=\tfrac12\|h_L-y\|^2$). Now substitute the two derivatives from Setup, one at a time, into the chain above:

$$ \frac{dh_\ell}{dz_\ell} = \operatorname{diag}\big(\sigma'(z_\ell)\big), \qquad \frac{dz_\ell}{dh_{\ell-1}} = W_\ell. $$

$$ g_{\ell-1} = g_\ell\,\operatorname{diag}\big(\sigma'(z_\ell)\big)\,W_\ell. $$

A row times a diagonal matrix is an elementwise product. Splitting the product there, left factor first, names a reusable intermediate:

$$ \underbrace{g_{z_\ell}}_{1\times n_\ell} := g_\ell\odot\sigma'(z_\ell)^T, \qquad \underbrace{g_{\ell-1}}_{1\times n_{\ell-1}} = \underbrace{g_{z_\ell}}_{1\times n_\ell}\ \underbrace{W_\ell}_{n_\ell\times n_{\ell-1}}. $$

Intuition: one chain-rule equation, two substitutions; $g_{z_\ell}$ is that same equation evaluated halfway through, and it is reused below for the parameter gradients too.

Two steps at once

The recursion above is this one equation applied repeatedly. Writing it out two layers back, straight from $h_{\ell-2}$ to $\mathcal{L}$, makes that explicit, before any substitution:

$$ \underbrace{\frac{d\mathcal{L}}{dh_{\ell-2}}}_{1\times n_{\ell-2}} = \underbrace{\frac{d\mathcal{L}}{dh_\ell}}_{1\times n_\ell}\cdot \underbrace{\frac{dh_\ell}{dz_\ell}}_{n_\ell\times n_\ell}\cdot \underbrace{\frac{dz_\ell}{dh_{\ell-1}}}_{n_\ell\times n_{\ell-1}}\cdot \underbrace{\frac{dh_{\ell-1}}{dz_{\ell-1}}}_{n_{\ell-1}\times n_{\ell-1}}\cdot \underbrace{\frac{dz_{\ell-1}}{dh_{\ell-2}}}_{n_{\ell-1}\times n_{\ell-2}}. $$

The first three factors are exactly $g_{\ell-1}$ from above; the equation is just $g_{\ell-1}$ taking one more step back, with its own two substitutions:

$$ g_{\ell-2} = \underbrace{\Big(g_\ell\cdot\frac{dh_\ell}{dz_\ell}\cdot\frac{dz_\ell}{dh_{\ell-1}}\Big)}_{=\;g_{\ell-1}}\cdot\frac{dh_{\ell-1}}{dz_{\ell-1}}\cdot\frac{dz_{\ell-1}}{dh_{\ell-2}} = g_{\ell-1}\,\operatorname{diag}\big(\sigma'(z_{\ell-1})\big)\,W_{\ell-1}. $$

Intuition: going back $k$ layers multiplies $k$ more of these small factors into the same row vector; backprop never forms this long product in one shot, it computes $g_\ell$, then $g_{\ell-1}$, then $g_{\ell-2}$, each one reusing the row vector the previous step already built.

Parameter gradients, from scratch

Take $\nabla_{W_{\ell-2}}\mathcal{L}$ directly, a weight matrix two layers back. Everything needed is already built: $g_{\ell-2}$ came from the two-step chain above, and substituting it through the same nonlinearity step used for $g_{z_\ell}$ gives:

$$ \underbrace{g_{z_{\ell-2}}}_{1\times n_{\ell-2}} = g_{\ell-2}\odot\sigma'(z_{\ell-2})^T. $$

$g_{z_{\ell-2}}$ is a row. Package it as a column instead, the standard convention for this last step, and give it a name:

$$ \underbrace{\delta_{\ell-2}}_{n_{\ell-2}\times1} \ :=\ g_{z_{\ell-2}}^T. $$

$z_{\ell-2}=W_{\ell-2}h_{\ell-3}+b_{\ell-2}$ is linear in $W_{\ell-2}$ and $b_{\ell-2}$ too, exactly as it is in $h_{\ell-3}$. The two vectors already in reach at this point in the chain are $\delta_{\ell-2}$ ($n_{\ell-2}\times1$, how much $z_{\ell-2}$ should change) and $h_{\ell-3}$ ($n_{\ell-3}\times1$, what fed into $z_{\ell-2}$ alongside $W_{\ell-2}$). Whatever $\nabla_{W_{\ell-2}}\mathcal{L}$ is, it must have the same shape as $W_{\ell-2}$ itself, $n_{\ell-2}\times n_{\ell-3}$ (the shape rule, Gradients, Jacobians, and Hessians). An outer product of these two vectors is the only way to build a matrix of that exact shape, and only one order works: $\delta_{\ell-2}h_{\ell-3}^T$ is $n_{\ell-2}\times n_{\ell-3}$, while the other order, $h_{\ell-3}\delta_{\ell-2}^T$, is $n_{\ell-3}\times n_{\ell-2}$, the wrong shape entirely. That pins the formula down:

$$ \nabla_{W_{\ell-2}}\mathcal{L} = \delta_{\ell-2}\,h_{\ell-3}^T. $$

Shapes

$$ \underbrace{\delta_{\ell-2}}_{n_{\ell-2}\times1}\ \underbrace{h_{\ell-3}^T}_{1\times n_{\ell-3}} = \underbrace{\nabla_{W_{\ell-2}}\mathcal{L}}_{n_{\ell-2}\times n_{\ell-3}}. $$

$\nabla_{b_{\ell-2}}\mathcal{L}=\delta_{\ell-2}$ needs no such argument: $b_{\ell-2}$ enters $z_{\ell-2}$ directly, with no other vector multiplying it, so the gradient is $\delta_{\ell-2}$ itself, unchanged.

By the same shape rule used everywhere above (output shape times input shape, Gradients, Jacobians, and Hessians), the raw derivative $d\mathcal{L}/dW_{\ell-2}$ is really $1\times n_{\ell-2}\times n_{\ell-3}$: the loss's own scalar shape, $1$, prepended to $W_{\ell-2}$'s own shape. That leading $1$ carries no information, exactly like the $1$ in $g_\ell$'s $1\times n_\ell$, so it is dropped by convention, identifying the gradient with a matrix of exactly $W_{\ell-2}$'s own shape: precisely $\nabla_{W_{\ell-2}}\mathcal{L}$ as just derived.

Intuition: the gradient for the weight connecting input $k$ to output $j$ is (how much output $j$ should change) times (how large input $k$ actually was); an input that was $0$ contributes no gradient to any weight reading from it, no matter how wrong the output was.

The same pattern, at every layer

Nothing above depended on being exactly two layers back. The identical argument applies verbatim at every $\ell$, with $\delta_\ell := g_{z_\ell}^T$ supplied by the recursion however far $\ell$ is from $L$:

$$ \nabla_{W_\ell}\mathcal{L} = \delta_\ell\,h_{\ell-1}^T, \qquad \nabla_{b_\ell}\mathcal{L} = \delta_\ell. $$

Intuition: the formula never referenced $\ell$'s distance from $L$, only $\delta_\ell$; the activation recursion is what supplies that vector at whatever depth it's needed, one step or several.

A consequence: rank 1, per example

$\delta_\ell h_{\ell-1}^T$ is exactly the outer-product construction from Rank, subspaces, and projections: a column times a row. For a single training example, $\nabla_{W_\ell}\mathcal{L}$ is therefore always rank $1$ (rank $0$ only if $\delta_\ell$ or $h_{\ell-1}$ is exactly $0$), no matter how large $W_\ell$ is. A minibatch of $B$ examples sums $B$ such rank-$1$ terms, $\frac1B\sum_{b=1}^B\delta_\ell^{(b)}\big(h_{\ell-1}^{(b)}\big)^T$, and a sum of $B$ rank-$1$ matrices has rank at most $B$, generically exactly $\min(B,n_\ell,n_{\ell-1})$.

Intuition: one example's gradient only ever says "nudge $W_\ell$ along this one row direction, scaled along this one column direction"; batching is what turns that single rank-$1$ nudge into an update with real structure across the weight matrix.

One backward pass, every parameter at once

Starting from a single row $g_L$ and running the recursion down to $\ell=1$ computes $\nabla_{W_\ell}\mathcal{L}$ and $\nabla_{b_\ell}\mathcal{L}$ for every layer, in one pass, at roughly the cost of one forward pass.

A JVP-based (forward-mode) approach would instead need to seed a separate pass per parameter, since forward mode propagates one input direction $v$ at a time (Gradients, Jacobians, and Hessians): with millions of parameters, that is millions of passes. Reverse mode wins precisely because there is one scalar output to seed from and, typically, far more parameters than that.

Intuition: forward mode answers "how does the output move if this one input moves"; a network has one loss but millions of weights, so asking the reverse question once, "how should every weight move to move this one loss", is the only question worth asking that many times cheaper.

Everything above is this general algorithm (Automatic differentiation) specialized to one particular computation graph: a plain chain of matmuls and elementwise nonlinearities, where every node has exactly one child.