Jacobians and the Chain Rule for Neural Networks

Welcome back. In our last lesson, we developed the tools to find the gradient of a scalar function, like a loss function, with respect to its vector or matrix parameters. We used a coordinate-free differential approach, which allowed us to handle complex expressions elegantly. This gave us the derivative at the very end of a computational graph.

Today, we will generalize this concept to handle the intermediate stages. Most layers in a neural network are vector-to-vector transformations. The derivative of such a function is no longer a vector (the gradient) but a matrix, which we call the Jacobian. We will then introduce the vector chain rule, showing how the Jacobians of individual layers can be multiplied together to find the gradient of the entire network. This is the mathematical foundation of backpropagation, the engine that powers deep learning.

1. From Gradients to Jacobians

In the previous lesson, we considered functions that map a vector to a scalar, $f : R^{n} \to R$ . The derivative, the gradient $\nabla_{x} f$ , was a vector containing the partial derivatives $\frac{\partial f}{\partial x _{i}}$ .

Now, consider a function that maps a vector to another vector, $y = f (x)$ , where $f : R^{n} \to R^{m}$ . We have $m$ output components, each a scalar function of the $n$ input components:

y_{1} y_{2} ⋮ y_{m} = f_{1} (x_{1}, x_{2}, \dots, x_{n}) = f_{2} (x_{1}, x_{2}, \dots, x_{n}) = f_{m} (x_{1}, x_{2}, \dots, x_{n})

To capture how the entire output vector $y$ changes with respect to the input vector $x$ , we need to compute all $m \times n$ possible partial derivatives, $\frac{\partial y _{i}}{\partial x _{j}}$ . The Jacobian matrix is the object that neatly organizes these partials.

[PDF] Matrix Calculus You Need For Deep Learning - arXiv

This paper provides a rigorous yet practical introduction. We'll use it to formalize the definition of the Jacobian.

Please read the section titled "Matrix calculus" from the beginning. You can start at the introduction and continue through the general definition up to the end of the large matrix equation defining the Jacobian, ending with this sentence. Focus on how the Jacobian is constructed by stacking the gradients of each scalar output function f_i(x). We will adhere to the numerator layout convention presented here, where the Jacobian has shape m \times n.

As defined in the paper, for a function $y = f (x)$ with $f : R^{n} \to R^{m}$ , the Jacobian matrix $\frac{\partial y}{\partial x}$ is an $m \times n$ matrix where the entry at row $i$ and column $j$ is $\frac{\partial y _{i}}{\partial x _{j}}$ .

\frac{\partial y}{\partial x} = \frac{\partial y _{1}}{\partial x _{1}} \frac{\partial y _{2}}{\partial x _{1}} ⋮ \frac{\partial y _{m}}{\partial x _{1}} \frac{\partial y _{1}}{\partial x _{2}} \frac{\partial y _{2}}{\partial x _{2}} ⋮ \frac{\partial y _{m}}{\partial x _{2}} \dots \dots ⋱ \dots \frac{\partial y _{1}}{\partial x _{n}} \frac{\partial y _{2}}{\partial x _{n}} ⋮ \frac{\partial y _{m}}{\partial x _{n}}

Notice that the $i$ -th row of the Jacobian is simply the gradient of the scalar function $y_{i}$ with respect to $x$ , $\nabla_{x} y_{i}$ . So, the Jacobian is a stack of gradient vectors.

Example: The Jacobian of a Linear Layer

Consider a fundamental operation in neural networks: an affine transformation $y = A x + b$ , where $A \in R^{m \times n}$ and $b \in R^{m}$ . This function maps $x \in R^{n}$ to $y \in R^{m}$ . The $i$ -th component of $y$ is:

y_{i} = j = 1 \sum n A_{ij} x_{j} + b_{i}

Let's find the partial derivative $\frac{\partial y _{i}}{\partial x _{k}}$ for some arbitrary index $k$ .

\frac{\partial y _{i}}{\partial x _{k}} = \frac{\partial}{\partial x _{k}} (j = 1 \sum n A_{ij} x_{j} + b_{i}) = j = 1 \sum n A_{ij} \frac{\partial x _{j}}{\partial x _{k}} + 0

