Matrix Calculus and Automatic Differentiation

Dive into Deep Learning · §24.3

Why the backward pass is cheap
Jacobians · the chain rule · forward- and reverse-mode autodiff.

From a slope to a whole matrix

Motivation

A network layer maps vectors to vectors, \mathbf f:\mathbb R^n\to\mathbb R^m, so its derivative is a matrix: the Jacobian.

It is the best local linear map: up close, \mathbf f is a matrix.
The chain rule becomes Jacobian multiplication, one long product per network.
That product is associative, and the cheap order to evaluate it is backpropagation.
We then build both flavours of autodiff in a few dozen lines of Python.

The Jacobian

the best local linear map

Vectors in, vectors out

The Jacobian

\mathbf f is differentiable at \mathbf x if a matrix \mathbf J tracks every first-order change:

\mathbf f(\mathbf x+\boldsymbol\delta)=\mathbf f(\mathbf x)+\mathbf J\boldsymbol\delta+o(\|\boldsymbol\delta\|), \qquad [\mathbf J]_{ij}=\frac{\partial f_i}{\partial x_j}.

Row i collects the partials of output i; column j says how every output reacts to nudging input j.

A small circle of inputs lands on (nearly) an ellipse, its image under \mathbf J. The leftover bend is the o(\|\boldsymbol\delta\|) remainder.

One object, three familiar faces

The Jacobian

The Jacobian quietly contains everything we already know about derivatives:

Scalar map (m=1): \mathbf J is one row, the gradient \partial f/\partial\mathbf x.

Vector map: the full m\times n matrix of partials.

Hessian: the Jacobian of the gradient field, \mathbf H=\mathbf J_{\nabla f}.

One construction recovers the slope, the gradient, and the matrix of second derivatives.

The numerical signature of a derivative

The Jacobian

The linear model must beat its own error: halve \boldsymbol\delta and the error should fall fourfold. Check \mathbf f(x,y)=(x^2y,\ \sin(x+y)) against a finite difference, prediction vs truth:

linear approx: [0.499    0.997424]
exact        : [0.498996 0.997424]
error / |delta|: 0.0015819444529976167

The error relative to \|\boldsymbol\delta\| is tiny, and shrinking \boldsymbol\delta tenfold shrinks it a hundredfold: the o(\|\boldsymbol\delta\|) fingerprint.

The chain rule is Jacobian composition

multiplication of matrices, in any order

Composition multiplies Jacobians

Chain rule

The multivariate chain rule (a sum over paths) is, in matrix form, just a product of Jacobians:

\mathbf J_{\mathbf g\circ\mathbf f}=\mathbf J_{\mathbf g}\,\mathbf J_{\mathbf f}, \qquad \mathbf J = \mathbf J_L\,\mathbf J_{L-1}\cdots\mathbf J_1.

A depth-L network is a composition \mathbf f_L\circ\cdots\circ\mathbf f_1, so its end-to-end derivative is one long matrix product, one local Jacobian per layer.

Two shapes dominate the product

Chain rule

Real networks alternate just two kinds of factor:

Dense \mathbf x\mapsto\mathbf W\mathbf x

Its own best linear map, so \mathbf J=\mathbf W.

Elementwise \boldsymbol\varphi

