Gradients of Vector-Valued Functions

Hello! Welcome back to our journey through the essential calculus for machine learning.

In the past few lessons, we've built a strong foundation. We learned to find the derivative of a function with respect to a single variable, extended this to multivariable functions using partial derivatives, and saw how to organize these into a gradient vector ( $\nabla f$ ) for scalar-valued functions. Most recently, we mastered the chain rule, which is crucial for differentiating nested functions.

Today, we will combine these concepts. We'll move beyond functions that output a single number (scalars) to those that output a whole vector of numbers. These are called vector-valued functions, and they are everywhere in machine learning, from a single layer in a neural network to complex data transformations.

Our goal for this lesson is to learn how to compute gradients of vector-valued functions. This will introduce us to a new and powerful mathematical object: the Jacobian matrix. Understanding the Jacobian is the final piece of the puzzle we need before we can tackle the core optimization algorithm of machine learning: gradient descent.

1. From Gradient Vectors to Jacobian Matrices

So far, we've dealt with functions like $f (x_{1}, x_{2}) = x_{1}^{2} + sin (x_{2})$ . This function takes a vector of inputs, $x = [x_{1}, x_{2}]$ , but produces a single scalar output. Its gradient, $\nabla f = [\frac{\partial f}{\partial x _{1}}, \frac{\partial f}{\partial x _{2}}]$ , is a vector that tells us the direction of steepest ascent.

But what if our function itself produces a vector? For example:

f (x) = [f_{1} (x_{1}, x_{2}) f_{2} (x_{1}, x_{2})] = [x_{1}^{2} + sin (x_{2}) x_{1} x_{2}]

This is a vector-valued function. It maps an input vector to an output vector. How do we find its "gradient"? Since there are multiple output functions, we need to find the partial derivatives for each of them. We organize all these partial derivatives into a matrix. This matrix is called the Jacobian.

The following video provides an excellent introduction to what the Jacobian is and why it's a natural extension of the derivatives we've already learned.

The Jacobian : Data Science Basics

Watch this video from ritvikmath to see how the Jacobian matrix is constructed when we move from single-output to multi-output functions.

Watch the segment from 00:51 to 03:39. Focus on the core idea: when you have multiple functions and multiple variables, you need to calculate all possible partial derivatives and collect them in a matrix.

As the video explained, the Jacobian matrix, denoted $J$ , is the collection of all first-order partial derivatives of a vector-valued function. If we have a function $f$ that maps a vector $x \in R^{n}$ to a vector $y \in R^{m}$ , the Jacobian is an $m \times n$ matrix.

Each row of the Jacobian is simply the gradient of one of the output functions.

J = \nabla f_{1} \nabla f_{2} ⋮ \nabla f_{m} = \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}}

To dig a bit deeper into the formal definition and an important note on notation, let's turn to a more detailed article.

The Matrix Calculus You Need For Deep Learning

This article, 'The Matrix Calculus You Need For Deep Learning', provides a rigorous definition of the Jacobian. It also introduces an important concept: layout conventions.

Read the section titled 'Matrix calculus'. Pay close attention to how the Jacobian is constructed by stacking gradient vectors. Also, note the difference between the 'numerator layout' (which we are using) and the 'denominator layout'. Being aware of this difference is very helpful when reading different machine learning papers.

2. Computing the Jacobian: A Step-by-Step Example

Let's solidify this by computing the Jacobian for a concrete example. Consider the function $f : R^{2} \to R^{2}$ defined as:

f (x_{1}, x_{2}) = [f_{1} f_{2}] = [x_{1}^{2} sin (x_{2}) x_{1} + 3 x_{2}^{3}]

Step 1: Compute the gradient of the first output function, $f_{1}$ .
This will be the first row of our Jacobian.

\nabla f_{1} = [\frac{\partial f _{1}}{\partial x _{1}}, \frac{\partial f _{1}}{\partial x _{2}}] = [2 x_{1} sin (x_{2}), x_{1}^{2} cos (x_{2})]

Step 2: Compute the gradient of the second output function, $f_{2}$ .
This will be the second row of our Jacobian.

\nabla f_{2} = [\frac{\partial f _{2}}{\partial x _{1}}, \frac{\partial f _{2}}{\partial x _{2}}] = [1, 9 x_{2}^{2}]

Step 3: Assemble the Jacobian matrix.
We stack the gradient vectors from Step 1 and Step 2.

J_{f} = [2 x_{1} sin (x_{2}) 1 x_{1}^{2} cos (x_{2}) 9 x_{2}^{2}]

This 2x2 matrix represents the complete "derivative" of our vector-valued function. Each entry tells you how one of the outputs changes with respect to a small change in one of the inputs.

