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

羿阁 萧箫 发自 凹非寺

量子位 | 公众号 QbitAI

什么,AI竟然能自己改进矩阵乘法,提升计算速度了?!

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

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

但自从德国数学家沃尔克·施特拉森(Volker Strassen)在1969年提出“施特拉森算法”后,矩阵乘法的计算速度一直进步甚微。

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

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

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

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

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

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

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

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

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

那么,这只“游戏”AI究竟是怎么打破50年前人类创下的纪录的?

从最强棋类AI进化而来

AlphaTensor,从DeepMind的最强通用棋类AI“AlphaZero”进化而来。

所以,矩阵乘法棋类有什么关系?

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

在训练时,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算法)。

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

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

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

现在,AlphaTensor的出现刷新了这一纪录:

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

例如,MIT计算机科学家Virginia Williams就表示:

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

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

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

AlphaTensor的研究团队都来自DeepMind。

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

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

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

他曾在AlphaGo和李世石大战中,担当AlphaGo的“人肉臂”(顺便把棋输入电脑),也是AlphaGo论文的共同一作。

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

对于这只AI达成的新成就,有网友调侃:

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

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

论文地址:
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