Nature封面：DeepMind推出AlphaTensor，用AI發(fā)現(xiàn)矩陣乘法算法

DeepMind 的 Alpha 系列 AI 智能體家族又多了一個(gè)成員——AlphaTensor，這次是用來(lái)發(fā)現(xiàn)算法。

數(shù)千年來(lái)，算法一直在幫助數(shù)學(xué)家們進(jìn)行基本運(yùn)算。早在很久之前，古埃及人就發(fā)明了一種不需要乘法表就能將兩個(gè)數(shù)字相乘的算法。希臘數(shù)學(xué)家歐幾里得描述了一種計(jì)算最大公約數(shù)的算法，這種算法至今仍在使用。在伊斯蘭的黃金時(shí)代，波斯數(shù)學(xué)家 Muhammad ibn Musa al-Khwarizmi 設(shè)計(jì)了一種求解線性方程和二次方程的新算法，這些算法都對(duì)后來(lái)的研究產(chǎn)生了深遠(yuǎn)的影響。

事實(shí)上，算法一詞的出現(xiàn)，有這樣一種說(shuō)法：波斯數(shù)學(xué)家 Muhammad ibn Musa al-Khwarizmi 名字中的 al-Khwarizmi 一詞翻譯為拉丁語(yǔ)為 Algoritmi 的意思，從而引出了算法一詞。不過(guò)，雖然今天我們對(duì)算法很熟悉，可以從課堂中學(xué)習(xí)、在科研領(lǐng)域也經(jīng)常遇到，似乎整個(gè)社會(huì)都在使用算法，然而發(fā)現(xiàn)新算法的過(guò)程是非常困難的。

現(xiàn)在，DeepMind 用 AI 來(lái)發(fā)現(xiàn)新算法。

在最新一期 Nature 封面論文《Discovering faster matrix multiplication algorithms with reinforcement learning》中，DeepMind 提出了 AlphaTensor，并表示它是第一個(gè)可用于為矩陣乘法等基本任務(wù)發(fā)現(xiàn)新穎、高效且可證明正確的算法的人工智能系統(tǒng)。簡(jiǎn)單來(lái)說(shuō)，使用 AlphaTensor 能夠發(fā)現(xiàn)新算法。這項(xiàng)研究揭示了 50 年來(lái)在數(shù)學(xué)領(lǐng)域一個(gè)懸而未決的問(wèn)題，即找到兩個(gè)矩陣相乘最快方法。

AlphaTensor 建立在 AlphaZero 的基礎(chǔ)上，而 AlphaZero 是一種在國(guó)際象棋、圍棋和將棋等棋盤(pán)游戲中可以打敗人類(lèi)的智能體。這項(xiàng)工作展示了 AlphaZero 從用于游戲到首次用于解決未解決的數(shù)學(xué)問(wèn)題的一次轉(zhuǎn)變。

矩陣乘法

矩陣乘法是代數(shù)中最簡(jiǎn)單的運(yùn)算之一，通常在高中數(shù)學(xué)課上教授。但在課堂之外，這種不起眼的數(shù)學(xué)運(yùn)算在當(dāng)代數(shù)字世界中產(chǎn)生了巨大的影響，在現(xiàn)代計(jì)算中無(wú)處不在。

兩個(gè) 3x3 矩陣相乘的例子。

你可能沒(méi)注意到，我們生活中處處隱藏著矩陣相乘，如智能手機(jī)中的圖像處理、識(shí)別語(yǔ)音命令、為電腦游戲生成圖形等都有它在背后進(jìn)行運(yùn)算。遍布世界各地的公司都愿意花費(fèi)大量的時(shí)間和金錢(qián)開(kāi)發(fā)計(jì)算硬件以有效地解決矩陣相乘。因此，即使是對(duì)矩陣乘法效率的微小改進(jìn)也會(huì)產(chǎn)生廣泛的影響。

