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

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}}$

3. Custom implementations and validation

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