How the Jacobian Matrix Works

A scalar derivative answers "how fast does one number change when I wiggle one other number?" Almost every function we care about in machine learning takes a vector or a matrix in and returns a vector or a matrix out. The derivative you actually need in that case is the Jacobian matrix: one stacked entry per input–output pair, evaluated at a point. Once you can read the Jacobian of a single linear layer, you can read the gradient of every linear layer, the backward of every attention block, and the contract of every kernel in a runtime like C-Kernel-Engine.

In the Matrix Wx+b: From Scalars To Transformers post, we followed the forward path: scalars become vectors, vectors become matrices, and matrix projections become transformer layers. This post is the backward-path sequel: if Y = XW + b, how does every output depend on every input and every weight? That dependency table is the Jacobian.

The Problem With One Number

For f(x) = x^2 the derivative f'(x) = 2x is a single number at every point. The picture is a one-dimensional slope on a one-dimensional line. As soon as the function has many inputs and many outputs — a linear layer, a softmax, a multi-head attention block — a single number is not enough. The slope depends on which input you wiggle, and you also need to know which output you are watching.

With n inputs and m outputs, that means m times n partial derivatives. The clean way to package them all in one place is the Jacobian.

The Definition

Let f be a vector-valued function with n inputs and m outputs:

\(f : \mathbb{R}^n \to \mathbb{R}^m\)

The Jacobian J of f at a point x is an m-by-n matrix. Its row index is the output being watched. Its column index is the input being wiggled. The entry at row i and column j is:

\(J_{ij} \;=\; \dfrac{\partial f_i}{\partial x_j}\)

The shape is always outputs-by-inputs. Each row is the gradient of one output. Each column is one input's sensitivity, viewed from every output. That shape rule will be the single most useful fact in this post, because every layer in a neural net maps an input vector space to an output vector space and therefore has a Jacobian of exactly that shape.

Set Up: One Token Through One Linear Layer

Take the smallest useful case. A single token, a single linear layer, a single forward pass. We use the row-vector convention throughout this post: X is a row, W has one column per output, Y is a row.

\(X \;=\; \begin{bmatrix} x_1 & x_2 & x_3 & x_4 \end{bmatrix} \quad \text{shape } \bigl[1 \times 4\bigr]\)

\(W \;=\; \begin{bmatrix} w_{11} & w_{12} & w_{13} & w_{14} & w_{15} & w_{16} \\ w_{21} & w_{22} & w_{23} & w_{24} & w_{25} & w_{26} \\ w_{31} & w_{32} & w_{33} & w_{34} & w_{35} & w_{36} \\ w_{41} & w_{42} & w_{43} & w_{44} & w_{45} & w_{46} \end{bmatrix} \quad \text{shape } \bigl[4 \times 6\bigr]\)

\(b \;=\; \begin{bmatrix} b_1 & b_2 & b_3 & b_4 & b_5 & b_6 \end{bmatrix} \quad \text{shape } \bigl[1 \times 6\bigr]\)

\(Y \;=\; X W + b \quad \text{shape } \bigl[1 \times 6\bigr]\)

This is the same orientation as the previous post: rows are tokens, columns of W are output directions. The transpose convention y = W^T x + b is equivalent — it just rotates the page. We will stay with the row-vector form because it is the one that makes "column j of W belongs to output j" obvious.

Forward, Spelled Out: The Six Output Equations

Write the matrix product out longhand. Each output y_j is a dot product of the input row X with column j of W, plus the bias b_j:

\(\begin{aligned} y_1 &= x_1 w_{11} + x_2 w_{21} + x_3 w_{31} + x_4 w_{41} + b_1 \\ y_2 &= x_1 w_{12} + x_2 w_{22} + x_3 w_{32} + x_4 w_{42} + b_2 \\ y_3 &= x_1 w_{13} + x_2 w_{23} + x_3 w_{33} + x_4 w_{43} + b_3 \\ y_4 &= x_1 w_{14} + x_2 w_{24} + x_3 w_{34} + x_4 w_{44} + b_4 \\ y_5 &= x_1 w_{15} + x_2 w_{25} + x_3 w_{35} + x_4 w_{45} + b_5 \\ y_6 &= x_1 w_{16} + x_2 w_{26} + x_3 w_{36} + x_4 w_{46} + b_6 \end{aligned}\)

