Computing the Jacobian Matrix

Hello! Welcome back to your course on the mathematical foundations of AI.

Introduction

In our last lesson, we formally derived the multivariable chain rule. We saw that the key to a clean, scalable formulation was to reinterpret the derivative of a multivariable function not as a single value, but as a linear map represented by a matrix. We named this the Jacobian matrix. You learned that the chain rule then elegantly becomes a product of these matrices, a concept that underpins the backpropagation algorithm.

Today's lesson drills down into the Jacobian itself. While you understand its role, the goal now is to master the mechanics of its construction. We will formally define the structure of the Jacobian and, most importantly, practice the skill this lesson is centered on: computing the Jacobian matrix for any given vector-valued function. We will also explore the different notational conventions you may encounter in academic literature.

This lesson solidifies a core computational skill required for analyzing and optimizing complex models.

From Gradient to Jacobian: A Natural Extension

Let's begin by solidifying the connection between the concepts you've already mastered—the derivative and the gradient—and the Jacobian. The video below provides an excellent conceptual bridge, progressing from simple to more complex functions.

As the video explains, we can think of this progression in terms of function mappings:

1. Derivative: For a function $f:\mathbb{R}\to \mathbb{R}$, the derivative $f^{\prime }(x)$ is a scalar that tells us the local rate of change.
2. Gradient: For a function $f:{\mathbb{R}}^n\to \mathbb{R}$, the gradient $\nabla f(\mathbf{x})$ is a vector. Each component tells us the rate of change of the scalar output with respect to one of the input variables.
3. Jacobian: For a function $\mathbf{f}:{\mathbb{R}}^n\to {\mathbb{R}}^m$, the Jacobian is a matrix. It captures how each of the $m$ output components changes with respect to each of the $n$ input components.

Formal Definition and Computation

Let's formalize the structure highlighted in the video. A vector-valued function $\mathbf{f}:{\mathbb{R}}^n\to {\mathbb{R}}^m$ is essentially a stack of $m$ scalar-valued functions, where each function takes an $n$-dimensional vector $\mathbf{x}$ as input:

$$\mathbf{f}(\mathbf{x})=\left[\begin{array}{c}{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)\end{array}\right]$$

The Jacobian matrix of $\mathbf{f}$, denoted as $J_{\mathbf{f}}$ or $\frac{d\mathbf{f}}{d\mathbf{x}}$, is the $m\times n$ matrix of all first-order partial derivatives.

The entry in the $i$-th row and $j$-th column is given by:

$$(J_{\mathbf{f}}{)}_{ij}=\frac{\partial f_i}{\partial x_j}$$

Expanding this out gives the full matrix:

$$J_{\mathbf{f}}=\left[\begin{array}{cccc}\frac{\partial f_1}{\partial x_1}& \frac{\partial f_1}{\partial x_2}& \cdots & \frac{\partial f_1}{\partial x_n}\\ \frac{\partial f_2}{\partial x_1}& \frac{\partial f_2}{\partial x_2}& \cdots & \frac{\partial f_2}{\partial x_n}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{\partial f_m}{\partial x_1}& \frac{\partial f_m}{\partial x_2}& \cdots & \frac{\partial f_m}{\partial x_n}\end{array}\right]$$

Notice that the $i$-th row of the Jacobian is simply the transpose of the gradient of the $i$-th component function: ${\left(\nabla f_i\right)}^T$. This connects back to the "stack of gradients" idea.

For a rigorous, text-based definition and a concise example, please review the following resource.

A Concrete Computational Example

Now, let's see the computation in action.

The following image provides three more quick examples, which are excellent for checking your understanding of how the dimensions of the input and output vectors determine the shape of the resulting Jacobian matrix.

Test your understanding!

Given your experience with regression models, consider a function that maps a 2D parameter vector $\mathbf{\theta }=[{\theta }_1,{\theta }_2{]}^T$ to a 3D prediction vector $\mathbf{y}\in {\mathbb{R}}^3$. Let the input data points be fixed at ${\mathbf{x}}_1=1,{\mathbf{x}}_2=2,{\mathbf{x}}_3=3$.

The function $\mathbf{f}:{\mathbb{R}}^2\to {\mathbb{R}}^3$ is defined as:

$$\mathbf{f}(\mathbf{\theta })=\left[\begin{array}{c}{y}_1\\ y_2\\ y_3\end{array}\right]=\left[\begin{array}{c}{\theta }_1+{\theta }_2{x}_1\\ {\theta }_1+{\theta }_2{x}_2\\ {\theta }_1+{\theta }_2{x}_3\end{array}\right]=\left[\begin{array}{c}{\theta }_1+{\theta }_2\\ {\theta }_1+2{\theta }_2\\ {\theta }_1+3{\theta }_2\end{array}\right]$$

Compute the Jacobian matrix $\frac{d\mathbf{f}}{d\mathbf{\theta }}$. What are its dimensions?

Show answer

The function maps from ${\mathbb{R}}^2$ (input $\mathbf{\theta }$) to ${\mathbb{R}}^3$ (output $\mathbf{y}$). Therefore, the Jacobian will be a $3\times 2$ matrix. The components of $\mathbf{f}$ are:

• $f_1({\theta }_1,{\theta }_2)={\theta }_1+{\theta }_2$
• $f_2({\theta }_1,{\theta }_2)={\theta }_1+2{\theta }_2$
• $f_3({\theta }_1,{\theta }_2)={\theta }_1+3{\theta }_2$

The variables are ${\theta }_1$ and ${\theta }_2$. We compute the partial derivatives for each row:

Row 1: Gradient of $f_1$

• $\frac{\partial f_1}{\partial {\theta }_1}=1$
• $\frac{\partial f_1}{\partial {\theta }_2}=1$

Row 2: Gradient of $f_2$

• $\frac{\partial f_2}{\partial {\theta }_1}=1$
• $\frac{\partial f_2}{\partial {\theta }_2}=2$

Row 3: Gradient of $f_3$

• $\frac{\partial f_3}{\partial {\theta }_1}=1$
• $\frac{\partial f_3}{\partial {\theta }_2}=3$

Assembling these into the matrix gives the Jacobian:

$$J_{\mathbf{f}}=\left[\begin{array}{cc}\frac{\partial f_1}{\partial {\theta }_1}& \frac{\partial f_1}{\partial {\theta }_2}\\ \frac{\partial f_2}{\partial {\theta }_1}& \frac{\partial f_2}{\partial {\theta }_2}\\ \frac{\partial f_3}{\partial {\theta }_1}& \frac{\partial f_3}{\partial {\theta }_2}\end{array}\right]=\left[\begin{array}{cc}1& 1\\ 1& 2\\ 1& 3\end{array}\right]$$

This $3\times 2$ matrix is often called the design matrix in the context of linear regression. Each element tells you how sensitive a specific prediction $y_i$ is to a small change in a parameter ${\theta }_j$.

An Important Note on Notational Conventions

As you delve deeper into AI literature, you will find that not all authors define the Jacobian with the same dimensions. This can be a source of significant confusion. The two main conventions are the Numerator Layout and the Denominator Layout.

To summarize this important point:

• Numerator Layout (what we use): For $\mathbf{f}:{\mathbb{R}}^n\to {\mathbb{R}}^m$, the Jacobian is an $m\times n$ matrix. The number of rows equals the dimension of the output (numerator), and the number of columns equals the dimension of the input (denominator). The $i$-th row is $(\nabla f_i{)}^T$.
• Denominator Layout: The Jacobian is an $n\times m$ matrix (the transpose of the numerator layout). The $j$-th column is $\nabla f_j$.

Being aware of this distinction is crucial. If you encounter a Jacobian in a paper with unexpected dimensions, it is likely due to the author's choice of layout convention. Always check the author's definitions.

Conclusion

In this lesson, we focused on the practical and definitional aspects of the Jacobian matrix, a cornerstone of multivariate calculus for machine learning.

Key Takeaways:

• The Jacobian matrix of a function $\mathbf{f}:{\mathbb{R}}^n\to {\mathbb{R}}^m$ is the $m\times n$ matrix of all first-order partial derivatives, where the element $(J_{\mathbf{f}}{)}_{ij}=\frac{\partial f_i}{\partial x_j}$.
• Conceptually, it can be constructed by stacking the transposed gradients of the output component functions, ${\left(\nabla f_i\right)}^T$, as its rows.
• Each element $(J_{\mathbf{f}}{)}_{ij}$ represents the sensitivity of the $i$-th output to a change in the $j$-th input.
• There are two main notational conventions (numerator and denominator layouts) that result in transposed versions of the Jacobian. Being aware of this is vital for reading academic literature.

Preview of the Next Lesson:
We have now generalized the first derivative from a scalar (derivative) to a vector (gradient) and finally to a matrix (Jacobian). In the next lesson, we will apply the same thinking to the second derivative. You will learn to define and compute the Hessian matrix for a scalar-valued function. The Hessian is fundamental to second-order optimization methods and for classifying the nature of critical points as minima, maxima, or saddle points.