Matrix calculus and reverse-mode automatic differentiation

Before this lesson: linear algebra 15 (JVP and VJP, first view) and linear algebra 16 (matrix differentials, $d\log\det X$). From the calculus track: the chain rule and computational graphs, automatic differentiation and numerical differentiation. This is the first lesson of the track, so it repeats the key facts briefly. Time: about 75 minutes, plus the notebook.

---

Why this matters #

In ML for science you differentiate more than neural networks.

• A Gaussian process or a Kalman filter has a likelihood with $\log\det K$ and $K^{-1}r$. Fitting its hyperparameters needs their gradients (lesson 07).
• An implicit time step contains a linear solve (lesson 02).
• Adjoint methods for ODEs, 4D-Var in weather forecasting and full-waveform inversion in seismology are reverse mode applied to a simulator (lessons 03, 04 and 06).
• A hybrid climate model puts a network inside a physics code.

PyTorch and JAX know the VJPs of their built-in operations. A Fortran physics code, a sparse direct solver or a custom GPU kernel has none: you must write its VJP yourself, or you cannot train through it. The same rules also decide the cost. Reverse mode gives the gradient with respect to $10^9$ parameters for a few times the cost of the loss, but it must keep (or recompute) the intermediate values. This memory cost is the central topic of lesson 03.

1. Predict first #

2. The picture #

One variable. Let $y = \tanh(x)$. A small change $dx$ gives $dy = (1 - y^2)\,dx$. So the derivative is the linear map "multiply by $1 - y^2$". Suppose a scalar loss $L$ depends on $y$, and we know $\bar y = \partial L/\partial y$. Then

so $\bar x = (1 - y^2)\,\bar y$.

Two facts appear already. First, the rule uses the output $y$, so the forward pass must keep it. Second, the same factor $1 - y^2$ moves a small change forward and a sensitivity backward.

Many variables. Let $y = F(x)$ with $x \in \R^n$ and $y \in \R^m$. The derivative is the linear map $dy = J\,dx$, where $J$ is the $m\times n$ Jacobian. There are two ways to use it.

• Push a direction $v$ forward: $v \mapsto Jv$. This is a Jacobian-vector product (JVP). It lives in the output space $\R^m$.
• Pull a sensitivity $\bar y$ back: $\bar y \mapsto J^\top \bar y$. This is a vector-Jacobian product (VJP). It lives in the input space $\R^n$.

A small example. $F(x_1, x_2) = (x_1 x_2,\ x_1 + 3x_2)$ at $x = (1, 2)$ has

The JVP with $v = (1, -1)$ is $Jv = (1, -2)$. The VJP with $\bar y = (1, 1)$ is $J^\top \bar y = (3, 4)$. The two are linked by one identity: $\langle \bar y, Jv\rangle = 1 - 2 = -1$ and $\langle J^\top\bar y, v\rangle = 3 - 4 = -1$. They are equal, as they must be for every $v$ and $\bar y$. This identity is the definition of the VJP in section 3.

A program. A program is a chain of small operations $x \to v_1 \to \dots \to v_N = L$ with Jacobians $J_1, \dots, J_N$. By the chain rule, its gradient is

Multiply from the right: start with the number $1$, then apply $J_N^\top$, then $J_{N-1}^\top$, and so on. Each step is one VJP, and the vector stays small. Nobody builds a Jacobian. For a softmax over $n$ classes the Jacobian has $n^2$ entries, but its VJP costs $O(n)$.

The figure runs this on a small network, $L = \sum_i \tanh(Wx + b)_i^2$.

Interactive figure: loading…

3. Definitions and main results #

Derivatives as linear maps #

Arrays of any shape are vectors. For matrices we use the Frobenius inner product $\langle A, B\rangle = \sum_{ij} A_{ij}B_{ij} = \operatorname{tr}(A^\top B)$.

Definition (derivative, JVP, VJP). $F$ is differentiable at $x$ if

