Nature封面:DeepMind推出AlphaTensor,用AI發現矩陣乘法算法
DeepMind 的 Alpha 系列 AI 智能體家族又多了一個成員——AlphaTensor,這次是用來發現算法。
數千年來,算法一直在幫助數學家們進行基本運算。早在很久之前,古埃及人就發明了一種不需要乘法表就能將兩個數字相乘的算法。希臘數學家歐幾里得描述了一種計算最大公約數的算法,這種算法至今仍在使用。在伊斯蘭的黃金時代,波斯數學家 Muhammad ibn Musa al-Khwarizmi 設計了一種求解線性方程和二次方程的新算法,這些算法都對后來的研究產生了深遠的影響。
事實上,算法一詞的出現,有這樣一種說法:波斯數學家 Muhammad ibn Musa al-Khwarizmi 名字中的 al-Khwarizmi 一詞翻譯為拉丁語為 Algoritmi 的意思,從而引出了算法一詞。不過,雖然今天我們對算法很熟悉,可以從課堂中學習、在科研領域也經常遇到,似乎整個社會都在使用算法,然而發現新算法的過程是非常困難的。
現在,DeepMind 用 AI 來發現新算法。
在最新一期 Nature 封面論文《Discovering faster matrix multiplication algorithms with reinforcement learning》中,DeepMind 提出了 AlphaTensor,并表示它是第一個可用于為矩陣乘法等基本任務發現新穎、高效且可證明正確的算法的人工智能系統。簡單來說,使用 AlphaTensor 能夠發現新算法。這項研究揭示了 50 年來在數學領域一個懸而未決的問題,即找到兩個矩陣相乘最快方法。
AlphaTensor 建立在 AlphaZero 的基礎上,而 AlphaZero 是一種在國際象棋、圍棋和將棋等棋盤游戲中可以打敗人類的智能體。這項工作展示了 AlphaZero 從用于游戲到首次用于解決未解決的數學問題的一次轉變。
矩陣乘法
矩陣乘法是代數中最簡單的運算之一,通常在高中數學課上教授。但在課堂之外,這種不起眼的數學運算在當代數字世界中產生了巨大的影響,在現代計算中無處不在。
兩個 3x3 矩陣相乘的例子。
你可能沒注意到,我們生活中處處隱藏著矩陣相乘,如智能手機中的圖像處理、識別語音命令、為電腦游戲生成圖形等都有它在背后進行運算。遍布世界各地的公司都愿意花費大量的時間和金錢開發計算硬件以有效地解決矩陣相乘。因此,即使是對矩陣乘法效率的微小改進也會產生廣泛的影響。
幾個世紀以來,數學家認為標準矩陣乘法算法是效率最高的算法。但在 1969 年,德國數學家 Volken Strassen 通過證明確實存在更好的算法,這一研究震驚了整個數學界。
標準算法與 Strassen 算法對比,后者少進行了一次乘法運算,為 7 次,而前者需要 8 次,整體效率大幅提高。
通過研究非常小的矩陣(大小為 2x2),Strassen 發現了一種巧妙的方法來組合矩陣的項以產生更快的算法。之后數十年,研究者都在研究更大的矩陣,甚至找到 3x3 矩陣相乘的高效方法,都還沒有解決。
DeepMind 的最新研究探討了現代 AI 技術如何推動新矩陣乘法算法的自動發現。基于人類直覺(human intuition)的進步,對于更大的矩陣來說,AlphaTensor 發現的算法比許多 SOTA 方法更有效。該研究表明 AI 設計的算法優于人類設計的算法,這是算法發現領域向前邁出的重要一步。
算法發現自動化的過程和進展
首先將發現矩陣乘法高效算法的問題轉換為單人游戲。其中,board 是一個三維度張量(數字數組),用于捕捉當前算法的正確程度。通過一組與算法指令相對應的所允許的移動,玩家嘗試修改張量并將其條目歸零。
當玩家設法這樣做時,將為任何一對矩陣生成可證明是正確的矩陣乘法算法,并且其效率由將張量清零所采取的步驟數來衡量。
這個游戲非常具有挑戰性,要考慮的可能算法的數量遠遠大于宇宙中原子的數量,即使對于矩陣乘法這樣小的情況也是如此。與幾十年來一直是人工智能挑戰的圍棋游戲相比,該游戲每一步可能的移動數量要多 30 個數量級(DeepMind 考慮的一種設置是 10^33 以上。)
為了解決這個與傳統游戲明顯不同的領域所面臨的挑戰,DeepMind 開發了多個關鍵組件,包括一個結合特定問題歸納偏置的全新神經網絡架構、一個生成有用合成數據的程序以及一種利用問題對稱性的方法。
接著,DeepMind 訓練了一個利用強化學習的智能體 AlphaTensor 來玩這個游戲,該智能體在開始時沒有任何現有矩陣乘法算法的知識。通過學習,AlphaTensor 隨時間逐漸地改進,重新發現了歷史上的快速矩陣算法(如 Strassen 算法),并且發現算法的速度比以往已知的要快。
AlphaTensor 玩的單人游戲,目標是找到正確的矩陣乘法算法。游戲狀態是一個由數字組成的立方數組(灰色表示 0,藍色表示 1,綠色表示 - 1),它代表了要完成的剩余工作。
舉例而言,如果學校里教的傳統算法可以使用 100 次乘法完成 4x5 與 5x5 矩陣相乘,通過人類的聰明才智可以將這一數字降至 80 次。與之相比,AlphaTensor 發現的算法只需使用 76 次乘法即可完成相同的運算,如下圖所示。
除了上述例子之外,AlphaTensor 發現的算法還首次在一個有限域中改進了 Strassen 的二階算法。這些用于小矩陣相乘的算法可以當做原語來乘以任意大小的更大矩陣。
AlphaTensor 還發現了具有 SOTA 復雜性的多樣化算法集,其中每種大小的矩陣乘法算法多達數千,表明矩陣乘法算法的空間比以前想象的要豐富。
在這個豐富空間中的算法具有不同的數學和實用屬性。利用這種多樣性,DeepMind 對 AlphaTensor 進行了調整,以專門發現在給定硬件(如 Nvidia V100 GPU、Google TPU v2)上運行速度快的算法。這些算法在相同硬件上進行大矩陣相乘的速度比常用算法快了 10-20%,表明了 AlphaTensor 在優化任意目標方面具備了靈活性。
AlphaTensor 具有一個對應于算法運行時的目標。當發現正確的矩陣乘法算法時,它會在指定硬件上進行基準測試,然后反饋給 AlphaTensor,以便在指定硬件上學習更高效的算法。
對未來研究和應用的影響
從數學的角度來看,對于旨在確定解決計算問題的最快算法的復雜性理論而言,DeepMind 的結果可以指導它的進一步研究。通過較以往方法更高效地探索可能的算法空間,AlphaTensor 有助于加深我們對矩陣乘法算法豐富性的理解。
此外,由于矩陣乘法是計算機圖形學、數字通信、神經網絡訓練和科學計算等很多計算任務的核心組成部分,AlphaTensor 發現的算法可以顯著提升這些領域的計算效率。
雖然本文只專注于矩陣乘法這一特定問題,但 DeepMind 希望能夠啟發更多的人使用 AI 來指導其他基礎計算任務的算法發現。并且,DeepMind 的研究還表明,AlphaZero 這種強大的算法遠遠超出了傳統游戲的領域,可以幫助解決數學領域的開放問題。
未來,DeepMind 希望基于他們的研究,更多地將人工智能用來幫助社會解決數學和科學領域的一些最重要的挑戰。
參考鏈接:
https://www.nature.com/articles/s41586-022-05172-4
https://github.com/deepmind/alphatensor
https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor
*博客內容為網友個人發布,僅代表博主個人觀點,如有侵權請聯系工作人員刪除。