8.2. Gradients of Matrix-Vector Operations

More typically, the functions we’ll need to take the gradient of will themselves be defined in terms of matrix and vector operations. In all of these examples, remember that we’re working with vector-to-scalar functions.

Example: Dot Product¶

Let $\vec a \in \mathbb{R}^n$ be some fixed vector (the equivalent of a constant in this context). Let’s find the gradient of

f(\vec x) = \vec a \cdot \vec x

I find it helpful to think about $f(\vec x)$ in its expanded form,

f(\vec x) = \vec a \cdot \vec x = a_1 x_1 + a_2 x_2 + \cdots + a_n x_n

Remember, $\nabla f(\vec x)$ contains all of the partial derivatives of $f$ , which we now need to compute.

What is $\frac{\partial f}{\partial x_1}$ ? To me, that looks like $a_1$ , since the first term is $a_1 x_1$ and none of the other terms involve $x_1$ .
Similarly, $\frac{\partial f}{\partial x_2} = a_2$ .
In general, $\frac{\partial f}{\partial x_i} = a_i$ .

Putting these together, we get

\nabla f(\vec x) = \begin{bmatrix} \frac{\partial f}{\partial x_1} \\ \frac{\partial f}{\partial x_2} \\ \vdots \\ \frac{\partial f}{\partial x_n} \end{bmatrix} = \begin{bmatrix} a_1 \\ a_2 \\ \vdots \\ a_n \end{bmatrix} = \vec a

Example: Norm and Chain Rule¶

Here’s an extremely important example that shows up everywhere in machine learning. Find the gradients of:

$f(\vec x) = \lVert \vec x \rVert^2$
$f(\vec x) = \lVert \vec x \rVert$

Solution