The term $\frac{\partial x _{j}}{\partial x _{k}}$ is 1 if $j = k$ and 0 otherwise. Thus, the only term that survives the sum is when $j = k$ .

\frac{\partial y _{i}}{\partial x _{k}} = A_{ik}

This means the entry in the $i$ -th row and $k$ -th column of the Jacobian is just the matrix element $A_{ik}$ . Therefore, the Jacobian of the function $f (x) = A x + b$ is simply the matrix $A$ .

\frac{\partial ( A x + b )}{\partial x} = A

This is a beautifully simple and important result. The Jacobian of a linear transformation is the transformation matrix itself.

Example: The Jacobian of an Element-wise Function

What about activation functions, like ReLU or sigmoid, which are applied element-wise? Let $y = f (x)$ where $y_{i} = σ (x_{i})$ for some scalar function $σ$ .

\frac{\partial y _{i}}{\partial x _{j}} = \frac{\partial σ ( x _{i} )}{\partial x _{j}}

Since $y_{i}$ only depends on $x_{i}$ , this partial derivative is zero if $i \neq = j$ . If $i = j$ , it's simply the derivative of $σ$ evaluated at $x_{i}$ , which we denote $σ^{'} (x_{i})$ .

\frac{\partial y _{i}}{\partial x _{j}} = {σ^{'} (x_{i}) 0 if i = j if i \neq = j

This means the Jacobian is a diagonal matrix:

\frac{\partial y}{\partial x} = σ^{'} (x_{1}) 0 ⋮ 0 0 σ^{'} (x_{2}) ⋮ 0 \dots \dots ⋱ \dots 00 ⋮ σ^{'} (x_{n}) = diag (σ^{'} (x))

This diagonal structure is a key reason for the computational efficiency of backpropagation in modern deep learning frameworks. We rarely need to materialize the full Jacobian; we only need to know its effect, which for element-wise operations is just an element-wise multiplication.

2. The Vector Chain Rule

Neural networks are deeply nested composite functions. A 3-layer MLP is of the form $f_{3} (f_{2} (f_{1} (x)))$ . To find the derivative of the entire composition, we need a chain rule that works for vector-valued functions and their Jacobians.

The Chain Rule

This video provides an excellent, intuitive derivation of the vector chain rule, building up from a multi-input, single-output case to the general matrix form.

First, watch from the start of the derivation to understand the "trace every pathway" intuition. For an output W and input R, a change in R affects W through all intermediate variables (X, Y, Z), and the total change is the sum of the changes along each path. Then, watch the critical segment from here, where the mess of summed partials is elegantly refactored into a product of Jacobian matrices. This is the main takeaway. Finally, see how this generalizes to compositions of multiple functions from this point on.

The key insight from the video is that the chain rule for vector functions is a direct generalization of the scalar case. If we have a composition of functions $z = f (y)$ and $y = g (x)$ , the chain rule allows us to find the derivative of $z$ with respect to $x$ by multiplying their Jacobians:

\frac{\partial z}{\partial x} = \frac{\partial z}{\partial y} \frac{\partial y}{\partial x}

Let's check the dimensions. If $g : R^{n} \to R^{m}$ and $f : R^{m} \to R^{p}$ , then:

$\frac{\partial y}{\partial x}$ is an $m \times n$ Jacobian.
$\frac{\partial z}{\partial y}$ is a $p \times m$ Jacobian.
The product $\frac{\partial z}{\partial y} \frac{\partial y}{\partial x}$ is a $(p \times m) \times (m \times n)$ matrix multiplication, resulting in a $p \times n$ matrix, which is the correct dimension for $\frac{\partial z}{\partial x}$ .

The order of multiplication is crucial and works "backwards" from the final output. For a deep neural network $y = f_{L} (\dots f_{2} (f_{1} (x)) \dots)$ , the Jacobian is the product of the Jacobians of each layer, starting from the last one.

This diagram illustrates the vector chain rule for a composite function F'(x). The derivative of the entire function is the product of the Jacobian matrices of each intermediate function, applied in reverse order.

3. Backpropagation: The Chain Rule on a Computational Graph

While the chain rule tells us what to compute (a product of Jacobians), backpropagation is an efficient algorithm for how to compute it on a computational graph.

At its core, backpropagation is about calculating the gradient of a final scalar loss, $L$ , with respect to every parameter and intermediate variable in the network. For any variable $u$ in the graph, we want to compute $\frac{\partial L}{\partial u}$ .

Let's consider a single "gate" or operation in the graph that computes $z = f (x, y)$ . During the forward pass, we compute and store the value of $z$ . During the backward pass, we assume we have already received the gradient from the subsequent parts of the network. This is the upstream gradient, $\frac{\partial L}{\partial z}$ . Our goal is to compute the downstream gradients, $\frac{\partial L}{\partial x}$ and $\frac{\partial L}{\partial y}$ .

The chain rule gives us the answer directly:

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial z} \frac{\partial z}{\partial x}

