Backpropagation — the algorithm that trains every deep network — is a single line of matrix calculus applied recursively. The Jacobian tells you "how does each output coordinate change when each input coordinate moves?" Backprop just stacks these Jacobians from the loss back to the weights. PyTorch's .backward() is doing exactly what's in this lesson, just with bigger matrices.
Learning Objectives
After this lesson, you will be able to:
Compute the Jacobian of a vector-valued function and apply the vector-valued chain rule J_h = J_g · J_f to chain transformations
Use the Hessian to reason about curvature, convexity, and the geometry of a loss landscape — and explain why deep-net Hessians are huge but Hessian-vector products are cheap
Apply the load-bearing matrix-calculus identities (∂(Ax)/∂x = A, ∂(x^T A x)/∂x = (A+A^T)x, ∂tr(AB)/∂A = B^T, ∂(log det A)/∂A = (A^{-1})^T) without flipping notation
Distinguish forward-mode (JVP) from reverse-mode (VJP) autodiff and explain why every ML framework defaults to reverse-mode for scalar losses
Derive the softmax + cross-entropy gradient ∂L/∂z = p − y_one_hot from scratch — the most beautiful formula in deep learning
Backprop is not magic. It is matrix calculus, applied recursively, with a clever trick. The previous lesson taught you what a derivative is. This one teaches you what a derivative of a vector function with respect to a vector is — and why the chain rule, generalised properly, is exactly what makes neural networks trainable.
The single-variable derivative f'(x) is a number. The gradient ∇f(x) (when f returns one number) is a vector. The Jacobian (when f returns a vector of m numbers) is an m × n matrix. The Hessian (second derivatives of a scalar function) is an n × n matrix. As you generalise from one input to many, and one output to many, derivatives become matrices and tensors — but the chain rule survives, dressed up as matrix multiplication.
Try it! Open the Python REPL (bottom-right of the screen: click Quick Actions, then Python) and type these lines yourself.
For a function f: R^n → R^m that takes an n-vector and returns an m-vector, the Jacobian J is the m × n matrix whose entry (i, j) is the partial derivative of the i-th output with respect to the j-th input:
Take f(x, y) = (x²y, x + y). We want the Jacobian at the point (x, y) = (2, 3).
The four partials:
∂f₁/∂x = ∂(x²y)/∂x = 2xy. At (2, 3): 2·2·3 = 12.
∂f₁/∂y = ∂(x²y)/∂y = x². At (2, 3): 4.
∂f₂/∂x = ∂(x+y)/∂x = 1.
∂f₂/∂y = ∂(x+y)/∂y = 1.
So the Jacobian at (2, 3) is:
Jf(2,3)=[12141]
The rows of the Jacobian are gradients of individual outputs; the columns describe how each input affects the entire output vector. Both readings matter — different parts of backprop use different orientations.
Switch between vector fields and scalar fields to see the Jacobian (m × n derivatives of a vector field) and the gradient + Hessian (1 × n and n × n derivatives of a scalar field) computed live at any point you click.
Loading visualization...
What Do You Think?
A function f: R³ → R² takes a 3-vector and returns a 2-vector. What shape is its Jacobian J_f in numerator layout?
Quick check
For the linear map y = W x with W ∈ R^{m × n}, what is the Jacobian ∂y/∂x?
This is the formula that makes deep learning compile. For composed functions h(x) = g(f(x)), the Jacobian of the composition is the product of Jacobians:
Sanity check: h(x, y) = (x+y)² + xy, so ∂h/∂x = 2(x+y) + y and ∂h/∂y = 2(x+y) + x. At (2, 3): [2·5 + 3, 2·5 + 2] = [13, 12]. ✓
This is the entire mathematical content of backpropagation. A neural network is a long composition of functions; the gradient of the loss with respect to any parameter is just a long chain of Jacobian matrix products.
Watch the chain rule fire through a stack of layers — input flows forward through three function blocks, then the gradient flows back as a product of Jacobians at each block.
Loading visualization...
Now animate the full forward-then-backward sweep. The animation makes it visceral why people call it "backprop": the loss signal genuinely walks right-to-left across the graph.
For nice functions (continuous second partials, i.e. C² functions) the Hessian is symmetric: ∂²f/(∂x∂y) = ∂²f/(∂y∂x). This is Clairaut's theorem (also known as Schwarz's theorem). Effectively every ML loss function meets the conditions, so you can almost always treat H as symmetric — which means it has real eigenvalues and orthogonal eigenvectors (the spectral theorem from the eigenvalues lesson).
All eigenvalues ≥ 0 ⇒ Hessian is positive semi-definite (PSD) ⇒ the function curves up in every direction at this point ⇒ the point is a local minimum candidate. If H is PSD everywhere, the function is globally convex — gradient descent can't get stuck.
All eigenvalues ≤ 0 ⇒ negative semi-definite ⇒ local maximum candidate.
Mixed signs ⇒ saddle point ⇒ flat in some directions, downhill in others. Saddle points are the dominant trap in non-convex deep-net optimisation.
Quick check
At a point where the Hessian has eigenvalues {+3, +1, −0.2}, what kind of stationary point is this?
These are the matrix-calculus identities you will use weekly. Memorising the first three buys you the next decade of ML papers; the rest are useful but lookup-able.
Convention note: This entire lesson uses numerator layout — the output index of a derivative becomes the row of the resulting Jacobian. Most DL textbooks (Goodfellow, Bishop, Roberts/Yaida/Hanin) use this convention. The Matrix Cookbook (Petersen & Pedersen) uses denominator layout, which transposes everything. If a paper's shapes look "off by a transpose," check which convention it's using.
Suppose y = Wx, where W is m × n. Then for a scalar loss L = L(y), the gradient with respect to W is:
∂W∂L=∂y∂L⋅xT
This single identity is responsible for >50% of the FLOPs in transformer training — every linear layer's backward pass is exactly this outer product, batched.
What Do You Think?
You have y = Wx, with W ∈ R^{m×n}. What is the shape of ∂L/∂W?
#Forward-Mode and Reverse-Mode: Two Ways to Walk a Computational Graph
When you compose a long chain of functions and want their derivative, there are two ways to evaluate the chain rule's matrix product. Both compute the same answer; they just associate the multiplications differently.
Suppose h = f_3 ∘ f_2 ∘ f_1 and you want a derivative. By the chain rule, J_h = J_3 · J_2 · J_1. Forward mode walks left to right (from input to output), accumulating a tangent vector v:
v ← initial tangent (a small input perturbation)
v ← J_1 · v # propagate through f_1
v ← J_2 · v # through f_2
v ← J_3 · v # through f_3 — done
Each step is a Jacobian-Vector product (JVP). Cost of one JVP: roughly the cost of one forward pass through the function.
When is this cheap? When the input has few dimensions — say, you want df/dx_1 for a scalar input. One JVP per input dimension. So forward mode costs O(n × forward) where n is the input dimension.
Reverse mode walks right to left (from output back to input), accumulating a cotangent vectorv^T:
v ← final cotangent (initialise as ∂L/∂L = 1 for scalar loss)
v ← v · J_3 # pull back through f_3
v ← v · J_2 # through f_2
v ← v · J_1 # back to input
Each step is a Vector-Jacobian product (VJP). Cost of one VJP: also roughly one forward pass. But you only need one VJP per scalar output. So reverse mode costs O(m × forward) where m is the output dimension.
In ML, the loss L is a scalar (m = 1) and there are millions to billions of parameters (n is huge). Reverse mode costs O(1 × forward) — a single backward pass gets you all n gradients. Forward mode would cost O(n × forward) — you'd run n full forward passes. For a 70B-parameter model, the difference is the difference between training completing in a week and not finishing this century.
Bridge: When you call loss.backward() in PyTorch, the framework is computing v^T · J for each parameter — pulling back the loss-gradient (a 1-vector for scalar losses) through every Jacobian in the network. That is reverse-mode autodiff, this lesson made concrete.
Bridge: JAX exposes both modes: jax.grad is reverse-mode (cheap for L: R^n → R), jax.jvp is forward-mode (cheap for L: R → R^m). Knowing which mode to pick comes down to the input/output dimension ratio you just saw.
Time for the payoff. We are going to derive the gradient of the most common ML loss — softmax + cross-entropy on a 1-layer network — using only the identities above. The derivation takes 5 lines. The result is breathtaking.
Bridge: The σ(x)(1−σ(x)) sigmoid derivative from the previous lesson is the chain rule's first link in every binary-classification backprop. The (p − y_one_hot) softmax+cross-entropy gradient derived here is the same idea generalised to multiclass — and it is why training one-hot classifiers feels "easy" compared to other losses. The clean, error-shaped gradient gives optimisers exactly the signal they need.
With ∂L/∂z in hand, the rest is the linear-layer identities from above:
∂W∂L=∂z∂LxT∂b∂L=∂z∂L∂x∂L=WT∂z∂L
That is the entire derivation. Five lines of math, three identities, one beautiful cancellation. In a 96-layer transformer, this exact pattern repeats 96 times — every linear layer's backward is outer(dz, x) for weights, W^T dz for the input — strung together by the chain rule. Backprop is just this, recursively.
Compute the same Jacobian symbolically (via sympy) and numerically (via manual finite differences), then verify they agree. This is what frameworks like JAX and PyTorch are doing under the hood — the symbolic path is what jax.jacfwd computes analytically; the numerical path is the sanity check (torch.autograd.gradcheck) that ML libraries use to catch bugs in their backward rules.
Loading visualization...
#Hessian-Vector Products: The Trick That Makes Second-Order DL Possible
The Hessian itself is unmanageable — for n parameters, it is n × n. But the Hessian-Vector Product Hv is just an n-vector, and it can be computed in O(forward) cost.
The trick: Hv is the gradient of (∇L · v) with respect to x. So you compute the gradient ∇L (one backward pass), take its inner product with v (cheap), then take another gradient with respect to x (one more backward pass). Total: two backward passes, no n × n matrix.
Hv=∇x(∇xL⋅v)
This is also called Pearlmutter's trick (Barak Pearlmutter, 1994). Before the deep-learning era it was a niche tool of Newton-method optimisers; in modern transformers it powers second-order training methods like Sophia (Liu et al. 2023) which use Hessian diagonal estimates to scale Adam-style updates.
Bridge: The Hessian-vector product Hv = ∂(∇L · v)/∂x is what Sophia, K-FAC, and Newton-Sketch use to approximate second-order steps without ever materialising the n²-sized Hessian. With n = 70B parameters, you cannot store H — but you can compute Hv.
Jacobian + chain rule = backprop. A neural network is a long composition of functions. Its gradient is one matrix product per layer, walked from output to input. Reverse-mode AD does this in time proportional to a single forward pass.
The matrix-calculus identity table buys you the world. Six identities — ∂(Ax)/∂x = A, ∂(x^T A x)/∂x = (A + A^T) x, ∂‖x‖²/∂x = 2x, ∂(x^T y)/∂x = y, ∂tr(AB)/∂A = B^T, ∂(log det A)/∂A = (A^{-1})^T — let you derive the gradient of essentially any standard ML loss in 5 minutes. Memorise them.
∂L/∂z = p − y_one_hot is the gradient of softmax + cross-entropy — and it is the cleanest, most useful gradient in supervised learning. It says: the gradient is the error vector. Push wrong probabilities down, push the correct probability up. This is the exact signal that makes deep learning trainable.
You now own the math behind every line of loss.backward(). The same machinery powers VAEs, diffusion models, transformer training, RLHF, normalizing flows, and second-order optimisers. The frameworks are sugar — this lesson is what is actually happening underneath.
For a function f: R^n → R^m, what is the shape of its Jacobian (in numerator layout)?
Next up: Probability and Bayes' Theorem. You will reason under uncertainty, watch a 99%-accurate test be wrong most of the time, and meet the foundation of every classification model.