矩阵乘法: 从 Strassen 到 Coppersmith–Winograd

这篇文章主要是对 MIT6.890 的 Lec19~23 的理解, 试图概括在 Strassen 之后一直到 (小) Coppersmith–Winograd 的矩阵乘法用到了一些什么技术. 本文的终点是证明 $\omega \le 2.404$.

张量描述

固定域 $F$, 我们称 $F$ 上一个双线性计算问题是对于

$$f_k=\sum t_{ijk}{x}_i{y}_j,$$

输入所有 $x_i,y_j$, 计算所有 $f_k$. 这里 $t\in F^{N\times M\times K}$ 是一个三阶张量.

线性算法与 rank

之前理解过转置原理的同学可能对这里更好理解.

我们希望完成这个双线性计算问题, 我们首先希望优化的是 $x,y$ 之间 的乘法问题, 换句话说, 我们认为 $t$ 是一个已经被固定的东西, 再次强调, 只有 $x,y$ 是输入.

当 $|F|$ 无限的时候, 加, 减, 乘所得到的结果无非是一个关于全体 $x_i,y_i$ 的多项式, 而且当它们不定的时候, 我们没有特殊关系可以约化它们, 所以正确的结果必须最后只剩下 $x_i{y}_j$ 的线性组合. 直觉上说, 我们的计算一定可以表示成, 先做一些乘法

$$P_{\lambda }=(u_{\lambda }^{\mathsf{T}}x)\cdot (v_{\lambda }^{\mathsf{T}}y).$$

然后每个输出是一个线性组合 $f_k=w_k^{\mathsf{T}}P$. 在 "bilinear complexity" 的意义下, 我们只关心这个向量 $P$ 的维数.

注意到当我们考虑多项式

$$\sum _{i,j,k}{t}_{ijk}{x}_i{y}_j{z}_k$$

的时候, 上面的写法就可以写作

$$\sum _{\lambda }(u_{\lambda }^{\mathsf{T}}x)(v_{\lambda }^{\mathsf{T}}y)(w_{\lambda }^{\mathsf{T}}z).$$

也就是说, 输入和输出的地位其实是相等的.

我们称可行的最小的 $\lambda$ 的长度为 $R(t)$, 这个张量的秩 (rank). 显然上述对称的写法是很好理解的, 而且对二阶张量, 也即矩阵来说, 和我们对矩阵定义的秩是一致的.

当 $|F|$ 无限的时候, 我们记 $C(t)$ 是前文提到的所需的 "乘法" 次数, 那么经过一些讨论, 我们其实可以证明 $R(t)\le C(t)\le 2R(t)$. 所以研究 $R(t)$, 在很大程度上就相当于研究矩阵乘法的效率问题.

但这并不对所有问题都完全奏效, 比如对多项式乘法来说, 在无限域上我们可以先进行求值, 然后做 $n$ 次乘法, 再插值回来. 在 "bilinear complexity" 意义下, 这个算法有 $R(t)=n$, 但它粗暴地忽略了求值插值的所需时间!

我们记 $t^{K\times M\times N}\oplus t^{\prime K^{\prime }\times M^{\prime }\times N^{\prime }}$ 是一个 $(K+K^{\prime })\times (M+M^{\prime })\times (N+N^{\prime })$ 维张量, 具体来说就是将 $t$ 放在立方体的 "左下角", $t^{\prime }$ 放在 "右上角".

$t^{K\times M\times N}\otimes t^{\prime K^{\prime }\times M^{\prime }\times N^{\prime }}$ 是一个 $(K{K}^{\prime })\times (M{M}^{\prime })\times (N{N}^{\prime })$ 维张量, 不意外地, 我们有

$$(t\otimes t^{\prime }{)}_{i{i}^{\prime }j{j}^{\prime }k{k}^{\prime }}=t_{ijk}{t}_{i^{\prime }{j}^{\prime }{k}^{\prime }}^{\prime }.$$

容易知道, $R(t\oplus t^{\prime })\le R(t)+R(t^{\prime })$, 而 $R(t\otimes t^{\prime })\le R(t)R(t^{\prime })$.

矩阵乘法

