参考文献:
- Discovering faster matrix multiplication algorithms with reinforcement learning:https://doi.org/10.1038/s41586-022-05172-4
Tensor
矩阵
A
∈
F
n
×
m
A \in \mathbb F^{n \times m}
A∈Fn×m和
B
∈
F
m
×
p
B \in \mathbb F^{m \times p}
B∈Fm×p,矩阵乘积
(
A
,
B
)
↦
C
:
=
A
B
(A,B) \mapsto C:=AB
(A,B)↦C:=AB 是双线性的(bilinear),因此可以表示为
3
3
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
T_{n,m,p}
Tn,m,p可以指定读写位置。先把
A
,
B
,
C
A,B,C
A,B,C展开为一维向量
A
′
,
B
′
,
C
′
A',B',C'
A′,B′,C′,如果
t
a
b
c
=
1
t_{abc}=1
tabc=1,那么
C
c
′
C'_c
Cc′的值的加和项中就有
A
a
′
⋅
B
b
′
A'_a \cdot B'_b
Aa′⋅Bb′
我们把
T
n
,
m
,
p
T_{n,m,p}
Tn,m,p做分解,得到
U
∈
Z
n
m
×
R
,
V
∈
Z
m
p
×
R
,
W
∈
Z
n
p
×
R
U \in \mathbb Z^{nm \times R},V \in \mathbb Z^{mp \times R},W \in \mathbb Z^{np \times R}
U∈Znm×R,V∈Zmp×R,W∈Znp×R,使得
T
n
,
m
,
p
=
∑
r
=
1
R
u
(
r
)
⊗
v
(
r
)
⊗
w
(
r
)
T_{n,m,p} = \sum_{r=1}^R u^{(r)} \otimes v^{(r)} \otimes w^{(r)}
Tn,m,p=r=1∑Ru(r)⊗v(r)⊗w(r)
其中的 ⊗ \otimes ⊗是列向量外积(outer product)或者叫张量积(tensor product)
可以写出根据 T n , m , p T_{n,m,p} Tn,m,p的分解 U , V , W U,V,W U,V,W控制的矩阵乘法:
如果分解后,列数 R R R越小,那么运算时的小矩阵乘法的数量就不大于 R R R。 著名的 Strassen’s algorithm,它就是找到了张量 T 2 , 2 , 2 T_{2,2,2} T2,2,2的 R = 7 R=7 R=7的分解。
AlphaTensor
这篇文章利用 Neural network 搭建了一个模型 AlphaTensor,发现了一系列的对张量 T n , m , p T_{n,m,p} Tn,m,p的更优分解:
具体的 U , V , W U,V,W U,V,W分解数据在文章附录,诸位可以自己去找。
另外,对于特殊的矩阵也有更优的算法。比如反对称阵(skew-symmetric matrix)和向量的乘法,仅需要 ( n − 1 ) ( n + 2 ) / 2 (n-1)(n+2)/2 (n−1)(n+2)/2个乘法运算:
上述的张量分解,还可以被应用于任意的双线性运算上,比如:循环卷积(多项式乘法,DFT),AlphaTensor 找到了对 n = 2 , 4 , 8 n=2,4,8 n=2,4,8的加速。