羿閣 蕭簫 發自 凹非寺
量子位 | 公眾號 QbitAI
什麼，AI竟然能自己改進矩陣乘法，提升計算速度了？！
還是直接打破人類50年前創下的最快紀錄的那種。
要知道，矩陣乘法可是電腦科學中最基礎的數學演算法之一，也是各種AI計算方法的基石，如今計算機處理圖像語音、壓縮資料等全都離不開它。
但自從德國數學家沃爾克·施特拉森（Volker Strassen）在1969年提出「施特拉森演算法」後，矩陣乘法的計算速度一直進步甚微。
現在，這隻新出爐的AI不僅改進了目前最優的4×4矩陣解法（50年前由施特拉森提出），還進一步提升了其他70餘種不同大小矩陣的計算速度。
這是DeepMind最新研究成果AlphaTensor，一經發出就登上了Nature封面。
有意思的是，AlphaTensor並非一開始就是專攻理論研究的，它的前身AlphaZero其實是個用來下下圍棋、國際象棋的「棋類AI」。
這項研究發佈後，一名在DeepMind工作6年的老員工表示：
我在DeepMind幹了這麼些年，能讓我吃驚的東西確實不多了，但這項研究確實讓我倒吸一口涼氣。
前Google大腦工程師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和GoogleTPU 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