对于 $K\times M\times N$ 的矩阵乘法来说, 问题就是一个 $F^{KM\times MN\times KN}$ 的张量, 其中

$$t_{ij,j^{\prime }k,i^{\prime }{k}^{\prime }}={\delta }_{i{i}^{\prime }}{\delta }_{j{j}^{\prime }}{\delta }_{k{k}^{\prime }}.$$

我们今后记这个张量为 $⟨K,M,N⟩$.

我们的第一个观察是 $⟨K,M,N⟩$ 同构于 $⟨M,N,K⟩$, 这是因为贡献关系可以看做一个完全三分图的所有三角形, 将顶点重标号并不影响什么. 所以我们有

$$R(⟨K,M,N⟩)=R(⟨M,N,K⟩)=R(⟨N,K,M⟩).$$

第二个观察就是 $⟨K,M,N⟩\otimes ⟨K^{\prime },M^{\prime },N^{\prime }⟩=⟨K{K}^{\prime },M{M}^{\prime },N{N}^{\prime }⟩$, 这无非就是分块矩阵乘法.

现在, 我们可以来定义什么是 $\omega$ 了. 我们记

$$\omega =\underset{n>1}{inf}\frac{\log R(⟨n,n,n⟩)}{\log n},$$

接下来, 我们声称: 对任意 $ϵ>0$, 存在 $O(n^{\omega +ϵ})$ 的矩阵乘法算法. 这是因为我们总存在一个常数 $T$ 使得 $\frac{\log R(⟨T,T,T⟩)}{\log T}\le \omega +ϵ$, 然后我们不断调用 $R(⟨T,T,T⟩)$ 给出的分块矩阵乘法递归下去, 主定理就保证了复杂度为 $O(n^{\omega +ϵ})$.

从这一步我们也能预感到矩阵乘法接下来的算法为什么现实中都不 work 了: 因为他们需要的 $T$ 可能都是天文数字!

Strassen 和 Pan 的构造

最初 Strassen 的构造无非指出, $R(⟨2,2,2⟩)\le 7$. 因此 $\omega \le {\log }_{2}7\approx 2.8074$.

然后, Pan (1980) 构造出了 $R(⟨70,70,70⟩)\le 143640$, 这导出了 $\omega \le 2.796$.

到目前为止, 远超我们理解的奇迹还没有出现.

Border rank

事情的转机来源于一个奇怪的现象. 我们考虑一些可以做逼近的域, 让一系列秩不超过 $r$ 的矩阵 $A_t$ 趋于 $A$, 我们有 $A$ 的秩仍然不超过 $r$. 这可以通过考察所有 $r+1$ 阶的余子式确定. 但奇怪的是, 对于三阶张量来说, 这件事是不成立的. 对于二次型

$$g_0=a_0{b}_0,g_1=a_0{b}_1+a_1{b}_0,$$

我们可以用

$$t(ϵ)=(1,ϵ)\otimes (1,ϵ)\otimes (0,1/ϵ)+(1,0)\otimes (1,0)\otimes (1,-1/ϵ)$$

来逼近它, 但这个二次型本身的秩是 $3$.

我们考虑多项式环 $F[ϵ]$, 我们定义 $R_h(t)$ 为最小的 $L$, 使得有一组解 $u_{\lambda }(ϵ),v_{\lambda }(ϵ),w_{\lambda }(ϵ)$ 使得

$$\sum _{\lambda }(u_{\lambda }^{\mathsf{T}}x)(v_{\lambda }^{\mathsf{T}}y)(w_{\lambda }^{\mathsf{T}}z)={ϵ}^{h}t+O({ϵ}^{h+1}).$$

显然按照这个定义, 我们有 $R_0(t)\ge R_1(t)\ge \dots$, 因为 $R_h(t)$ 总是整数, 我们有下界 $\underset{―}{R}(t)=\underset{h}{min}{R}_h(t)$.

第一个观察是, $R_0(t)$ 依然是可以被 $R_h(t)$ 控制的, 这是因为我们提取 $(u_{\lambda }^{\mathsf{T}}x)(v_{\lambda }^{\mathsf{T}}y)(w_{\lambda }^{\mathsf{T}}z)$ 的 ${ϵ}^h$ 次项系数, 这只需要做 $O(h^2)$ 次乘法.

