AlphaTensor 横空出世！打破矩阵乘法计算速度 50 年纪录，DeepMind 新研究再刷 Nature 封面，详细算法已开源

什么，AI 竟然能自己改进矩阵乘法，提升计算速度了？！

还是直接打破人类 50 年前创下的最快纪录的那种。

要知道，矩阵乘法可是计算机科学中最基础的数学算法之一，也是各种 AI 计算方法的基石，如今计算机处理图像语音、压缩数据等全都离不开它。

但自从德国数学家沃尔克・施特拉森（Volker Strassen）在 1969 年提出"施特拉森算法"后，矩阵乘法的计算速度一直进步甚微。

现在，这只新出炉的 AI 不仅改进了目前最优的 4×4 矩阵解法（50 年前由施特拉森提出），还进一步提升了其他 70 余种不同大小矩阵的计算速度。

这是 DeepMind 最新研究成果 AlphaTensor，一经发出就登上了 Nature 封面。

有意思的是，AlphaTensor 并非一开始就是专攻理论研究的，它的前身 AlphaZero 其实是个用来下下围棋、国际象棋的"棋类 AI"。

这项研究发布后，一名在 DeepMind 工作 6 年的老员工表示：

我在 DeepMind 干了这么些年，能让我吃惊的东西确实不多了，但这项研究确实让我倒吸一口凉气。

前谷歌大脑工程师 Eric Jang 也激动转发：干得好！

那么，这只"游戏"AI 究竟是怎么打破 50 年前人类创下的纪录的？

从最强棋类 AI 进化而来

AlphaTensor，从 DeepMind 的最强通用棋类 AI"AlphaZero"进化而来。

所以，矩阵乘法和棋类有什么关系？

和棋盘一样，矩阵看起来也是方方正正的，每一格可以用对应的数据表示。

因此研究人员突发奇想，能不能直接把 AI 做矩阵乘法，当成是 AI 在棋盘上下棋？

其中棋盘代表要解决的乘法问题，下棋步骤代表解决问题的步骤，对应的规则被命名为 TensorGame，一种新的"3D 棋类游戏"。

但与棋类 AI 略有不同的是，AlphaZero 要找到的是做矩阵乘法的最佳算法 -- 即通过尽可能少的步骤，来"赢"得比赛，也就是计算出最终结果。

在了解 AlphaTensor 具体如何训练之前，先来简单回顾一下矩阵乘法的计算。

以计算最简单的 2×2 矩阵乘法为例：

正常来说，我们需要计算 8 次乘法，再通过 4 次加法来获得最终的结果：

但在矩阵乘法运算中，乘法的复杂度是 O (n³)，而加法的复杂度只有 O (n²)，n 越大时此方法的收益就越大。

因此，如果能想办法降低做乘法的步骤，就能进一步加速矩阵乘法的运算速度。

例如根据经典的 Strassen 算法，两个 2×2 的矩阵相乘只需做 7 次乘法，时间复杂度也会进一步下降。

当然，这只是最简单的矩阵乘法之一。

对于更大、更复杂的矩阵乘法来说，计算出最终结果的可能性只会越来越多 --

甚至对于两个矩阵相乘的方法来说，最终可能性比宇宙中的原子还要多（数量级达到 10 的 33 次方）。

与 AlphaZero 之前搞定的围棋游戏相比，AlphaTensor 的计算量还要更大，因为矩阵乘法比围棋可能的步骤还要多出 30 倍左右。

它同样采用强化学习训练，并在训练之前先学习了一些人类计算矩阵乘法的方法，避免在过程中"无脑乱猜"，浪费不必要的计算量。

在训练时，AlphaTensor 每一步都会从一个可选择的操作集（包含下一步可以做的所有计算动作）选择要完成的下一个动作，最终训练自己通过更少的步骤达成计算目标。

具体在选择的过程中，AlphaTensor 采取了树搜索（Tree Search）的方法，即基于现有游戏结果预测下一个最可能降低步骤的动作。

出乎研究者们意料的是，AlphaTensor 发现的计算矩阵乘法的方法真的挺有效。

例如在英伟达 V100 GPU 和谷歌 TPU v2 这两种硬件上，使用 AlphaTensor 发现的算法计算矩阵乘法，比常用算法要快上 10~20% 左右。

（当然研究者们也表示，其他处理器还得看硬件逻辑，计算方法不一定针对每个处理器都有这么好的加速作用）

具体而言，AlphaTensor 一共改进了 70 多种不同大小矩阵的计算方法。

效率超越 70 + 现有计算方法