幾個(gè)世紀(jì)以來(lái)，數(shù)學(xué)家認(rèn)為標(biāo)準(zhǔn)矩陣乘法算法是效率最高的算法。但在 1969 年，德國(guó)數(shù)學(xué)家 Volken Strassen 通過(guò)證明確實(shí)存在更好的算法，這一研究震驚了整個(gè)數(shù)學(xué)界。

標(biāo)準(zhǔn)算法與 Strassen 算法對(duì)比，后者少進(jìn)行了一次乘法運(yùn)算，為 7 次，而前者需要 8 次，整體效率大幅提高。

通過(guò)研究非常小的矩陣（大小為 2x2），Strassen 發(fā)現(xiàn)了一種巧妙的方法來(lái)組合矩陣的項(xiàng)以產(chǎn)生更快的算法。之后數(shù)十年，研究者都在研究更大的矩陣，甚至找到 3x3 矩陣相乘的高效方法，都還沒(méi)有解決。

DeepMind 的最新研究探討了現(xiàn)代 AI 技術(shù)如何推動(dòng)新矩陣乘法算法的自動(dòng)發(fā)現(xiàn)。基于人類(lèi)直覺(jué)（human intuition）的進(jìn)步，對(duì)于更大的矩陣來(lái)說(shuō)，AlphaTensor 發(fā)現(xiàn)的算法比許多 SOTA 方法更有效。該研究表明 AI 設(shè)計(jì)的算法優(yōu)于人類(lèi)設(shè)計(jì)的算法，這是算法發(fā)現(xiàn)領(lǐng)域向前邁出的重要一步。

算法發(fā)現(xiàn)自動(dòng)化的過(guò)程和進(jìn)展

首先將發(fā)現(xiàn)矩陣乘法高效算法的問(wèn)題轉(zhuǎn)換為單人游戲。其中，board 是一個(gè)三維度張量（數(shù)字?jǐn)?shù)組），用于捕捉當(dāng)前算法的正確程度。通過(guò)一組與算法指令相對(duì)應(yīng)的所允許的移動(dòng)，玩家嘗試修改張量并將其條目歸零。

當(dāng)玩家設(shè)法這樣做時(shí)，將為任何一對(duì)矩陣生成可證明是正確的矩陣乘法算法，并且其效率由將張量清零所采取的步驟數(shù)來(lái)衡量。

這個(gè)游戲非常具有挑戰(zhàn)性，要考慮的可能算法的數(shù)量遠(yuǎn)遠(yuǎn)大于宇宙中原子的數(shù)量，即使對(duì)于矩陣乘法這樣小的情況也是如此。與幾十年來(lái)一直是人工智能挑戰(zhàn)的圍棋游戲相比，該游戲每一步可能的移動(dòng)數(shù)量要多 30 個(gè)數(shù)量級(jí)（DeepMind 考慮的一種設(shè)置是 10^33 以上。）

為了解決這個(gè)與傳統(tǒng)游戲明顯不同的領(lǐng)域所面臨的挑戰(zhàn)，DeepMind 開(kāi)發(fā)了多個(gè)關(guān)鍵組件，包括一個(gè)結(jié)合特定問(wèn)題歸納偏置的全新神經(jīng)網(wǎng)絡(luò)架構(gòu)、一個(gè)生成有用合成數(shù)據(jù)的程序以及一種利用問(wèn)題對(duì)稱(chēng)性的方法。

接著，DeepMind 訓(xùn)練了一個(gè)利用強(qiáng)化學(xué)習(xí)的智能體 AlphaTensor 來(lái)玩這個(gè)游戲，該智能體在開(kāi)始時(shí)沒(méi)有任何現(xiàn)有矩陣乘法算法的知識(shí)。通過(guò)學(xué)習(xí)，AlphaTensor 隨時(shí)間逐漸地改進(jìn)，重新發(fā)現(xiàn)了歷史上的快速矩陣算法（如 Strassen 算法），并且發(fā)現(xiàn)算法的速度比以往已知的要快。

