Gradients of Matrix-Matrix Multiplication in Deep Learning

Understanding Artificial Neural Networks with Hands-on Experience - Part 1. Matrix Multiplication, Its Gradients and Custom Implementations
https://coolgpu.github.io/coolgpu_blog/github/pages/2020/09/22/matrixmultiplication.html

coolgpu_blog
https://github.com/coolgpu/coolgpu_blog/tree/master/_posts

We talk about matrix multiplication because it can be used in Convolution and ConvTranspose operations to make things simpler. In this post, we will focus on derivation of the gradients of matrix multiplication. While the derivation process may seem complex, the final results will be in a pretty simple form and are easy to remember.

1. Matrix multiplication

The definition of matrix /ˈmeɪtrɪks/ multiplication /ˌmʌltɪplɪˈkeɪʃn/ can be found in every linear algebra /ˈældʒɪbrə/ book. Let’s use the definition from Wikipedia. Given a $m\times k$ matrix $\mathbf{A}$ and a $k\times n$ matrix $\mathbf{B}$

$$\begin{array}{r}\mathbf{A}=\left[\begin{array}{cccc}{a}_{11}& a_{12}& \cdots & a_{1k}\\ a_{21}& a_{22}& \cdots & a_{2k}\\ ⋮& ⋮& \ddots & ⋮\\ a_{m1}& a_{m2}& \cdots & a_{mk}\end{array}\right]\end{array}\text{\hspace{1em}(1)}$$

and

$$\begin{array}{r}\mathbf{B}=\left[\begin{array}{cccc}{b}_{11}& b_{12}& \cdots & b_{1n}\\ b_{21}& b_{22}& \cdots & b_{2n}\\ ⋮& ⋮& \ddots & ⋮\\ b_{k1}& b_{k2}& \cdots & b_{kn}\end{array}\right]\end{array}\text{\hspace{1em}(1)}$$

their matrix product $\mathbf{C}=\mathbf{A}\mathbf{B}$ is defined as

$$\begin{array}{r}\mathbf{C}=\left[\begin{array}{cccc}{c}_{11}& c_{12}& \cdots & c_{1n}\\ c_{21}& c_{22}& \cdots & c_{2n}\\ ⋮& ⋮& c_{ij}& ⋮\\ c_{m1}& c_{m2}& \cdots & c_{mn}\end{array}\right]\end{array}\text{\hspace{1em}(2)}$$

where its element $c_{ij}$ is given by

$$c_{ij}=\sum _{t=1}^k{a}_{it}{b}_{tj}\text{\hspace{1em}(3)}$$

for $i=1,\dots ,m$ and $j=1,\dots ,n$. In other words, $c_{ij}$ is the dot product of the $i$th row of $\mathbf{A}$ and the $j$th column of $\mathbf{B}$.

2. Derivation of the gradients

2.1. Dimensions of the gradients

If we are considering an isolated matrix multiplication, the partial derivative matrix $\mathbf{C}$ with respect to either matrix $\mathbf{A}$ and matrix $\mathbf{B}$ would be a 4-D hyper-space relationship, referred to as Jacobian Matrix. You will also find that there will be many zeros in the 4-D Jacobian Matrix because, as shown in Equation (3), $c_{ij}$ is a function of only the $i$th row of $\mathbf{A}$ and the $j$th column of $\mathbf{B}$, and independent of other rows of $\mathbf{A}$ and other columns of $\mathbf{B}$.

Jacobian matrix and determinant
https://en.wikipedia.org/wiki/Jacobian_matrix_and_determinant

isolate /ˈaɪsəleɪt/
vt. 使孤立；使绝缘；使隔离
n. 隔离种群
vi. 孤立；隔离
adj. 隔离的；孤立的

What we are considering here is not an isolated matrix multiplication. Instead, we are talking about matrix multiplication inside a neural network that will have a scalar loss function. For example, consider a simple case where the loss $L$ is the mean of matrix $\mathbf{C}$:

