矩阵乘法: 从 Strassen 到 Coppersmith–Winograd

这篇文章主要是对 MIT6.890 的 Lec19~23 的理解, 试图概括在 Strassen 之后一直到 (小) Coppersmith–Winograd 的矩阵乘法用到了一些什么技术. 本文的终点是证明 \(\omega \leq 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) \leq C(t) \leq 2R(t)\). 所以研究 \(R(t)\), 在很大程度上就相当于研究矩阵乘法的效率问题.

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

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

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

\[(t\otimes t')_{ii' jj' kk'} = t_{ijk} t'_{i'j'k'}. \]

容易知道, \(R(t\oplus t') \leq R(t) + R(t')\), 而 \(R(t\otimes t') \leq R(t)R(t')\).

矩阵乘法

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

\[t_{ij, j'k, i'k'} = \delta_{ii'} \delta_{jj'} \delta_{kk'}. \]

我们今后记这个张量为 \(\langle K, M, N \rangle\).

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

\[R(\langle K, M, N \rangle) = R(\langle M, N, K \rangle) = R(\langle N, K, M \rangle). \]

第二个观察就是 \(\langle K, M, N \rangle \otimes \langle K', M', N' \rangle = \langle KK', MM', NN' \rangle\), 这无非就是分块矩阵乘法.

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

\[\omega = \inf_{n > 1} \frac{\log R(\langle n,n,n \rangle)}{\log n}, \]

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

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

Strassen 和 Pan 的构造

最初 Strassen 的构造无非指出, \(R(\langle 2,2,2 \rangle) \leq 7\). 因此 \(\omega \leq \log_2 7 \approx 2.8074\).

然后, Pan (1980) 构造出了 \(R(\langle 70, 70, 70 \rangle) \leq 143640\), 这导出了 \(\omega \leq 2.796\).

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

Border rank

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

\[g_0 = a_0b_0, \quad g_1 = a_0 b_1 + a_1 b_0, \]

我们可以用

\[ t(\epsilon) = (1, \epsilon) \otimes (1, \epsilon)\otimes (0, 1/\epsilon) + (1, 0) \otimes (1, 0) \otimes (1, -1/\epsilon) \]

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

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

\[ \sum_{\lambda} (u_\lambda^{\mathsf T} x) (v_\lambda^{\mathsf T} y) (w_\lambda^{\mathsf T} z) = \epsilon^h t + O(\epsilon^{h+1}). \]

显然按照这个定义, 我们有 \(R_0(t) \geq R_1(t) \geq \dots\), 因为 \(R_h(t)\) 总是整数, 我们有下界 \(\underline{R}(t) = \min_h R_h(t)\).

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

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

\[\omega \leq \frac{3\log \underline{R}(\langle K,M,N \rangle)}{\log (KMN)}. \]

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

\[\langle T,T, T \rangle = \langle K, M, N \rangle \otimes \langle M, N, K \rangle \otimes \langle N, K, M \rangle, \]

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

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

\[R_{3st}(\langle T^s,T^s, T^s \rangle) \leq h^{3s}. \]

也就有

\[ \omega \leq \frac{\log \left( O(st^2) \cdot h^{3s} \right)}{s\log (NMK)} = \frac{3\log h}{\log(NMK)} + O\left( \frac{\log s}s \right). \]

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

Bini, Capovani, Romani (1979) 得到了 \(\underline{R}(\langle 2, 2, 3 \rangle)\leq 10\). 进而得到了 \(\omega \leq 2.780\).

Schönhage \(\tau\)-定理

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

Schönhage \(\bm \tau\)-定理: 对于 \(r>p\) 以及

\[\underline{R}\left( \bigoplus_{i=1}^p \langle k_i, m_i, n_i \rangle \right)\leq r, \]
令 \(\tau\) 是下述等式

\[\sum_{i=1}^p (k_im_in_i)^\tau = r \]
的解, 我们有 \(\omega \leq 3\tau\).

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

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

\[R(t^{\otimes s}) \leq r^s \cdot \operatorname{poly}(h\cdot s), \]

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

\[\bigotimes_i \langle k_i, m_i, n_i \rangle^{\otimes s_i} = \left\langle \prod_i k_i^{s_i}, \prod_i m_i^{s_i}, \prod_i n_i^{s_i} \right\rangle, \]

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

\[\binom{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\leq \binom{s}{s_1, \dots, s_p} (KMN)^\tau \cdot \binom{s+p-1}{p-1}\) 来进行控制, 也就得到了

\[R\left( \binom{s}{s_1, \dots, s_p} \odot \langle K, M, N\rangle \right) \leq \binom{s}{s_1, \dots, s_p} (KMN)^\tau \cdot \binom{s+p-1}{p-1} \operatorname{poly}(hs). \]

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

\[\color{red} {R\left( \langle K, M, N\rangle \right) \leq (KMN)^\tau \cdot \binom{s+p-1}{p-1} \operatorname{poly}(hs)} \]

当 \(s\to \infty\) 时, 我们有

\[\frac{3\log \left[(KMN)^\tau \cdot \binom{s+p-1}{p-1} \operatorname{poly}(hs)\right]}{\log KMN} \to 3\tau. \]

所以 \(\omega \leq 3\tau\).

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

Coppersmith–Winograd 张量

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

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

我们这里主要讨论小 CW 张量 \(\mathfrak{cw}_q \in F^{(q+1)\times (q+1) \times (q+1)}\), 其中 \(q\geq 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}\) 就等同于 \(\langle 1, 1, q \rangle\), \(T_{101}\) 等同于 \(\langle q, 1, 1 \rangle\), \(T_{110}\) 等同于 \(\langle 1, q, 1\rangle\).

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

\[\mathfrak{cw}_q^{\otimes k} = \sum_{\substack{I,J,K\\ I_i + J_i + K_i = 2}} T_{IJK}. \]

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

\[t = \sum_{\substack{I,J,K\\ I_i + J_i + K_i = 2\\|I|=|J|=|K| = 2N}} T_{IJK}, \]

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

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

首先我们注意到:

如果我们可以对某个常数 \(\phi\) 构造出一族

\[|L\cap S^X\times S^Y\times S^Z| \geq \phi^{N-o(N)}, \]
就可以给出控制

\[\omega \leq \frac{\log \frac{c^3}{\phi}}{\log q}. \]

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

\[\phi^{N-o(N)} q^{3N \cdot \tau} = c^{3N+o(N)} \]

的解是

\[\tau = (1+o(1)) \frac{\log \frac{c^3}{\phi}}{3\log q}, \]

所以有

\[\omega \leq \frac{\log \frac{c^3}{\phi}}{\log q}. \]

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

\(\omega \leq 2.404\) 的构造

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

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

第 1 步: 哈希与极值组合

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

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

由于 \(p\) 充分大, 所以是奇素数, \(\frac 12\) 是良定义的.

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

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

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

\[|S| \geq \frac{N}{e^{O(\sqrt {\log N})}}. \]

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

那么接下来, 我们令

\[S_X = h_X^{-1}(S), \quad S_Y = h_Y^{-1}(S), \quad 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| = \binom{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} \geq \frac{1}{p^{1+\epsilon}} \binom{3N}{N,N,N}. \]

第 2 步

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

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

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

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

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

\[\binom{3N}{N,N,N} \cdot \frac 32 \left( \binom{2N}{N} - 1 \right). \]

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

\[\mathbb E[| L \cap S^X\times S^Y\times S^Z |] \geq \frac 1{8p^\epsilon} \binom{3N}{N}. \]

根据 Stirling 近似, 我们有

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

因此 \(\phi = \frac{27}{4}\) 满足要求.

结合前面的结果, 我们有

\[\omega \leq \frac{\log \frac{4 c^3}{27}}{\log q}. \]

\(c\leq q+2\) 的构造, 以及剩下的?

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

\[ \begin{aligned} & \sum_{i=1}^q \epsilon (x_0+\epsilon x_i)(y_0+\epsilon y_i)(z_0+\epsilon z_i) \\ & \quad - (x_0 + \epsilon^2 X_1)(y_0 + \epsilon^2 Y_1)(z_0 + \epsilon^2 Z_1) \\ & \quad + (1 - q\epsilon)x_0y_0z_0 \\ & = \epsilon^3 \mathfrak{cw}_q + O(\epsilon^4), \end{aligned} \]

有 \(c\leq q+2\). 当取 \(q=8\) 的时候, 就有

\[\omega \leq \log_8 \left( \frac{4000}{27} \right) < 2.40364. \]

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

EntropyIncreaser

“征服世界的将是这样一些人：开始的时候，他们试图找到梦想中的乐园，最终，当他们无法找到时，就亲自创造了它。”