接下来的观察就更为有趣了, 我们不难验证: $R_h(t\oplus t^{\prime })\le R_h(t)+R_h(t^{\prime })$ 和 $R_{h+h^{\prime }}(t\oplus t^{\prime })\le R_h(t)+R_{h^{\prime }}(t^{\prime })$. 后者的巧妙指出就在于, $h$ 的次数贡献是相加而不是相乘. 这直接导出了如下结果:

$$\omega \le \frac{3\log \underset{―}{R}(⟨K,M,N⟩)}{\log (KMN)}.$$

这是因为, 首先我们将它划归到方阵, 也即

$$⟨T,T,T⟩=⟨K,M,N⟩\otimes ⟨M,N,K⟩\otimes ⟨N,K,M⟩,$$

我们有 $T=\log (KNM)$, 此时设 $h=R_t(⟨K,M,N⟩)$, 那么我们有 $R_{3t}(⟨T,T,T⟩)\le h^3$.

接下来要用到我们前面的观察了, 我们将它再张量 $s$ 次, 得到

$$R_{3st}(⟨T^s,T^s,T^s⟩)\le h^{3s}.$$

也就有

$$\omega \le \frac{\log (O(s{t}^2)\cdot h^{3s})}{s\log (NMK)}=\frac{3\log h}{\log (NMK)}+O\left(\frac{\log s}{s}\right).$$

因此, 我们不需要控制 rank, 而是只要能够控制 border rank, 我们就可以得到更快的矩阵乘法了.

Bini, Capovani, Romani (1979) 得到了 $\underset{―}{R}(⟨2,2,3⟩)\le 10$. 进而得到了 $\omega \le 2.780$.

Schönhage $\tau$-定理

接下来, 矩阵乘法迎来了指数上最大的单次改进, 这来源于 Schönhage 的观察: 我们未必要直接控制某一个矩阵乘法的 border rank, 而是控制一些几个并行的矩阵乘法的 border rank.

Schönhage $\mathit{\tau }$-定理: 对于 $r>p$ 以及

$$\underset{―}{R}(\underset{i=1}{\overset{p}{⨁}}⟨k_i,m_i,n_i⟩)\le r,$$

令 $\tau$ 是下述等式

$$\sum _{i=1}^p(k_i{m}_i{n}_i{)}^{\tau }=r$$

的解, 我们有 $\omega \le 3\tau$.

这个定理的证明涉及一些琐碎的计算, 这里我们主要展示其主要思想.

将上面的这个张量记为 $t$, 有 $R_h(t)=s$, 我们考虑它的 $s$ 阶张量积 $t^{\otimes s}$, 我们可以控制

$$R(t^{\otimes s})\le r^s\cdot \mathrm{poly}(h\cdot s),$$

再考虑 $t^{\otimes s}$ 的展开式中, 有一个直和因子

$$\underset{i}{⨂}⟨k_i,m_i,n_i{⟩}^{\otimes s_i}=⟨\prod _i{k}_i^{s_i},\prod _i{m}_i^{s_i},\prod _i{n}_i^{s_i}⟩,$$

它总共出现了 $(\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p})$ 次, 如果我们选取 $s_i$ 最大化

$$(\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p}){\left(\prod _i{k}_i^{s_i}\prod _i{m}_i^{s_i}\prod _i{n}_i^{s_i}\right)}^{\tau },$$

记三部分乘积分别为 $K,M,N$, 就可以用 $r^s\le (\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p})(KMN{)}^{\tau }\cdot (\genfrac{}{}{0ex}{}{s+p-1}{p-1})$ 来进行控制, 也就得到了

$$R((\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p})\odot ⟨K,M,N⟩)\le (\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p})(KMN{)}^{\tau }\cdot (\genfrac{}{}{0ex}{}{s+p-1}{p-1})\mathrm{poly}(hs).$$

这是同时计算常数 $a=(\genfrac{}{}{0ex}{}{s}{s_1,\dots ,s_p})$ 个 $⟨K,M,N⟩$ 时候得到的界, 但我们可以通过继续做张量积 $⟨K,M,N⟩$, 也就是在算更大的矩阵乘法的时候把这个东西摊掉, 也就是可以当做下述事实成立来控制 $\omega$, 具体证明就是做张量积的时候进行放缩, 细节略去.

$$R(⟨K,M,N⟩)\le (KMN{)}^{\tau }\cdot (\genfrac{}{}{0ex}{}{s+p-1}{p-1})\mathrm{poly}(hs)$$

当 $s\to \mathrm{\infty }$ 时, 我们有

$$\frac{3\log [(KMN{)}^{\tau }\cdot (\genfrac{}{}{0ex}{}{s+p-1}{p-1})\mathrm{poly}(hs)]}{\log KMN}\to 3\tau .$$

所以 $\omega \le 3\tau$.

基于一些构造, Schönhage (1981) 在此之上确定了 $\omega \le 2.522$.

Coppersmith–Winograd 张量

先前的构造都是考虑了一些矩阵乘法问题的直和的 border rank 被如何控制. Coppersmith 和 Winograd 对这一想法做出了如下变通:

1. 我们先将一些矩阵乘法问题求和成为一个张量, 但它们未必是直和的形式.
2. 我们将这个张量做很高次的张量积, 然后将一些变量 "置零" 来去除一些纠缠在一起的项, 然后我们数一数剩下来的东西, 发现完全分解成很多个矩阵乘法问题的直和.
3. 应用 Schönhage $\tau$-定理来控制 $\omega$.

我们这里主要讨论小 CW 张量 ${\mathfrak{cw}}_q\in F^{(q+1)\times (q+1)\times (q+1)}$, 其中 $q\ge 2$. 我们记变量为 $X_0=\{x_0\},X_1=\{x_1,\dots ,x_q\}$ 以及类似的 $Y,Z$. 令 $T_{ijk}=X_i{Y}_j{Z}_k$. 我们定义

$${\mathfrak{cw}}_q=T_{110}+T_{011}+T_{101}=\sum _{i+j+k=2}{T}_{ijk}=\sum _{i=1}^q(x_0{y}_i{z}_i+x_i{y}_0{z}_i+x_i{y}_i{z}_0).$$

可以验证, $T_{011}$ 就等同于 $⟨1,1,q⟩$, $T_{101}$ 等同于 $⟨q,1,1⟩$, $T_{110}$ 等同于 $⟨1,q,1⟩$.

对于 $k$ 位的数 $I,J,K\in \{0,1{\}}^k$, 我们可以自然地定义 $T_{IJK}$, 并且有

$${\mathfrak{cw}}_q^{\otimes k}=\sum _{\begin{array}{c}I,J,K\\ I_i+J_i+K_i=2\end{array}}{T}_{IJK}.$$

现在我们考虑 ${\mathfrak{cw}}_q^{\otimes 3N}$, 为了方便分析, 我们先让 $X_I$ 中只保留 $I$ 有 $2N$ 个 $1$ 的, $Y,Z$ 类似处理, 记

$$t=\sum _{\begin{array}{c}I,J,K\\ I_i+J_i+K_i=2\\ |I|=|J|=|K|=2N\end{array}}{T}_{IJK},$$

如果 $\underset{―}{R}({\mathfrak{cw}}_q)=c$, 那么我们知道 $R(t)\le R({\mathfrak{cw}}_q^{\otimes 3N})\le c^{3N+o(N)}$.

我们记 $t$ 中保留的指标集为 $L$, 我们的问题就是: 如何选取三个子集 $S^X,S^Y,S^Z$, 使得 $|L\cap S^X\times S^Y\times S^Z|$ 中任意两个点都没有相同的坐标 $I\ne I^{\prime },J\ne J^{\prime },K\ne K^{\prime }$ (下称独立), 且交集尽量大?

首先我们注意到:

如果我们可以对某个常数 $\varphi$ 构造出一族

$$|L\cap S^X\times S^Y\times S^Z|\ge {\varphi }^{N-o(N)},$$

就可以给出控制

$$\omega \le \frac{\log \frac{c^3}{\varphi }}{\log q}.$$

这是因为我们考虑使用 Schönhage $\tau$-定理, 现在我们有 ${\varphi }^{N-o(N)}$ 个大小为 $q^{3N}$ 的矩阵乘法, 那么

$${\varphi }^{N-o(N)}{q}^{3N\cdot \tau }=c^{3N+o(N)}$$

的解是

$$\tau =(1+o(1))\frac{\log \frac{c^3}{\varphi }}{3\log q},$$

所以有

$$\omega \le \frac{\log \frac{c^3}{\varphi }}{\log q}.$$

接下来就是考虑如何构造出一组合适的选取方法了.

$\omega \le 2.404$ 的构造

对于这个 $S_X,S_Y,S_Z$ 的构造大概分两步.

1. 我们首先通过一个随机的方式将所有的指标 $I,J,K$ 映到一些桶里 (有的被清零), 我们保证不同的桶之间两个点两两独立.
2. 接下来, 我们对每个桶内部进行一些处理, 使得每个桶内只保留一些独立的点.

第 1 步: 哈希与极值组合

首先让我们来构造一族哈希函数. 选取一个充分大的质数 $p$ 介于 $[3(\genfrac{}{}{0ex}{}{2N}{2}),4(\genfrac{}{}{0ex}{}{2N}{2})]$ 之间, 我们考虑一个随机数 $w_0$ 以及随机向量 $w$, 都是在模 $p$ 下均匀随机的, 然后考虑 $modp$ 下的哈希函数

$$h_X(I)=w^{\mathsf{T}}I,h_Y(J)=h_X(J)+w_0,h_Z(K)=\frac{1}{2}{h}_Y(2-K).$$

由于 $p$ 充分大, 所以是奇素数, $\frac{1}{2}$ 是良定义的.

显然, 由于这三个函数都是仿射变换, 所以 $I\ne I^{\prime }$ 的时候, $h_X(I),h_X(I^{\prime })$ 是独立随机的, $h_Y,h_Z$ 有类似性质. 同时, 注意到在 $L$ 里 $I+J+K=2$ 总是成立, 我们还有 $h_X(I)+h_Y(J)=2h_Z(K)(\mathrm{mod}p)$, 也就是说, $h_X(I),h_Z(K),h_Y(J)$ 总是等差数列.

之所以构造这样的哈希函数, 是因为极值组合学中有如下构造:

(Behrend, 1946) 对于 $N$, 存在 $[N]$ 的子集 $S$ 中不存在非平凡的长为 $3$ 的等差数列, 同时

$$|S|\ge \frac{N}{e^{O(\sqrt{\log N})}}.$$

我们取 $p=N/2$, 就得到了 $modp$ 下无等差数列的构造. 对此, 我们其实是想要结果: 对于任意 $ϵ$, 对充分大的 $p$, 我们希望 $|S|\ge p^{1-ϵ}$.

那么接下来, 我们令

$$S_X=h_X^{-1}(S),S_Y=h_Y^{-1}(S),S_Z=h_Z^{-1}(S).$$

由于 $S$ 是无等差数列的, 那么 $h_X(I),h_Y(J),h_Z(K)\in S$ 的情况下, 如果 $(I,J,K)\in L$, 我们就有它们必须是平凡等差数列, 也即 $h_X(I)=h_Y(J)=h_Z(K)$.

因此对每一个 $s\in S$, 我们可以把映到 $s$ 的 $I,J,K$ 称作一个桶 (bucket), 此时三个下标都被映到了 $s$, 这也很显然地说明了桶和桶之间是坐标不交的.

接下来让我们先算一下我们还剩下多少. 注意我们 $L$ 中为了简化分析加了条件, 我们考虑每一列 $(I_i,J_i,K_i)$, 它出现 $110,101,011$ 必须各 $N$ 次, 因此

$$|L|=(\genfrac{}{}{0ex}{}{3N}{N,N,N}).$$

由于 $h_X(I),h_Y(J)$ 的取值独立, 它们又唯一确定了 $h_Z(K)$, 我们有

$$Pr[h(I,J,K)=s]=\frac{1}{p^2}.$$

因此

$$\mathbb{E}[|L\cap S^X\times S^Y\times S^Z|]=\frac{|L||S|}{p^2}\ge \frac{1}{p^{1+ϵ}}(\genfrac{}{}{0ex}{}{3N}{N,N,N}).$$

第 2 步

接下来我们考虑通过一个简单的策略来实现 "清零".

我们随便按照某个顺序考虑一个桶里的所有元素 $B_s$, 如果一个元素和排在它前面的某个元素有一个相同坐标, 我们就选取一个不同的坐标, 将当前元素的这个坐标清零 (注意 $L$ 中元素任意两者最多一个相同坐标). 这个操作显然至少会删掉这个元素自身, 但也会同时删掉其他元素. 我们考虑用某种方法来估计一个上界.

我们考虑两个元素如果不独立就连一条边, 当我们进行一次清零操作的时候, 假设这个坐标此时对应了 $\ell$ 个元素, 那么会一次性删去 $(\genfrac{}{}{0ex}{}{\ell }{2})$ 条边, 由于这个清零操作引发的时候还有一个元素和其他元素有边, 所以总共删去至少 $(\genfrac{}{}{0ex}{}{\ell }{2})+1\ge \ell$ 条边. 也就是说, 我们可以用边数来控制删去的元素数量.

那么边数的期望就比较好求了. 容易验证, 当 $I=I^{\prime }$ 的时候, $h_X(I),h_Y(J),h_Y(J^{\prime })$ 三者两两独立 (其他情况类似), 也就是说一条边有 $|S|/p^3=1/p^{2+ϵ}$ 的概率继续存在.

通过一些简单的计数我们可以知道, 一开始 $L$ 中的边数是

$$(\genfrac{}{}{0ex}{}{3N}{N,N,N})\cdot \frac{3}{2}((\genfrac{}{}{0ex}{}{2N}{N})-1).$$

最后我们利用选取的 $p$ 的大小, 就可以得到上界, 经过了最终的清零后, 我们有

$$\mathbb{E}[|L\cap S^X\times S^Y\times S^Z|]\ge \frac{1}{8p^{ϵ}}(\genfrac{}{}{0ex}{}{3N}{N}).$$

根据 Stirling 近似, 我们有

$$\mathbb{E}[|L\cap S^X\times S^Y\times S^Z|]\ge \mathrm{\Omega }\left(\frac{1}{p^{ϵ}\sqrt{N}}{\left(\frac{27}{4}\right)}^N\right).$$

因此 $\varphi =\frac{27}{4}$ 满足要求.

结合前面的结果, 我们有

$$\omega \le \frac{\log \frac{4c^3}{27}}{\log q}.$$

$c\le q+2$ 的构造, 以及剩下的?

我们还需要知道 $c=\underset{―}{R}({\mathfrak{cw}}_q)$ 能做到多好. 实际上, 由于

$$\begin{array}{rl}& \sum _{i=1}^qϵ(x_0+ϵx_i)(y_0+ϵy_i)(z_0+ϵz_i)\\ & -(x_0+{ϵ}^2{X}_1)(y_0+{ϵ}^2{Y}_1)(z_0+{ϵ}^2{Z}_1)\\ & +(1-qϵ)x_0{y}_0{z}_0\\ & ={ϵ}^3{\mathfrak{cw}}_q+O({ϵ}^4),\end{array}$$

有 $c\le q+2$. 当取 $q=8$ 的时候, 就有

$$\omega \le {\log }_8\left(\frac{4000}{27}\right)<2.40364.$$

这里有一个尚未解决的问题, 就是是否有 $c\le q+1$, 但是如果 $q=2$ 的时候 $c\le q+1$ 可以构造出来, 我们就有 $\omega =2$ 了.

更多?

这里其实并没有介绍 Coppersmith–Winograd 的最后结果 ($\omega \le 2.3755$), 为了达到最后的结果, 他们还考虑了更加复杂的张量, 被称作大 CW 张量.

据说目前最优的结果 ($\omega \le 2.3728596$) 又发明了被称作 laser 法, 我还没有了解过, 按下不表.

此外, Coppersmith–Winograd 的方法似乎不是唯一的路径, Cohn, Kleinberg, Szegedy 和 Umans (2003) 通过表示论另辟蹊径, 达到了和 CW 很接近的结果. 我之后有时间的话可能会尝试解读一下.