Backpropagation › Jacobian Matrices and Gradient Flow

Before diving into the full backward pass, it's worth understanding Jacobian matrices carefully, because they formalize how gradients flow through vector-valued operations.

In neural network backpropagation, the loss $\mathcal{L}$ is a scalar, and each operation in the network transforms a vector to another vector. At each node, we don't technically need the full Jacobian. We only need the vector-Jacobian product (VJP), which maps the upstream gradient vector to a downstream gradient vector.

To see why the full Jacobian is unnecessary, consider a layer that maps $\mathbb{R}^{512} \to \mathbb{R}^{512}$. The full Jacobian would be a $512 \times 512$ matrix of 262,144 entries. But we only need to propagate a gradient vector of length 512 backward. The VJP computes this directly in $O(512^2)$ time, the same cost as forming the Jacobian, but without ever explicitly materializing it. For larger layers, the savings are more dramatic, and specialized operations like batch normalization have structured Jacobians that admit even faster VJP computations.

In practice, backpropagation never explicitly forms full Jacobian matrices. Instead, it computes vector-Jacobian products. If the upstream gradient is a row vector $\mathbf{v}^T$ (with the same shape as the output), the downstream gradient is:

This produces a row vector with the same shape as the input. The key insight is that computing this product is much cheaper than computing $\mathbf{J}$ explicitly, then multiplying. PyTorch's autograd system is built entirely around this VJP abstraction.

This VJP perspective also clarifies what backpropagation is computing at each step. The upstream gradient $\mathbf{v}$ represents "how the loss changes when the output of this node changes." The Jacobian $\mathbf{J}$ represents "how the output of this node changes when the input changes." Their product gives "how the loss changes when the input changes," which is exactly the downstream gradient we need to propagate.

The power of computational graphs is that each operation only needs to define its own local Jacobian (or VJP). The full gradient is then assembled by chaining these local gradients via the multivariate chain rule.

Here are the local gradients for the key operations in a neural network:

Addition: $\mathbf{z} = \mathbf{x} + \mathbf{y}$. The Jacobian with respect to each input is the identity matrix $\mathbf{I}$. Gradients pass through unchanged: $\frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}}$ and $\frac{\partial \mathcal{L}}{\partial \mathbf{y}} = \frac{\partial \mathcal{L}}{\partial \mathbf{z}}$. This is why residual connections are so effective: the skip path provides a direct gradient highway back to earlier layers. The gradient flows through without any attenuation from the addition node itself.

Multiplication by scalar $w$: $\mathbf{z} = w \cdot \mathbf{x}$. The gradient with respect to $\mathbf{x}$ is $w \cdot \frac{\partial \mathcal{L}}{\partial \mathbf{z}}$ (scale up or down), and the gradient with respect to $w$ is $\mathbf{x}^T \frac{\partial \mathcal{L}}{\partial \mathbf{z}}$.

Matrix multiplication: $\mathbf{z} = \mathbf{W}\mathbf{x}$. The Jacobian with respect to $\mathbf{x}$ is $\mathbf{W}^T$, giving $\frac{\partial \mathcal{L}}{\partial \mathbf{x}} = \mathbf{W}^T \frac{\partial \mathcal{L}}{\partial \mathbf{z}}$. The gradient with respect to $\mathbf{W}$ is the outer product $\frac{\partial \mathcal{L}}{\partial \mathbf{z}} \mathbf{x}^T$. Notice the transpose: gradients flow backward through the weight matrix transposed, which is why the backward pass through a linear layer involves $\mathbf{W}^T$ while the forward pass involves $\mathbf{W}$.

Elementwise nonlinearity: $\mathbf{a} = \sigma(\mathbf{z})$. The Jacobian is a diagonal matrix with $\sigma'(z_i)$ on the diagonal, so $\frac{\partial \mathcal{L}}{\partial \mathbf{z}} = \frac{\partial \mathcal{L}}{\partial \mathbf{a}} \odot \sigma'(\mathbf{z})$, where $\odot$ denotes elementwise multiplication. Because the Jacobian is diagonal, we never form a full matrix; we just do an elementwise multiply.