with a linear map $F'(x)$. The JVP is the map $v \mapsto F'(x)[v]$. The VJP is the adjoint map $\bar y \mapsto F'(x)^*[\bar y]$, defined by

If you flatten $x$ and $y$ into vectors, $F'(x)$ is the Jacobian $J$ and $F'(x)^* = J^\top$. The inner-product form is better for matrices: it gives $\bar A$ directly as a matrix of the shape of $A$, with no flattening.

Notation. For a scalar loss $L$, a bar marks the gradient of $L$ with respect to a variable: $\bar y = \partial L/\partial y$. It has the shape of $y$. It is also called the cotangent or the adjoint of $y$.

Lemma (a gradient through a map is a VJP). If $L = \ell(F(x))$ and $\bar y = \nabla \ell(F(x))$, then $\bar x = F'(x)^*[\bar y]$.

Proof.

1. $dL = \langle \bar y, dy\rangle$, by the definition of the gradient of $\ell$.
2. $dy = F'(x)[dx]$, by the definition of the derivative.
3. So $dL = \langle \bar y, F'(x)[dx]\rangle = \langle F'(x)^*[\bar y], dx\rangle$ for every $dx$. The gradient is the vector that represents $dL$ in this way, so $\bar x = F'(x)^*[\bar y]$. $\blacksquare$

Recipe for any VJP.

1. Write the differential $dy$ in terms of $dx$.
2. Write $\langle \bar y, dy\rangle$.
3. Move everything away from $dx$ with transposes and the cyclic rule $\operatorname{tr}(ABC) = \operatorname{tr}(CAB)$, until it reads $\langle G, dx\rangle$. Then $\bar x = G$.

$G$ must have the shape of $x$. This is a free check of your algebra.

The VJPs of the primitives #

Here $\odot$ is the elementwise product, $\mathbf 1$ is the all-ones vector and $A^{-\top} = (A^\top)^{-1}$.

Primitive	Differential	VJP
$y = \tanh(x)$, elementwise	$dy = (1 - y^2)\odot dx$	$\bar x = (1 - y^2)\odot\bar y$
$y = e^x$; $y = \log x$	$dy = y\odot dx$; $dy = dx/x$	$\bar x = y\odot\bar y$; $\bar x = \bar y/x$
$y = \operatorname{sum}(x)$ over some axes	$dy = \operatorname{sum}(dx)$	$\bar x = \bar y$ broadcast back to the shape of $x$
$Y = X + \mathbf 1 b^\top$ (bias $b$ added to every row)	$dY = dX + \mathbf 1\,db^\top$	$\bar X = \bar Y$; $\bar b$ = sum of the rows of $\bar Y$
$s = \operatorname{softmax}(z)$	$ds = s\odot(dz - \langle s, dz\rangle\mathbf 1)$	$\bar z = s\odot(\bar s - \langle s, \bar s\rangle\mathbf 1)$
$C = AB$	$dC = dA\,B + A\,dB$	$\bar A = \bar C B^\top$, $\bar B = A^\top \bar C$
$X = A^{-1}B$ (linear solve)	$A\,dX = dB - dA\,X$	$\bar B = A^{-\top}\bar X$, $\bar A = -\bar B X^\top$
$y = \log\lvert\det A\rvert$	$dy = \operatorname{tr}(A^{-1}dA)$	$\bar A = \bar y\,A^{-\top}$

We derive the important rows now. Each derivation follows the recipe.

Sum and broadcasting are adjoints of each other. $\langle \bar y, \operatorname{sum}(dx)\rangle = \langle \operatorname{broadcast}(\bar y), dx\rangle$. This is why the gradient of a bias is a sum over the batch: the forward pass broadcasts $b$ to every row, so the backward pass sums the rows.

Matrix product $C = AB$.

1. $dC = dA\,B + A\,dB$ (product rule).
2. First term, with the cyclic rule:

3. Second term:

4. So $\bar A = \bar C B^\top$ and $\bar B = A^\top \bar C$. $\blacksquare$