Stop here. This is the picture you need to hold in your head before any derivative makes sense. Each output is a sum of four products, and each product uses one entry of X and one entry of one column of W.

Which Weights Touch Which Output?

The dependency structure is what makes the Jacobian tractable. Read the six equations column by column:

• y_1 depends only on column 1 of W, i.e. on W[:,1] = (w_11, w_21, w_31, w_41)^T.
• y_2 depends only on column 2 of W, i.e. on W[:,2] = (w_12, w_22, w_32, w_42)^T.
• y_3 depends only on column 3 of W.
• y_4 depends only on column 4 of W.
• y_5 depends only on column 5 of W.
• y_6 depends only on column 6 of W.

So y_1 does not depend on w_12, w_22, w_32, or w_42 at all. Those weights contribute to y_2, not to y_1. If you wiggle w_12 by a small amount, y_1 does not move. That is the single most important observation in this whole post, because every partial derivative that is "not on the diagonal" is exactly zero.

The Local Partials, One Output At A Time

Now differentiate. The bias disappears because b_j is constant with respect to every weight. Take the first output and write out every partial derivative explicitly.

Derivatives Of y_1

Look at y_1 = x_1 w_11 + x_2 w_21 + x_3 w_31 + x_4 w_41 + b_1. Each weight in column 1 of W touches y_1 once, with coefficient being the matching input:

\(\begin{aligned} \dfrac{\partial y_1}{\partial w_{11}} &= x_1 \\ \dfrac{\partial y_1}{\partial w_{21}} &= x_2 \\ \dfrac{\partial y_1}{\partial w_{31}} &= x_3 \\ \dfrac{\partial y_1}{\partial w_{41}} &= x_4 \end{aligned}\)

The weights in every other column are nowhere in the equation for y_1. So their partial derivatives are exactly zero:

\(\begin{aligned} \dfrac{\partial y_1}{\partial w_{12}} &= 0 \\ \dfrac{\partial y_1}{\partial w_{22}} &= 0 \\ \dfrac{\partial y_1}{\partial w_{32}} &= 0 \\ \dfrac{\partial y_1}{\partial w_{42}} &= 0 \end{aligned}\)

And the same pattern repeats for every other column of W — there are 24 weights in W and 23 of them do not affect y_1.

Derivatives Of y_2

Now y_2 = x_1 w_12 + x_2 w_22 + x_3 w_32 + x_4 w_42 + b_2. By the same rule, the non-zero partials are the ones in column 2 of W:

\(\begin{aligned} \dfrac{\partial y_2}{\partial w_{12}} &= x_1 \\ \dfrac{\partial y_2}{\partial w_{22}} &= x_2 \\ \dfrac{\partial y_2}{\partial w_{32}} &= x_3 \\ \dfrac{\partial y_2}{\partial w_{42}} &= x_4 \end{aligned}\)

The non-zero rule for column 1 of W inverts: the weights that fed y_1 do not feed y_2:

\(\begin{aligned} \dfrac{\partial y_2}{\partial w_{11}} &= 0 \\ \dfrac{\partial y_2}{\partial w_{21}} &= 0 \\ \dfrac{\partial y_2}{\partial w_{31}} &= 0 \\ \dfrac{\partial y_2}{\partial w_{41}} &= 0 \end{aligned}\)

Derivatives Of y_3, y_4, y_5, y_6

Same shape again. For each output y_j, the only non-zero partials are the four entries of column j of W, and each of those four partials is the matching input:

\(\dfrac{\partial y_j}{\partial w_{i, k}} \;=\; \begin{cases} x_i & \text{if } k = j \\ 0 & \text{if } k \ne j \end{cases}\)

Read the indices carefully:

• i is the input index (the row of W).
• k is the weight column index, which is also the output index because of the orientation we chose.
• The partial is non-zero only when the weight's column matches the output you are differentiating, and in that case the value is the input at row i.

