The chain rule as matrix products — Linear algebra for LLMs

The Jacobian

\displaystyle y = f(x),\ x \in \mathbb{R}^n,\ y \in \mathbb{R}^m: \quad J_{ij} = \frac{\partial y_i}{\partial x_j}

For a function from vectors to vectors, the derivative is a matrix: row $i$ holds the slopes of output $i$ with respect to every input. Shape $m \times n$ , outputs by inputs.

Write

Jacobian

J_{ij} = \partial y_i / \partial x_j

, shape (outputs × inputs).

A linear map is its own Jacobian

\displaystyle y = Wx \;\Rightarrow\; J = W

$y_i = \sum_j W_{ij} x_j$ , so $\partial y_i / \partial x_j = W_{ij}$ . The best linear approximation of a linear map is the map itself.

Composing functions multiplies Jacobians

\displaystyle z = g(f(x)) \;\Rightarrow\; J_{z,x} = J_{z,y}\, J_{y,x}

The chain rule for vectors is a matrix product, in the same right-to-left order as composing maps (Part 3). A network is a long composition, so its derivative is a long product of Jacobians.

Write

Chain rule:

J_{g \circ f} = J_g\, J_f

.

Elementwise functions have diagonal Jacobians

\displaystyle y = \sigma(x) \text{ entrywise} \;\Rightarrow\; J = \operatorname{diag}(\sigma'(x))

Output $i$ depends only on input $i$ , so the Jacobian is diagonal, and multiplying by it is just $\odot\, \sigma'(x)$ (Part 2). Code never builds it.

Backprop: start from the loss

\displaystyle \frac{\partial L}{\partial x} = \frac{\partial L}{\partial z}\, J_{z,y}\, J_{y,x}

\displaystyle (1 \times m)(m \times k)(k \times n)

The loss is a scalar, so $\partial L / \partial z$ is a single row. Multiplying from the left keeps a row vector at every step: a vector–Jacobian product, cheap. Multiplying the Jacobians together first would build big matrices for nothing. That’s why gradients flow backward, from the loss to the input.

Write

Backprop = vector–Jacobian products, from the loss backward:

\delta_x = \delta_y J_{y,x}

.

✏️ Your turn

One backward step

\displaystyle y = Wx, \quad W = \begin{bmatrix} 1 & 2 \\ 3 & 4 \end{bmatrix}, \quad \frac{\partial L}{\partial y} = [1, -1]

Write $\partial L / \partial x$ .

Work it out on paper, then go to the next card to check.

Answer

\displaystyle \frac{\partial L}{\partial x} = [1, -1]\begin{bmatrix} 1 & 2 \\ 3 & 4 \end{bmatrix} = [-2, -2]

The upstream row times the Jacobian $W$ . In column layout the same thing is written $W^\top \frac{\partial L}{\partial y}$ .

In PyTorch

torch.autograd.functional.jacobian(f, x)     # the full J, for small checks
torch.autograd.grad(L, x)                    # what backward computes: a VJP

In a transformer

Autograd stores what each operation needs for its vector–Jacobian product during the forward pass (the activations), then runs the products in reverse. That stored data is why training uses far more memory than inference.

📌 Recap

The chain rule in one card

Jacobian: $J_{ij} = \partial y_i / \partial x_j$ ; for $y = Wx$ , $J = W$ .
Composition multiplies Jacobians; elementwise functions give diagonals ( $\odot$ ).
Backprop multiplies a row vector by Jacobians, from the loss backward.

Write

The three lines above.