论文算法发现

Discovering faster matrix multiplication algorithms with reinforcement learning (AlphaTensor)

AlphaTensor:用强化学习发现更快的矩阵乘法算法

Alhussein Fawzi, Matej Balog, Aja Huang, Thomas Hubert, Bernardino Romera-Paredes, et al. (DeepMind) · Google DeepMind · Nature · 2022-10

代码主页

一句话DeepMind 把"找矩阵乘法算法"建模成单人游戏 TensorGame,用 Sampled AlphaZero 搜索张量分解,在 Z2(模 2)算术下把 4x4 矩阵乘法从 Strassen 两层递归的 49 次乘法降到 47 次(1969 年以来首次),标准算术下也刷新了 (4,5,5) 76 vs 80 等多个尺寸,并为 V100/TPU v2 定制出比 cuBLAS 快 8.5%-23.9% 的分块算法。

这是什么

矩阵乘法的乘法次数下界是算法领域的经典开放问题:Strassen 1969 年证明 2x2 矩阵乘法只需 7 次乘法(而非 8 次),递归应用可把复杂度降到 O(N^2.81)。此后 50 多年,4x4 的最优纪录一直停留在 Strassen 递归两层的 49 次。找更优分解等价于求矩阵乘法张量 T_n 的低秩分解,而张量秩分解是 NP-hard,4x4 的动作空间比 3x3 大约 10^10 倍,人工和 SAT/组合搜索都够不到。

这篇 2022 年 10 月的 Nature 论文提出 AlphaTensor:把"减掉一个 rank-1 张量"当成游戏的一步,玩家从目标张量出发,每步选一组因子 (u,v,w) 做减法,减到零张量即得到一个可证明正确的算法,步数越少奖励越高。用 AlphaZero 的神经网络 + MCTS 玩这个游戏,同一个 agent 同时训练分解 n,m,p<=5 的全部尺寸,并支持在不同 ring(Z2 / 标准算术)和不同目标(乘法次数 / 实际硬件运行时间)下搜索。

结果分三层:理论上打破多个尺寸的最优乘法次数纪录;数学上发现 4x4 存在 14,236 个互不等价的 rank-49 分解,说明解空间远比已知的丰富;工程上把奖励换成实测运行时间,为特定硬件搜出实用加速。它是"AI 发现底层计算算法"路线在 LLM 时代之前的代表作,后续 AlphaDev(汇编排序)、AlphaEvolve 都沿这条线走。

主结果表(论文 Fig 3):各尺寸 (n,m,p) 矩阵乘法此前已知最优乘法次数 vs AlphaTensor 找到的(模 2 / 标准算术两列),红色为改进项——4x4x4 模 2 下 47(原 49)、(4,5,5) 标准算术 76(原 80)等。右侧散点是用 recombination 拼装后在 n,m,p<=12 范围刷新的 70 多个尺寸,越靠上改进越大。

机制与做法

TensorGame:把算法发现变成单人游戏

n x m 与 m x p 矩阵乘法对应一个 nm x mp x pn 的三维 0/1 张量 T_{n,m,p};任何把 T 写成 R 个 rank-1 张量之和的分解,都严格对应一个用 R 次标量乘法的算法(Strassen 算法就是 T_2 的 rank-7 分解)。游戏状态 S_t 初始为目标张量,每步选因子三元组 (u,v,w)(元素取自小离散集合,如 {-2,-1,0,1,2}),更新 S_t <- S_{t-1} - u⊗v⊗w,减到零张量游戏结束,正确性由构造自动保证。每步奖励 -1(压步数),超过步数上限则按残差张量的秩上界罚分;把奖励换成实测运行时间就能做硬件定制搜索。

硬件定制算法的实测加速(论文 Fig 5):以运行时间为奖励、为 8,192 尺寸优化的 4x4 分块算法,在 V100 GPU(a)和 TPU v2(b)上相对标准 matmul(GPU 上即 cuBLAS)的中位加速,蓝色 AlphaTensor、红色 Strassen-square 对照。注意 c 图:为 GPU 优化的算法在 TPU 上只剩 2.8%,说明它学到的是对特定硬件+编译器栈的适配,不是普适更优。