$$L=\frac{1}{m\times n}\sum _{i=1}^m\sum _{j=1}^n{c}_{ij}\text{\hspace{1em}(4)}$$

our focus is the partial derivatives of scalar $L$ w.r.t. the input matrix $\mathbf{A}$ and $\mathbf{B}$, $\frac{\partial L}{\partial \mathbf{A}}$ and $\frac{\partial L}{\partial \mathbf{B}}$, respectively. Therefore, $\frac{\partial L}{\partial \mathbf{A}}$ has the same dimension as $\mathbf{A}$, which is another $m\times k$ matrix, and $\frac{\partial L}{\partial \mathbf{B}}$ has the same dimension as $\mathbf{B}$, which is another $k\times n$ matrix.

2.2. The chain rule

We will use the chain rule to do backpropagation of gradients. For such an important tool in neural networks, it doesn’t hurt to briefly summarize the chain rule just like in the previous post for one more time. Given a function $L\left(x_1,x_2,\dots x_N\right)$ as

$$L\left(x_1,\dots x_N\right)=L\left(f_1\left(x_1,\dots x_N\right),f_2\left(x_1,\dots x_N\right),\dots ,f_M\left(x_1,\dots x_N\right)\right)\text{\hspace{1em}(5)}$$

Then the gradient of $L$ w.r.t $x_i$ can be computed as

$$\frac{\partial L}{\partial x_i}=\frac{\partial L}{\partial f_1}\frac{\partial f_1}{\partial x_i}+\frac{\partial L}{\partial f_2}\frac{\partial f_2}{\partial x_i}+\cdots +\frac{\partial L}{\partial f_M}\frac{\partial f_M}{\partial x_i}=\sum _{m=1}^M\frac{\partial L}{\partial f_m}\frac{\partial f_m}{\partial x_i}\text{\hspace{1em}(6)}$$

Equation (6) can be understood from two perspectives:

• Summation means that all possible paths through which $x_i$ contributes to $L$ should be included
• Product means that, along each path $m$, the output gradient equals the upstream passed in, $\frac{\partial L}{\partial f_m}$, times the local gradient, $\frac{\partial f_m}{\partial x_i}$.

2.3. Derivation of the gradient $\frac{\partial L}{\partial \mathbf{A}}$

In this section, we will use a $2\times 4$ matrix $\mathbf{A}$ and a $4\times 3$ matrix $\mathbf{B}$ as an example to step-by-step derive the partial derivative of $\frac{\partial L}{\partial \mathbf{A}}$. Please note that the same derivation can be performed on a general $m\times k$ matrix $\mathbf{A}$ and $k\times n$ matrix $\mathbf{B}$. A specific example is used here purely for the purpose of making it more straightforward.

Let’s start with writing the matrix $\mathbf{A}$, $\mathbf{B}$ and their matrix product $\mathbf{C}=AB$ in expanded format.

expand /ɪkˈspænd/
vt. 扩张；使膨胀；详述
vi. 张开，展开；发展

$$\mathbf{A}=\left[\begin{array}{cccc}{a}_{11}& a_{12}& a_{13}& a_{14}\\ a_{21}& a_{22}& a_{23}& a_{24}\end{array}\right]\text{\hspace{1em}(7)}$$

and

$$\mathbf{B}=\left[\begin{array}{ccc}{b}_{11}& b_{12}& b_{13}\\ b_{21}& b_{22}& b_{23}\\ b_{31}& b_{32}& b_{33}\\ b_{41}& b_{42}& b_{43}\end{array}\right]\text{\hspace{1em}(7)}$$

