Vector-Jacobian Products in Practice

Every layer in a neural network is a function that maps an input vector to an output vector, and in principle its full derivative is a Jacobian matrix — one row per output, one column per input. For a layer with 1,000 inputs and 1,000 outputs that matrix has a million entries, and a deep network has many such layers chained together. Computing and multiplying full Jacobians at every layer would be prohibitively slow and memory-hungry, yet this is exactly what training a network by gradient descent seems to require.

Backpropagation sidesteps this by never forming the Jacobian at all. Instead, it computes a vector-Jacobian product (VJP): given the gradient vector flowing backward from later layers, it directly produces the gradient vector that should flow into earlier layers, using a formula specific to the layer's operation (matrix multiply, ReLU, softmax, and so on) that costs about the same as one forward pass through the layer — not the far larger cost of building and multiplying a full matrix.

Concretely, for a layer $y = f(x)$, the chain rule needs $v^\top J$, where $J$ is the Jacobian of $f$ and $v$ is the incoming gradient vector. Frameworks like PyTorch and JAX implement a hand-written "backward" rule for each operation that computes $v^\top J$ directly — for a linear layer $y = Wx$, the VJP is simply $v^\top W$, a matrix-vector product, never requiring the Jacobian $W$ itself to be treated as anything bigger than the weight matrix already stored in memory.

For a linear layer mapping 1,000 inputs to 1,000 outputs, forming the Jacobian explicitly would need a million stored numbers; the VJP instead reuses the existing 1,000×1,000 weight matrix and does one matrix-vector multiply, exactly matching the cost of the forward pass.