可以用爱因斯坦求和替代的那些矩阵运算

技术背景

在前面的几篇文章中我们分别介绍过numpy中的爱因斯坦求和函数Einsum和MindSpore框架中的爱因斯坦求和算子Einsum的基本用法。而我们需要知道，爱因斯坦求和其实还可以实现非常多的功能，甚至可以替代大部分的矩阵运算，比如常见的点乘、元素乘、求和等等这些都是可以的。那我们就逐一看一下可以用爱因斯坦求和来替代的那些函数和方法。

案例演示

在numpy、Jax框架和MindSpore框架中都是支持爱因斯坦求和算符的，那么这里为了方便演示，我们采用的是numpy来做一些参考案例:

In [1]: import numpy as np

In [2]: x = np.arange(3)

In [3]: x
Out[3]: array([0, 1, 2])

In [4]: y = np.arange(3, 6)

In [5]: y
Out[5]: array([3, 4, 5])

In [6]: P = np.arange(1, 10).reshape(3,3)

In [7]: P
Out[7]: 
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

矩阵转置

矩阵转置，或者是调换矩阵的某两个维度，这个功能用爱因斯坦求和来做是非常清晰的，我们先看一下相应的公式：

\[P^T=\left[ \begin{matrix} P_{00}&P_{01}&P_{02}\\ P_{10}&P_{11}&P_{12}\\ P_{20}&P_{21}&P_{22} \end{matrix} \right]^T= \left[ \begin{matrix} P_{00}&P_{10}&P_{20}\\ P_{01}&P_{11}&P_{21}\\ P_{02}&P_{12}&P_{22} \end{matrix} \right] \]

一般矩阵转置我们如果用numpy来操作的话，只需要使用P=P.T就可以了，而这个功能用爱因斯坦求和算子也是可以实现的：

In [40]: np.allclose(P.T, np.einsum('kl->lk', P))
Out[40]: True

这里有一个比较有意思的事情是，如果不指定生成的序号，但是给定的爱因斯坦算符顺序如果前面的大于后面的，也可以实现矩阵转置的功能，比如下面的一个案例：

In [41]: np.allclose(P.T, np.einsum('ji', P))
Out[41]: True

元素乘

对应于两个矩阵（矢量、张量）之间的元素乘法，普通操作我们可以直接用\(x*y\)来实现（假定维度大小为3）：

\[x*y = \left[ \begin{matrix} x_0\\x_1\\x_2 \end{matrix} \right]* \left[ \begin{matrix} y_0\\y_1\\y_2 \end{matrix} \right]=\left[ \begin{matrix} x_0y_0\\x_1y_1\\x_2y_2 \end{matrix} \right] \]

对应于代码实现：

In [8]: np.allclose(x*y, np.einsum('k,k->k', x, y))
Out[8]: True

矩阵内求和

把矩阵中的所有元素相加：

\[SUM(x)=SUM(\left[ \begin{matrix} x_0\\x_1\\x_2 \end{matrix} \right])=x_0+x_1+x_2 \]

对应于Python代码实现为：

In [9]: np.allclose(np.sum(x), np.einsum('k->', x))
Out[9]: True

In [12]: np.allclose(np.sum(P), np.einsum('kl->', P))
Out[12]: True

In [13]: np.allclose(np.sum(P, axis=-1), np.einsum('kl->k', P))
Out[13]: True

In [14]: np.allclose(np.sum(P, axis=0), np.einsum('kl->l', P))
Out[14]: True

那么，既然求和能算，同样的平均值也是可以计算的，这里就不展开介绍了。

矩阵点乘

这个应用场景很多，比如当我们需要计算两个向量之间的夹角的时候，就会用到矩阵点乘。矩阵点乘的定义如下：

\[x\cdot y = \left[ \begin{matrix} x_0\\x_1\\x_2 \end{matrix} \right]\cdot \left[ \begin{matrix} y_0\\y_1\\y_2 \end{matrix} \right]=x_0y_0+x_1y_1+x_2y_2 \]

对应的Python代码实现如下所示：

In [15]: np.allclose(np.dot(x, y), np.einsum('k,k->', x, y))
Out[15]: True

矩阵向量乘

这个应用场景也非常多，比如我们经常所用到的向量的伸缩、旋转等，都可以用一系列的矩阵作用在一个向量上来表示，相关的计算公式为：

\[P\cdot x=\left[ \begin{matrix} P_{00}&P_{01}&P_{02}\\ P_{10}&P_{11}&P_{12}\\ P_{20}&P_{21}&P_{22} \end{matrix} \right]\cdot \left[ \begin{matrix} x_0\\x_1\\x_2 \end{matrix} \right]= \left[ \begin{matrix} P_{00}x_0+P_{01}x_1+P_{02}x_2\\P_{10}x_0+P_{11}x_1+P_{12}x_2\\P_{20}x_0+P_{21}x_1+P_{22}x_2 \end{matrix} \right] \]