矩阵乘法是计算机要做的最关键数学计算之一。

同时，它也是机器学习计算中不可或缺的基础，无论在 AI 处理手机图像、理解语音命令，还是渲染电脑游戏画面（计算机图形学）等方面，都能见到它的身影。

如今我们做矩阵乘法，很大程度上仍然离不开 50 年前的 Strassen 算法。

1969 年，德国数学家沃尔克・施特拉森（Volker Strassen）证明，将两个 2×2 的矩阵相乘，不一定需要进行 8 次乘法。

他巧妙的通过构造 7 个中间变量，用增加 14 次加法为代价省去了一次乘法，这种方法被称为"施特拉森算法"（Strassen 算法）。

基于 Strassen 算法逻辑，沃尔克・施特拉森改进了当时的一大批矩阵乘法。

50 多年来，尽管针对一些不容易适应计算机代码的地方进行了轻微改进，但该算法一直是大多数矩阵大小上最有效的方法。

现在，AlphaTensor 的出现刷新了这一纪录：

它发现了一种仅用 47 次乘法就能将两个 4×4 的矩阵相乘的算法，超过了施特拉森算法所需的 49 次乘法。

不仅如此，AlphaTensor 还发现了比以前想象的更丰富的矩阵乘法算法空间 -- 每种尺寸上多达数千个算法。

最终，它在 70 种不同大小矩阵的矩阵乘法中击败了现有的最佳算法。

举个例子，2 个 9×9 矩阵相乘所需的步骤数从 511 步减少到 498 步，2 个 11×11 矩阵相乘所需的步骤数从 919 步减少到 896 步……

所以在时间复杂度上，AlphaTensor 是否做出了对应的突破？

对此论文介绍称，目前最优的矩阵乘法时间复杂度，仍然是 2021 年 3 月 MIT & 哈佛大学研究中达成的这一数值（AlphaTensor 改善的时间复杂度并不比它更低）--

BUT，这个操作起来实在是太麻烦了，所以在实际计算中用处不大，除非计算的是天文数字大小的矩阵。

换而言之，即使 Strassen 算法的复杂度只达到 O (n^2.81)，但在大多数情况下，都要比上面那个时间复杂度更低的计算方法更实用。

嗯，更别提在不少特定矩阵乘法中还超过了 Strassen 算法的 AlphaTensor 了。

同时研究人员也表示，AlphaTensor 设计的算法具有一定的灵活性。

它不仅可能推进各种应用程序重新设计算法，还可能优化能源使用量和数值稳定性等指标，帮助在实际应用时防止算法运行时出现小的舍入误差（包括 Strassen 算法等计算矩阵乘法，都会出现一定的误差）。

此外，虽然目前这些突破还只是针对特定算法改进的，但也有科学家认为 AlphaTensor 的潜力不止于此。

例如，MIT 计算机科学家 Virginia Williams 就表示：

研究者们可以再尝试一下，去搞明白这些特定算法中有没有什么特殊规律。此外，也可以研究一下如果将这些特殊算法组合起来，是否能发现更多更优的计算方法。

目前 AlphaTensor 的相关代码已经开源。

共同一作也是 AlphaGo 关键"摆棋手"

AlphaTensor 的研究团队都来自 DeepMind。

5 位共同一作分别是 Alhussein Fawzi、Matej Balog、黄士杰、Thomas Hubert 和 Bernardino Romera-Paredes。

其中黄士杰来自中国台湾，本科毕业于台湾交通大学计算机与信息科学专业，在台湾师范大学获得研究生、博士学位，后前往加拿大阿尔伯塔大学攻读博士后，于 2012 年加入 DeepMind。

他曾在 AlphaGo 和李世石大战中，担当 AlphaGo 的"人肉臂"（顺便把棋输入电脑），也是 AlphaGo 论文的共同一作。

对于这只 AI 达成的新成就，有网友调侃：

有意思的是，这只 AI 竟然是基于旧的矩阵乘法运算规则，算出这个新矩阵乘法计算方法的。

论文地址：

• https://www.nature.com/articles/s41586-022-05172-4

参考链接：

• [1]https://www.technologyreview.com/2022/10/05/1060717/deepmind-uses-its-game-playing-ai-to-best-a-50-year-old-record-in-computer-science/

• [2]https://www.nature.com/articles/d41586-022-03166-w

• [3]https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor

• [4]https://twitter.com/DeepMind/status/1577677899108421633

本文来自微信公众号：量子位 （ID：QbitAI），作者：羿阁、萧箫