Backpropagation with Jacobian-Vector Products

1. Backpropagation and Computational Graphs

At its core, a neural network is a deeply nested function. Deriving gradients for such complex functions by manually expanding the entire expression on paper is tedious, error-prone, and not modular. A more systematic approach is to represent the network as a computational graph, where each node represents a basic operation (e.g., matrix multiplication, an activation function).

Backpropagation is an algorithm that efficiently computes gradients by traversing this graph backward, from the final loss to the input parameters. The process involves two passes:

1. Forward Pass: Data flows through the graph from left to right. Each node computes its output and caches any values needed for the backward pass.
2. Backward Pass: Gradient information flows from right to left. Each node receives a gradient from "upstream" and uses it to compute gradients for its own inputs, which are then passed "downstream".

The video below gives an excellent overview of this concept.

As the video explains, the chain rule at each node takes the form:

Downstream Gradient = Upstream Gradient × Local Gradient

Let's formalize this for the vector case.

2. The Vector Chain Rule and Jacobian-Vector Products

When we move from scalar operations to vector-valued functions, as in neural network layers, the "local gradient" becomes a Jacobian matrix. For a function $\mathbf{y}=f(\mathbf{x})$, where $\mathbf{x}\in {\mathbb{R}}^n$ and $\mathbf{y}\in {\mathbb{R}}^m$, the Jacobian ${\mathbf{J}}_f(\mathbf{x})$ is the $m\times n$ matrix of all partial derivatives $\frac{\partial y_i}{\partial x_j}$.

If we have a final scalar loss $\mathcal{L}$ that is a function of $\mathbf{y}$, the upstream gradient is ${\nabla }_{\mathbf{y}}\mathcal{L}$. The chain rule to find the downstream gradient ${\nabla }_{\mathbf{x}}\mathcal{L}$ becomes a matrix-vector multiplication:

This operation, ${\mathbf{J}}^T\mathbf{v}$, is the cornerstone of backpropagation. It's known as a Jacobian-vector product (JvP). Backpropagation is, in essence, a sequence of these JvP calculations, where the output of one JvP (the downstream gradient) becomes the input vector for the next JvP upstream.

A crucial insight, which we will see in action, is that we almost never need to explicitly construct the full Jacobian matrix. Doing so would be incredibly memory-intensive. Instead, for most common operations, we can derive a much more efficient, direct procedure to compute the result of the JvP. This is sometimes called the "VJP trick" (for Vector-Jacobian Product, referring to ${\mathbf{v}}^T\mathbf{J}$, which is dimensionally equivalent to the JvP for our purposes).

3. The Efficiency of Reverse Mode (Backpropagation)

The chain of JvPs can be computed from left-to-right (forward mode) or right-to-left (reverse mode). For training neural networks, reverse mode—which is what we call backpropagation—is vastly more efficient.

This is because training involves a function that maps many parameters (weights and biases) to a single scalar loss. Reverse mode starts with one gradient (${\nabla }_{\mathcal{L}}\mathcal{L}=1$) and propagates this single piece of information backward, performing computationally cheap matrix-vector products at each step. Forward mode, in contrast, would require passing a full Jacobian matrix forward at each step, involving expensive matrix-matrix products.

The following article provides a clear, compelling numerical example of this efficiency gain.

4. Backpropagation in a 2-Layer MLP

Let's now trace a full backward pass through a simple MLP. Our network will take a vector $\mathbf{x}\in {\mathbb{R}}^{d_{in}}$, pass it through one hidden layer with a ReLU activation, and then to an output layer producing logits ${\mathbf{z}}_2\in {\mathbb{R}}^{d_{out}}$.

Forward Pass:

1. ${\mathbf{z}}_1={\mathbf{W}}_1\mathbf{x}+{\mathbf{b}}_1$
2. ${\mathbf{a}}_1=\text{ReLU}({\mathbf{z}}_1)$
3. ${\mathbf{z}}_2={\mathbf{W}}_2{\mathbf{a}}_1+{\mathbf{b}}_2$
4. $\mathbf{p}=\text{softmax}({\mathbf{z}}_2)$
5. $\mathcal{L}=\text{CrossEntropy}(\mathbf{p},\mathbf{y})$

Where ${\mathbf{W}}_1,{\mathbf{b}}_1,{\mathbf{W}}_2,{\mathbf{b}}_2$ are the parameters and $\mathbf{y}$ is the one-hot encoded true label.

Backward Pass:
We start at the end and work our way back, computing the gradient of the loss $\mathcal{L}$ with respect to each variable. The initial upstream gradient is ${\nabla }_{\mathcal{L}}\mathcal{L}=1$.

Step 1: Gradient of Loss w.r.t. Logits (${\nabla }_{{\mathbf{z}}_2}\mathcal{L}$)

This is the first "real" gradient. As we derived in a previous lesson, the combination of Softmax and Cross-Entropy loss yields a remarkably simple gradient.

This vector is our first upstream gradient, which we'll denote ${\mathbf{g}}_{z2}$.

Step 2: Backprop through Layer 2 (${\mathbf{z}}_2={\mathbf{W}}_2{\mathbf{a}}_1+{\mathbf{b}}_2$)

We need to compute ${\nabla }_{{\mathbf{W}}_2}\mathcal{L}$, ${\nabla }_{{\mathbf{b}}_2}\mathcal{L}$, and the new upstream gradient ${\nabla }_{{\mathbf{a}}_1}\mathcal{L}$. The upstream gradient we have is ${\mathbf{g}}_{z2}$.