Sampled AlphaZero + 两个关键 trick

网络输入当前张量和历史动作,输出 policy(对巨大动作空间做采样而非枚举,即 Sampled AlphaZero)和 value(用 quantile regression 学回报分布,推理时取 75% 以上分位数均值,鼓励冒险)。MCTS 用网络引导,玩完的游戏回流训练。

两个对成败关键的设计:一是合成示范——张量分解难但反向构造容易,随机采样因子拼出 500 万个"张量-分解"对做监督信号,与 RL 目标混合训练,单用任何一种都明显更差;二是利用张量对称性做数据增广(随机基变换、随机置换因子顺序),等于把一个目标张量变成成千上万个等价目标。训练用 64 个 TPU v3 核做 learner、1,600 个 TPU v4 actor,60 万步约一周收敛。

从单个分解到实用加速

小尺寸分解可以递归/组合出大尺寸算法:论文用 recombination(拼装小张量的分解)刷新了 n,m,p<=12 范围内 70 多个张量的已知最优 rank。硬件实验则不做递归:把 8,192x8,192 矩阵切成 4x4 个 2,048 块,只在最外层用 AlphaTensor 搜出的 4x4 分块算法(奖励即 JAX/XLA 编译后的实测时间),块内乘法仍走标准 cuBLAS/TPU kernel。有趣的发现是:搜出来的算法加法比 Strassen-square 还多,但操作模式更容易被 XLA 融合,所以更快——它优化的其实是"对这套编译器栈友好",给 GPU 调的算法搬到 TPU 上就不灵(Fig 5c:GPU 版在 TPU 上只有 2.8% vs TPU 版 10.3%)。

关键结果

Z2(模 2)算术下 4x4 矩阵乘法 47 次乘法,低于 Strassen 两层递归的 49 次,是 1969 年以来该尺寸的首次改进;递归应用得到 Z2 上 O(N^2.778) 的算法(注意:仅限模 2,不适用于浮点)。
标准算术下也有真实改进:(4,5,5) 76 次(原纪录 80)、(4,4,5) 63(原 64)、(3,4,5) 47(原 48);(5,5,5) 在 Z2 下 96(原 98)。经 recombination 拼装,刷新 n,m,p<=12 范围内 70 多个张量的已知最优 rank。
为 4x4x4 标准算术找到 14,236 个互不等价的 rank-49 分解(此前只知道 Strassen-square 一个),全部数据随仓库发布,可用 Colab 验证不等价性。
以实测运行时间为奖励,为 8,192x8,192 矩阵搜出的 4x4 分块算法:V100 GPU 上比 cuBLAS 快 8.5%(Strassen-square 只快 4.3%),TPU v2 上快 10.3%;矩阵越大加速越明显,20,480 时 GPU 上达 23.9%。
超出矩阵乘法:为 skew-symmetric 矩阵-向量乘积发现了 (n-1)(n+2)/2 ≈ n²/2 次乘法的通用算法(此前已知 ~n²),渐近最优,是从小尺寸解归纳出的可证明一般规律。
训练成本不低:64 TPU v3 核 learner + 1,600 个 TPU v4 actor,跑约一周;引用数 S2 737 / Google Scholar 约 1,181(2026-08 查)。

实证核查

扎实核心产物(分解本身)全部公开且机器可验证,论文对适用范围的表述也诚实;水分主要在媒体转述层:47 次纪录仅限模 2 算术、10-20% 加速的对比基线和测法都有讲究。另外 RL 训练代码未开源,"用 RL 才能找到"这一点很快被便宜得多的经典搜索部分推翻。

打破 Strassen 保持 50 年的 4x4 矩阵乘法纪录(47 vs 49 次乘法)。