Shape check: $A$ is $m\times n$, $B$ is $n\times k$, so $C$ and $\bar C$ are $m\times k$. Then $\bar C B^\top$ is $(m\times k)(k\times n) = m\times n$, the shape of $A$. For a linear layer $y = Wx$: $\bar W = \bar y\, x^\top$ (an outer product, rank one) and $\bar x = W^\top \bar y$. The backward pass is two matrix products of the forward size, so it costs about two forward passes.

Softmax $s_i = e^{z_i}/\sum_j e^{z_j}$ (row by row for a batch).

1. $\log s_i = z_i - \log\sum_j e^{z_j}$, so $d\log s_i = dz_i - \langle s, dz\rangle$.
2. $ds_i = s_i\,d\log s_i$, so

3. The matrix $\operatorname{diag}(s) - ss^\top$ is symmetric, so the VJP applies the same matrix: $\bar z = s\odot\bar s - s\,\langle s, \bar s\rangle$. The cost is $O(n)$; the $n\times n$ matrix is never built. $\blacksquare$

Linear solve $X = A^{-1}B$, where $B$ is a vector or a matrix.

1. Differentiate $AX = B$: $dA\,X + A\,dX = dB$, so $dX = A^{-1}(dB - dA\,X)$.
2. Move $A^{-1}$ to the other side of the inner product as its transpose:

3. Name $\bar B = A^{-\top}\bar X$: this is one solve with $A^\top$. The right side becomes

4. Read off $\bar B = A^{-\top}\bar X$ and $\bar A = -\bar B X^\top$. $\blacksquare$

The forward pass factors $A$ (for example $PA = LU$). The same factors also solve with $A^\top$, because $A^\top = U^\top L^\top P$. So the backward solve is cheap, and no inverse is ever formed.

