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.