AlphaTensor：发现更快的矩阵乘法

参考文献：

1. Discovering faster matrix multiplication algorithms with reinforcement learning：https://doi.org/10.1038/s41586-022-05172-4

Tensor

矩阵$A\in {\mathbb{F}}^{n\times m}$和$B\in {\mathbb{F}}^{m\times p}$，矩阵乘积 $(A,B)\mapsto C:=AB$ 是双线性的（bilinear），因此可以表示为$3$维的张量（Tensor）：
T n , m , p = [ t a b c ] ∈ { 0 , 1 } n m , m p , n p T_{n,m,p} = [t_{abc}] \in \{0,1\}^{nm,mp,np} Tn,m,p​=[tabc​]∈{0,1}nm,mp,np
这个$T_{n,m,p}$可以指定读写位置。先把$A,B,C$展开为一维向量$A^{\prime },B^{\prime },C^{\prime }$，如果$t_{abc}=1$，那么$C_c^{\prime }$的值的加和项中就有$A_a^{\prime }\cdot B_b^{\prime }$

我们把$T_{n,m,p}$做分解，得到$U\in {\mathbb{Z}}^{nm\times R},V\in {\mathbb{Z}}^{mp\times R},W\in {\mathbb{Z}}^{np\times R}$，使得
$$T_{n,m,p}=\sum _{r=1}^R{u}^{(r)}\otimes v^{(r)}\otimes w^{(r)}$$

其中的$\otimes$是列向量外积（outer product）或者叫张量积（tensor product）

可以写出根据$T_{n,m,p}$的分解$U,V,W$控制的矩阵乘法：

如果分解后，列数$R$越小，那么运算时的小矩阵乘法的数量就不大于$R$。 著名的 Strassen’s algorithm，它就是找到了张量$T_{2,2,2}$的$R=7$的分解。

AlphaTensor

这篇文章利用 Neural network 搭建了一个模型 AlphaTensor，发现了一系列的对张量$T_{n,m,p}$的更优分解：

具体的$U,V,W$分解数据在文章附录，诸位可以自己去找。

另外，对于特殊的矩阵也有更优的算法。比如反对称阵（skew-symmetric matrix）和向量的乘法，仅需要$(n-1)(n+2)/2$个乘法运算：

上述的张量分解，还可以被应用于任意的双线性运算上，比如：循环卷积（多项式乘法，DFT），AlphaTensor 找到了对$n=2,4,8$的加速。