AlphaTensor 横空出世!打破矩阵乘法计算速度 50 年纪录,DeepMind 新研究再刷 Nature 封面,详细算法已开源
- 量子位
2022-10-06 16:27
什么,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),作者:羿阁、萧箫
广告声明:文内含有的对外跳转链接(包括不限于超链接、二维码、口令等形式),用于传递更多信息,节省甄选时间,结果仅供参考,IT之家所有文章均包含本声明。