AlphaTensor 玩的單人游戲，目標(biāo)是找到正確的矩陣乘法算法。游戲狀態(tài)是一個(gè)由數(shù)字組成的立方數(shù)組（灰色表示 0，藍(lán)色表示 1，綠色表示 - 1），它代表了要完成的剩余工作。

舉例而言，如果學(xué)校里教的傳統(tǒng)算法可以使用 100 次乘法完成 4x5 與 5x5 矩陣相乘，通過(guò)人類(lèi)的聰明才智可以將這一數(shù)字降至 80 次。與之相比，AlphaTensor 發(fā)現(xiàn)的算法只需使用 76 次乘法即可完成相同的運(yùn)算，如下圖所示。

除了上述例子之外，AlphaTensor 發(fā)現(xiàn)的算法還首次在一個(gè)有限域中改進(jìn)了 Strassen 的二階算法。這些用于小矩陣相乘的算法可以當(dāng)做原語(yǔ)來(lái)乘以任意大小的更大矩陣。

AlphaTensor 還發(fā)現(xiàn)了具有 SOTA 復(fù)雜性的多樣化算法集，其中每種大小的矩陣乘法算法多達(dá)數(shù)千，表明矩陣乘法算法的空間比以前想象的要豐富。

在這個(gè)豐富空間中的算法具有不同的數(shù)學(xué)和實(shí)用屬性。利用這種多樣性，DeepMind 對(duì) AlphaTensor 進(jìn)行了調(diào)整，以專(zhuān)門(mén)發(fā)現(xiàn)在給定硬件（如 Nvidia V100 GPU、Google TPU v2）上運(yùn)行速度快的算法。這些算法在相同硬件上進(jìn)行大矩陣相乘的速度比常用算法快了 10-20%，表明了 AlphaTensor 在優(yōu)化任意目標(biāo)方面具備了靈活性。

AlphaTensor 具有一個(gè)對(duì)應(yīng)于算法運(yùn)行時(shí)的目標(biāo)。當(dāng)發(fā)現(xiàn)正確的矩陣乘法算法時(shí)，它會(huì)在指定硬件上進(jìn)行基準(zhǔn)測(cè)試，然后反饋給 AlphaTensor，以便在指定硬件上學(xué)習(xí)更高效的算法。

對(duì)未來(lái)研究和應(yīng)用的影響

從數(shù)學(xué)的角度來(lái)看，對(duì)于旨在確定解決計(jì)算問(wèn)題的最快算法的復(fù)雜性理論而言，DeepMind 的結(jié)果可以指導(dǎo)它的進(jìn)一步研究。通過(guò)較以往方法更高效地探索可能的算法空間，AlphaTensor 有助于加深我們對(duì)矩陣乘法算法豐富性的理解。

此外，由于矩陣乘法是計(jì)算機(jī)圖形學(xué)、數(shù)字通信、神經(jīng)網(wǎng)絡(luò)訓(xùn)練和科學(xué)計(jì)算等很多計(jì)算任務(wù)的核心組成部分，AlphaTensor 發(fā)現(xiàn)的算法可以顯著提升這些領(lǐng)域的計(jì)算效率。

雖然本文只專(zhuān)注于矩陣乘法這一特定問(wèn)題，但 DeepMind 希望能夠啟發(fā)更多的人使用 AI 來(lái)指導(dǎo)其他基礎(chǔ)計(jì)算任務(wù)的算法發(fā)現(xiàn)。并且，DeepMind 的研究還表明，AlphaZero 這種強(qiáng)大的算法遠(yuǎn)遠(yuǎn)超出了傳統(tǒng)游戲的領(lǐng)域，可以幫助解決數(shù)學(xué)領(lǐng)域的開(kāi)放問(wèn)題。

未來(lái)，DeepMind 希望基于他們的研究，更多地將人工智能用來(lái)幫助社會(huì)解決數(shù)學(xué)和科學(xué)領(lǐng)域的一些最重要的挑戰(zhàn)。

參考鏈接：

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

https://github.com/deepmind/alphatensor

https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor