谷歌DeepMind發(fā)現(xiàn)更快排序算法，已集成到C++庫

接觸過基礎(chǔ)計(jì)算機(jī)科學(xué)課程的朋友們，肯定都曾親自動(dòng)手設(shè)計(jì)排序算法——也就是借助代碼將無序列表中的各個(gè)條目按升序或降序方式重新排列。這是個(gè)有趣的挑戰(zhàn)，可行的操作方法也多種多樣。人們?cè)度氪罅繒r(shí)間探索如何更高效地完成排序任務(wù)。

作為一項(xiàng)基礎(chǔ)操作，大多數(shù)編程語言的標(biāo)準(zhǔn)庫中都內(nèi)置有排序算法。世界各地的代碼庫中使用了許多不同的排序技術(shù)和算法來在線組織大量數(shù)據(jù)，但至少就與 LLVM 編譯器配套使用的 C++ 庫而言，排序代碼已經(jīng)有十多年沒有任何變化了。

近日，谷歌 DeepMind AI 小組如今開發(fā)出一種強(qiáng)化學(xué)習(xí)工具 AlphaDev，能夠在無需通過人類代碼示例做預(yù)訓(xùn)練的情況下，開發(fā)出極限優(yōu)化的算法。如今，這些算法已經(jīng)集成到 LLVM 標(biāo)準(zhǔn) C++ 排序庫中，這是十多年來排序庫部分第一次發(fā)生變化，也是第一次將通過強(qiáng)化學(xué)習(xí)設(shè)計(jì)的算法添加到該庫中。

把編程過程視為“游戲”

由于不必預(yù)先接觸任何人類游戲策略，DeepMind 系統(tǒng)往往能發(fā)現(xiàn)人類從未設(shè)想過的攻關(guān)思路。當(dāng)然，由于完全依靠自我對(duì)抗來學(xué)習(xí)經(jīng)驗(yàn)，DeepMind 在某些情況下也會(huì)形成可被人類利用的盲點(diǎn)。

這種方法跟編程其實(shí)非常相似。大語言模型之所以能夠編寫出有效代碼，就是因?yàn)樗鼈兛吹竭^大量人類代碼示例。但也正因?yàn)槿绱?，語言模型很難產(chǎn)出人類之前沒做過的東西。如果我們希望對(duì)普遍存在的現(xiàn)有算法（例如排序函數(shù)）做進(jìn)一步優(yōu)化，那么繼續(xù)依賴現(xiàn)有人類代碼將很難突破固有思路的束縛。那么，如何才能讓 AI 找到真正的新方向？

DeepMind 的研究人員采用了與國際象棋和圍棋相同的方法：把代碼優(yōu)化任務(wù)轉(zhuǎn)化成單人“組裝游戲”。AlphaDev 系統(tǒng)開發(fā)出一種 x86 匯編算法，會(huì)將代碼的運(yùn)行延遲視為一個(gè)分?jǐn)?shù)，在努力將該分?jǐn)?shù)最小化的同時(shí)確保代碼能夠順暢跑通。通過強(qiáng)化學(xué)習(xí)，AlphaDev 逐漸具備了編寫緊湊、高效代碼的能力。

AlphaDev 基于 AlphaZero。DeepMind 向來以開發(fā)能自學(xué)游戲規(guī)則的 AI 軟件而聞名。這種思路被證明效果拔群，也先后攻克了國際象棋、圍棋和《星際爭霸》等諸多游戲難題。雖然具體細(xì)節(jié)因所玩游戲而異，但 DeepMind 軟件確實(shí)能在重復(fù)游玩中不斷學(xué)習(xí)，持續(xù)探索能令得分最大化的辦法。

AlphaDev 的兩個(gè)核心組件是學(xué)習(xí)算法和表示函數(shù)。

AlphaDev 學(xué)習(xí)算法可以結(jié)合 DRL 和隨機(jī)搜索優(yōu)化算法來玩組裝游戲。AlphaDev 中的主要學(xué)習(xí)算法是 AlphaZero 33 的擴(kuò)展，AlphaZero 33 是一種著名的 DRL 算法，其中訓(xùn)練神經(jīng)網(wǎng)絡(luò)以指導(dǎo)搜索完成游戲。

表示函數(shù)，負(fù)責(zé)跟蹤代碼開發(fā)時(shí)的整體性能，其中包括算法的常規(guī)結(jié)構(gòu)以及對(duì) x86 寄存器和內(nèi)存的使用。該系統(tǒng)會(huì)單獨(dú)添加匯編指令，通過蒙特卡洛樹搜索（同樣是一種從游戲系統(tǒng)中借用的方法）進(jìn)行選擇。樹狀結(jié)構(gòu)允許系統(tǒng)快速將搜索范圍縮小至包含大量潛在指令的有限區(qū)域，而蒙特卡洛方法則以一定程度的隨機(jī)性從這個(gè)分支區(qū)域內(nèi)選擇具體指令。（請(qǐng)注意，這里所說的“指令”是為創(chuàng)建有效、完整程序集而選擇特定寄存器等操作。）

之后，系統(tǒng)會(huì)評(píng)估匯編代碼的延遲和有效性狀態(tài)，為其打分并與前一次得分進(jìn)行比較。而通過強(qiáng)化學(xué)習(xí)，系統(tǒng)會(huì)在給定的程序狀態(tài)之下保留樹結(jié)構(gòu)中不同分支的工作信息。隨著時(shí)間推移，系統(tǒng)將逐漸“學(xué)會(huì)”如何以最高得分（代表最低延遲）獲得游戲勝利（成功完成排序）。AlphaDev 的主要表示函數(shù)基于 Transformer。

為了訓(xùn)練 AlphaDev 發(fā)現(xiàn)新算法，AlphaDev 在每輪中都會(huì)觀察它生成的算法和中央處理器 (CPU) 中包含的信息，然后通過選擇要添加到算法中的指令完成游戲。AlphaDev 必須有效地搜索大量可能的指令組合，以找到可以排序的算法，并且還要比當(dāng)前最好的算法更快，同時(shí)代理模型可以根據(jù)算法的正確性和延遲獲得獎(jiǎng)勵(lì)。

圖 A：組裝游戲，圖 B：獎(jiǎng)勵(lì)計(jì)算

最終，AlphaDev 發(fā)現(xiàn)了新的排序算法，這些算法可以讓 LLVM libc++ 排序庫得到改進(jìn)：對(duì)于較短的序列，排序庫的速度提高了 70%；對(duì)于超過 250,000 個(gè)元素的序列，速度提高了約 1.7%。

具體而言，該算法的創(chuàng)新主要在于兩種指令序列：AlphaDev Swap Move（交換移動(dòng)）和 AlphaDev Copy Move（復(fù)制移動(dòng)），通過這兩個(gè)指令，AlphaDev 跳過了一個(gè)步驟，以一種看似錯(cuò)誤但實(shí)際上是捷徑的方式連接項(xiàng)目。

左圖：帶有 min(A,B,C) 的原始 sort3 實(shí)現(xiàn)。?

右圖：AlphaDev Swap Move - AlphaDev 發(fā)現(xiàn)你只需要 min(A,B)。

左：max (B, min (A, C)) 的原始實(shí)現(xiàn)用于對(duì)八個(gè)元素進(jìn)行排序的更大排序算法。

?右：AlphaDev 發(fā)現(xiàn)在使用其復(fù)制移動(dòng)時(shí)只需要 max (B, min (A, C))。

這套系統(tǒng)的主要優(yōu)勢在于，其訓(xùn)練過程不借助任何代碼示例。相反，系統(tǒng)會(huì)自主生成代碼示例，然后對(duì)其做出評(píng)估。過程當(dāng)中，它也就逐漸掌握了關(guān)于有效排序的指令組合信息。

從排序到散列

在發(fā)現(xiàn)更快的排序算法后，DeepMind 測試了 AlphaDev 是否可以概括和改進(jìn)不同的計(jì)算機(jī)科學(xué)算法：散列。

哈希是計(jì)算中用于檢索、存儲(chǔ)和壓縮數(shù)據(jù)的基本算法。就像使用分類系統(tǒng)來定位某本書的圖書管理員一樣，散列算法可以幫助用戶知道他們正在尋找什么以及在哪里可以找到它。這些算法獲取特定密鑰的數(shù)據(jù)（例如用戶名“Jane Doe”）并對(duì)其進(jìn)行哈希處理——這是一個(gè)將原始數(shù)據(jù)轉(zhuǎn)換為唯一字符串（例如 1234ghfty）的過程。計(jì)算機(jī)使用此散列來快速檢索與密鑰相關(guān)的數(shù)據(jù)，而不是搜索所有數(shù)據(jù)。

DeepMind 將 AlphaDev 應(yīng)用于數(shù)據(jù)結(jié)構(gòu)中最常用的散列算法之一，以嘗試發(fā)現(xiàn)更快的算法。當(dāng)將其應(yīng)用于散列函數(shù)的 9-16 字節(jié)范圍時(shí)，AlphaDev 發(fā)現(xiàn)的算法速度提高了 30%。

今年，AlphaDev 的新哈希算法被發(fā)布到開源 Abseil 庫中，可供全球數(shù)百萬開發(fā)人員使用，該庫現(xiàn)在每天被數(shù)萬億次使用。