Log-determinant (Jacobi's formula, from linear algebra 16).

1. $A + dA = A(I + E)$ with $E = A^{-1}dA$, so $\det(A + dA) = \det A\,\det(I + E)$.
2. $\det(I + E) = 1 + \operatorname{tr}E + O(\lVert E\rVert^2)$.
3. So $d\log\lvert\det A\rvert = \operatorname{tr}(A^{-1}dA) = \langle A^{-\top}, dA\rangle$, and $\bar A = \bar y\,A^{-\top}$. $\blacksquare$

The formula also holds when $\det A < 0$. Compute the value from an LU or Cholesky factorization (slogdet in NumPy), never as $\log(\det A)$: $\det A$ overflows or underflows long before $\log\lvert\det A\rvert$ does. For a symmetric positive definite $K = LL^\top$ (Cholesky), $\log\det K = 2\sum_i \log L_{ii}$ and $\bar K = \bar y\,K^{-1}$.

Reverse mode is back substitution #

A program computes $v_0 = x$, then $v_k = \varphi_k(\text{parents of } k)$ for $k = 1, \dots, N$, in this order. $L = v_N$ is a scalar. Each $v_k$ can be an array. The tape (also called the Wengert list) stores the nodes in this order, with their parents and the values their VJPs need.

Theorem (the reverse sweep gives the gradient). Set $\bar v_N = 1$, and for $i = N - 1, \dots, 0$

where $\partial_i\varphi_k$ is the derivative of $\varphi_k$ with respect to its input $v_i$. Then $\bar v_0 = \nabla_x L$.

Why this is true, in one picture. Write all the linearized steps as one system of equations. Each step depends only on earlier steps, so the system is triangular. The gradient is the solution of the transposed system, and a transposed lower-triangular system is solved from the last unknown to the first: that is the backward sweep.

Proof of the theorem

1. Linearize every node: $dv_k = \sum_i M_{ki}\,dv_i$, where $M_{ki} = \partial_i\varphi_k$ if $i$ is a parent of $k$, and $M_{ki} = 0$ otherwise. Parents come earlier, so $M_{ki} = 0$ for $i \ge k$: $M$ is strictly lower triangular (block by block).
2. With $dv_0 = dx$, all these equations together read $(I - M)\,dv = E\,dx$, where $E$ puts $dx$ into block 0. $I - M$ is unit lower triangular, so it is invertible.
3. $dL = dv_N = e_N^\top (I - M)^{-1}E\,dx$, where $e_N$ selects block $N$. So $\nabla_x L = E^\top w$ with $w = (I - M)^{-\top}e_N$.
4. $(I - M)^\top$ is unit upper triangular. Solve $(I - M)^\top w = e_N$ by back substitution, last block first: $w_N = 1$ and $w_i = \sum_k M_{ki}^\top w_k$. This is the sweep of the theorem with $w_i = \bar v_i$, and $E^\top w = w_0 = \bar v_0$. $\blacksquare$

Three consequences:

• Order. The sweep visits the nodes in reverse tape order. A node is processed only after all its users have added their terms.
• Fan-out adds. A value used by several nodes collects one term per use. A parameter shared by many layers, or by many time steps of a simulation, gets the sum of all terms.
• Forward mode solves the same system the other way: forward substitution of $(I - M)\,dv = E\,v$ for one input direction $v$, in program order, together with the forward pass.

Forward or reverse? #

Let $F : \R^n \to \R^m$.

	forward mode (JVP)	reverse mode (VJP)
one pass gives	$Jv$ (column $j$ for $v = e_j$)	$J^\top u$ (row $i$ for $u = e_i$)
full Jacobian	$n$ passes	$m$ passes
gradient of a scalar loss	$n$ passes	1 pass
time of one pass	a small constant times $\operatorname{cost}(F)$	a small constant times $\operatorname{cost}(F)$
extra memory	one tangent per live value; nothing is kept	the tape: every value a VJP needs, for the whole program

The small constant is the cheap gradient principle. For programs made of arithmetic operations, the Baur–Strassen theorem bounds the cost of the gradient by 5 times the cost of the function. For a network dominated by matrix products, the backward pass costs about two forward passes.

Interactive figure: loading…

• Reverse wins when $m \ll n$: a loss, a data misfit, a log-likelihood. Training, 4D-Var, seismic full-waveform inversion and fitting the hyperparameters of a Gaussian process all have $m = 1$.
• Forward wins when $n \ll m$, or when memory is the limit. Examples: how a whole simulated field changes with three source parameters; the Jacobian of a small ODE right-hand side for an implicit solver. A sparse Jacobian (a PDE stencil) needs only a few JVPs with grouped ("colored") seeds.
• Memory is the price of reverse mode. For a long time loop the tape grows with the number of steps. Checkpointing stores only some values and recomputes the rest (lesson 03).

Checking a gradient #

A gradient code needs a test. Pick a random direction $v$ and compare the directional derivative $\langle\nabla f(x), v\rangle$ with the central difference

Two errors compete.

• The truncation error comes from the Taylor series. For central differences the $h^2$ terms cancel, so it is $O(h^2)$. (For the forward difference $(f(x + hv) - f(x))/h$ it is only $O(h)$.)
• The rounding error comes from the subtraction of two nearly equal numbers. It is about $\varepsilon\lvert f\rvert/h$, where $\varepsilon \approx 2.2\cdot 10^{-16}$ is the machine epsilon of float64.

The sum $C h^2 + \varepsilon\lvert f\rvert/h$ is smallest when both parts have similar size, at $h \approx \varepsilon^{1/3} \approx 6\cdot 10^{-6}$ (times the scale of $x$). There the relative error is about $10^{-10}$. One random direction tests all entries of the gradient at once for 2 function calls; a full finite-difference gradient needs $2n$ calls.

Interactive figure: loading…

Answers to section 1.

1. No inverse. The backward pass needs one solve with $A^\top$ ($\bar b = A^{-\top}\bar x$), and it reuses the LU factors of the forward pass. Then $\bar A = -\bar b\,x^\top$ is an outer product.
2. No difference. Since $\sum_i s_i = 1$, $\langle s, \bar s + c\mathbf 1\rangle = \langle s, \bar s\rangle + c$, and the shift $c$ cancels in $\bar s - \langle s, \bar s\rangle\mathbf 1$. In matrix form, $(\operatorname{diag}(s) - ss^\top)\mathbf 1 = 0$: the softmax ignores a common shift of its inputs, so its VJP ignores a common shift of its cotangent.
3. Add them. $\bar x = \sin x + x\cos x$: the first term comes through the factor $x$, the second through $\sin x$. This is the product rule, and it is fan-out in the theorem above.

4. Worked example #

(a) A Gaussian negative log-likelihood. Gaussian process regression, Kalman filters and 4D-Var all contain the function

with a covariance matrix $K$ and a residual $b$ (the constant $\tfrac n2 \log 2\pi$ is dropped). Take

Forward pass (the tape). The inputs are $K$ and $b$.

1. $v_1 = \alpha = \operatorname{solve}(K, b) = K^{-1}b$, so $\alpha = \tfrac15(3 - 2,\ -1 + 4) = (0.2,\ 0.6)$.
2. $v_2 = q = b^\top\alpha = 1 \cdot 0.2 + 2 \cdot 0.6 = 1.4$.
3. $v_3 = \log\det K = \log 5 \approx 1.6094$.
4. $L = \tfrac12 v_2 + \tfrac12 v_3 = 0.7 + 0.8047 \approx 1.50472$.

Backward pass, from $\bar L = 1$.

1. $L = \tfrac12 v_2 + \tfrac12 v_3$, so $\bar v_2 = \tfrac12$ and $\bar v_3 = \tfrac12$.
2. Log-determinant: $\bar K \mathrel{+}= \bar v_3\,K^{-\top} = \tfrac12 K^{-1} = \begin{pmatrix} 0.3 & -0.1 \\ -0.1 & 0.2 \end{pmatrix}$.
3. Inner product $q = b^\top\alpha$: $\bar b \mathrel{+}= \bar v_2\,\alpha = (0.1,\ 0.3)$ and $\bar\alpha = \bar v_2\,b = (0.5,\ 1)$.
4. Solve: $\bar b_{\text{s}} = K^{-\top}\bar\alpha$, so $\bar b_{\text{s}} = \tfrac15(1.5 - 1,\ -0.5 + 2) = (0.1,\ 0.3)$. Then $\bar b \mathrel{+}= \bar b_{\text{s}}$, and $\bar K \mathrel{+}= -\bar b_{\text{s}}\alpha^\top = -\begin{pmatrix} 0.02 & 0.06 \\ 0.06 & 0.18 \end{pmatrix}$.

Result.

$b$ is used twice (by the solve and by the inner product), so its cotangent has two terms: this is fan-out. The closed form $\tfrac12(K^{-1} - \alpha\alpha^\top)$ is the classic gradient of a Gaussian negative log-likelihood with respect to its covariance matrix.

(b) Softmax. $z = (0, \log 3)$ gives $e^z = (1, 3)$ and $s = (1/4,\ 3/4)$. Take $L = s_1$, so $\bar s = (1, 0)$ and $\langle s, \bar s\rangle = 1/4$. Then

Check with the known derivatives: $\partial s_1/\partial z_1 = s_1(1 - s_1) = 3/16$ and $\partial s_1/\partial z_2 = -s_1 s_2 = -3/16$. The entries sum to zero, as the shift invariance requires.

5. In code #

NumPy checks the worked example. The backward pass below follows the four backward steps of section 4 line by line; a random-direction central difference tests it.

import numpy as np

# L(K, b) = 1/2 b^T K^{-1} b + 1/2 log det K
def loss(K, b):
    return 0.5 * b @ np.linalg.solve(K, b) + 0.5 * np.linalg.slogdet(K)[1]

def loss_vjp(K, b):
    alpha = np.linalg.solve(K, b)              # forward: one solve
    Kbar = 0.5 * np.linalg.inv(K).T            # logdet VJP (fine for 2x2; use a factorization at scale)
    bbar = 0.5 * alpha                         # q = b^T alpha, first use of b
    abar = 0.5 * b                             # q = b^T alpha, cotangent of alpha
    bs = np.linalg.solve(K.T, abar)            # solve VJP: one solve with K^T
    return Kbar - np.outer(bs, alpha), bbar + bs   # fan-out: b gets two terms

K = np.array([[2.0, 1.0], [1.0, 3.0]])
b = np.array([1.0, 2.0])
Kbar, bbar = loss_vjp(K, b)
print(loss(K, b))        # 1.5047189562170504
print(Kbar)              # [[ 0.28 -0.16]  [-0.16  0.02]]
print(bbar)              # [0.2 0.6]

# Gradient check in one random direction (central differences, float64)
rng = np.random.default_rng(0)
dK, db = rng.standard_normal((2, 2)), rng.standard_normal(2)
h = 1e-5
fd = (loss(K + h * dK, b + h * db) - loss(K - h * dK, b - h * db)) / (2 * h)
exact = np.sum(Kbar * dK) + bbar @ db
print(abs(fd - exact) / abs(exact) < 1e-8)   # True

# Jacobi's formula d log|det X| = tr(X^{-1} dX), also when det X < 0
X = np.array([[1.0, 2.0, 0.0], [3.0, 1.0, 1.0], [0.0, 1.0, 2.0]])
dX = rng.standard_normal((3, 3))
fd = (np.linalg.slogdet(X + h * dX)[1] - np.linalg.slogdet(X - h * dX)[1]) / (2 * h)
exact = np.trace(np.linalg.solve(X, dX))
print(round(np.linalg.det(X), 6), abs(fd - exact) / abs(exact) < 1e-8)   # -11.0 True

In PyTorch the same check is torch.autograd.gradcheck, and a new primitive with a hand-written VJP is a torch.autograd.Function (in JAX: jax.custom_vjp). Its backward method is one row of the table in section 3. The notebook builds a small tape in NumPy and checks every rule of the table.

6. Where it shows up #

• Every autodiff framework. loss.backward() is the sweep of section 3. To add your own primitive (a fused kernel, a call into a C++ or Fortran physics library, a sparse solver), you register its VJP.
• Gaussian likelihoods. Gaussian process hyperparameters, Kalman-filter likelihoods and the background term of 4D-Var (lessons 06 and 07) need solves and log-determinants. By the chain rule, the gradient with respect to a kernel parameter $\theta$ is $\langle \bar K, \partial K/\partial\theta\rangle = \tfrac12\operatorname{tr}\big((K^{-1} - \alpha\alpha^\top)\,\partial K/\partial\theta\big)$ with $\alpha = K^{-1}b$: section 4 plus one more inner product.
• Normalizing flows and change of variables. The term $\log\lvert\det A\rvert$ with the Jacobian $A = \partial T/\partial x$ of the flow has the VJP $\bar y\,A^{-\top}$, which needs a factorization of $A$. This is one reason why flows use triangular or coupling layers, whose determinants are cheap.
• Seismic imaging. Full-waveform inversion fits a velocity model with millions of cells to recorded seismograms. The misfit is one number, so reverse mode (there called the adjoint-state method) gives the whole gradient for about the cost of one more wave simulation (lesson 04).
• Weather forecasting. 4D-Var gets the gradient of its cost function from an adjoint model of the forecast model (lesson 06). In the other direction, a JVP tells how a whole simulated field moves with each of a few physical parameters.
• Training at scale. Activation memory is the tape: it grows with depth, batch size and sequence length. Activation checkpointing trades it for recomputation.

7. Common mistakes #

• Building Jacobians or inverses. An $n\times n$ softmax Jacobian, or $A^{-1}$ for a solve. Use the VJP formulas: $O(n)$ for softmax, one transposed solve for solve.
• $A^{-1}$ instead of $A^{-\top}$ for a non-symmetric $A$, in the solve and log-det VJPs. The shapes match, so only a test with a non-symmetric matrix catches this bug.
• Overwriting instead of adding cotangents at fan-out. A related bug is aliasing: if two cotangents are the same array object, an in-place += on one changes the other.
• Forgetting to unbroadcast. The cotangent of a bias is summed over the batch. A shape check catches this.
• Symmetric parameters. One off-diagonal parameter gets $\bar K_{12} + \bar K_{21}$ (section 4).
• A weak gradient check. Forward differences, float32, $h = 10^{-12}$, or one coordinate only. Use float64, central differences, $h \approx 10^{-6}$ to $10^{-5}$ and a random direction.
• "Reverse mode is always better." For $n \le m$, or with a tight memory budget, forward mode can be the better choice.

8. Check yourself #

1. $C = AB$ with $A \in \R^{m\times n}$ and $B \in \R^{n\times k}$. Give $\bar A$ and $\bar B$ and check their shapes.
   Answer

   $\bar A = \bar C B^\top$: $(m\times k)(k\times n) = m\times n$. $\bar B = A^\top\bar C$: $(n\times m)(m\times k) = n\times k$. Both have the shapes of $A$ and $B$.
2. Scalar solve: $ax = b$ with $a = 2$, $b = 4$ and $\bar x = 1$. Compute $\bar b$ and $\bar a$ with the solve VJP, and check $\bar a$ against $\partial(b/a)/\partial a$.
   Answer

   $x = 2$. $\bar b = \bar x/a = 0.5$ and $\bar a = -\bar b\,x = -1$. Directly: $\partial(b/a)/\partial a = -b/a^2 = -1$.
3. Why is the VJP of $\log\lvert\det A\rvert$ equal to $\bar y\,A^{-\top}$ and not $\bar y\,A^{-1}$? When do the two agree?
   Answer

   $dy = \operatorname{tr}(A^{-1}dA)$, and $\langle G, dA\rangle = \operatorname{tr}(G^\top dA)$, so $G^\top = A^{-1}$, that is, $G = A^{-\top}$. They agree when $A$ is symmetric (for example a covariance matrix).
4. A climate emulator maps 4 forcing parameters to a field of $10^6$ values, and a scalar loss compares the field with observations. (a) How do you get the full Jacobian of the emulator, and with how many passes? (b) And the gradient of the loss with respect to the 4 parameters?
   Answer

   (a) Forward mode, 4 JVPs (reverse mode would need $10^6$ VJPs). (b) One reverse pass, or 4 forward passes: both are cheap. Forward mode keeps no tape, so it wins if the emulator's intermediate values do not fit in memory; otherwise one reverse pass is the cheapest.
5. In the sweep, why must the nodes be visited in reverse tape order? What goes wrong if a node is processed before all its users?
   Answer

   The sweep is back substitution with an upper triangular matrix: $\bar v_i$ needs the final $\bar v_k$ of every node $k$ that uses $v_i$, and those come later on the tape. If node $i$ is processed too early, it passes on an incomplete cotangent, and the paths through the missing users are lost.
6. Your gradient check in float64 with $h = 10^{-10}$ gives a relative error of $10^{-4}$. Bug, or a bad check?
   Answer

   Probably the check. The rounding error is about $\varepsilon\lvert f\rvert/h \approx 2\cdot10^{-6}\lvert f\rvert$, which is large when $\lvert f\rvert$ is much larger than the directional derivative. Repeat with $h \approx 10^{-5}$ (scaled to $x$). If the error stays above about $10^{-7}$ for every $h$, look for a bug: a bug gives an error floor that does not move with $h$.

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. Define the vector-Jacobian product $\bar x = F'(x)^*[\bar y]$ with inner products.

   Show answer

2. $C = AB$. Give the VJPs $\bar A$ and $\bar B$.

   Show answer

3. $X = A^{-1}B$ (a linear solve). Give the VJPs, and say what the backward pass computes.

   Show answer

4. Give the VJP of $y = \log\lvert\det A\rvert$, and the formula behind it.

   Show answer

5. A value $v$ is used by two later nodes. How is $\bar v$ computed in reverse mode?

   Show answer

6. How many passes does each mode need for the full Jacobian of $F: \R^n \to \R^m$?

   Show answer

7. What is the price of reverse mode, compared with forward mode?

   Show answer

8. Why do central differences beat forward differences in a gradient check?

   Show answer

9. What is the best step $h$ for central differences in float64, and why?

   Show answer

10. Why is the softmax VJP $O(n)$?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Implement backprop for this small graph." — Record the nodes in execution order; seed $\bar L = 1$; visit the nodes in reverse order, call each local VJP, and add into the parents' cotangents.
• "Forward or reverse mode: costs, memory, when?" — $n$ JVPs or $m$ VJPs for an $m\times n$ Jacobian, each a few times the cost of $F$; reverse mode stores the tape; reverse for losses, forward for few inputs or tight memory.
• "What is the gradient of $\log\det K$? Of $b^\top K^{-1}b$ with respect to $K$?" — $K^{-\top}$; and $-K^{-\top}bb^\top K^{-\top}$, which is $-\alpha\alpha^\top$ with $\alpha = K^{-1}b$ for a symmetric $K$.
• "Why is the softmax backward pass $O(n)$?" — $\bar z = s\odot(\bar s - \langle s, \bar s\rangle\mathbf 1)$; the Jacobian is never formed.
• "How do you add a sparse solver written in C++ to PyTorch?" — A torch.autograd.Function: forward solves and saves the factorization; backward does one transposed solve for $\bar b$ and forms $\bar A = -\bar b x^\top$ (only on the sparsity pattern of $A$, if $A$ is sparse).
• "What is the VJP of a VJP?" — The JVP. A VJP is linear in its cotangent, and the adjoint of an adjoint is the original map (Blondel and Roulet, §8.4).

10. Practice #

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

1. Write the VJP of the matrix product and check it with central differences in a random direction.
2. Write the VJP of a linear solve without forming $A^{-1}$, and test it with a non-symmetric matrix.
3. Check Jacobi's formula $d\log\lvert\det X\rvert = \operatorname{tr}(X^{-1}dX)$ by finite differences, also for $\det X < 0$.
4. Build a 40-line tape with backward for add, mul, tanh and sum; get the gradient of a two-layer network and check it against finite differences.
5. Plot the finite-difference error against $h$ (the V-curve) and find the best step for the forward and the central difference.

11. Watch and read #

What is Automatic Differentiation?Ari Seff · 14 min

Forward and reverse mode on a small computational graph, and why reverse mode wins for a scalar loss.

Lecture 1 Part 2: Derivatives as Linear OperatorsMIT OpenCourseWare · 18.S096 Matrix Calculus · 48 min

The derivative as a linear map, the view used in section 3 of this lesson.

The spelled-out intro to neural networks and backpropagation: building microgradAndrej Karpathy · 2 h 26 min

A scalar tape with a backward pass, built line by line; watch how the gradients of a reused value are added.

• Bright, Edelman, Johnson, Matrix Calculus (for Machine Learning and Beyond): ch. 2 (derivatives as linear operators), §5.1 (the matrix inner product), §6.3 (reverse-mode "adjoint" differentiation), ch. 7 (derivatives of the determinant and the inverse), §8.3–8.4 (forward and reverse mode on computational graphs). The lectures are on MIT OpenCourseWare 18.S096.
• Blondel and Roulet, The Elements of Differentiable Programming: §2.3 (linear maps, adjoints, JVPs and VJPs), §8.1–8.3 (forward and reverse mode, their costs, the Baur–Strassen theorem), §8.4 (VJPs of primitives).
• Giles, An extended collection of matrix derivative results for forward and reverse mode algorithmic differentiation (2008): §2.2–2.3 (product, inverse, determinant, solve), §2.4 (the Gaussian likelihood).
• autodidact: a small, readable version of the HIPS autograd library. Read its tape and VJP table after you finish exercise 4.