That single rule, written out, is the entire Jacobian of Y = XW + b with respect to W.

Where The Jacobian Appears

The full object that contains every partial y_j / partial w_i, k is the Jacobian of Y with respect to W. It is the conceptual object that holds the entire dependency table between every output and every weight. If you wrote it out literally as a 2D matrix, you would have to flatten either the outputs or the weights, and you would lose the structure. It is cleaner to keep three indices: output j, weight row i, weight column k.

The non-zero pattern is the thing that matters. Using the same notation as above, the rule for the Jacobian of Y with respect to W is:

\(\dfrac{\partial y_j}{\partial W[:, j]} \;=\; X\)

\(\dfrac{\partial y_j}{\partial W[:, k]} \;=\; 0 \quad \text{when } k \ne j\)

In words: every output y_j "sees" exactly one column of W, namely column j, and the sensitivity along that column is the full input row X. Every other column contributes zero. That is the key mental model.

The Jacobian is what you would write down if you wanted to do forward-mode autodiff through this layer, or if you wanted to compute the Jacobian-vector product for many different upstream vectors one at a time. You almost never actually do that for a linear layer. You almost always do the reverse: take the upstream gradient, multiply by the local Jacobian, return the input gradient. That is the next section.

Backprop: Multiply The Local Derivative By The Parent Gradient

Now flip direction. Suppose some loss L depends on Y. We are given the upstream gradient

\(G \;=\; \dfrac{\partial \mathcal{L}}{\partial Y} \;=\; \begin{bmatrix} g_1 & g_2 & g_3 & g_4 & g_5 & g_6 \end{bmatrix} \quad \text{shape } \bigl[1 \times 6\bigr]\)

where g_j = partial L / partial y_j. We want the gradient of the loss with respect to the weights, dW. Apply the chain rule column by column.

Column 1 Of dW

dW[:,1] collects the contribution of every weight in column 1 of W to the loss. The loss flows in through y_1 only, with weight g_1. Each weight in column 1 has local derivative equal to the matching input:

\(\begin{aligned} \mathrm{d}w_{11} &= \dfrac{\partial \mathcal{L}}{\partial w_{11}} \;=\; \dfrac{\partial \mathcal{L}}{\partial y_1} \cdot \dfrac{\partial y_1}{\partial w_{11}} \;=\; g_1 \cdot x_1 \\ \mathrm{d}w_{21} &= \dfrac{\partial \mathcal{L}}{\partial w_{21}} \;=\; g_1 \cdot x_2 \\ \mathrm{d}w_{31} &= \dfrac{\partial \mathcal{L}}{\partial w_{31}} \;=\; g_1 \cdot x_3 \\ \mathrm{d}w_{41} &= \dfrac{\partial \mathcal{L}}{\partial w_{41}} \;=\; g_1 \cdot x_4 \end{aligned}\)

Stack those four partials as a column:

\(\mathrm{d}W[:,1] \;=\; X^{\mathsf T} \cdot g_1 \;=\; \begin{bmatrix} x_1 \\ x_2 \\ x_3 \\ x_4 \end{bmatrix} \cdot g_1 \;=\; \begin{bmatrix} x_1 g_1 \\ x_2 g_1 \\ x_3 g_1 \\ x_4 g_1 \end{bmatrix}\)

Column 2 Of dW

Same shape, but the upstream gradient is now g_2:

\(\mathrm{d}W[:,2] \;=\; X^{\mathsf T} \cdot g_2 \;=\; \begin{bmatrix} x_1 g_2 \\ x_2 g_2 \\ x_3 g_2 \\ x_4 g_2 \end{bmatrix}\)

Columns 3 Through 6

Same again, with g_3, g_4, g_5, g_6:

\(\mathrm{d}W[:,k] \;=\; X^{\mathsf T} \cdot g_k \quad \text{for } k = 1, 2, 3, 4, 5, 6\)

Stacking The Columns

Glue the six columns together. The full weight gradient is the 4 x 6 matrix:

\(\mathrm{d}W \;=\; \begin{bmatrix} x_1 g_1 & x_1 g_2 & x_1 g_3 & x_1 g_4 & x_1 g_5 & x_1 g_6 \\ x_2 g_1 & x_2 g_2 & x_2 g_3 & x_2 g_4 & x_2 g_5 & x_2 g_6 \\ x_3 g_1 & x_3 g_2 & x_3 g_3 & x_3 g_4 & x_3 g_5 & x_3 g_6 \\ x_4 g_1 & x_4 g_2 & x_4 g_3 & x_4 g_4 & x_4 g_5 & x_4 g_6 \end{bmatrix}\)

The Compact Formula

The pattern is now obvious: each column of dW is the input row X transposed and scaled by one entry of G. That is exactly an outer product, and an outer product between a column and a row is a matrix product:

\(\mathrm{d}W \;=\; X^{\mathsf T} \cdot G\)

Shape check:

\(\bigl[4 \times 1\bigr] \cdot \bigl[1 \times 6\bigr] \;=\; \bigl[4 \times 6\bigr]\)

Right shape, right contents. This is the backward of one token through one linear layer.

Backprop For The Other Inputs

The same chain rule gives the gradient with respect to the input row X and the bias b.

Input Gradient dX

Each input x_i touches every output y_j through the weight w_i, j. The chain rule says

\(\dfrac{\partial \mathcal{L}}{\partial x_i} \;=\; \sum_{j=1}^{6} \dfrac{\partial \mathcal{L}}{\partial y_j} \cdot \dfrac{\partial y_j}{\partial x_i} \;=\; \sum_{j=1}^{6} g_j \cdot w_{i, j}\)

Stack those four sums back into a row:

\(\mathrm{d}X \;=\; G \cdot W^{\mathsf T}\)

Shape check:

\(\bigl[1 \times 6\bigr] \cdot \bigl[6 \times 4\bigr] \;=\; \bigl[1 \times 4\bigr]\)

Bias Gradient db

Every y_j depends on b_j with derivative 1, and on no other b_k. So

\(\dfrac{\partial \mathcal{L}}{\partial b_j} \;=\; \dfrac{\partial \mathcal{L}}{\partial y_j} \cdot 1 \;=\; g_j\)

Stacking:

\(\mathrm{d}b \;=\; G\)

That is the entire backward pass for the single-token linear layer. Three formulas, three shape checks:

\(\begin{aligned} \mathrm{d}W &= X^{\mathsf T} G & \bigl[4 \times 1\bigr] \cdot \bigl[1 \times 6\bigr] &= \bigl[4 \times 6\bigr] \\ \mathrm{d}X &= G W^{\mathsf T} & \bigl[1 \times 6\bigr] \cdot \bigl[6 \times 4\bigr] &= \bigl[1 \times 4\bigr] \\ \mathrm{d}b &= G & \bigl[1 \times 6\bigr] &= \bigl[1 \times 6\bigr] \end{aligned}\)

Many Tokens: The Batch Extension

Stack T tokens as rows of a matrix X. The forward is now a matrix-matrix multiply:

\(X \in \mathbb{R}^{T \times N}, \quad W \in \mathbb{R}^{N \times K}, \quad b \in \mathbb{R}^{K}, \quad Y = X W + \mathbf{1}\, b^{\mathsf T} \in \mathbb{R}^{T \times K}\)

with the upstream gradient G = partial L / partial Y also of shape T x K. The three backward formulas extend with no extra work:

\(\begin{aligned} \mathrm{d}W &= X^{\mathsf T} G & \bigl[N \times T\bigr] \cdot \bigl[T \times K\bigr] &= \bigl[N \times K\bigr] \\ \mathrm{d}X &= G W^{\mathsf T} & \bigl[T \times K\bigr] \cdot \bigl[K \times N\bigr] &= \bigl[T \times N\bigr] \\ \mathrm{d}b &= \sum_{t=1}^{T} G[t, :] & \bigl[1 \times K\bigr] &= \bigl[1 \times K\bigr] \end{aligned}\)

The bias gradient is now a sum over the token axis, because every token's bias contribution goes into the same vector b. That is the only place the batch dimension shows up in the backward.

The Jacobian Of Matrix Multiplication

