Key matrix derivatives

0. $f(A) = \sigma_i(A)$, the i-th singular value of $A$

Assume all singular values are unique. The trick is to find a closed form equation for $\sigma_i(A)$ in terms of $A$. Calculating the gradient from that point is the easy part. Given the SVD, $A = UDV^T$, with $u_i, v_i$ as i-th columns of $U, V$ respectively, note that:

This gives us:

Thus

the gradient of $\sigma_i(A)$ is given by $\nabla f = u_iv_i^T$.

Verification using Python

import torch
x = torch.randn(1000,1000).double()
dx = torch.randn_like(x).mul(1e-7)

u1,d1,v1 = torch.linalg.svd(x)
u2,d2,v2 = torch.linalg.svd(x+dx)

index = 0
df_true = d2[index] - d1[index]
df_pred = ((u1[:,index].view(-1,1) @ v1.t()[:,index].view(1,-1)) * dx).sum()

torch.allclose(df_true, df_pred)

'''
Outputs:
True
'''

1. Jacobian of elementwise functions

Jacobian of an elementwise function $f: R^n \rightarrow R^n$ is a diagonal matrix of shape $n \times n$. Eg consider the function $\sigma(x) = \frac{1}{1+e^{-x}}$. Its derivative is $\sigma'(x) = \sigma(x)(1 - \sigma(x))$. Applying it to each element of a vector we get:

Each diagonal entry $\textcolor{#1d4ed8}{\sigma'(x_1)}$, $\textcolor{#be185d}{\sigma'(x_2)}$, $\textcolor{#047857}{\sigma'(x_3)}$ only depends on its own input $x_i$, since $\sigma(x_i)$ does not depend on $x_j$ for $j \neq i$; that is exactly why every off-diagonal entry is $0$.

Practically it is easier to store the diagonal matrix as a vector.

2. Softmax function

The softmax function ($\text{softmax}: R^n \rightarrow R^n$) is defined as:

We calculate the partial derivative $\frac{\partial y_i}{\partial x_j}$ separately for the case $i = j$ (an output coordinate with respect to its own input) and $i \neq j$ (an output coordinate with respect to a different input):

This expression forms the entry $i,j$ of the Jacobian matrix. Collecting the $i = j$ entries into $\textcolor{#1d4ed8}{\text{diag}(y)}$ and the $-y_iy_j$ entries into $-\textcolor{#be185d}{yy^T}$, the Jacobian takes a succinct form:

3. Row-wise vector functions

This can be seen as an extension of the idea of elementwise functions. If a function $f: R^{\text{in}} \rightarrow R^{\text{out}}$ is applied to each row of a matrix of shape $(n, \text{in})$, the output is a matrix of shape $(n, \text{out})$. The Jacobian for each row is a matrix of shape $(\text{out}, \text{in})$ and there are $n$ of them.

$f$ acts on each row of the matrix with shape $(n, \text{in})$.

Another point of view is that it can be seen as a function acting on the entire matrix. From that point of view, the Jacobian has a shape $(n \times \text{out}, n \times \text{in})$. Writing $\textcolor{#1d4ed8}{J_1} = \frac{\partial \textcolor{#1d4ed8}{y_1}}{\partial \textcolor{#1d4ed8}{x_1}}$, $\textcolor{#be185d}{J_2} = \frac{\partial \textcolor{#be185d}{y_2}}{\partial \textcolor{#be185d}{x_2}}$, $\textcolor{#047857}{J_3} = \frac{\partial \textcolor{#047857}{y_3}}{\partial \textcolor{#047857}{x_3}}$ for the per-row Jacobians, this whole-matrix Jacobian is:

Each $J_i$ has shape $(\text{out}, \text{in})$, one per row, and only $x_i$ affects $y_i$, so every off-diagonal block is $0$: exactly the same block-diagonal reasoning as the elementwise Jacobian in Section 1, just with matrix blocks in place of scalars.

The Jacobian is a block diagonal matrix, with $n$ blocks running down the diagonal and zeros everywhere else. Each block is the per-row Jacobian, of shape $(\text{out}, \text{in})$.

It is clear that storing the entire block diagonal matrix is a lot more expensive than storing the individual Jacobian matrices. This idea is explored later on in the article on VJP and JVP.

← Jacobian of a function · Function is a DAG →