Strassen 这个名字大多数接触过算法的人都不陌生用分治把矩阵乘法的复杂度从 (O(n^3)) 压到 (O(n^{\log_2 7}))。(\log_2 7 \approx 2.807)意思是矩阵规模翻倍时朴素算法的工作量变成 8 倍Strassen 只变成约 7 倍。这个差距看着不大但阶次差异决定了它在大规模下会越拉越开。我最初在教材上看到那七个中间矩阵时只觉得推导像某种线性代数魔术。直到真的用 Python 实现了一版跑通、对比、调阈值、踩浮点精度的坑之后才意识到算法设计、工程实现、性能权衡这三件事在 Strassen 身上是绑在一起的。这篇文章打算完整记录一次实现过程会给出可直接复现的代码会把分块、补零、递归终止这些容易被忽略的决策讲明白也会聊聊为什么在实际跑分里它经常输给 NumPy 的底层 BLAS——这本身就是一个很有意思的话题。如果你正学算法、想在面试里把这个算法讲得不像背课本或者已经写过矩阵运算但搞不清为何自己实现比 NumPy 慢这篇文章应该能帮上忙。1. 矩阵乘法为什么值得较真先从 O(n^3) 的三重循环说起1.1 朴素乘法的三次方循环先看最熟悉的版本。两个 (n \times n) 矩阵相乘结果矩阵 (C) 的每个元素都是 (A) 的一行和 (B) 的一列做内积[ C[i][j] \sum_{k0}^{n-1} A[i][k] \cdot B[k][j] ]写成 Python 就是三层循环def multiply_naive(A, B): n len(A) C [[0] * n for _ in range(n)] for i in range(n): for j in range(n): for k in range(n): C[i][j] A[i][k] * B[k][j] return C把循环次数数一遍(i, j, k) 各取 (n) 个值一共执行 (n^3) 次乘法另外还有 (n^3) 次加法。所以朴素矩阵乘法的复杂度是 (O(n^3))。(n1000) 时乘法次数是 (10^9)。这个量级放在现代 CPU 上如果用他人写好的高度优化库来完成很快但如果用上面这段纯 Python 跑慢得会让你对人生怀疑。问题不在于 Python 语法慢而在于整个算法的乘法次数本身就是按三次方增长的。1.2 分块之后为什么复杂度没变主定理视角“既然 (O(n^3)) 大那用分治不行吗”这是很自然的想法。把矩阵 (A, B) 切成四块[ A \begin{bmatrix} A_{11} A_{12} \ A_{21} A_{22} \end{bmatrix}, \quad B \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} ]这需要 8 个大小为 (n/2) 的子矩阵乘法加上 4 次子矩阵加法。递归式为[ T(n) 8T(n/2) \Theta(n^2) ]主定理里 (a8, b2)于是 (\log_b a \log_2 8 3)(f(n) n^2 O(n^{3-\varepsilon}))属于第一种情况所以[ T(n) \Theta(n^3) ]也就是说普通分块只是把同样的 (O(n^3)) 换了个写法一点忙都没帮上。关键瓶颈在于 8 次子乘法它们在递归树里占主导地位。Strassen 的全部要点就是把 8 变成 7。2. Strassen 怎么把“8 次乘法”变成“7 次”核心推导与验证2.1 七个中间矩阵Strassen 的奇妙之处在于2×2 分块后并不直接算那 8 个乘积而是先构造 7 个“组合乘积”。这里我用 (P_1) 到 (P_7) 表示不同教材的字母不一样有的写 (M_1)有的写 (S_1)但数学本质相同[ P_1 A_{11} \cdot (B_{12} - B_{22}) ] [ P_2 (A_{11} A_{12}) \cdot B_{22} ] [ P_3 (A_{21} A_{22}) \cdot B_{11} ] [ P_4 A_{22} \cdot (B_{21} - B_{11}) ] [ P_5 (A_{11} A_{22}) \cdot (B_{11} B_{22}) ] [ P_6 (A_{12} - A_{22}) \cdot (B_{21} B_{22}) ] [ P_7 (A_{11} - A_{21}) \cdot (B_{11} B_{12}) ]注意看每个 (P) 都是一次子矩阵乘法。原来 8 次现在只有 7 次。代价是加减法变多了但加减法只是 (O(n^2)) 量级的操作在递归主项里低于 (n^{\log_2 7})所以渐进意义上可被高阶吸收掉。2.2 重组公式与验证算完 (P_1 \sim P_7) 后用以下方式重组出 (C) 的四个分块[ 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 ]很多第一次看的人都会问这公式是硬凑的吗其实验证一次就能放心。拿 (C_{11}) 举例完整展开[ P_5 A_{11}B_{11} A_{11}B_{22} A_{22}B_{11} A_{22}B_{22} ] [ P_4 A_{22}B_{21} - A_{22}B_{11} ] [ -P_2 -(A_{11}B_{22} A_{12}B_{22}) ] [ P_6 A_{12}B_{21} A_{12}B_{22} - A_{22}B_{21} - A_{22}B_{22} ]逐项相加(A_{11}B_{22})、(A_{22}B_{11})、(A_{22}B_{22})、(A_{12}B_{22})、(A_{22}B_{21}) 这些项全部抵消最终剩下[ C_{11} A_{11}B_{11} A_{12}B_{21} ]和普通分块的结果一模一样。(C_{12}, C_{21}, C_{22}) 也是同样机制。所以说Strassen 不是近似而是精确的重组只是通过更多加减法换来了更少的乘法。得到新的递归式[ T(n) 7T(n/2) \Theta(n^2) ]这里 (\log_2 7 \approx 2.807)所以复杂度变成 (\Theta(n^{\log_2 7}))。教科书里经常把 (\log_2 7) 写成 (\log 7)就是这个来源。3. 编码前的三个决定矩阵表示、补零策略与递归阈值算法的数学核心清楚了但真正写 Python 时会发现摆在面前的有三个工程决策任何一个选错代码都会慢到让你误以为算法本身是错的。3.1 用 NumPy 视图切分矩阵而不是拷贝子矩阵第一种自然想法是用 Python 嵌套列表表示矩阵然后递归时用列表切片。比如A[:m][:m]看起来挺方便但列表切片会把整块数据复制出来。问题在于Strassen 每一层都要切出 8 个子块、做 7 次递归每切一次就复制 (O(n^2)) 的元素。这些复制成本层层叠加很容易把节省下来的乘法优势全部吃光。更好的选择是用 NumPy 数组。NumPy 的切片返回的是原数组的“视图”view比如A11 A[:m, :m]这一步并没有把数据拷出来只是记录了起始位置、步长和形状。真正发生数据复制是在做A11 A12这类运算时但那也正是必须计算的加法没法避免。所以我倾向于直接使用 NumPy 写核心递归这样代码干净也能直观展示算法本身的分治结构。如果非要纯 Python 复刻不是不行但性能表现会差一个量级这篇文章就不展开纯列表版了。3.2 奇数尺寸怎么处理补零到下一个 2 的幂递归切分时遇到奇数尺寸最麻烦。如果 (n) 是奇数n // 2取整后左上和右下分块尺寸不一致代码会复杂很多。常用做法是先把矩阵补零到下一个 2 的幂。比如 (n100) 补到 128(n65) 也补到 128。这会让递归时每个子块大小始终是偶数直到到达递归终点。def _next_power_of_two(n): return 1 (n - 1).bit_length()实现入口时把原始矩阵放进一个更大的全零矩阵左上角def _pad_to_power_of_two(mat): n mat.shape[0] m _next_power_of_two(n) if m n: return mat, n padded np.zeros((m, m), dtypemat.dtype) padded[:n, :n] mat return padded, n计算完成后只取结果矩阵左上角 (n \times n) 的部分。补零的代价是矩阵尺寸可能被放大但通常不会超过变成原来的 2 倍而且换来的是递归逻辑的极大简化。也有另一种思路允许奇数尺寸在递归时把一个矩阵切成尺寸分别为n//2和n - n//2的两块。这能省去补零的浪费但每个子矩阵尺寸不再是整齐的 2 的幂递归代码里到处都是需要处理的不整齐边界代码量和出错概率都上升。我见过不少课程作业为了省那点补零开销最后 debug 了一整晚。个人建议初版先用补零跑通后再考虑做非均匀分块。3.3 递归阈值别老老实实递归到 1×1如果递归一直走到1×1矩阵每次递归都要做一次乘法还要承担海量的 Python 函数调用和 NumPy 小矩阵调度开销。实测结果是当矩阵已经小到某个程度继续用 Strassen 递归是纯粹的负优化。解决办法是设一个阈值threshold。当矩阵尺寸小于等于阈值时直接用最简单的A BNumPy 点积或普通三重循环不再继续递归。标准实现里这个阈值通常取 64 或 128但它非常依赖具体机器和库版本。我习惯把阈值做成函数参数跑几组数据再定。比如def _strassen_core(A, B, threshold): n A.shape[0] if n threshold: return A B ...这样改起来方便也好比较不同阈值下的表现。4. 完整 Python 实现与正确性验证4.1 完整代码以下是一份可以复制的完整实现。用 NumPy 表示矩阵入口函数负责补零和还原核心递归函数负责 Strassen 计算。import numpy as np def _next_power_of_two(n): 返回不小于 n 的最小 2 的幂。 return 1 (n - 1).bit_length() def _pad_to_power_of_two(mat): 把方阵补零到 2 的幂尺寸返回填充后矩阵和原始尺寸。 n mat.shape[0] m _next_power_of_two(n) if m n: return mat, n padded np.zeros((m, m), dtypemat.dtype) padded[:n, :n] mat return padded, n def _strassen_core(A, B, threshold): Strassen 递归核心要求 A、B 是同尺寸方阵且边长是 2 的幂。 n A.shape[0] if n threshold: return A B m n // 2 # 切分矩阵这里用的是视图不复制数据 A11 A[:m, :m] A12 A[:m, m:] A21 A[m:, :m] A22 A[m:, m:] B11 B[:m, :m] B12 B[:m, m:] B21 B[m:, :m] B22 B[m:, m:] # 七个 Strassen 中间乘积 P1 _strassen_core(A11, B12 - B22, threshold) P2 _strassen_core(A11 A12, B22, threshold) P3 _strassen_core(A21 A22, B11, threshold) P4 _strassen_core(A22, B21 - B11, threshold) P5 _strassen_core(A11 A22, B11 B22, threshold) P6 _strassen_core(A12 - A22, B21 B22, threshold) P7 _strassen_core(A11 - A21, B11 B12, threshold) # 重组结果 C11 P5 P4 - P2 P6 C12 P1 P2 C21 P3 P4 C22 P5 P1 - P3 - P7 # 组装为大矩阵。先分配空的 C再往四个象限填。 C np.empty((n, n), dtypeA.dtype) C[:m, :m] C11 C[:m, m:] C12 C[m:, :m] C21 C[m:, m:] C22 return C def strassen_matmul(A, B, threshold64): 对外入口。自动补零、递归计算、裁剪还原。 A、B 是任意同尺寸方阵。 assert A.shape B.shape and A.ndim 2 and A.shape[0] A.shape[1] n A.shape[0] A_pad, orig_n _pad_to_power_of_two(A) B_pad, _ _pad_to_power_of_two(B) result _strassen_core(A_pad, B_pad, threshold) return result[:orig_n, :orig_n].copy()代码里值得注意的地方有几处。首先是C np.empty((n, n), dtypeA.dtype)这里用empty而不是zeros因为我们会在每个分块位置都赋值不需要先把整个矩阵清零。其次是所有的B12 - B22、A11 A12这类表达式都会创建临时数组这正是 Strassen 在 Python 里常数比较大的原因之一后面第六部分会专门说怎么优化。4.2 核心函数怎么读_strassen_core的逻辑对应前面第七部分的公式基线如果 (n \le \text{threshold})直接用普通矩阵乘。否则分成 4 个子块计算 7 个乘积每个乘积都是规模减半的递归调用。最后按公式组装回四个象限。代码顺序和公式完全一一对应方便对照。用A[:m, :m]得到的是视图所以切割操作几乎不花时间。入口函数strassen_matmul则处理两件脏活把任意尺寸矩阵补成 2 的幂尺寸算完再裁掉补零区域并返回副本。如果没有入口这层封装调用者就得自己保证矩阵尺寸是 2 的幂这在日常使用里很烦人。4.3 正确性验证用随机矩阵对照朴素乘法写完了别急着信。最靠谱的验证方式是把 Strassen 结果和 NumPy 标准乘法做随机对比覆盖各种边长包括奇数和不是 2 的幂的数def verify(threshold64): for n in [1, 2, 3, 7, 15, 16, 31, 32, 65, 100, 127, 129, 200]: for _ in range(3): A np.random.rand(n, n) B np.random.rand(n, n) expected A B got strassen_matmul(A, B, threshold) if not np.allclose(got, expected, rtol1e-8, atol1e-10): raise ValueError(fn{n} 验证失败) print(全部通过)第一次跑这个验证时我最常遇到的问题是补零函数没写对比如padded[:n, :n]写成了padded[:n, n:]结果错误发生在角落位置普通测试看不出来必须靠随机矩阵和allclose才能暴露。另外np.allclose的 rtol 和 atol 要稍微放宽一点因为 Strassen 的中间加减法和朴素乘法的运算顺序不同浮点误差会有微小差异但通常量级在 (10^{-12}) 左右。5. 跑分对比什么规模下才值得用 Strassen5.1 理论乘法数先算一笔账先看理论。假设矩阵边长为 2 的幂朴素乘法需要的标量乘法次数是 (n^3)Strassen 递归则是 (7^{\log_2 n}) 次忽略阈值层具体数字如表n朴素乘法次数Strassen 乘法次数压缩倍数1282,097,152823,5432.5525616,777,2165,764,8012.91512134,217,72840,353,6073.3310241,073,741,824282,475,2493.8020488,589,934,5921,977,326,7434.35从表里能看出随着规模增大Strassen 的优势在扩大但绝不夸张。关键是它把指数从 3 降到了 2.807规模越大优势越明显。注意这还只是乘法次数Strassen 的加法次数比朴素多实际压缩倍数没有表格这么理想。5.2 实测观察教学版 Strassen 经常输给 NumPy如果你看完理论就开始期待strassen_matmul在跑分时碾压A B那我要先泼一盆冷水。我自己在不同机器上对比过一个纯 Python NumPy 实现的 Strassen在 (n256) 到 (n2048) 区间内通常比 NumPy 直接乘法要慢有时慢 23 倍甚至更多。原因不是 Strassen 算法没用而是竞争对象太强。NumPy 的底层调用的 BLAS 类库是高度优化的 C 程序里面有缓存分块、SIMD 指令、指令级并行等一大堆优化。而 Python 版 Strassen 每一层都要创建大量中间数组Python 函数调用开销、NumPy 对象调度开销全部叠加在一起常数因子大得吓人。到了 (n1024)理论上的 3.8 倍乘法优势根本盖不住常数差距。那么 Strassen 在什么时候能赢如果你把比较对象换成纯 Python 三层循环写出的朴素乘法那它几乎从几百阶尺度开始就能明显胜出。但这样的对比本质上是“用 C 库的教学实现”对“用纯 Python 的朴素实现”不公平却也能说明问题Strassen 的价值是减少乘法次数而不是减少 Python 层的调度开销。如果你想看到 Strassen 真正发挥威力通常需要把它实现在底层数据结构上用 C、C 或 CUDA并且优化掉中间矩阵的反复分配。那已经超出了“Python 实现”的范围但能解释为什么学术界和工业界对大矩阵乘法依然重视这类算法。5.3 想看到 Strassen 赢需要满足哪些条件结合我的测试经验Python 里想让 Strassen 跑出接近理论的收益至少要满足几个条件阈值要选对。太小则递归过深太大则退化成朴素分块。64 或 128 是常见起点但最好针对自己的数据规模调一调。矩阵要大。到了 (n \ge 4096) 左右乘法次数优势才有机会抵抗 Python 层常数但前提是你已经把临时数组分配减到最少。不能用纯列表。列表切片的复制成本会彻底毁掉分治优势。所以我的建议是如果你只是为了理解算法Python 版跑通就是胜利如果你为了追求性能直接研究如何用高性能库或扩展语言里的 tile/分块实现不要在 Python 层和 BLAS 较劲。6. 实现中容易踩的坑与后续优化方向6.1 浮点精度问题Strassen 的浮点误差比朴素乘法大这是个经典结论。原因在于它把同样的元素用不同的加减顺序组合特别是 (P_5 P_4 - P_2 P_6) 这类式子会有正负项互相抵消数值上容易放大舍入误差。实际影响有多大对随机矩阵、几百阶以内np.allclose默认容差通常能过。但如果矩阵元素量级差异巨大或者矩阵接近病态误差可能明显变大。如果你在写需要高精度结果的代码建议保留一份朴素乘法的对照结果如果误差大到不可接受Strassen 就不适合继续向下递归这时候可以把阈值调大让更多层使用普通乘法减少 Strassen 的应用深度。6.2 内存分配是最容易被忽略的杀手在_strassen_core里A11 A22这样的操作会为中间结果分配新数组B12 - B22也会。7 个乘积每个都有 1 到 2 个临时矩阵再加上重组阶段的加减一层递归能产生十几块临时数组。如果矩阵很大这些临时数组会频繁触发内存分配回收带来不小的开销。一个可行的优化办法是预先分配一块足够大的临时缓冲然后用np.subtract(B12, B22, outtmp1)这类带out参数的运算把结果写进缓冲而不是新开数组。这样可以明显降低分配压力。另外一个容易踩的坑是补零后忘了裁剪。如果入口函数返回的是m \times m矩阵而不是原始的 (n \times n)后续下游代码会出现形状不匹配。我在第一个版本里就遇到过最后在验证函数里用各种异形边长才揪出来。6.3 加法次数可以进一步压缩Winograd 变体标准 Strassen 一层的矩阵加减法次数是 19 次左右也就是 7 次乘法之前有 11 次加减重组阶段有 8 次加减。每一层都这么多常数不小。Winograd 变体通过重新组合加减顺序把加减法数从 19 降到 15乘法仍然是 7 次。原理不复杂就是把一些重复的加减合并。实际工程里很多追求性能的矩阵乘法库采用的就是这种变体。想深入的人可以搜一下 Winograd variant of Strassen阅读那些实现时会发现它们在底层刻意避免任何不必要的数组分配。如果只是想跑通算法标准版本已经够了如果后续想对比性能可以尝试改成 Winograd 形式并测试加减法次数降低后带来的收益。6.4 最后再说几句实践心得写 Strassen 最大的收获不是背下那七个公式而是理解一个道理算法复杂度的改进往往是用更复杂的常数因子换来的。8 次乘法变成 7 次听起来只是少了 12.5%但因为递归每层都在少乘才最终改变了指数。这个“每层少一次指数就变”的视角比单纯背公式有用得多。另外调试递归函数时一定要用形状多样、带奇数的测试矩阵。只测 (2^n) 尺寸很容易漏掉补零和裁剪的问题比如我遇到过 (n128) 全对但 (n100) 结果错一半的情况就是因为裁剪边界写错。把验证函数直接放进测试用例里每次改完代码跑一遍能省很多时间。如果你想继续扩展还可以试试不同阈值对运行时间的影响曲线或者把_strassen_core改成非递归版本用迭代式分块来减少 Python 递归开销。每一层改动都会让你对这个算法的理解更深一层。
