Key matrix derivatives

0. , the i-th singular value of

Assume all singular values are unique. The trick is to find a closed form equation for in terms of . Calculating the gradient from that point is the easy part. Given the SVD, , with as i-th columns of respectively, note that:

This gives us:

Thus

the gradient of is given by .

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 is a diagonal matrix of shape . Eg consider the function . Its derivative is . Applying it to each element of a vector we get:

Each diagonal entry , , only depends on its own input , since does not depend on for ; that is exactly why every off-diagonal entry is .

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

---

2. Softmax function

The softmax function () is defined as:

We calculate the partial derivative separately for the case (an output coordinate with respect to its own input) and (an output coordinate with respect to a different input):

This expression forms the entry of the Jacobian matrix. Collecting the entries into and the entries into , 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 is applied to each row of a matrix of shape , the output is a matrix of shape . The Jacobian for each row is a matrix of shape and there are of them.

acts on each row of the matrix with shape .

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 . Writing , , for the per-row Jacobians, this whole-matrix Jacobian is:

Each has shape , one per row, and only affects , so every off-diagonal block is : 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 blocks running down the diagonal and zeros everywhere else.

Each block is the per-row Jacobian, of shape .

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 →