Now zoom out. The same dependency story, with different letters, is the Jacobian of plain matrix multiplication. Let

\(A \in \mathbb{R}^{M \times N}, \quad B \in \mathbb{R}^{N \times K}, \quad C = A \cdot B \in \mathbb{R}^{M \times K}\)

Each output element is one dot product:

\(C[i, j] \;=\; \sum_{r=0}^{N-1} A[i, r] \cdot B[r, j]\)

C[i, j] depends on exactly one row of A and exactly one column of B. That is the same structure we just saw in the single-token case — column j of W played the role of column j of B, and the input row played the role of row i of A. The Jacobian is huge but regular, and the non-zero partials are slices of the other matrix:

\(\dfrac{\partial C[i, j]}{\partial A[p, q]} \;=\; \begin{cases} B[q, j] & \text{if } p = i \\ 0 & \text{otherwise} \end{cases}, \qquad \dfrac{\partial C[i, j]}{\partial B[p, q]} \;=\; \begin{cases} A[i, p] & \text{if } q = j \\ 0 & \text{otherwise} \end{cases}\)

The picture below makes that dependency structure visual. The row of A in purple and the column of B in purple are the only entries that flow into the highlighted cell C[2, 3] in yellow.

Backprop Through A Matrix Multiplication

Given the upstream gradient G = partial L / partial C, the chain rule for matrix functions gives the two compact formulas:

\(\begin{aligned} \dfrac{\partial \mathcal{L}}{\partial A} &= G \cdot B^{\mathsf T} \\ \dfrac{\partial \mathcal{L}}{\partial B} &= A^{\mathsf T} \cdot G \end{aligned}\)

These are vector-Jacobian products. Conceptually you are multiplying G by the transpose of the local Jacobian of the matmul layer, layer by layer, in reverse order. In practice the runtimes never build the full Jacobian — they call a specialized backward kernel that returns G * B^T directly.

Why Backprop Does Not Build The Full Jacobian

It is worth saying out loud just how big these Jacobians get. Walk the ladder of three concrete sizes. Each one is real, each one is a count of entries the runtime would have to materialize if it actually built the full Jacobian. Nobody builds any of them. The point of backpropagation is that you do not have to.

Push that matrix through a single linear layer for one token, and you get 4096 output partials, one per output direction. The full Jacobian of the 4096 outputs with respect to all 16.7 million weights would be a 4096 x 16.7M table.

For a full 4096-by-4096 matrix multiplication, the output is 4096 x 4096 entries and one of the inputs alone is 4096 x 4096 entries. Stack the partials of every output with respect to every input entry and the table becomes a four-dimensional object.

The reason backpropagation works in practice is that the loss is a scalar at the top of the network, so what you actually need at every layer is a vector-Jacobian product, not a Jacobian. Two equivalent ways to phrase it:

• VJP (reverse mode). Given an upstream vector g in R^m, return J_f^T g in R^n. This is what backpropagation computes. PyTorch, JAX, and the CKE backward kernels all implement VJPs.
• JVP (forward mode). Given a tangent vector v in R^n, return J_f v in R^m. Cheap when the input space is small.

For a model with one loss and millions of parameters, reverse mode is dramatically cheaper — one backward pass through the whole graph, instead of one per parameter. That is why every modern training stack is reverse mode under the hood, and why every specialized backward kernel is secretly a VJP.

How This Maps To C-Kernel-Engine

Once you see that the Jacobian is the local derivative object, the architecture of a kernel runtime becomes obvious. Each kernel only needs to know its own Jacobian. The orchestrator stitches the kernels into a forward and a backward pass.

For the Y = XW + b layer the kernel contract looks like this:

\(\begin{aligned} \text{Forward:} &\quad Y = X W + b \\ \text{Backward:} &\quad \mathrm{d}X = G W^{\mathsf T} \\ \text{Backward:} &\quad \mathrm{d}W = X^{\mathsf T} G \\ \text{Backward:} &\quad \mathrm{d}b = \sum_{\text{tokens}} G \end{aligned}\)

The kernel does not build the full Jacobian. It uses the structure of the Jacobian to compute the compact GEMM / GEMV forms directly. The Jacobian is the conceptual object. The kernel implements the computational shortcut.

• Forward kernel. Takes X, W, b. Computes Y = XW + b. Saves the inputs and the output so the backward kernel can read them.
• Backward kernel. Takes the upstream gradient G. Reads the saved X, W. Returns dX = G W^T, dW = X^T G, and db = sum G.
• Orchestrator. Walks the graph, calls each forward kernel in order, then walks the graph in reverse calling each backward kernel, accumulating gradients where multiple downstream paths meet.

The same three-role structure generalizes beyond the linear layer. A LayerNorm kernel, a softmax kernel, a flash-attention kernel, a quantization kernel: each one ships with its own forward and backward. None of them know about the others. The orchestrator is the only thing that needs to know the whole graph, and the orchestrator itself does not need to understand any of the local math — it just has to know the shapes and the kernel names.

The Toy Scalar Example, For Sanity

Drop the matrices and look at a 2D-to-2D toy map:

\(f(x, y) \;=\; \bigl( x^{2},\; x y \bigr)\)

The Jacobian is the 2 x 2 matrix

\(J(x, y) \;=\; \begin{bmatrix} 2x & 0 \\ y & x \end{bmatrix}\)

At the point (1, 2),

\(J(1, 2) \;=\; \begin{bmatrix} 2 & 0 \\ 2 & 1 \end{bmatrix}\)

This map is not conformal: the singular values of J(1, 2) are sigma_1 ~ 2.92 and sigma_2 ~ 0.68, with ratio sigma_1 / sigma_2 ~ 4.27. Push a small circle of points around (1, 2) through J(1, 2). The output is a stretched ellipse whose principal axes are the right singular vectors of J and whose axis lengths are sigma_1 and sigma_2:

This is the same story as Y = XW + b, just with a matrix you can see on one page. The Jacobian is the local linear map, and the local linear map is what every backward kernel implements.

What To Take Away

• The Jacobian is the matrix of every partial derivative of a vector-valued function, with shape outputs-by-inputs.
• For Y = XW + b with the row convention X in R^(1 x N) and W in R^(N x K), each output y_j depends only on column j of W.
• The non-zero partials are partial y_j / partial w_i, k = x_i when k = j and zero otherwise. Every off-diagonal weight derivative is exactly zero.
• Backprop through that layer is three formulas: dW = X^T G, dX = G W^T, db = G.
• For matrix multiplication C = AB, the Jacobian is huge but extremely regular: every output C[i, j] depends on exactly one row of A and one column of B, and the partials are slices of the other matrix.
• Backprop never forms the full Jacobian. It computes vector-Jacobian products, which is why a specialized backward kernel can replace the full derivative.
• A kernel runtime like C-Kernel-Engine treats every kernel as a local math contract. Forward kernel, backward kernel, orchestrator. The Jacobian lives inside the kernel; the runtime stitches the kernels.
• The Jacobian is the conceptual object. The kernel implements the computational shortcut.

Read Next: Jacobians Inside C-Kernel-Engine

The important production point is not the diagram-generation pipeline. It is that Jacobians explain why a runtime such as C-Kernel-Engine needs explicit forward kernels, backward kernels, shape contracts, and memory ownership. The runtime is not trying to materialize giant derivative matrices. It is trying to compile the local derivative structure into kernels that can run predictably on real CPU hardware.

• Start with What Is the C-Kernel-Engine? if you want the broader runtime framing: graph, lowering, memory plan, generated C, and measurable parity.
• Read DL/d-LLM: The Full Backward Pass for the backprop side of this same idea: gradients flow backward through local operations instead of building one enormous Jacobian for the whole model.
• Read SIMD Deep Dive for the hardware side: once the derivative structure is known, the work still has to become vectorized loads, dot products, reductions, and cache-friendly loops.
• Read Linux System Programming For AI Kernels for the systems side: Jacobian math becomes useful only when the runtime can control memory, threads, NUMA placement, and repeatable measurement.

That is the bridge from calculus to CKE: the Jacobian is the mathematical object, vector-Jacobian products are the training shortcut, and C kernels are the concrete execution artifact.