The gradient handles functions with many inputs but a single scalar output (like a loss). The Jacobian generalizes this one step further — it handles functions with many inputs and many outputs, which is exactly what every intermediate layer of a neural network is.
Definition
For a function \(\mathbf{f}: \mathbb{R}^n \to \mathbb{R}^m\) (taking \(n\) inputs and producing \(m\) outputs), the Jacobian is an \((m,n)\) matrix — row \(i\) is the gradient of output \(f_i\) with respect to all \(n\) inputs. Each entry \(J_{ij} = \frac{\partial f_i}{\partial x_j}\) captures how much output \(i\) changes when input \(j\) changes.
Numerical Example
At \((x,y)=(2,1)\): \(\mathbf{J} = \begin{bmatrix}4 & 4\\1 & 2\end{bmatrix}\) — a \((2,2)\) matrix, since there are 2 outputs and 2 inputs.
Why a Neural Network Layer Needs a Jacobian, Not Just a Gradient
A single layer \(\mathbf{y} = \mathbf{W}\mathbf{x}+\mathbf{b}\) maps a vector input to a vector output — it has multiple inputs and multiple outputs. When backpropagation passes a gradient signal backward through this layer, it needs to know how each output was affected by each input — that's the layer's Jacobian. In fact, for this specific linear layer, the Jacobian with respect to \(\mathbf{x}\) is simply \(\mathbf{W}\) itself — which is exactly why backpropagation through a linear layer multiplies the incoming gradient by \(\mathbf{W}^\top\) (see Matrix Transpose).
Code
import torch
def f(v):
x, y = v[0], v[1]
return torch.stack([x**2 * y, x + y**2])
x = torch.tensor([2.0, 1.0], requires_grad=True)
J = torch.autograd.functional.jacobian(f, x)
print(J)
# tensor([[4., 4.],
# [1., 2.]])Common Mistakes
- Assuming the Jacobian is only a theoretical construct — automatic differentiation libraries compute Jacobian-vector products implicitly and efficiently every time you call
.backward()through a multi-output intermediate layer, even though you rarely form the full Jacobian matrix explicitly (it can be huge). - Confusing the Jacobian's shape convention — different textbooks transpose it differently; always check whether rows are outputs or inputs for a given source.
Interview Relevance
Q: "How does the Jacobian relate to a single neural network layer?" A layer computing \(\mathbf{y}=\mathbf{W}\mathbf{x}+\mathbf{b}\) maps a vector to a vector, so its derivative with respect to \(\mathbf{x}\) is a full Jacobian matrix — which for this linear case is exactly \(\mathbf{W}\). Backpropagation through the layer multiplies the incoming gradient by this Jacobian (via \(\mathbf{W}^\top\)) to compute the gradient with respect to the layer's inputs.
Practice Question
For \(\mathbf{f}(x,y) = [3x+y,\ x-2y]\), write out the \((2,2)\) Jacobian matrix. Is it constant (the same everywhere) or does it depend on \(x\) and \(y\)? What does that tell you about whether \(\mathbf{f}\) is linear?