實(shí)際可用的代碼

復(fù)雜程序中的排序機(jī)制能夠處理大量任意條目的集合。但在標(biāo)準(zhǔn)庫層面來看，這種能力源自一系列高度限定的具體函數(shù)。這些函數(shù)各自只能處理一種或幾種情況。例如，某些單獨(dú)算法只能對(duì) 3、4 或 5 個(gè)條目做排序。我們也可以同時(shí)使用一組函數(shù)對(duì)任意數(shù)量的條目作排序，但原則上每一次函數(shù)調(diào)用最多只能對(duì) 4 個(gè)條目做排序。

DeepMind 在每個(gè)函數(shù)上都設(shè)置了 AlphaDev，其實(shí)際運(yùn)行方式有著很大區(qū)別。對(duì)于負(fù)責(zé)處理特定數(shù)量條目的函數(shù)，可以編寫出不含任何分支的代碼，即根據(jù)變量狀態(tài)執(zhí)行不同的代碼。因此代碼性能往往與所涉及的指令數(shù)量成反比。

AlphaDev 已經(jīng)成功將 sort-3、sort-5 和 sort-8 的指令數(shù)量各減一，在 sort-6 和 sort-7 中的指令削減量甚至更多。只有 sort-4 上沒能找到改進(jìn)現(xiàn)有代碼的方法。而在實(shí)際系統(tǒng)上的重復(fù)運(yùn)行測試證明，更少的指令確實(shí)帶來了更好的性能。

至于對(duì)可變數(shù)量條目進(jìn)行排序，則要求代碼中包含分支，而不同處理器專用于處理這些分支的元件數(shù)量也有區(qū)別。

對(duì)于這類情況，研究人員在 100 臺(tái)不同的計(jì)算設(shè)備上對(duì)代碼性能做出了評(píng)估。AlphaDev 在這類場景下同樣找到了進(jìn)一步榨取性能的方法，下面我們以一次最多排序 4 個(gè)條目的函數(shù)為例，看看它到底是怎么操作的。

在 C++ 庫的現(xiàn)有實(shí)現(xiàn)中，代碼需要進(jìn)行一系列測試來確認(rèn)具體需要對(duì)多少個(gè)條目做排序，再根據(jù)條目數(shù)量調(diào)用相應(yīng)的排序函數(shù)。

而 AlphaDev 修改后的代碼則采取更加“神奇”的思路：它先測試是不是 2 個(gè)條目，如果是則調(diào)用相應(yīng)函數(shù)立即做排序。如果數(shù)量大于 2 個(gè)，則代碼會(huì)先對(duì)前 3 個(gè)條目做排序。這樣如果確實(shí)只有 3 個(gè)條目，則返回排序結(jié)果。由于實(shí)際是有 4 個(gè)條目要做排序，所以 AlphaDev 會(huì)運(yùn)行專門代碼，以非常高效的方式將第 4 個(gè)條目插入到前 3 個(gè)已經(jīng)排序完成的條目中的適當(dāng)位置。

這種辦法聽起來有點(diǎn)怪異，但事實(shí)證明其性能確實(shí)始終優(yōu)于現(xiàn)有代碼。

由于 AlphaDev 確實(shí)生成了更高效的代碼，所以研究團(tuán)隊(duì)打算把這些成果重新合并到 LLVM 標(biāo)準(zhǔn) C++ 庫中。但問題是這些代碼為匯編格式，而非 C++。所以他們必須通過逆向計(jì)算找到能夠生成相同程序集的 C++ 代碼。

現(xiàn)如今，代碼成果已經(jīng)被合并至 LLVM 工具鏈內(nèi)，成為十多年來這部分代碼的首次更新。研究人員估計(jì)，AlphaDev 生成的新代碼正每天被執(zhí)行數(shù)萬億次。

結(jié)束語

“太棒了！將我們程序員很早就學(xué)會(huì)的這種基本排序任務(wù)的速度提高了 70%?？吹?AI 在我們都依賴的算法和庫中提供重大加速，真是令人興奮?！庇?a target="_blank">開發(fā)者對(duì)谷歌 DeepMind 的成果表示振奮。

但也有開發(fā)者并不買賬：“相當(dāng)令人失望……1.7% 的改善？5 個(gè)元素的序列 70%？可能是最不受歡迎、最不切實(shí)際的應(yīng)用研究……”也有開發(fā)者表示：“說發(fā)現(xiàn)了新算法是不是有點(diǎn)誤導(dǎo)人？似乎更像是算法優(yōu)化。無論如何這仍然很酷?！?/p>