Guide

Math toolkit

Jacobians and vector-Jacobian products

When a layer maps a vector to a vector, its slopes form a matrix called the Jacobian. Training only needs one vector times that matrix, which a single backward sweep gives.

12 of 187 · 3 min read

A layer that takes a vector and returns a vector has a slope for every pair of input and output entries. Those slopes form a matrix, the Jacobian. Full Jacobians are huge, and training almost never needs one. It needs a single vector times the Jacobian, which is a vector-Jacobian product. Backprop computes exactly that.

Concept

The Jacobian of f: ℝⁿ → ℝᵐ is an m × n matrix with entries Jᵢⱼ = ∂fᵢ/∂xⱼ. Row i says how output i responds to every input. Column j says how every output responds to input j.

A linear layer

For a linear layer y = W x, the Jacobian with respect to x is just W. The slope of each output with respect to each input is the corresponding weight.

Take the weight matrix from the linear layer example, W = [[0.5, 0, −1], [1, −1, 0]]. Its Jacobian has shape 2 × 3. Suppose the gradient flowing into the output is v = [1, 0]. The vector-Jacobian product is:

vᵀ J = [1, 0] · W = [0.5, 0, −1]

Notice that the result is the first row of W. Backprop through a linear layer multiplies by the transpose of the weights in the direction of the gradient, which is why you see Wᵀ throughout backprop.

Why products, not full Jacobians

A single reverse sweep computes a vector times the Jacobian, one row of work per output. The loss is one scalar, so the upstream vector is just 1, and one sweep yields the gradient for every parameter. Materializing the full Jacobian of a neural network is usually impossible. Frameworks ask you to provide the vector and compute the product directly.

Tip

Forward-mode versus reverse-mode. Forward mode computes J u for an input direction u, which is cheap for few inputs and many outputs. Reverse mode computes vᵀ J, which is cheap for one scalar output and many inputs. Training is the second case.

Two rules you use often

  • Sum rule: (f + g)' = f' + g'. Bias terms and weighted sums add their slopes.
  • Product rule: (f · g)' = f' · g + f · g'. You see it less in nets, but it appears in many derivations.

Scaling a logit by a constant also changes the slope. Softmax with temperature uses z / T, which puts a factor of 1 / T on the chain. A lower temperature sharpens the probabilities and scales the gradients.

Common mistakes

  • Forgetting the transpose. The vector-Jacobian product of y = W x uses Wᵀ, not W.
  • Trying to build the full Jacobian. For a big layer it has too many entries. Use the product instead.
  • Confusing J u with vᵀ J. The first is forward mode, the second is reverse mode. Backprop uses the second.

Check yourself

  1. For y = W x with W of shape 2 × 3, what is the Jacobian with respect to x? It is W itself, with shape 2 × 3.
  2. Why do frameworks use vector-Jacobian products? The full Jacobian is too large to build. One product gives the gradient the loss needs in one reverse pass.

Key takeaways

  • The Jacobian holds every slope of a vector-valued function.
  • Training needs a vector times the Jacobian, not the whole matrix.
  • For a linear layer, the vector-Jacobian product is the transpose of the weights applied to the upstream gradient.
  • Reverse mode is efficient when the output is a single scalar, as the loss is.

Further learning

Next, the slopes of the functions that appear in every network: derivatives of common functions.

Calculus and gradients

Derivatives, partial derivatives, the chain rule and gradients. This is the math that turns a loss number into parameter updates.

The chain rule on a graph

When one quantity depends on another through a chain of operations, their slopes multiply. Applied on a graph, the chain rule gives every gradient in one backward pass.

Backpropagation from scratch

How training computes the gradient of the loss for every weight in one backward pass. Covers computational graphs, reverse-mode autodiff and the bugs that break training.