$$\begin{array}{rl}\mathbf{C}& =\left[\begin{array}{ccc}{c}_{11}& c_{12}& c_{13}\\ c_{21}& c_{22}& c_{23}\end{array}\right]=\left[\begin{array}{cccc}{a}_{11}& a_{12}& a_{13}& a_{14}\\ a_{21}& a_{22}& a_{23}& a_{24}\end{array}\right]\left[\begin{array}{ccc}{b}_{11}& b_{12}& b_{13}\\ b_{21}& b_{22}& b_{23}\\ b_{31}& b_{32}& b_{33}\\ b_{41}& b_{42}& b_{43}\end{array}\right]\\ & =\left[\begin{array}{ccc}{a}_{11}{b}_{11}+a_{12}{b}_{21}+a_{13}{b}_{31}+a_{14}{b}_{41}& a_{11}{b}_{12}+a_{12}{b}_{22}+a_{13}{b}_{32}+a_{14}{b}_{42}& a_{11}{b}_{13}+a_{12}{b}_{23}+a_{13}{b}_{33}+a_{14}{b}_{43}\\ a_{21}{b}_{11}+a_{22}{b}_{21}+a_{23}{b}_{31}+a_{24}{b}_{41}& a_{21}{b}_{12}+a_{22}{b}_{22}+a_{23}{b}_{32}+a_{24}{b}_{42}& a_{21}{b}_{13}+a_{22}{b}_{23}+a_{23}{b}_{33}+a_{24}{b}_{43}\end{array}\right]\end{array}\text{\hspace{1em}(8)}$$

Consider an arbitrary element of $\mathbf{A}$, for example $a_{23}$, we have the local partial derivative of $\mathbf{C}$ w.r.t. $a_{23}$ based on Equation (8).

$$\frac{\partial L}{\partial \mathbf{A}}=\frac{\partial L}{\partial \mathbf{C}}\frac{\partial \mathbf{C}}{\partial \mathbf{A}}$$

$$\begin{array}{rl}\frac{\partial c_{11}}{\partial a_{23}}& =0\\ \frac{\partial c_{12}}{\partial a_{23}}& =0\\ \frac{\partial c_{13}}{\partial a_{23}}& =0\\ \frac{\partial c_{21}}{\partial a_{23}}& =\frac{\partial }{\partial a_{23}}\left(a_{21}{b}_{11}+a_{22}{b}_{21}+a_{23}{b}_{31}+a_{24}{b}_{41}\right)=0+0+\frac{\partial }{\partial a_{23}}\left(a_{23}{b}_{31}\right)+0=b_{31}\\ \frac{\partial c_{22}}{\partial a_{23}}& =\frac{\partial }{\partial a_{23}}\left(a_{21}{b}_{12}+a_{22}{b}_{22}+a_{23}{b}_{32}+a_{24}{b}_{42}\right)=0+0+\frac{\partial }{\partial a_{23}}\left(a_{23}{b}_{32}\right)+0=b_{32}\\ \frac{\partial c_{23}}{\partial a_{23}}& =\frac{\partial }{\partial a_{23}}\left(a_{21}{b}_{13}+a_{22}{b}_{23}+a_{23}{b}_{33}+a_{24}{b}_{43}\right)=0+0+\frac{\partial }{\partial a_{23}}\left(a_{23}{b}_{33}\right)+0=b_{33}\end{array}\text{\hspace{1em}(9)}$$

Using the chain rule, we have the partial derivative of the loss $L$ w.r.t. $a_{23}$

$$\begin{array}{rl}\frac{\partial L}{\partial a_{23}}& =\frac{\partial L}{\partial c_{11}}\frac{\partial c_{11}}{\partial a_{23}}+\frac{\partial L}{\partial c_{12}}\frac{\partial c_{12}}{\partial a_{23}}+\frac{\partial L}{\partial c_{13}}\frac{\partial c_{13}}{\partial a_{23}}+\frac{\partial L}{\partial c_{21}}\frac{\partial c_{21}}{\partial a_{23}}+\frac{\partial L}{\partial c_{22}}\frac{\partial c_{22}}{\partial a_{23}}+\frac{\partial L}{\partial c_{23}}\frac{\partial c_{23}}{\partial a_{23}}\\ & =0+0+0+\frac{\partial L}{\partial c_{21}}{b}_{31}+\frac{\partial L}{\partial c_{22}}{b}_{32}+\frac{\partial L}{\partial c_{23}}{b}_{33}\\ & =\frac{\partial L}{\partial c_{21}}{b}_{31}+\frac{\partial L}{\partial c_{22}}{b}_{32}+\frac{\partial L}{\partial c_{23}}{b}_{33}\end{array}\text{\hspace{1em}(10)}$$

The second line in Equation (10) used the results from Equation (9).

Following a similar manner, we can derive the other elements of $\frac{\partial L}{\partial \mathbf{A}}$ as below

$$\begin{array}{rl}\frac{\partial L}{\partial \mathbf{A}}& =\left[\begin{array}{cccc}\frac{\partial L}{\partial a_{11}}& \frac{\partial L}{\partial a_{12}}& \frac{\partial L}{\partial a_{13}}& \frac{\partial L}{\partial a_{14}}\\ \frac{\partial L}{\partial a_{21}}& \frac{\partial L}{\partial a_{22}}& \frac{\partial L}{\partial a_{23}}& \frac{\partial L}{\partial a_{24}}\end{array}\right]\\ & =\left[\begin{array}{cccc}\frac{\partial L}{\partial c_{11}}{b}_{11}+\frac{\partial L}{\partial c_{12}}{b}_{12}+\frac{\partial L}{\partial c_{13}}{b}_{13}& \frac{\partial L}{\partial c_{11}}{b}_{21}+\frac{\partial L}{\partial c_{12}}{b}_{22}+\frac{\partial L}{\partial c_{13}}{b}_{23}& \frac{\partial L}{\partial c_{11}}{b}_{31}+\frac{\partial L}{\partial c_{12}}{b}_{32}+\frac{\partial L}{\partial c_{13}}{b}_{33}& \frac{\partial L}{\partial c_{11}}{b}_{41}+\frac{\partial L}{\partial c_{12}}{b}_{42}+\frac{\partial L}{\partial c_{13}}{b}_{43}\\ \frac{\partial L}{\partial c_{21}}{b}_{11}+\frac{\partial L}{\partial c_{22}}{b}_{12}+\frac{\partial L}{\partial c_{23}}{b}_{13}& \frac{\partial L}{\partial c_{21}}{b}_{21}+\frac{\partial L}{\partial c_{22}}{b}_{22}+\frac{\partial L}{\partial c_{23}}{b}_{23}& \frac{\partial L}{\partial c_{21}}{b}_{31}+\frac{\partial L}{\partial c_{22}}{b}_{32}+\frac{\partial L}{\partial c_{23}}{b}_{33}& \frac{\partial L}{\partial c_{21}}{b}_{41}+\frac{\partial L}{\partial c_{22}}{b}_{42}+\frac{\partial L}{\partial c_{23}}{b}_{43}\end{array}\right]\end{array}\text{\hspace{1em}(11)}$$

Equation (11) can be equivalently rewritten as a matrix product.

$$\begin{array}{r}\frac{\partial L}{\partial \mathbf{A}}=\left[\begin{array}{ccc}\frac{\partial L}{\partial c_{11}}& \frac{\partial L}{\partial c_{12}}& \frac{\partial L}{\partial c_{13}}\\ \frac{\partial L}{\partial c_{21}}& \frac{\partial L}{\partial c_{22}}& \frac{\partial L}{\partial c_{23}}\end{array}\right]\left[\begin{array}{cccc}{b}_{11}& b_{21}& b_{31}& b_{41}\\ b_{12}& b_{22}& b_{32}& b_{42}\\ b_{13}& b_{23}& b_{33}& b_{43}\end{array}\right]\end{array}\text{\hspace{1em}(12)}$$

In fact, the first matrix is the upstream derivative $\frac{\partial L}{\partial \mathbf{C}}$ and the second matrix is the transpose of $\mathbf{B}$. Then we have

$$\frac{\partial L}{\partial \mathbf{A}}=\frac{\partial L}{\partial \mathbf{C}}{\mathbf{B}}^T\text{\hspace{1em}(13)}$$

