前言

Strassen最早于1968年由Volker Srassen发表于论文《Gaussian Elimination is not Optimal》，将矩阵乘法的算法复杂度首次从 $O(n^3),where\space log_2(8)=3$ 降低到 $O(n^2.807),where\space log_2(7)≈2.807$ 。后来有不同的方法基于此算法进行改进。

MNN Strassenconv1x1采用的是2008年发表的这边文章中的算法《Memory efficient scheduling of Strassen-Winograd’s matrix multiplication algorithm》，主要针对memory访问上进行优化。

原理

普通矩阵乘

以 $C=AB$ 为例，假设 $A$ 和 $B$ 均为 $n × n$ 矩阵，则 $C$ 中每个元素需要通过以下方式计算：

$c_{ij}=\displaystyle\sum_{k=1}^na_{ik}b_{kj}$

即每个元素计算需要 $n$ 次乘法， $n-1$ 次加法，C共有 $n^2$ 个元素，故一共需要 $n^3$ 次乘法和 $n^2(n-1)$ 次加法，故算法复杂度为 $O(n^3)$ 。

分治法

Strassen算法基于分治的思想，因此我们首先考虑一个简单的分治策略。

假设矩阵 $A$ 和矩阵 $B$ 都是 $n×n,(n=2^k)$ 的方阵，求 $C=AB$ ，则每个 $n×n$ 的矩阵都可以分割为四个 $(n/2) ×( n/2)$ 的矩阵：

$A=\begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix} \space B=\begin{bmatrix} B_{11} & B_{12} \\ B_{21} & B_{22} \end{bmatrix} \space C=\begin{bmatrix} C_{11} & C_{12} \\ C_{21} & C_{22} \end{bmatrix}$

于是 $C=AB$ 可以改写为：

$\begin{bmatrix} C_{11} & C_{12} \\ C_{21} & C_{22} \end{bmatrix}=\begin{bmatrix} A_{11} & A_{12} \\ A_{21} & A_{22} \end{bmatrix}\begin{bmatrix} B_{11} & B_{12} \\ B_{21} & B_{22} \end{bmatrix}$

展开有：

$C_{11} = A_{11}B_{11}+A_{12}B_{21} \\ C_{12} = A_{11}B_{12}+A_{12}B_{22} \\ C_{21} = A_{21}B_{11}+A_{22}B_{21} \\ C_{22} = A_{21}B_{12}+A_{22}B_{22}$

每个等式需要两次矩阵乘法和一次矩阵加法。若用 $T(n)$ 表示两个 $n×n$ 矩阵之间的乘法，则有如下递归式：

$T(n)=8T(\frac{n}{2})+\varTheta(n^2)$

其中：

1. $8T(\frac{n}{2})$ 表示8次子矩阵乘法，子矩阵规模为 $(n/2) ×( n/2)$ ；
1. $\varTheta(n^2)$ 表示4次子矩阵加法以及合并 $C$ 矩阵的时间复杂度。

根据递归式求解，得 $T(n)=\varTheta(n^{log_2^8})=\varTheta(n^3)$ ，与普通矩阵乘的时间复杂度相同。分治法并不能起到加速的效果。

原始版本Strassen

分治法包含了8次矩阵相乘和4次矩阵相加，相比矩阵加法，矩阵乘法是非常慢的。这8次矩阵相乘正是瓶颈的来源。于是我们想到能不能减少矩阵相乘的次数，哪怕代价是更多的矩阵相加。Strassen正是利用了这一点。

1. 同样还是每个矩阵分成4份，然后创建如下10个中间矩阵（10次矩阵加法 $10×\frac{n}{2}×\frac{n}{2}$ ，时间复杂度为 $\varTheta(n^2)$ ）：

$S_1=B_{12}-B_{22} \\ S_2=A_{11}+A_{12} \\ S_3=A_{21}+A_{22} \\ S_4=B_{21}-B_{11} \\ S_5=A_{11}+A_{22} \\ S_6=B_{11}+B_{22} \\ S_7=A_{12}-A_{22} \\ S_8=B_{21}+B_{22} \\ S_9=A_{11}-A_{21} \\ S_{10}=B_{11}+B_{12}$

1. 而后是执行7次矩阵乘法（时间复杂度 $7T(\frac{n}{2})=\varTheta(n^{log_2^7})=\varTheta(n^{2.807})$ ）：

$P_1=A_{11}S_1 \\ P_2=S_2B_{22} \\ P_3=S_3B_{11} \\ P_4=A_{22}S_4 \\ p_5=S_5S_6 \\ P_6=S_7S_8 \\ P_7=S_9S_{10}$

1. 最后通过这7个矩阵计算得到C矩阵（8次矩阵加法 $8×\frac{n}{2}×\frac{n}{2}$ ，时间复杂度 $\varTheta(n^2)$ ）：

$C_{11}=P_5+P_4-P_2+P_6 \\ C_{12}=P_1+P_2 \\ C_{21}=P_3+P_4 \\ C_{22}=P_5+P_1-P_3-P_7$

综合可得如下递归式：

$T(n)=\begin{cases} \varTheta(1) &\text{if } n=1 \\ 7T(\frac{n}{2})\varTheta(n^2) &\text{if } n>1 \end{cases}$

进而总的时间复杂度为 $T(n)=\varTheta(n^{log_2^7})=\varTheta(n^{2.807})$

改进版本Strassen

1. 8次矩阵加法：

$\begin{align*} & S_1=A_{21}+A_{22} \quad T_1=B_{12}-B_{11} \\ & S_2=S_1-A_{11} \quad\enspace T_2=B_{22}-T_1 \\ & S_3=A_{11}-A_{21} \quad T_3=B_{22}-B_{12} \\ & S_4=A_{12}-S_2 \quad\enspace T_4=T_2-B_{21} \end{align*}$

1. 7次矩阵乘法：

$\begin{align*} & P_1=A_{11}B_{11} \quad & P_5=S_1T_1 \\ & P_2=A_{12}B_{21} \quad & P_6=S_2T_2 \\ & P_3=S_4B_{22} \quad & P_7=S_3T_3 \\ & P_4=A_{22}T_4 \end{align*}$

1. 7次矩阵加法：

$\begin{align*} &U_1=P_1+P_2 \quad U_5=U_4+P_3 \\ &U_2=P_1+P_6 \quad U_6=U_3-P_4 \\ &U_3=U_2+P_7 \quad U_7=U_3+P_5 \\ &U_4=U_2+P_5 \end{align*}$

1. 最终结果：

$C=\begin{bmatrix} U_1 & U_5 \\ U_6 & U_7 \end{bmatrix}$

改进版本的Strassen相比原始版本矩阵加法次数由18次降低到15次。同时该算法同时分析各中间矩阵块之间的依赖关系，充分考虑到了计算中的内存复用，每次递归只需要动态分配两块临时内存 $X_1$ ， $X_2$ 即可， $C_{11}$ ， $C_{12}$ ， $C_{21}$ ， $C_{22}$ 都可以在计算中复用。