As we did in the previous example, we can expand $f(\vec x) = \lVert \vec x \rVert^2$ to get
$f(\vec x) = \vec x \cdot \vec x = x_1^2 + x_2^2 + \cdots + x_n^2$
For each $i$ , $\frac{\partial f}{\partial x_i} = 2x_i$ . So,
$\nabla f(\vec x) = \begin{bmatrix} 2x_1 \\ 2x_2 \\ \vdots \\ 2x_n \end{bmatrix} = \boxed{2\vec x}$
Think of this as the equivalent of the “power rule” for vectors.
There are two ways to find the gradient of $f(\vec x) = \lVert \vec x \rVert$ : directly, or by using the chain rule. It’s not immediately obvious how the chain rule should work here, so we’ll start with the direct method and reason about how the chain rule may arise.
Direct method: Let’s start by expanding $f(\vec x) = \lVert \vec x \rVert$ like we did above.
$f(\vec x) = \sqrt{\vec x \cdot \vec x} = (\vec x \cdot \vec x)^{1/2} = (x_1^2 + x_2^2 + \cdots + x_n^2)^{1/2}$
For each $i$ , the (regular, scalar-to-scalar) chain rule tells us that
$\frac{\partial f}{\partial x_i} = \frac{1}{2}(x_1^2 + x_2^2 + \cdots + x_n^2)^{-1/2} \cdot 2x_i = \frac{x_i}{\sqrt{x_1^2 + x_2^2 + \cdots + x_n^2}} = \frac{x_i}{\lVert \vec x \rVert}$
So,
$\nabla f(\vec x) = \begin{bmatrix} \frac{x_1}{\lVert \vec x \rVert} \\ \frac{x_2}{\lVert \vec x \rVert} \\ \vdots \\ \frac{x_n}{\lVert \vec x \rVert} \end{bmatrix} = \boxed{\frac{\vec x}{\lVert \vec x \rVert}}$
Chain rule method: Let me start by writing $f(\vec x)$ in terms of a composition of two functions.
$f(\vec x) = \lVert \vec x \rVert = \sqrt{\lVert \vec x \rVert^2} = h(g(\vec x))$
where $g(\vec x) = \lVert \vec x \rVert^2$ and $h(x) = \sqrt{x}$ . Note that $g: \mathbb{R}^n \to \mathbb{R}$ is the vector-to-scalar function we found the gradient of above, and $h: \mathbb{R} \to \mathbb{R}$ is a scalar-to-scalar function.
Then, generalizing the calculation we did with the first method, we have a “chain rule” for a function $h(g(\vec x))$ (where $h$ is scalar-to-scalar and $g$ is vector-to-scalar):
$\nabla f(\vec x) = \underbrace{\left(\frac{\text{d}h}{\text{d}x}(g (\vec x))\right)}_{h' (g(\vec x))} \nabla g(\vec x)$
Remember that $h(x) = \sqrt{x}$ , so $\frac{\text{d}h}{\text{d}x}(x) = \frac{1}{2\sqrt{x}}$ and $\frac{\text{d}h}{\text{d}x}(g(\vec x)) = \frac{1}{2\sqrt{g(\vec x)}} = \frac{1}{2\lVert \vec x \rVert}$ . This means
$f(\vec x) = \left(\frac{\text{d}h}{\text{d}x}(g (\vec x))\right) \nabla g(\vec x) = \left( \frac{1}{2\lVert \vec x \rVert} \right) 2 \vec x = \frac{\vec x}{\lVert \vec x \rVert}$
which is what we found earlier.

Example: Norm to an Exponent¶

Find the gradient of $f(\vec x) = \lVert \vec x \rVert^p$ , where $p$ is some real number.

Solution

We can treat this as a composition of two functions, $g(\vec x) = \lVert \vec x \rVert$ and $h(x) = x^p$ , and use the chain rule introduced in the solution to the previous example.

$\frac{\text{d}h}{\text{d}x}(x) = p x^{p-1}$ and $\nabla g(\vec x) = \frac{\vec x}{\lVert \vec x \rVert}$ . Putting these together yields

\begin{align*} \nabla f(\vec x) &= \left( \frac{\text{d}h}{\text{d}x}(g (\vec x)) \right) \nabla g(\vec x) \\ &= p\, g(\vec x)^{p-1} \frac{\vec x}{\lVert \vec x \rVert} \\ &= p\, \lVert \vec x \rVert^{p-1} \frac{\vec x}{\lVert \vec x \rVert} \\ &= p\, \lVert \vec x \rVert^{p-2} \vec x \end{align*}

Example: Log Sum Exp¶

If $\vec x \in \mathbb{R}^n$ , we can define the log sum exp function as

f(\vec x) = \log \left( \sum_{i=1}^n e^{x_i} \right)

What is $\nabla f(\vec x)$ ? (The answer is called the softmax function, and comes up all the time in machine learning, when we want our models to output predicted probabilities in a classification problem.)

Solution

Let’s look at the partial derivatives with respect to each $x_i$ .

\frac{\partial f}{\partial x_i} = \frac{\partial}{\partial x_i} \left( \log \left( \sum_{j=1}^n e^{x_j} \right) \right) = \left(\frac{1}{\sum_{j=1}^n e^{x_j}} \right) \frac{\partial}{\partial x_i} \left( \sum_{j=1}^n e^{x_j} \right) = \frac{e^{x_i}}{\sum_{j=1}^n e^{x_j}}

Then,

\nabla f(\vec x) = \begin{bmatrix} \frac{e^{x_1}}{\sum_{j=1}^n e^{x_j}} \\ \frac{e^{x_2}}{\sum_{j=1}^n e^{x_j}} \\ \vdots \\ \frac{e^{x_n}}{\sum_{j=1}^n e^{x_j}} \end{bmatrix} = \frac{1}{\sum_{j=1}^n e^{x_j}} \begin{bmatrix} e^{x_1} \\ e^{x_2} \\ \vdots \\ e^{x_n} \end{bmatrix}

There isn’t really a way to simplify the expression using matrix-vector operations, so I’ll leave it as-is. As mentioned above, the gradient we’re looking at is called the softmax function. The softmax function maps $\mathbb{R}^n \to \mathbb{R}^n$ , meaning its a vector-to-vector function.

Let’s suppose we have the matrix $\begin{bmatrix} 3 \\ 5 \\ -1 \end{bmatrix}$ . What does passing it through the softmax function yield?

\text{softmax}\left(\begin{bmatrix} 3 \\ {\color{orange}5} \\ -1 \end{bmatrix}\right) = \frac{1}{e^3 + e^{5} + e^{-1}}\begin{bmatrix} e^3 \\ e^5 \\ e^{-1} \end{bmatrix} \approx \begin{bmatrix} 0.119 \\ {\color{orange}0.879} \\ 0.0002 \end{bmatrix}

The output vector has the same number of elements as the input vector, but each element is between 0 and 1, and the sum of elements is 1, meaning that we can interpret the outputted vector as a probability distribution. Larger values in the output correspond to larger values in the input, and almost all of the “mass” is concentrated at the maximum element of the input vector (position 2), hence the name “soft” max. (The “hard” max might be $\begin{bmatrix} 0 \\ 1 \\ 0 \end{bmatrix}$ in this case.)

Example: Quadratic Forms¶

Suppose $x \in \mathbb{R}^n$ and $A$ is an $n \times n$ matrix. The function

f(\vec x) = \vec x^T A \vec x

is called a quadratic form, and its gradient is given by

\nabla f(\vec x) = (A + A^T) \vec x

We won’t directly cover the proof of this formula here; one place to find it is here. Instead, we’ll focus our energy on understanding how it works, since it’s extremely important.

Let $A = \begin{bmatrix} a & b \\ c & d \end{bmatrix}$ . Expand out $f(\vec x) = \vec x^T A \vec x$ and compute $\nabla f(\vec x)$ directly by computing partial derivatives, and verify that the result you get matches the formula above.
In quadratic forms, we typically assume that $A$ is symmetric, meaning that $A = A^T$ . Why do you think this assumption is made (what does it help with)?
- Hint: Let $A = \begin{bmatrix} 3 & 2 \\ 6 & 1 \end{bmatrix}$ and $B = \begin{bmatrix} 3 & 4 \\ 4 & 1 \end{bmatrix}$ . Compute $\nabla (\vec x^T A \vec x)$ and $\nabla (\vec x^T B \vec x)$ .
If $A$ is any symmetric $n \times n$ matrix, what is $\nabla f(\vec x)$ ?
Suppose $A$ is symmetric and $n \times n$ , $\vec b \in \mathbb{R}^n$ , and $c \in \mathbb{R}$ . Find the gradient of
$f(\vec x) = \vec x^T A \vec x + \vec b \cdot \vec x + c$

Solution

If $A = \begin{bmatrix} a & b \\ c & d \end{bmatrix}$ , then
$\begin{align*} f(\vec x) &= \vec x^T A \vec x \\ &= \begin{bmatrix} x_1 & x_2 \end{bmatrix} \begin{bmatrix} a & b \\ c & d \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \end{bmatrix} \\ &= \begin{bmatrix} x_1 & x_2 \end{bmatrix} \begin{bmatrix} a x_1 + b x_2 \\ c x_1 + d x_2 \end{bmatrix} \\ &= a x_1^2 + (b + c) x_1 x_2 + d x_2^2 \end{align*}$
Then,
$\frac{\partial f}{\partial x_1} = 2a {\color{#3d81f6}x_1} + (b + c) {\color{#3d81f6}x_2}, \quad \frac{\partial f}{\partial x_2} = (b + c) {\color{#3d81f6}x_1} + 2d {\color{#3d81f6}x_2}$
$\nabla f(\vec x) = \begin{bmatrix} 2a {\color{#3d81f6}x_1} + (b + c) {\color{#3d81f6}x_2} \\ (b + c) {\color{#3d81f6}x_1} + 2d {\color{#3d81f6}x_2} \end{bmatrix} = \begin{bmatrix} 2a & b + c \\ b + c & 2d \end{bmatrix} \begin{bmatrix} {\color{#3d81f6}x_1} \\ {\color{#3d81f6}x_2} \end{bmatrix} = (A + A^T) \vec x$
since $A^T = \begin{bmatrix} a & c \\ b & d \end{bmatrix}$ , meaning $A + A^T = \begin{bmatrix} 2a & b + c \\ b + c & 2d \end{bmatrix}$ .
For a particular quadratic form, there are infinitely many choices of matrices $A$ that represent it. To illustrate, let’s look at $A = \begin{bmatrix} 3 & 2 \\ 6 & 1 \end{bmatrix}$ and $B = \begin{bmatrix} 3 & 4 \\ 4 & 1 \end{bmatrix}$ as provided in the hint.
$\vec x^T A \vec x = \begin{bmatrix} x_1 & x_2 \end{bmatrix} \begin{bmatrix} 3 & 2 \\ 6 & 1 \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \end{bmatrix} = 3 x_1^2 + (2 + 6) x_1 x_2 + x_2^2$
$\vec x^T B \vec x = \begin{bmatrix} x_1 & x_2 \end{bmatrix} \begin{bmatrix} 3 & 4 \\ 4 & 1 \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \end{bmatrix} = 3 x_1^2 + (4 + 4) x_1 x_2 + x_2^2$
Note that both $\vec x^T A \vec x$ and $\vec x^T B \vec x$ are equal to the expression $3x_1^2 + 8x_1x_2 + x_2^2$ . In fact, any matrix of the form $\begin{bmatrix} 3 & b \\ c & 1 \end{bmatrix}$ where $b + c = 8$ would produce the same quadratic form.
So, to avoid this issue of having infinitely many choices of the matrix $A$ , we pick the symmetric matrix $A$ , where $A = A^T$ . As we’re about to see, this choice of $A$ simplifies the calculation of the gradient.
If $A$ is any symmetric $n \times n$ matrix, then $A = A^T$ , and $A + A^T = 2A$ . So,
$\nabla (\vec x^T A \vec x) = (A + A^T) \vec x = 2A \vec x$
This is also an important rule; don’t forget it.
Think of $f(\vec x) = \vec x^T A \vec x + \vec b \cdot \vec x + c$ as the matrix-vector equivalent of a quadratic function, $ax^2 + bx + c$ . The derivative of $ax^2 + bx + c$ is $2ax + b$ . Check out what the gradient of $f(\vec x)$ ends up being!

\begin{align*} f(\vec x) &= \vec x^T A \vec x + \vec b \cdot \vec x + c \\ \nabla f(\vec x) &= (A + A^T) \vec x + \vec b \\ &= 2A \vec x + \vec b \qquad \text{(since $A$ is symmetric)} \end{align*}

Summary of Important Gradient Rules¶

These are the core rules you need to know moving forward, not just because we’re about to use them in an important proof, but because they’ll come up repeatedly in your future machine learning work.

Function	Name	Gradient
$f(\vec x) = \vec a \cdot \vec x$	dot product	$\nabla f(\vec x) = \vec a$
$f(\vec x) = \lVert \vec x \rVert^2$	squared norm	$\nabla f(\vec x) = 2\vec x$
$f(\vec x) = \vec x^T A \vec x$	quadratic form	$\nabla f(\vec x) = (A + A^T) \vec x$ if $A$ is symmetric, $\nabla f(\vec x) = 2A \vec x$

We now apply the gradient concepts to examples and matrix-vector operations.

Optimization¶

In the calculus of scalar-to-scalar functions, we have a well-understood procedure for finding the extrema of a function. The general strategy is to take the derivative, set it to zero, and solve for the inputs (called critical points) that satisfy that condition. To be thorough, we’d perform a second derivative test to check whether each critical point is a maximum, minimum, or neither.

In the land of vector-to-scalar functions, the equivalent is to solve for where the gradient is zero, which corresponds to finding where all partial derivatives are zero. Assessing whether we’ve arrived at a maximum or minimum is more difficult to do in the vector-to-scalar case, and we will save a discussion of this for Chapter 8.5.

As an example, consider

f(\vec x) = \vec x^T \begin{bmatrix} 3 & 4 \\ 4 & 1 \end{bmatrix} \vec x + \begin{bmatrix} 1 \\ 2 \end{bmatrix} \cdot \vec x + 3

As we computed earlier, the gradient of $f(\vec x) = \vec x^T A \vec x + \vec b \cdot \vec x + c$ is $\nabla f(\vec x) = 2A \vec x + \vec b$ for symmetric $A$ . So,

\nabla f(\vec x) = 2 \begin{bmatrix} 3 & 4 \\ 4 & 1 \end{bmatrix} \vec x + \begin{bmatrix} 1 \\ 2 \end{bmatrix} = \begin{bmatrix} 6x_1 + 8x_2 + 1 \\ 8x_1 + 2x_2 + 2 \end{bmatrix}

To find the critical points, we set the gradient to zero and solve the resulting system. We can also accomplish this by using the inverse of $A$ , if we happen to have it:

\nabla f(\vec x) = 0 \implies 2 A \vec x + \vec b = 0 \implies \vec x^* = -\frac{1}{2}A^{-1} \vec b

Either way, we find that $\vec x^* = \begin{bmatrix} -7/26 \\ 1/13 \end{bmatrix}$ satisfies $\nabla f(\vec x^*) = 0$ , which corresponds to a local minimum.

import numpy as np
import plotly.graph_objs as go

# Define the function
def f(x, y):
    return 3 * x ** 2 + 8 * x * y + y ** 2 + x + 2 * y + 3

fig = plot_gradient_on_surface(
    f = f,
    lim = 5,
    xaxis_title = 'x₁',
    yaxis_title = 'x₂',
    zaxis_title = 'f(x₁, x₂)',
    title='',
    dfx1 = lambda x, y: 6 * x + 8 * y + 1,
    dfx2 = lambda x, y: 8 * x + 2 * y + 2,
    point = np.array([-7/26, 1/13]),
)

fig.update_layout(title='', scene_camera=dict(eye=dict(x=1, y=2, z=2)))

# Annotate the local minimum point
x_star = -7/26
y_star = 1/13
z_star = f(x_star, y_star)
fig.add_trace(
    go.Scatter3d(
        x=[x_star],
        y=[y_star],
        z=[z_star],
        mode='markers+text',
        marker=dict(size=12, color='gold', symbol='circle'),
        text=["local minimum"],
        textposition="top center",
        textfont=dict(color='gold', size=14),
        name="local minimum"
    )
)

Minimizing Mean Squared Error¶

Remember, the goal of this section is to minimize mean squared error,

R_\text{sq}(\vec w) = \frac{1}{n} \lVert \vec y - X \vec w \rVert^2

In the general case, $X$ is an $n \times (d + 1)$ matrix, $y \in \mathbb{R}^n$ , and $\vec w \in \mathbb{R}^{d+1}$ .

We’re now equipped with the tools to minimize $R_\text{sq}(\vec w)$ by taking its gradient and setting it to zero. Hopefully, we end up with the same conditions on $\vec w^*$ that we derived in Chapter 6.3.

In the most recent example we saw, the optimal vector $\vec x^*$ corresponded to a local minimum. We know that we won’t run into such an issue here since $R_\text{sq}(\vec w)$ cannot output a negative number (it is the average of squared losses), so its minimum possible output is 0, meaning that there will be some global minimizer $\vec w^*$ .

Let’s start by rewriting the squared norm as a dot product and eventually matrix multiplication.

\begin{align*}R_\text{sq}(\vec w) = \frac{1}{n} \lVert \vec y - X \vec w \rVert^2 &= \frac{1}{n} (\vec y - X \vec w) \cdot (\vec y - X \vec w) \\ &= \underbrace{\frac{1}{n} (\vec y - X \vec x)^T (\vec y - X \vec w)}_{\text{since } \vec u \cdot \vec v = \vec u^T \vec v} \\ &= \frac{1}{n} \left( \vec y^T - (X \vec w)^T \right) (\vec y - X \vec w) \\ &= \frac{1}{n} \left( \vec y^T \vec y - {\color{orange}\vec y^T X \vec w} - {\color{orange}(X \vec w)^T \vec y} + (X \vec w)^T X \vec w \right)\end{align*}

Let’s focus on the two terms in orange. They are both equal: they are both the dot product of $\vec y$ and $X \vec w$ . Ideally, I want to express each term as a dot product of $\vec w$ with something, since I’m taking the gradient with respect to $\vec w$ . Remember, the dot product is a scalar, and the transpose of a scalar is just that same scalar. So,

\vec y^T X \vec w = (\vec y^T X \vec w)^T = \vec w^T X^T \vec y = \vec w^T (X^T \vec y)

so, performing this substitution in for both orange terms gives us

\begin{align*}R_\text{sq}(\vec w) &= \frac{1}{n} \left( \vec y^T \vec y - {\color{orange}\vec w^T (X^T \vec y)} - {\color{orange}\vec w^T X^T \vec y} + \vec w^T (X^T X) \vec w \right) \\ &= \frac{1}{n} \left( \vec y^T \vec y - 2 \vec w^T (X^T \vec y) + \vec w^T (X^T X) \vec w \right)\end{align*}

Now, we’re ready to take the gradient, which we’ll do term by term.

$\nabla \left( \vec y^T \vec y \right) = \vec 0$ , since $\vec y^T \vec y$ is a constant with respect to $\vec w$
$\nabla \left( 2 \vec w^T (X^T \vec y) \right) = 2 X^T \vec y$ using the dot product rule, since this is the dot product between $2X^T \vec y$ (a vector) and $\vec w$ (a vector)
$\nabla \left( \vec w^T (X^T X) \vec w \right) = 2X^T X \vec w$ , using the quadratic form rule, since $X^T X$ is a symmetric matrix

Plugging these terms in gives us

\begin{align*}R_\text{sq}(\vec w) &= \frac{1}{n} \left( \vec y^T \vec y - 2 \vec w^T (X^T \vec y) + \vec w^T (X^T X) \vec w \right) \\ \nabla R_\text{sq}(\vec w) &= \frac{1}{n} \left( \nabla \left(\vec y^T \vec y \right) - \nabla \left( 2 \vec w^T (X^T \vec y) \right) + \nabla \left( \vec w^T (X^T X) \vec w \right) \right) \\ &= \frac{1}{n} \left( 0 - 2 X^T \vec y + 2X^T X \vec w \right) \\ &= \boxed{\frac{2}{n} (X^T X \vec w - X^T \vec y)} \end{align*}

Finally, to find the minimizer $\vec w^*$ , we set the gradient to zero and solve.

\frac{2}{n} (X^T X \vec w^* - X^T \vec y) = 0 \implies X^TX \vec w^* = X^T \vec y

Stop me if this feels familiar... these are the normal equations once again! It shouldn’t be a surprise that we ended up with the same conditions on $\vec w^*$ that we derived in Chapter 6.3, since we were solving the same problem.

We’ve now shown that the minimizer of

R_\text{sq}(\vec w) = \frac{1}{n} \lVert \vec y - X \vec w \rVert^2

is given by solving $X^TX \vec w^* = X^T \vec y$ . These equations have a unique solution if $X^TX$ is invertible, and infinitely many solutions otherwise. If $\vec w^*$ satisfies the normal equations, then $X \vec w^*$ is the vector in $\text{colsp}(X)$ that is closest to $\vec y$ . All of that interpretation from Chapter 6.3 and Chapter 7 carry over; we’ve just introduced a new way of finding the solution.

Heads up: In Homework 9, you’ll follow similar steps to minimize a new objective function, that resembles $R_\text{sq}(\vec w)$ but involves another term. There, you’ll minimize

R_\text{ridge}(\vec w) = \lVert \vec y - X \vec w \rVert^2 + \lambda \lVert \vec w \rVert^2

where $\lambda > 0$ is a constant, called the regularization hyperparameter. (Notice the missing $\frac{1}{n}$ .) A good way to practice what you’ve learned (and to get a head start on the homework) is to compute the gradient of $R_\text{ridge}(\vec w)$ and set it to zero. We’ll walk through what the significance of $R_\text{ridge}(\vec w)$ is in the homework.