Equation (13) shows that, for a matrix multiplication $\mathbf{C}=\mathbf{A}\mathbf{B}$ in a neural network, the derivative of the loss $L$ w.r.t matrix $\mathbf{A}$ equals the upstream derivative $\frac{\partial L}{\partial \mathbf{C}}$ times the transpose of matrix $\mathbf{B}$.

Let’s check the dimensions. On the left hand side of Equation (13), $\frac{\partial L}{\partial \mathbf{A}}$ has a dimension of $m\times k$, the same as $\mathbf{A}$. On the right hand side, $\frac{\partial L}{\partial \mathbf{C}}$ has a dimension of $m\times n$ and ${\mathbf{B}}^T$ has a dimension of $n\times k$; therefore, their matrix product has a dimension of $m\times k$ and matches that of $\frac{\partial L}{\partial \mathbf{A}}$.

2.4. Derivation of the gradient $\frac{\partial L}{\partial \mathbf{B}}$

Similarly, for $\frac{\partial L}{\partial \mathbf{B}}$, let’s consider an arbitrary element of $\mathbf{B}$, for example $b_{12}$, we have the local partial derivative of $\mathbf{C}$ w.r.t. $b_{12}$ based on Equation (8) above.

$$\begin{array}{rl}\frac{\partial c_{11}}{\partial b_{12}}& =0\\ \frac{\partial c_{12}}{\partial b_{12}}& =\frac{\partial }{\partial b_{12}}\left(a_{11}{b}_{12}+a_{12}{b}_{22}+a_{13}{b}_{32}+a_{14}{b}_{42}\right)=a_{11}\\ \frac{\partial c_{13}}{\partial b_{12}}& =0\\ \frac{\partial c_{21}}{\partial b_{12}}& =0\\ \frac{\partial c_{22}}{\partial b_{12}}& =\frac{\partial }{\partial b_{12}}\left(a_{21}{b}_{12}+a_{22}{b}_{22}+a_{23}{b}_{32}+a_{24}{b}_{42}\right)=a_{21}\\ \frac{\partial c_{23}}{\partial b_{12}}& =0\end{array}\text{\hspace{1em}(14)}$$

Using the chain rule, we have the partial derivative of the loss $L$ w.r.t. $b_{12}$

$$\begin{array}{rl}\frac{\partial L}{\partial b_{12}}& =\frac{\partial L}{\partial c_{11}}\frac{\partial c_{11}}{\partial b_{12}}+\frac{\partial L}{\partial c_{12}}\frac{\partial c_{12}}{\partial b_{12}}+\frac{\partial L}{\partial c_{13}}\frac{\partial c_{13}}{\partial b_{12}}+\frac{\partial L}{\partial c_{21}}\frac{\partial c_{21}}{\partial b_{12}}+\frac{\partial L}{\partial c_{22}}\frac{\partial c_{22}}{\partial b_{12}}+\frac{\partial L}{\partial c_{23}}\frac{\partial c_{23}}{\partial b_{12}}\\ & =0+\frac{\partial L}{\partial c_{12}}{a}_{11}+0+0+\frac{\partial L}{\partial c_{22}}{a}_{21}+0\\ & =a_{11}\frac{\partial L}{\partial c_{12}}+a_{21}\frac{\partial L}{\partial c_{22}}\end{array}\text{\hspace{1em}(15)}$$

The second line in Equation (15) used the results from Equation (14). Following a similar manner again, we can derive the other elements of $\frac{\partial L}{\partial \mathbf{B}}$ as shown below