Test your understanding!

Compute the Jacobian matrix for the function $f : R^{3} \to R^{2}$ defined by:

f (x_{1}, x_{2}, x_{3}) = [x_{1}^{2} + e^{x_{2}} x_{2} cos (x_{3})]

What are the dimensions of the resulting Jacobian?

Show answer

The function maps from $R^{3}$ to $R^{2}$ , so the Jacobian will be a 2x3 matrix.

Gradient of $f_{1} = x_{1}^{2} + e^{x_{2}}$ :
$\nabla f_{1} = [\frac{\partial f _{1}}{\partial x _{1}}, \frac{\partial f _{1}}{\partial x _{2}}, \frac{\partial f _{1}}{\partial x _{3}}] = [2 x_{1}, e^{x_{2}}, 0]$
Gradient of $f_{2} = x_{2} cos (x_{3})$ :
$\nabla f_{2} = [\frac{\partial f _{2}}{\partial x _{1}}, \frac{\partial f _{2}}{\partial x _{2}}, \frac{\partial f _{2}}{\partial x _{3}}] = [0, cos (x_{3}), - x_{2} sin (x_{3})]$
Assemble the Jacobian: $J_{f} = [2 x_{1} 0 e^{x_{2}} cos (x_{3}) 0 - x_{2} sin (x_{3})]$

3. The Power of the Jacobian: The Vector Chain Rule

The true power of the Jacobian emerges when we combine it with the chain rule. In the previous lesson, we saw that for scalar composite functions $h (x) = g (f (x))$ , the derivative is $h^{'} (x) = g^{'} (f (x)) \cdot f^{'} (x)$ .

How does this work for vector-valued functions? Consider a neural network, where the output of layer 1, $a^{(1)}$ , becomes the input to layer 2, $a^{(2)}$ . This is a composition of vector-valued functions: $a^{(2)} = layer_{2} (layer_{1} (x))$ .

The chain rule extends beautifully to this scenario: the Jacobian of the composite function is simply the matrix product of the individual Jacobians.

If $h (x) = g (f (x))$ , then:
$J_{h} (x) = J_{g} (f (x)) \cdot J_{f} (x)$

This is the vector chain rule. It's not just an abstract formula; it's the engine behind backpropagation in neural networks. It allows us to calculate how a change in a parameter deep inside the network affects the final output by multiplying a series of matrices.

The same ritvikmath video gives a fantastic and intuitive explanation of this process.

The Jacobian : Data Science Basics

Watch the second part of the video to see a practical application of the vector chain rule in the context of a simple neural network.

Watch from 03:39 to 09:45. You don't need to memorize the specifics of the neural network layers. Instead, focus on the main idea: to find the effect of the input x on the final output p, you can simply multiply the Jacobian matrices of each step in the process.

The video demonstrates that what could be a messy process of tracing every possible path through the network becomes a clean and efficient matrix multiplication. Given your background in computer science, this should resonate with how composing well-defined modules can simplify complex systems. The vector chain rule provides the mathematical framework for this composition.

Conclusion

Today, we generalized our understanding of derivatives to handle vector-valued functions. This led us to the Jacobian matrix, a fundamental tool in multivariable calculus and machine learning.

Key Takeaways:

A vector-valued function maps an input vector to an output vector (e.g., $f : R^{n} \to R^{m}$ ).
The Jacobian matrix is the generalization of the derivative for vector-valued functions. It's an $m \times n$ matrix where each row is the gradient of one of the output functions.
The gradient of a scalar function ( $m = 1$ ) is a special case of the Jacobian, resulting in a $1 \times n$ row matrix.
The vector chain rule states that the Jacobian of a composite function is the matrix product of the Jacobians of its parts. This is the mathematical basis for the backpropagation algorithm used to train neural networks.

Preview of the Next Lesson:
We have now assembled all the necessary mathematical machinery! We understand gradients and Jacobians, and we know how to compute them for complex, nested functions. In the next lesson, we will finally put it all into practice by learning about Gradient Descent. We'll see how the gradient vector, which we now know how to compute, tells us the exact direction to "step" to minimize a function, allowing our machine learning models to learn from data.

Contour Plot of f(x, y) with Gradient Descent Paths — This image provides a sneak peek of what we'll be doing next. The lines show how we can use the gradient to iteratively 'walk' down the surface of a function to find its minimum point. This is the core idea of gradient descent.

Recommended resources

1. From Gradient Vectors to Jacobian Matrices

2. Computing the Jacobian: A Step-by-Step Example

3. The Power of the Jacobian: The Vector Chain Rule

Conclusion