论文正文明确限定在 Z2(模 2)算术,标准/浮点算术下 4x4 仍是 49 次,这一改进对实际浮点 matmul 没有直接用处。repo issue #5('4x4x4 algorithm for real arithmetic?')中社区确认此点,并指出模 2 下 Stothers 博士论文早有 48 次的结果;图形程序员 Fabian Giesen 的评论文章(fgiesen.wordpress.com, 2022-10-06)进一步指出媒体报道普遍略去了这个限定。

发现的算法在 GPU/TPU 上比常用算法快 10-20%。

数据本身真实(Fig 5:V100 上 8,192 尺寸 8.5%、20,480 尺寸 23.9%),但按 Giesen 的分析:对比基线是普通 matmul 而非 Strassen(比 Strassen-square 只多几个点);只测了 >=8,192 的超大矩阵、单层 4x4 分块、走 JAX/XLA 而非手工调优 kernel;且浮点下这类快速算法有数值稳定性代价,论文未在正文展开。另外 repo issue #13 指出 benchmark 脚本把 V100 时钟锁在 1530 MHz——高于 V100 的最大 boost 频率 1380 MHz,该 issue 至今无官方回应。

开源仓库配套论文(github.com/google-deepmind/alphatensor,2.8k stars)。

仓库只含四样东西:发现的分解数据(factorizations_r.npz / factorizations_f2.npz)、14,236 个不等价分解及验证 Colab、V100 benchmark 脚本、recombination 代码;RL 训练代码(Sampled AlphaZero)未开源,issue #4 中作者明确答复 'not open sourced',网络结构只有 Supplementary 里的伪代码。好处是结果本身可独立验证(分解正确性是机器可查的),坏处是搜索过程无法复现。仓库最后一次 push 是 2024-04,基本处于只读状态。

(隐含叙事)这类发现需要 AlphaZero 级别的 RL + 数千 TPU。

论文发布数周后,Kauers & Moosbauer 用普通工作站上的经典随机游走搜索(flip graph,arXiv:2212.01175)把 Z2 下 (5,5,5) 从 AlphaTensor 的 96 降到 95、(4,4,5) 从 63 降到 62——起点正是 AlphaTensor 公布的分解。后续 Adaptive Flip Graph(Arai et al., ISSAC 2024)和 Flip Graphs with Symmetry(Moosbauer & Poole, 2025)继续刷新多个尺寸。说明 RL 不是唯一路径,但 AlphaTensor 的结果作为搜索起点确实推动了这条线。

与我们方向的关系

这是'AI 发现 AI 底层计算'方向的奠基工作:目标(matmul)正是深度学习自身最烧的算子,形成了自举闭环的雏形。方法论上最值得借鉴的是问题形式化——把'找算法'压成'张量分解游戏',正确性由构造保证(减到零张量即正确),这样 RL 只需要优化效率而不用管对错;凡是能写成这种'解空间离散 + 正确性可机器验证 + 目标可度量'形式的算法发现问题(kernel 调度、量化方案、通信模式)都可以套这个框架。合成示范(逆向构造监督数据)和对称性增广两个 trick 在其他搜索问题上也通用。

对照后续工作看更有意思:flip graph 系列证明便宜的经典搜索在同一问题上能追平甚至反超,而 DeepMind 自己的后续(AlphaDev、FunSearch、AlphaEvolve)转向了 LLM+进化搜索,其中 AlphaEvolve 2025 年在标准算术下把 4x4 降到 48 次复数乘法——补上了 AlphaTensor 没做到的事。读这篇时值得想清楚:专用 RL 搜索 vs 经典组合搜索 vs LLM 引导搜索,各自适合什么样的解空间。

阅读笔记

读的是 Nature 开放获取全文(无 arXiv 版)。benchmark 若要自己跑,repo 的 benchmarking/ 只支持 V100 + 特定 CUDA 版本,issue #6 要求支持自定义 GPU 被关闭未实现。论文中 modular arithmetic 均指 Z2;'rank' 即乘法次数。媒体报道(含 DeepMind 官方博客)对 mod-2 限定和加速基线的表述比论文正文宽松,引用时以论文数字为准。