$$\begin{gathered} \frac{\partial L}{\partial \mathbf{B}}=\left[\begin{array}{ccc}\frac{\partial L}{\partial b_{11}}& \frac{\partial L}{\partial b_{12}}& \frac{\partial L}{\partial b_{13}}\\ \frac{\partial L}{\partial b_{21}}& \frac{\partial L}{\partial b_{22}}& \frac{\partial L}{\partial b_{23}}\\ \frac{\partial L}{\partial b_{31}}& \frac{\partial L}{\partial b_{32}}& \frac{\partial L}{\partial b_{33}}\\ \frac{\partial L}{\partial b_{41}}& \frac{\partial L}{\partial b_{42}}& \frac{\partial L}{\partial b_{43}}\end{array}\right] \\ =\left[\begin{array}{ccc}{a}_{11}\frac{\partial L}{\partial c_{11}}+a_{21}\frac{\partial L}{\partial c_{21}}& a_{11}\frac{\partial L}{\partial c_{12}}+a_{21}\frac{\partial L}{\partial c_{22}}& a_{11}\frac{\partial L}{\partial c_{13}}+a_{21}\frac{\partial L}{\partial c_{23}}\\ a_{12}\frac{\partial L}{\partial c_{11}}+a_{22}\frac{\partial L}{\partial c_{21}}& a_{12}\frac{\partial L}{\partial c_{12}}+a_{22}\frac{\partial L}{\partial c_{22}}& a_{12}\frac{\partial L}{\partial c_{13}}+a_{22}\frac{\partial L}{\partial c_{23}}\\ a_{13}\frac{\partial L}{\partial c_{11}}+a_{23}\frac{\partial L}{\partial c_{21}}& a_{13}\frac{\partial L}{\partial c_{12}}+a_{23}\frac{\partial L}{\partial c_{22}}& a_{13}\frac{\partial L}{\partial c_{13}}+a_{23}\frac{\partial L}{\partial c_{23}}\\ a_{14}\frac{\partial L}{\partial c_{11}}+a_{24}\frac{\partial L}{\partial c_{21}}& a_{14}\frac{\partial L}{\partial c_{12}}+a_{24}\frac{\partial L}{\partial c_{22}}& a_{14}\frac{\partial L}{\partial c_{13}}+a_{24}\frac{\partial L}{\partial c_{23}}\end{array}\right]\text{\hspace{1em}(16)} \end{gathered}$$

This can be rewritten as a matrix product.

$$\frac{\partial L}{\partial \mathbf{B}}=\left[\begin{array}{cc}{a}_{11}& a_{21}\\ a_{12}& a_{22}\\ a_{13}& a_{23}\\ a_{14}& a_{24}\end{array}\right]\left[\begin{array}{ccc}\frac{\partial L}{\partial c_{11}}& \frac{\partial L}{\partial c_{12}}& \frac{\partial L}{\partial c_{13}}\\ \frac{\partial L}{\partial c_{21}}& \frac{\partial L}{\partial c_{22}}& \frac{\partial L}{\partial c_{23}}\end{array}\right]\text{\hspace{1em}(17)}$$

In fact, the first matrix is the transpose of $\mathbf{A}$ and the second matrix is the upstream derivative $\frac{\partial L}{\partial \mathbf{C}}$. Then we have

$$\frac{\partial L}{\partial \mathbf{B}}={\mathbf{A}}^T\frac{\partial L}{\partial \mathbf{C}}\text{\hspace{1em}(18)}$$

Equation (18) shows that, for a matrix multiplication $\mathbf{C}=\mathbf{A}\mathbf{B}$ in a neural network, the derivative of the loss $L$ w.r.t matrix $\mathbf{B}$ equals the transpose of matrix $\mathbf{A}$ times the upstream derivative $\frac{\partial L}{\partial \mathbf{C}}$. Let’s check the dimensions once more. On the left hand side of Equation (18), $\frac{\partial L}{\partial \mathbf{B}}$ has a dimension of $k\times n$, the same as $\mathbf{B}$. On the right hand side, ${\mathbf{A}}^T$ has a dimension of $k\times m$ and $\frac{\partial L}{\partial \mathbf{C}}$ has a dimension of $m\times n$; therefore, their matrix product has a dimension of $k\times n$ and matches that of $\frac{\partial L}{\partial \mathbf{B}}$.

Again, the above derivations can be generalized to any matrix multiplication. If you have time, you can derive it by yourself, just make sure the subscript indices are correct.

3. Custom implementations and validation