• Gradient w.r.t. ${\mathbf{a}}_1$:
  The local Jacobian is $\frac{\partial {\mathbf{z}}_2}{\partial {\mathbf{a}}_1}={\mathbf{W}}_2$. The JvP gives us the new upstream gradient for the next layer back:

  $${\mathbf{g}}_{a1}={\nabla }_{{\mathbf{a}}_1}\mathcal{L}={\left(\frac{\partial {\mathbf{z}}_2}{\partial {\mathbf{a}}_1}\right)}^T{\nabla }_{{\mathbf{z}}_2}\mathcal{L}={\mathbf{W}}_2^T{\mathbf{g}}_{z2}$$

• Gradient w.r.t. ${\mathbf{b}}_2$:
  The local Jacobian is $\frac{\partial {\mathbf{z}}_2}{\partial {\mathbf{b}}_2}=\mathbf{I}$. The JvP is trivial:

  $${\nabla }_{{\mathbf{b}}_2}\mathcal{L}={\left(\frac{\partial {\mathbf{z}}_2}{\partial {\mathbf{b}}_2}\right)}^T{\nabla }_{{\mathbf{z}}_2}\mathcal{L}={\mathbf{I}}^T{\mathbf{g}}_{z2}={\mathbf{g}}_{z2}$$

  (Note: If using batching, this requires a sum over the batch dimension).

• Gradient w.r.t. ${\mathbf{W}}_2$:
  This is the most interesting part. Explicitly forming the Jacobian $\frac{\partial {\mathbf{z}}_2}{\partial {\mathbf{W}}_2}$ is cumbersome as it would be a 3rd-order tensor. Instead, we derive the JvP directly. The result is a simple outer product:

  $${\nabla }_{{\mathbf{W}}_2}\mathcal{L}={\mathbf{g}}_{z2}{\mathbf{a}}_1^T$$

  This is a highly efficient operation that completely avoids large intermediate matrices. For a detailed derivation of this crucial step, the following resources are excellent.

Step 3: Backprop through ReLU (${\mathbf{a}}_1=\text{ReLU}({\mathbf{z}}_1)$)

The upstream gradient is now ${\mathbf{g}}_{a1}={\nabla }_{{\mathbf{a}}_1}\mathcal{L}$. Since ReLU is an element-wise function, its Jacobian is diagonal. The JvP simplifies to an element-wise product (Hadamard product, $\odot$).

The derivative ${\text{ReLU}}^{\prime }({\mathbf{z}}_1)$ is a vector containing 1 for elements where $z_{1,i}>0$ and 0 otherwise. This operation effectively acts as a "gradient router," only allowing gradients to pass through neurons that were active during the forward pass.

Step 4: Backprop through Layer 1 (${\mathbf{z}}_1={\mathbf{W}}_1\mathbf{x}+{\mathbf{b}}_1$)

This step is identical in form to Step 2, but with different variables. The upstream gradient is ${\mathbf{g}}_{z1}$.

• ${\nabla }_{{\mathbf{b}}_1}\mathcal{L}={\mathbf{g}}_{z1}$
• ${\nabla }_{{\mathbf{W}}_1}\mathcal{L}={\mathbf{g}}_{z1}{\mathbf{x}}^T$

At this point, we have computed the gradients for all parameters (${\mathbf{W}}_1,{\mathbf{b}}_1,{\mathbf{W}}_2,{\mathbf{b}}_2$), and the backpropagation pass is complete. An optimization algorithm like SGD or Adam can now use these gradients to update the parameters.

5. From Theory to Code: Modular Implementation

This node-by-node, JvP-based view of backpropagation maps directly to how modern deep learning frameworks are designed. In PyTorch, you can define custom operations by subclassing torch.autograd.Function and implementing forward and backward static methods.

• The forward method performs the operation and uses ctx.save_for_backward to cache any tensors (like inputs or outputs) needed for the gradient calculation.
• The backward method receives the upstream gradient (grad_output) and must return the gradients with respect to each of the function's inputs. This backward method is precisely an implementation of the JvP for that operation.

The Michigan Online video provides a fantastic explanation of this concept.

Conclusion

In this lesson, we demystified backpropagation by framing it as an efficient application of the chain rule on a computational graph. We saw that the core mathematical operation is the Jacobian-vector product, and that the algorithm's efficiency comes from avoiding the explicit creation of large Jacobian matrices.

Key Takeaways:

• Backpropagation is reverse-mode automatic differentiation on a computational graph.
• The chain rule for vector functions is applied sequentially via Jacobian-vector products (JvPs).
• For a typical neural network (many parameters to a single loss), reverse-mode is far more efficient than forward-mode differentiation.
• The JvP for a linear layer's weights simplifies to an outer product (${\nabla }_{\mathbf{W}}\mathcal{L}={\mathbf{g}}_{out}{\mathbf{a}}_{in}^T$).
• The JvP for an element-wise activation simplifies to an element-wise product with the local derivative (${\mathbf{g}}_{in}={\mathbf{g}}_{out}\odot f^{\prime }({\mathbf{z}}_{in})$).
• This modular, JvP-based structure is directly mirrored in the torch.autograd.Function API, allowing for custom, differentiable layers.

We have focused on first-order derivatives (gradients and Jacobians). In our next lesson, we will extend these ideas to second-order derivatives, exploring the Hessian matrix and its approximations, which are critical for understanding the curvature of the loss landscape and for advanced optimization techniques.