对应的Python代码如下所示：

In [16]: np.allclose(np.dot(P, x), np.einsum('kl,l->k', P, x))
Out[16]: True

In [25]: np.allclose(np.dot(P, x[:, None]), np.einsum('kl,lm->km', P, x[:, None]))
Out[25]: True

In [31]: np.allclose(np.dot(P, P.T), np.einsum('kl,lm->km', P, P.T))
Out[31]: True

在上述案例中我们还包含了矩阵跟矩阵之间的乘法，这些基本运算都是可以通用的。

克罗内克积

克罗内克积，又叫张量积，比如两个矢量或者矩阵之间没有耦合关系，那么可以用一个克罗内克积来总体表示这两个矢量或者矩阵组成的矢量或者矩阵，该运算被定义为：

\[x\otimes y^{T}=\left[ \begin{matrix} x_0\\x_1\\x_2 \end{matrix} \right]\otimes \left[y_0, y_1, y_2\right]=\left[ \begin{matrix} x_0y_0&x_0y_1&x_0y_2\\ x_1y_0&x_1y_1&x_1y_2\\ x_2y_0&x_2y_1&x_2y_2 \end{matrix} \right] \]

对应Python代码实现如下所示：

In [36]: np.allclose(np.kron(x[:, None], y), np.einsum('kl,l->kl', x[:, None], y))
Out[36]: True

In [37]: np.allclose(np.kron(x, y), np.einsum('kl,l->kl', x[:, None], y).reshape(9))
Out[37]: True

需要注意的是，爱因斯坦求和运算只能减少总的维度数量，但是不可改变维度大小，因此有时候会需要用到reshape的功能配合使用。

取对角元

这个应用也好理解，就是把矩阵的每一个对角元素取出来，用公式描述就是：

\[diag(P)=diag(\left[ \begin{matrix} P_{00}&P_{01}&P_{02}\\ P_{10}&P_{11}&P_{12}\\ P_{20}&P_{21}&P_{22} \end{matrix} \right])=\left[P_{00}, P_{11}, P_{22}\right] \]

相关的Python代码实现如下所示：

In [46]: np.allclose(np.diag(P), np.einsum('ii->i', P))
Out[46]: True

求矩阵迹

矩阵的迹（Trace），就是对所有的对角元进行求和，那么有了上一步使用爱因斯坦求和函数提取所有的对角元之后，其实我们可以稍微调整一下，就能得到求矩阵迹的方法。首先看下矩阵迹的公式定义：

\[Tr(P) = Tr(\left[ \begin{matrix} P_{00}&P_{01}&P_{02}\\ P_{10}&P_{11}&P_{12}\\ P_{20}&P_{21}&P_{22} \end{matrix} \right])=P_{00}+P_{11}+P_{22} \]

相关的Python代码实现如下所示：

In [47]: np.allclose(np.trace(P), np.einsum('ii->', P))
Out[47]: True

多重运算

有时候会涉及到一系列的矩阵按照顺序作用在一个向量上，如果从张量的角度来考虑的话，其中的维度还可以非常灵活的变化，不一定全都是方阵。应该说，这也是爱因斯坦求和算子的重大意义所在。如果不使用爱因斯坦求和算子，那么要计算\(A\cdot B\cdot C\cdot x\)这样的一个过程，可以多次嵌套使用numpy的dot点乘函数。但是这样比较麻烦，一般推荐可以使用numpy中的另外一个函数：multi_dot，相关的Python代码实现如下所示：

In [39]: np.allclose(np.linalg.multi_dot((P, P, P, x)), np.einsum('ij,jk,kl,l->i', P, P, P, x))
Out[39]: True

在这种多重运算的过程中，可以使用einsum_path去找到一条更好的归并路径，以达到提升算法性能的效果。

总结概要

本文主要基于Python的Numpy库，介绍一些爱因斯坦求和算子Einsum的应用场景，包括求和、求内外积、求转置等等。我们需要明确的是，爱因斯坦求和算子的意义主要在于矩阵的多重运算时，可以通过爱因斯坦求和约定将这种复杂问题定义成一个张量网络，通过图模型去找到一个更好的缩并路径，以得到更好的算法复杂度。而如果只是普通的点乘求和之类的运算，其实并不是Einsum的主要功能。但是这些功能也可以用爱因斯坦求和的形式来实现，也说明了这个约定的先进性。当然，也有众多的矩阵运算功能是无法直接通过爱因斯坦求和算子来实现的，比如矩阵求逆、求本征值、矩阵扩维、矩阵重构还有向量叉乘等等。只有在合适的地方使用Einsum，才能体现它的真正价值。