With the derived Equations (13) and (18), it is in fact pretty easy to implement the backward pass of matrix multiplication. Please see the example implementation on GitHub for a network that simply takes the mean of the matrix product $\mathbf{C}=\mathbf{A}\mathbf{B}$ as the loss. The core part is just a 3-line code as demonstrated below.

grad_C_manual = (torch.ones(C.shape, dtype=torch.float64) / C.numel())

grad_A_manual = grad_C_manual.mm(B.t())
grad_B_manual = A.t().mm(grad_C_manual)

The first line calculate the derivative of the loss w.r.t $\mathbf{C}$ for the mean operation in Equation (4), which serves as the upstream gradient for $\frac{\partial L}{\partial \mathbf{A}}$ and $\frac{\partial L}{\partial \mathbf{B}}$.

The second and third lines compute $\frac{\partial L}{\partial \mathbf{A}}$ and $\frac{\partial L}{\partial \mathbf{B}}$ using the chain rule based on Equations (13) and (18), respectively. The function $t()$ is a matrix transpose operation.

To validate our derivations and implementation, we compared these results with those from Torch built-in implementation via $loss.backward()$ and they matched.

Demo_MatrixMultiplication_backward.py
https://github.com/coolgpu/Demo_Matrix_Multiplication_backward/blob/master/Demo_MatrixMultiplication_backward.py

#!/usr/bin/env python
# coding=utf-8

import matplotlib

import torch

print(matplotlib.__version__)

# A is a (MxP) matrix and B is a (PxN) matrix, so C=AxB is a (MxN) matrix

M, P, N = 2, 3, 4

# torch.randint(low=0, high, size, \*, generator=None, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False) -> Tensor
# Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive).
# requires_grad (bool, optional) - If autograd should record operations on the returned tensor. Default: False.
A = torch.randint(0, 100, (M, P), requires_grad=True, dtype=torch.float64)
B = torch.randint(0, 100, (P, N), requires_grad=True, dtype=torch.float64)

# https://pytorch.org/docs/stable/generated/torch.mm.html
# torch.mm(input, mat2, *, out=None) -> Tensor
# Performs a matrix multiplication of the matrices input and mat2.
# If `input` is a (n×m) tensor, `mat2` is a (m×p) tensor, `out` will be a (n×p) tensor.
# This function does not broadcast. For broadcasting matrix products, see torch.matmul().
C = A.mm(B)

# Tensor.retain_grad() -> None
# Enables this Tensor to have their grad populated during backward(). This is a no-op for leaf tensors.
C.retain_grad()

# calculate the loss simply as the mean of C
# torch.mean(input, *, dtype=None) -> Tensor
# Returns the mean value of all elements in the input tensor. Input must be floating point or complex.
loss = C.mean()
print(f"\nloss = {loss.item()}")

# perform build-in backpropagation
# Tensor.backward(gradient=None, retain_graph=None, create_graph=False, inputs=None)
# Computes the gradient of current tensor wrt graph leaves.
loss.backward()

print('\nA=\n', A)
print('B=\n', B)
print('C=\n', C)

print('\nbuilt-in dL/dC=\n', C.grad)
print('built-in dL/dA=\n', A.grad)
print('built-in dL/dB=\n', B.grad)
# Tensor.grad
# This attribute is None by default and becomes a Tensor the first time a call to backward() computes gradients for self.
# The attribute will then contain the gradients computed and future calls to backward() will accumulate (add) gradients into it.

# Now perform maunal calculation of the gradients dL/dC, dL/dA and dL/dB
grad_C_manual = (torch.ones(C.shape, dtype=torch.float64) / C.numel())
grad_A_manual = grad_C_manual.mm(B.t())
grad_B_manual = A.t().mm(grad_C_manual)

print('\nmanual dL/dC=\n', grad_C_manual)
print('manual dL/dA=\n', grad_A_manual)
print('manual dL/dA=\n', grad_B_manual)

diff_grad_C = grad_C_manual - C.grad
diff_grad_A = grad_A_manual - A.grad
diff_grad_B = grad_B_manual - B.grad

