Jacobians and the Multivariate Chain Rule

In our previous lesson, we established how to compute derivatives of scalar-valued functions using the gradient and the Hessian. We saw that the gradient points in the direction of steepest ascent, and the Hessian describes the function's local curvature. These are fundamental tools for optimisation. Now, we will extend these concepts to functions that output vectors, not just single scalars.

This lesson will introduce you to the Jacobian matrix, the generalization of the derivative for vector-valued functions. We will then learn how to apply the multivariate chain rule, a powerful technique for differentiating compositions of such functions. Mastering this is crucial because many advanced machine learning models, most notably neural networks, are built by chaining together multiple vector-to-vector transformations. Understanding how to differentiate through this chain is the key to training them.

1. The Jacobian Matrix: Derivative of a Vector Function

In the last lesson, we took the derivative of a function $f : R^{n} \to R$ and got a vector (the gradient). What happens if our function maps a vector to another vector, for example, $f : R^{n} \to R^{m}$ ? How can we describe its "derivative"?

The derivative in this case should represent the best linear approximation of the function at a point. For a function mapping $R^{n} \to R^{m}$ , this linear approximation is described by an $m \times n$ matrix. This matrix is called the Jacobian.

The following video from Khan Academy provides an excellent geometric intuition for what the Jacobian represents. It shows how the transformation $f$ acts on a small neighborhood around a point and how the Jacobian matrix captures this local, linear behavior.

The Jacobian matrix

This video builds the Jacobian matrix from its geometric interpretation as a local linear approximation of a vector-valued function.

Please watch the video from the beginning to end (The Jacobian matrix). Pay close attention to how the columns of the Jacobian matrix are formed by observing how the output vector changes in response to small steps along each input axis.

As the video explains, the Jacobian matrix $J_{f}$ (often denoted as $D f$ or $\frac{\partial f}{\partial x}$ ) contains all the first-order partial derivatives of the function $f$ . If $f$ has component functions $(f_{1}, f_{2}, \dots, f_{m})$ and the input vector is $x = (x_{1}, x_{2}, \dots, x_{n})$ , the Jacobian is constructed by stacking the gradients of each component function.

The following reading provides the formal definition.

The Matrix Calculus You Need For Deep Learning

This article section formally introduces the Jacobian as a stack of gradient vectors.

Please read the section titled "Matrix calculus". Note how the Jacobian is built by taking the gradient of each output function (f_1, f_2, etc.) with respect to the input vector and stacking these gradients as rows. This article uses the "numerator layout," which we will adhere to in this course.

So, the Jacobian matrix of $f$ at a point $x$ is defined as:

J_{f} (x) = \frac{\partial f}{\partial x} = \frac{\partial f _{1}}{\partial x _{1}} \frac{\partial f _{2}}{\partial x _{1}} ⋮ \frac{\partial f _{m}}{\partial x _{1}} \frac{\partial f _{1}}{\partial x _{2}} \frac{\partial f _{2}}{\partial x _{2}} ⋮ \frac{\partial f _{m}}{\partial x _{2}} \dots \dots ⋱ \dots \frac{\partial f _{1}}{\partial x _{n}} \frac{\partial f _{2}}{\partial x _{n}} ⋮ \frac{\partial f _{m}}{\partial x _{n}} = \nabla f_{1} (x)^{T} \nabla f_{2} (x)^{T} ⋮ \nabla f_{m} (x)^{T}

Each row of the Jacobian is the gradient (transposed into a row vector) of one of the output components of $f$ . The entry at row $i$ , column $j$ is $\frac{\partial f _{i}}{\partial x _{j}}$ , telling us how the $i$ -th output changes with respect to the $j$ -th input.

Example Calculation

Let's compute the Jacobian for the function $f : R^{2} \to R^{2}$ defined by:

f (x, y) = (f_{1} (x, y) f_{2} (x, y)) = (x^{2} y 5 x + sin (y))

Find the gradient of the first component function, $f_{1} = x^{2} y$ :
$\nabla f_{1} = (\frac{\partial f _{1}}{\partial x} \frac{\partial f _{1}}{\partial y}) = (2 x y x^{2})$
Find the gradient of the second component function, $f_{2} = 5 x + sin (y)$ :
$\nabla f_{2} = (\frac{\partial f _{2}}{\partial x} \frac{\partial f _{2}}{\partial y}) = (5 cos (y))$
Assemble the Jacobian by stacking the transposed gradients:
$J_{f} (x, y) = ((\nabla f_{1})^{T} (\nabla f_{2})^{T}) = (2 x y 5 x^{2} cos (y))$

This is a $2 \times 2$ matrix, as expected for a function from $R^{2}$ to $R^{2}$ .

2. The Multivariate Chain Rule