\frac{\partial L}{\partial y} = \frac{\partial L}{\partial z} \frac{\partial y}{\partial y}

Here, $\frac{\partial L}{\partial z}$ is the upstream gradient (which is a row vector since L is a scalar), and $\frac{\partial z}{\partial x}$ is the local Jacobian of the gate. The downstream gradient is found by multiplying the upstream gradient by the local Jacobian.

This diagram shows how gradients are propagated backward through a computational node. The upstream gradient $\partial L / \partial z$ is multiplied by the local Jacobians ($\partial z / \partial x$ and $\partial z / \partial y$) to compute the gradients for the inputs, $\partial L / \partial x$ and $\partial L / \partial y$. Notice the result always has the same shape as the variable it's a gradient of.

This process starts at the very end of the network, where the gradient of the loss with respect to itself, $\frac{\partial L}{\partial L}$ , is simply the scalar 1. We then recursively apply this rule, moving backward through the graph, gate by gate, until we have the gradients for all parameters.

CS231n Winter 2016: Lecture 4: Backpropagation, Neural Networks 1

Andrej Karpathy's classic CS231n lecture provides a definitive explanation of this process.

First, watch from this segment to understand the local/global gradient passing mechanism. Each 'gate' in the computational graph receives an upstream (global) gradient and multiplies it by its own local gradient to pass downstream. Then, watch the crucial discussion on vector gradients from this point. Karpathy explains that while the derivatives are technically Jacobian matrices, we almost never form them explicitly because they are often very large and sparse (e.g., diagonal for element-wise operations like ReLU). This practical insight is key to understanding how autograd systems are implemented efficiently.

As Karpathy emphasizes, although the theory involves Jacobian-matrix products, in practice, the backward pass for most common operations (element-wise functions, matrix multiplication, summation) can be implemented much more efficiently without ever constructing the full, dense Jacobian matrix. This is how libraries like PyTorch and TensorFlow achieve their performance.

Conclusion

In this lesson, we have built the complete mathematical machinery for automatic differentiation.

Key Takeaways:

The Jacobian matrix $\frac{\partial y}{\partial x}$ is the generalization of the gradient for vector-valued functions $f : R^{n} \to R^{m}$ , with dimensions $m \times n$ .
For an affine layer $y = A x + b$ , the Jacobian is simply $A$ . For an element-wise activation function, the Jacobian is a diagonal matrix.
The vector chain rule expresses the Jacobian of a composite function as the product of the Jacobians of its constituent functions: $\frac{\partial z}{\partial x} = \frac{\partial z}{\partial y} \frac{\partial y}{\partial x}$ .
Backpropagation is a practical algorithm that applies the chain rule recursively on a computational graph. At each node, it computes downstream gradients by multiplying the upstream gradient by the local Jacobian.

We are now equipped to derive the gradients for any part of a neural network. In our next lesson, we will put this into practice by deriving the gradients for the most common loss functions in deep learning: Mean Squared Error (for regression) and Cross-Entropy (for classification). We'll see how combining the Jacobian of the final activation function with the gradient of the loss often leads to a surprisingly simple and intuitive result.