print('\nDifference between custom implementation and Torch built-in:')
print('diff_grad_C max difference:', diff_grad_C.abs().max().detach().numpy())
print('diff_grad_A max difference:', diff_grad_A.abs().max().detach().numpy())
print('diff_grad_B max difference:', diff_grad_B.abs().max().detach().numpy())

print('\nDone!')

/home/yongqiang/miniconda3/bin/python /home/yongqiang/stable_diffusion_work/stable_diffusion_diffusers/yongqiang.py 
3.7.1

loss = 6445.375

A=
 tensor([[ 2., 75., 68.],
        [ 1., 44.,  7.]], dtype=torch.float64, requires_grad=True)
B=
 tensor([[31., 37., 26., 41.],
        [72., 37., 76., 47.],
        [74., 76., 89., 75.]], dtype=torch.float64, requires_grad=True)
C=
 tensor([[10494.,  8017., 11804.,  8707.],
        [ 3717.,  2197.,  3993.,  2634.]], dtype=torch.float64,
       grad_fn=<MmBackward0>)

built-in dL/dC=
 tensor([[0.1250, 0.1250, 0.1250, 0.1250],
        [0.1250, 0.1250, 0.1250, 0.1250]], dtype=torch.float64)
built-in dL/dA=
 tensor([[16.8750, 29.0000, 39.2500],
        [16.8750, 29.0000, 39.2500]], dtype=torch.float64)
built-in dL/dB=
 tensor([[ 0.3750,  0.3750,  0.3750,  0.3750],
        [14.8750, 14.8750, 14.8750, 14.8750],
        [ 9.3750,  9.3750,  9.3750,  9.3750]], dtype=torch.float64)

manual dL/dC=
 tensor([[0.1250, 0.1250, 0.1250, 0.1250],
        [0.1250, 0.1250, 0.1250, 0.1250]], dtype=torch.float64)
manual dL/dA=
 tensor([[16.8750, 29.0000, 39.2500],
        [16.8750, 29.0000, 39.2500]], dtype=torch.float64,
       grad_fn=<MmBackward0>)
manual dL/dA=
 tensor([[ 0.3750,  0.3750,  0.3750,  0.3750],
        [14.8750, 14.8750, 14.8750, 14.8750],
        [ 9.3750,  9.3750,  9.3750,  9.3750]], dtype=torch.float64,
       grad_fn=<MmBackward0>)

Difference between custom implementation and Torch built-in:
diff_grad_C max difference: 0.0
diff_grad_A max difference: 0.0
diff_grad_B max difference: 0.0

Done!

Process finished with exit code 0

4. Summary

In this post, we demonstrated how to derive the gradients of matrix multiplication in neural networks. While the derivation steps seem complex, the final equations of the gradients are pretty simple and easy to implement:

$$\frac{\partial L}{\partial \mathbf{A}}=\frac{\partial L}{\partial \mathbf{C}}{\mathbf{B}}^T$$

$$\frac{\partial L}{\partial \mathbf{B}}={\mathbf{A}}^T\frac{\partial L}{\partial \mathbf{C}}$$

In real neural networks applications, the matrix $\mathbf{A}$ and $\mathbf{B}$ typically come from the outputs of other layers. In those scenarios, the gradients $\frac{\partial L}{\partial \mathbf{A}}$ and $\frac{\partial L}{\partial \mathbf{B}}$ can serve as the upsteam gradients of those layers in backpropagation computing.

References

[1] Yongqiang Cheng, https://yongqiang.blog.csdn.net/
[2] Understanding Artificial Neural Networks with Hands-on Experience - Part 1. Matrix Multiplication, Its Gradients and Custom Implementations, https://coolgpu.github.io/coolgpu_blog/github/pages/2020/09/22/matrixmultiplication.html
[3] U-Net: Convolutional Networks for Biomedical Image Segmentation, https://arxiv.org/abs/1505.04597
[4] Deep Residual Learning for Image Recognition, https://arxiv.org/abs/1512.03385