AlphaTensor橫空出世！打破矩陣乘法計算速度50年紀(jì)錄，DeepMind新研究再刷Nature封面，詳細(xì)算法已開源

羿閣 2022-10-06 13:41:53 來源：量子位

羿閣 蕭簫 發(fā)自 凹非寺

量子位 | 公眾號 QbitAI

什么，AI竟然能自己改進(jìn)矩陣乘法，提升計算速度了？！

還是直接打破人類50年前創(chuàng)下的最快紀(jì)錄的那種。

要知道，矩陣乘法可是計算機(jī)科學(xué)中最基礎(chǔ)的數(shù)學(xué)算法之一，也是各種AI計算方法的基石，如今計算機(jī)處理圖像語音、壓縮數(shù)據(jù)等全都離不開它。

但自從德國數(shù)學(xué)家沃爾克·施特拉森（Volker Strassen）在1969年提出“施特拉森算法”后，矩陣乘法的計算速度一直進(jìn)步甚微。

現(xiàn)在，這只新出爐的AI不僅改進(jìn)了目前最優(yōu)的4×4矩陣解法（50年前由施特拉森提出），還進(jìn)一步提升了其他70余種不同大小矩陣的計算速度。

這是DeepMind最新研究成果AlphaTensor，一經(jīng)發(fā)出就登上了Nature封面。

有意思的是，AlphaTensor并非一開始就是專攻理論研究的，它的前身AlphaZero其實是個用來下下圍棋、國際象棋的“棋類AI”。

這項研究發(fā)布后，一名在DeepMind工作6年的老員工表示：

我在DeepMind干了這么些年，能讓我吃驚的東西確實不多了，但這項研究確實讓我倒吸一口涼氣。

前谷歌大腦工程師Eric Jang也激動轉(zhuǎn)發(fā)：干得好！

那么，這只“游戲”AI究竟是怎么打破50年前人類創(chuàng)下的紀(jì)錄的？

從最強(qiáng)棋類AI進(jìn)化而來

AlphaTensor，從DeepMind的最強(qiáng)通用棋類AI“AlphaZero”進(jìn)化而來。

所以，矩陣乘法和棋類有什么關(guān)系？

和棋盤一樣，矩陣看起來也是方方正正的，每一格可以用對應(yīng)的數(shù)據(jù)表示。

因此研究人員突發(fā)奇想，能不能直接把AI做矩陣乘法，當(dāng)成是AI在棋盤上下棋？

其中棋盤代表要解決的乘法問題，下棋步驟代表解決問題的步驟，對應(yīng)的規(guī)則被命名為TensorGame，一種新的“3D棋類游戲”。

但與棋類AI略有不同的是，AlphaZero要找到的是做矩陣乘法的最佳算法——即通過盡可能少的步驟，來“贏”得比賽，也就是計算出最終結(jié)果。

在了解AlphaTensor具體如何訓(xùn)練之前，先來簡單回顧一下矩陣乘法的計算。

以計算最簡單的2×2矩陣乘法為例：

正常來說，我們需要計算8次乘法，再通過4次加法來獲得最終的結(jié)果：

但在矩陣乘法運算中，乘法的復(fù)雜度是O(n3)，而加法的復(fù)雜度只有O(n2)，n越大時此方法的收益就越大。

因此，如果能想辦法降低做乘法的步驟，就能進(jìn)一步加速矩陣乘法的運算速度。

例如根據(jù)經(jīng)典的Strassen算法，兩個2×2的矩陣相乘只需做7次乘法，時間復(fù)雜度也會進(jìn)一步下降。

當(dāng)然，這只是最簡單的矩陣乘法之一。

對于更大、更復(fù)雜的矩陣乘法來說，計算出最終結(jié)果的可能性只會越來越多——

甚至對于兩個矩陣相乘的方法來說，最終可能性比宇宙中的原子還要多（數(shù)量級達(dá)到10的33次方）。

與AlphaZero之前搞定的圍棋游戲相比，AlphaTensor的計算量還要更大，因為矩陣乘法比圍棋可能的步驟還要多出30倍左右。

它同樣采用強(qiáng)化學(xué)習(xí)訓(xùn)練，并在訓(xùn)練之前先學(xué)習(xí)了一些人類計算矩陣乘法的方法，避免在過程中“無腦亂猜”，浪費不必要的計算量。

在訓(xùn)練時，AlphaTensor每一步都會從一個可選擇的操作集（包含下一步可以做的所有計算動作）選擇要完成的下一個動作，最終訓(xùn)練自己通過更少的步驟達(dá)成計算目標(biāo)。

具體在選擇的過程中，AlphaTensor采取了樹搜索（Tree Search）的方法，即基于現(xiàn)有游戲結(jié)果預(yù)測下一個最可能降低步驟的動作。

出乎研究者們意料的是，AlphaTensor發(fā)現(xiàn)的計算矩陣乘法的方法真的挺有效。

例如在英偉達(dá)V100 GPU和谷歌TPU v2這兩種硬件上，使用AlphaTensor發(fā)現(xiàn)的算法計算矩陣乘法，比常用算法要快上10~20%左右。