Couples no coordinates, so \mathbf J=\operatorname{diag}(\varphi').

So \mathbf J=\operatorname{diag}(\boldsymbol\varphi_L')\,\mathbf W_L\cdots\operatorname{diag}(\boldsymbol\varphi_1')\,\mathbf W_1: weights interleaved with cheap diagonal masks. For ReLU that mask is just zeros and ones, which is why a backward pass through an activation costs O(n).

Associativity: two brackets, two modes

Chain rule

Matrix multiplication is associative, so the chain \mathbf J_L\cdots\mathbf J_1 may be bracketed in any order:

Multiply right-to-left \Rightarrow propagate inputs forward. Forward mode.
Multiply left-to-right \Rightarrow propagate sensitivities backward. Reverse mode.

Choosing forward vs reverse mode is nothing but choosing the cheaper way to multiply a chain of matrices.

Identities, derived not tabulated

differentiate one component, reassemble

Four identities, one method

Identities

Differentiate one component, reassemble, then sanity-check the 1\times1 scalar collapse:

\nabla_{\mathbf x}\,\mathbf a^\top\mathbf x = \mathbf a

\nabla_{\mathbf x}\,\mathbf x^\top\mathbf A\mathbf x = (\mathbf A+\mathbf A^\top)\mathbf x

\nabla_{\mathbf W}\|\mathbf W\mathbf x-\mathbf y\|^2 = 2(\mathbf W\mathbf x-\mathbf y)\mathbf x^\top

softmax + cross-entropy, w.r.t. the logits =\mathbf p-\mathbf y

No table needed: the least-squares gradient is the rank-one outer product a linear layer returns in backprop.

Derive, then verify: autograd audits the first two at a random point:

grad a^T x   equals a        : True
grad x^T A x equals (A+A^T)x : True

The cancellation that fuses softmax with cross-entropy

Identities

The softmax Jacobian \operatorname{diag}(\mathbf p)-\mathbf p\mathbf p^\top, composed with cross-entropy, collapses to \mathbf p-\mathbf y whatever the Jacobian; the fused softmax–cross-entropy primitive also avoids \log of an underflowed probability via log-sum-exp. Verify against autograd:

z = torch.tensor([1.0, -0.5, 2.0], requires_grad=True)
p = torch.softmax(z, dim=0)
J = torch.diag(p) - torch.outer(p, p)          # our formula
J_ad = torch.autograd.functional.jacobian(lambda z: torch.softmax(z, 0), z)
y = torch.tensor([0.0, 1.0, 0.0])              # one-hot target (class 1)
loss = -(y * torch.log(torch.softmax(z, 0))).sum()
loss.backward()
print('softmax Jacobian matches autograd:', torch.allclose(J, J_ad, atol=1e-6))
print('d loss / d z :', z.grad.detach().numpy().round(6))
print('p - y        :', (p.detach() - y).numpy().round(6))

softmax Jacobian matches autograd: True
d loss / d z : [ 0.253716 -0.943388  0.689672]
p - y        : [ 0.253716 -0.943388  0.689672]

Forward mode and dual numbers

carry the derivative alongside the value

An algebra that carries derivatives

Forward mode

Adjoin a symbol \varepsilon\neq0 with \varepsilon^2=0, the exact encoding of discard second-order terms:

(a+b\varepsilon)(c+d\varepsilon)=ac+(ad+bc)\varepsilon.

The \varepsilon-coefficient ad+bc is the product rule. So running a program on x+1\cdot\varepsilon yields

f(x+\varepsilon)=f(x)+f'(x)\,\varepsilon.

Differentiation is free: seed the input’s \varepsilon-part with 1, read the output’s \varepsilon-part.

A 15-line class differentiates for free

Forward mode

A handful of overloaded operators is forward-mode AD: the \varepsilon-part returns f'(x)=2x\cos(x^2)+e^x in one pass, no formula for f' written. Differentiate f(x)=\sin(x^2)+e^x at x=1.3:

def f(x):
    return (x * x).sin() + x.exp()                  # sin(x^2) + e^x, in dual algebra

x0 = 1.3
out = f(Dual(x0, 1.0))                              # seed the derivative with 1
exact = 2 * x0 * np.cos(x0**2) + np.exp(x0)         # f'(x) by hand
print('dual result :', out)
print('value f(x)  :', np.sin(x0**2) + np.exp(x0))
print("derivative  : dual %.6f  vs  exact %.6f" % (out.b, exact))

dual result : 4.662200 + 3.360101 eps
value f(x)  : 4.662200318713363
derivative  : dual 3.360101  vs  exact 3.360101

One pass per input: cheap when tall

Forward mode

With a vector tangent, one forward pass carries a Jacobian–vector product \mathbf J\mathbf v, a linear combination of columns, never forming \mathbf J. Seeding \mathbf e_j reads off literal column j.

Forward mode costs one pass per input. Cheap for tall Jacobians (m\gg n), and exactly wrong for a deep net, whose loss is one scalar over millions of inputs.

Reverse mode, the tape, and backprop

record forward, replay backward

Why reverse mode is the right cost model

Reverse mode

Reverse mode multiplies left-to-right, carrying a vector–Jacobian product \mathbf u^\top\mathbf J, one pass per output.

A scalar loss is a single-row Jacobian (m=1), so one backward sweep yields the entire gradient.

Gradient of a scalar in one reverse sweep at a small constant multiple of the primal cost: the cheap-gradient principle (Baur–Strassen).

One row or one column per pass

Reverse mode

The cost rule, drawn as shapes. A pass buys one row (reverse, a VJP) or one column (forward, a JVP); the cheap mode is the one whose unit matches your Jacobian’s short side.

A scalar loss is one row: one VJP, one backward pass. The full m\times n matrix costs \min(m,n) passes either way, which is why it is never formed.

The tape: a diamond, not a chain

Reverse mode

Record each op as a node with a backward closure; replay in reverse topological order, seeding the output adjoint with 1.

For r=uv+u, y=r^2, the value r feeds both arguments of y, so its adjoint arrives twice.

A node’s adjoint is the sum over its outgoing edges, the chain rule’s “sum over paths”, which is why the tape accumulates with +=, not =.

Thirty lines reproduce the backward pass

Reverse mode

One forward pass records the tape; one backward pass yields both partials, matching the framework’s own autograd. Real engines change only the details (tensor nodes, VJP backwards, a larger primitive set), but the skeleton is exactly this:

u, v = Var(2.0), Var(-3.0)
r = u * v + u                                       # shared intermediate
y = r * r                                           # r feeds both arguments
backprop(y)
ut = torch.tensor(2.0, requires_grad=True)
vt = torch.tensor(-3.0, requires_grad=True)
rt = ut * vt + ut
(rt * rt).backward()                                # framework autograd
print('our tape   : dy/du = %.4f  dy/dv = %.4f' % (u.grad, v.grad))
print('torch      : dy/du = %.4f  dy/dv = %.4f' % (ut.grad, vt.grad))

our tape   : dy/du = 16.0000  dy/dv = -16.0000
torch      : dy/du = 16.0000  dy/dv = -16.0000

The asymmetry, counted: 200 passes vs 1

Reverse mode

With both engines on the table, the cost claim stops being an assertion. Take a scalar function of n = 200 inputs: forward mode assembles the gradient one input direction per pass; the tape runs one backward sweep.

forward mode : 200 passes (one per input)
reverse mode : 1 backward sweep for all 200 partials
max |forward - reverse| over the gradient: 0.0e+00

Identical gradients, a factor-n gap in pass counts; per-pass costs match to a small constant, so the work ratio is \Theta(n). Replace 200 by the 10^8 parameters of a modern network: that factor is why no framework trains in forward mode.

Never form the Jacobian

Reverse mode

A dense m\times n Jacobian costs \min(m,n) passes and \Theta(mn) storage; what you want costs one pass and holds only a vector.

If you are assembling a Jacobian, you have probably written down a matrix you could have multiplied through. Autograd exposes the JVP and the VJP; you compose them.

One order up, the Hessian–vector product \mathbf H\mathbf v is the directional derivative of the gradient map: curvature in a direction without ever forming \mathbf H.

Curvature for the price of a gradient

Reverse mode

Differentiate the gradient once, in one direction. For L(\mathbf x) = \tfrac12\mathbf x^\top\mathbf A\mathbf x the Hessian is \mathbf A itself, so the right answer is known in advance, and no 2\times2 matrix of second derivatives is ever assembled:

# Hv for L(x) = x^T A x / 2, where H = A exactly
A = torch.tensor([[3.0, 1.0], [1.0, 2.0]])
L = lambda x: 0.5 * x @ A @ x
x, v = torch.tensor([1.0, -2.0]), torch.tensor([1.0, 0.5])
_, Hv = torch.func.jvp(torch.func.grad(L), (x,), (v,))  # forward over reverse
print('forward-over-reverse Hv:', Hv)
print('A v                    :', A @ v)

forward-over-reverse Hv: tensor([3.5000, 2.0000])
A v                    : tensor([3.5000, 2.0000])

Cost: a small constant multiple of one gradient, rather than n gradients. The absolute cost still follows the primal computation. This is what makes Newton and conjugate-gradient methods, and curvature diagnostics, tractable at scale.

Differentiating through an equation

Reverse mode

Some outputs are defined implicitly: \mathbf x^\star(\boldsymbol\theta) solves \mathbf g(\mathbf x^\star, \boldsymbol\theta) = \mathbf 0 (a fixed point, an equilibrium, an inner argmin). No need to tape the solver.

Implicit function theorem. Differentiating the identity \mathbf g(\mathbf x^\star(\boldsymbol\theta), \boldsymbol\theta) \equiv \mathbf 0 gives \;\frac{\partial\mathbf x^\star}{\partial\boldsymbol\theta} = -\bigl(\frac{\partial\mathbf g}{\partial\mathbf x}\bigr)^{-1}\frac{\partial\mathbf g}{\partial\boldsymbol\theta}, both Jacobians at the solution only.

One linear solve per gradient, however many iterations the solver ran. Checked on a 2\times2 Newton solve, against a finite difference of the solver’s output:

implicit-function gradient: [-0.11311451 -0.69322129]
finite-difference check   : [-0.11311451 -0.69322129]

The equation carries the derivative, not the algorithm that solved it: the seed of the adjoint method, deep equilibrium models, and bilevel optimization.

The price of the tape, and buying it back

Reverse mode

Reverse mode must keep every forward intermediate alive until the backward sweep consumes it: O(L) activation memory, often the binding constraint on batch size and depth.

Checkpointing stores only every \sqrt{L}-th activation and recomputes each segment when the sweep reaches it: O(\sqrt{L}) memory for roughly one extra forward pass.

Recap

Wrap-up

The Jacobian is the best local linear map; the gradient (m=1) and the Hessian are special cases.
The chain rule is Jacobian composition; a deep net is one long matrix product.
Its associativity is what yields the two autodiff modes.

Forward mode (dual numbers, JVP): one column per pass, cheap when tall.
Reverse mode (the tape, VJP): one row per pass, cheap when wide. A loss is maximally wide, so backprop is reverse-mode AD, paid for in stored activations (checkpointing buys them back).
Compose, never assemble: \mathbf H\mathbf v by differentiating the gradient; implicit outputs by one linear solve.

Two products, JVP and VJP, compose into everything a training loop needs.