The true power of the Jacobian becomes apparent when we differentiate composite functions. Many complex systems, including neural networks, are built by chaining simpler functions together. For example, if we have $y = f (u)$ and $u = g (x)$ , the final output $y$ is a function of the initial input $x$ through the composition $y = (f \circ g) (x)$ .

The multivariate chain rule states that the Jacobian of the composite function is the matrix product of the individual Jacobians.

The Chain Rule of Calculus - Eli Bendersky's website

This section from Eli Bendersky's website clearly states the general formula for the multivariate chain rule.

Read the section titled "Multivariate chain rule - general formulation". Pay close attention to the core formula: D(f \circ g)(a)=Df(g(a)) \cdot Dg(a). Notice that it's a matrix multiplication, and the Jacobian of the outer function f is evaluated at the output of the inner function, g(a).

The rule is:

J_{f \circ g} (x) = J_{f} (g (x)) \cdot J_{g} (x)

This is a beautiful analogy to the single-variable case $\frac{d y}{d x} = \frac{d y}{d u} \frac{d u}{d x}$ , but with scalar derivatives replaced by Jacobian matrices and multiplication replaced by matrix multiplication. The order matters!

Let's see this in action with a concrete example.

The Chain Rule of Calculus - Eli Bendersky's website

This section provides a detailed, step-by-step example of applying the chain rule.

Please read the entire section "Multivariate chain rule - examples". I encourage you to pause and perform the calculations yourself as you read. The example computes the rate of change of temperature along a parametric curve, which is a classic application.

Let's recap the example from the reading. We have a temperature field $T = f (x, y, z)$ and a path through this field given by $g (t) = (x (t), y (t), z (t))$ . We want to find how temperature changes with respect to time, $\frac{d T}{d t}$ . This is the derivative of the composition $f \circ g$ .

Function $g : R \to R^{3}$ takes time $t$ to a position in space. Its Jacobian is a $3 \times 1$ matrix (a column vector).
Function $f : R^{3} \to R$ takes a position to a scalar temperature. Its Jacobian is a $1 \times 3$ matrix (a row vector, which is just the transposed gradient $\nabla f^{T}$ ).

Applying the chain rule, $D (f \circ g) (t) = D f (g (t)) \cdot D g (t)$ , we multiply a $1 \times 3$ matrix by a $3 \times 1$ matrix, resulting in a $1 \times 1$ matrix—a scalar, which is exactly what we expect for $\frac{d T}{d t}$ .

3. Application: The Essence of Backpropagation

The multivariate chain rule is not just a mathematical curiosity; it is the engine that powers the training of virtually all modern neural networks. A deep neural network is nothing more than a long chain of composed functions.

A visual representation of the multivariate chain rule applied to a sequence of functions. The derivative of the overall function F'(x) is the product of the Jacobian matrices of each function in the chain.

Imagine the diagram above represents a simple neural network.

$x$ is the input data.
$a$ , $b$ , and $c$ are the outputs of successive layers (activations).
$y$ is the final output, from which we compute a scalar loss, $L$ .

Training this network means finding the gradient of the loss $L$ with respect to the parameters within each transformation (e.g., in the function that maps $x$ to $a$ ). To find $\frac{\partial L}{\partial x}$ , we apply the chain rule:

\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y} \frac{\partial y}{\partial c} \frac{\partial c}{\partial b} \frac{\partial b}{\partial a} \frac{\partial a}{\partial x}

Each term $\frac{\partial ( output )}{\partial ( input )}$ is a Jacobian matrix. The famous backpropagation algorithm is simply a computationally efficient procedure for calculating this long product of Jacobians, moving backward from the final loss to the earliest parameters in the network. We will revisit this in detail when we study neural networks, but now you understand its mathematical heart.

Conclusion

In this lesson, we have moved from scalar- to vector-valued functions and equipped ourselves with the necessary calculus tools.

Here are the key takeaways:

The Jacobian matrix is the generalization of the derivative for vector-valued functions. Its entries are all the first-order partial derivatives, capturing how each output component changes with respect to each input component.
Geometrically, the Jacobian defines the best local linear approximation of a vector function at a specific point.
The multivariate chain rule provides a way to differentiate compositions of functions by multiplying their Jacobian matrices in the correct order.
This rule is the mathematical foundation for the backpropagation algorithm, which is essential for training deep neural networks.

In the next lesson, we will focus on computing derivatives for specific types of matrix expressions that appear constantly in linear models, such as $x^{T} A x$ and $∥ Ax - b ∥^{2}$ . This will give you a practical "cookbook" for many common derivations in statistical machine learning.

Recommended resources

1. The Jacobian Matrix: Derivative of a Vector Function

Example Calculation

2. The Multivariate Chain Rule

3. Application: The Essence of Backpropagation

Conclusion