（當(dāng)然研究者們也表示，其他處理器還得看硬件邏輯，計算方法不一定針對每個處理器都有這么好的加速作用）

具體而言，AlphaTensor一共改進(jìn)了70多種不同大小矩陣的計算方法。

效率超越70+現(xiàn)有計算方法

矩陣乘法是計算機(jī)要做的最關(guān)鍵數(shù)學(xué)計算之一。

同時，它也是機(jī)器學(xué)習(xí)計算中不可或缺的基礎(chǔ)，無論在AI處理手機(jī)圖像、理解語音命令，還是渲染電腦游戲畫面（計算機(jī)圖形學(xué)）等方面，都能見到它的身影。

如今我們做矩陣乘法，很大程度上仍然離不開50年前的Strassen算法。

1969年，德國數(shù)學(xué)家沃爾克·施特拉森（Volker Strassen）證明，將兩個2×2的矩陣相乘，不一定需要進(jìn)行8次乘法。

他巧妙的通過構(gòu)造7個中間變量，用增加14次加法為代價省去了一次乘法，這種方法被稱為“施特拉森算法”（Strassen算法）。

基于Strassen算法邏輯，沃爾克·施特拉森改進(jìn)了當(dāng)時的一大批矩陣乘法。

50多年來，盡管針對一些不容易適應(yīng)計算機(jī)代碼的地方進(jìn)行了輕微改進(jìn)，但該算法一直是大多數(shù)矩陣大小上最有效的方法。

現(xiàn)在，AlphaTensor的出現(xiàn)刷新了這一紀(jì)錄：

它發(fā)現(xiàn)了一種僅用47次乘法就能將兩個4×4的矩陣相乘的算法，超過了施特拉森算法所需的49次乘法。

不僅如此，AlphaTensor還發(fā)現(xiàn)了比以前想象的更豐富的矩陣乘法算法空間——每種尺寸上多達(dá)數(shù)千個算法。

最終，它在70種不同大小矩陣的矩陣乘法中擊敗了現(xiàn)有的最佳算法。

舉個例子，2個9×9矩陣相乘所需的步驟數(shù)從511步減少到498步，2個11×11矩陣相乘所需的步驟數(shù)從919步減少到896步……

所以在時間復(fù)雜度上，AlphaTensor是否做出了對應(yīng)的突破？

對此論文介紹稱，目前最優(yōu)的矩陣乘法時間復(fù)雜度，仍然是2021年3月MIT&哈佛大學(xué)研究中達(dá)成的這一數(shù)值（AlphaTensor改善的時間復(fù)雜度并不比它更低）——

BUT，這個操作起來實在是太麻煩了，所以在實際計算中用處不大，除非計算的是天文數(shù)字大小的矩陣。

換而言之，即使Strassen算法的復(fù)雜度只達(dá)到O(n^2.81)，但在大多數(shù)情況下，都要比上面那個時間復(fù)雜度更低的計算方法更實用。

嗯，更別提在不少特定矩陣乘法中還超過了Strassen算法的AlphaTensor了。

同時研究人員也表示，AlphaTensor設(shè)計的算法具有一定的靈活性。

它不僅可能推進(jìn)各種應(yīng)用程序重新設(shè)計算法，還可能優(yōu)化能源使用量和數(shù)值穩(wěn)定性等指標(biāo)，幫助在實際應(yīng)用時防止算法運行時出現(xiàn)小的舍入誤差（包括Strassen算法等計算矩陣乘法，都會出現(xiàn)一定的誤差）。

此外，雖然目前這些突破還只是針對特定算法改進(jìn)的，但也有科學(xué)家認(rèn)為AlphaTensor的潛力不止于此。

例如，MIT計算機(jī)科學(xué)家Virginia Williams就表示：

研究者們可以再嘗試一下，去搞明白這些特定算法中有沒有什么特殊規(guī)律。此外，也可以研究一下如果將這些特殊算法組合起來，是否能發(fā)現(xiàn)更多更優(yōu)的計算方法。

目前AlphaTensor的相關(guān)代碼已經(jīng)開源。

共同一作也是AlphaGo關(guān)鍵“擺棋手”

AlphaTensor的研究團(tuán)隊都來自DeepMind。

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

其中黃士杰來自中國臺灣，本科畢業(yè)于臺灣交通大學(xué)計算機(jī)與信息科學(xué)專業(yè)，在臺灣師范大學(xué)獲得研究生、博士學(xué)位，后前往加拿大阿爾伯塔大學(xué)攻讀博士后，于2012年加入DeepMind。

他曾在AlphaGo和李世石大戰(zhàn)中，擔(dān)當(dāng)AlphaGo的“人肉臂”（順便把棋輸入電腦），也是AlphaGo論文的共同一作。

對于這只AI達(dá)成的新成就，有網(wǎng)友調(diào)侃：

有意思的是，這只AI竟然是基于舊的矩陣乘法運算規(guī)則，算出這個新矩陣乘法計算方法的。

論文地址：
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