DeepMind 新作 AlphaDev ---- 強化學(xué)習(xí)探索更優(yōu)排序算法
前言
DeepMind 最近在 Nature 發(fā)表了一篇論文 AlphaDev[2, 3]，一個利用強化學(xué)習(xí)來探索更優(yōu)排序算法的AI系統(tǒng)。
AlphaDev 系統(tǒng)直接從 CPU 匯編指令的層面入手去探索更優(yōu)的排序算法，因為相對于高級編程語言來說，在匯編指令層級對存儲和寄存器的操作可以更加的靈活，所以能發(fā)現(xiàn)更多潛在的調(diào)優(yōu)策略。
在 AlphaDev 的論文中，只關(guān)注探索短序列排序：
定長序列排序（比如 sort3 算法只能對長度為3的序列進(jìn)行排序） 變長序列排序（比如 ?variable sort5 算法可以對長度為1~5的變長序列進(jìn)行排序）
而對于長序列的排序，可以被分解為短序列的排序。
DeepMind 通過 AlphaDev 發(fā)現(xiàn)了比目前人工調(diào)優(yōu)算法更優(yōu)的定長短序列排序算法 sort3，sort4 和 sort5 ，并且已經(jīng)將代碼提交到了 LLVM 標(biāo)準(zhǔn) C++ 庫[4] 。
簡單來說，AlphaDev 將探索更高效排序算法的過程，建模為一個單玩家的匯編游戲（single-player game， AssemblyGame）。
游戲的過程就是玩家從 CPU 匯編指令集合中，選取一系列的指令組合得到一個新的排序算法。不過這個過程是非常有挑戰(zhàn)的，玩家需要考慮，匯編指令的組合空間并最終得得到一個正確和高效的算法。
該游戲主要包括以下難點：
匯編游戲的搜索空間和圍棋類似（10^700） 只要有一條指令沒弄對，可能就會導(dǎo)致整個算法錯誤
AlphaDev 系統(tǒng)詳解
將排序算法表示為 CPU 匯編指令
首先來看一個簡單的變長（variable sort2）短排序函數(shù)的 C 代碼實現(xiàn)，排序結(jié)果從小到大：
void?variable_sort_2(int?length,?int?*a)?{
????switch?(length)?{
????case?0:
????case?1:
????????return;
????case?2:
????????int?tmp?=?a[0];
????????//?a[0]?保存兩者之間的最小值
????????a[0]?=?(a[1]?0])???a[1]?:?a[0];
????????//?a[1]?保存兩者之間的最大值
????????a[1]?=?(a[1]?1];
????????return;
????}
}
通過 gcc 生成對應(yīng)的匯編代碼，我用的 gcc 版本是 11.3.0，命令 gcc -S -O1 -o sort2.s sort2.c
匯編代碼只保留了核心部分，生成的結(jié)果和論文中的示例有些許不同但是原理是一致的：
variable_sort_2:
.LFB0:
; %edi 寄存器保存參數(shù) length 的值
; cmpl 指令對比 %edi 和 常量 2
cmpl $2, %edi
; 相等就跳轉(zhuǎn)到 .L3 標(biāo)簽處，
; 對應(yīng) C 代碼的 case 2
je .L3
.L1:
; 不等于 2 就直接返回，
; 對應(yīng) C 代碼 case 0 和 1
ret
.L3:
; 將 a[0] 賦值給寄存器 %edx
movl (%rsi), %edx
; 將 a[1] 賦值給寄存器 %eax
movl 4(%rsi), %eax
; 對比 %edx 和 %eax
cmpl %edx, %eax
; 將 %edx 賦值給 %ecx
movl %edx, %ecx
; cmov 是條件移動指令根據(jù) cmpl
; 指令的結(jié)果判斷是否執(zhí)行
; 如果 %eax <= %edx
; 則將 %eax 賦值給 %ecx
cmovle %eax, %ecx
; 此時 %ecx 保存了最小值
; 將 %ecx 賦值給 a[0]
movl %ecx, (%rsi)
; 如果 %eax 小于 %edx
; 則將 %edx 賦值給 %eax
cmovl %edx, %eax
; 此時 %eax 保存了最大值
; 將 %eax 賦值給 a[1]
movl %eax, 4(%rsi)
jmp .L1
一般來說匯編程序所做的事情基本都是，將內(nèi)存的值復(fù)制到寄存器，然后對寄存器的值作修改，再將寄存器的值寫回到內(nèi)存中。
而 AlphaDev 系統(tǒng)只關(guān)注 x86 處理器架構(gòu)所支持的匯編指令集合的一個子集。
每條匯編指令的格式均為：操作碼<操作數(shù)A, 操作數(shù)B> 比如：
mov? ? 移動指令，表示將 A 的值賦值給 Bcmp? ? 比較指令，相當(dāng)于 執(zhí)行A - B操作，但是不會對 A 和 B 做修改，而是根據(jù)相減的結(jié)果設(shè)置特殊的 flag 寄存器，更多內(nèi)容可以參考[5]cmovX條件移動指令，根據(jù)X和 flag 寄存器的值判斷是否執(zhí)行將 A 賦值給 B 的操作，一般都是出現(xiàn)在cmp指令之后。X可以是L（是否滿足小于條件），G（是否滿足大于條件），LE（是否滿足小于或等于條件），GE（是否滿足大于等于條件）。jX? ? ? ? ? ?條件跳轉(zhuǎn)指令，根據(jù)X和 flag 寄存器的值判斷是否執(zhí)行跳轉(zhuǎn)到指定標(biāo)記位置操作，A 可以是匯編程序代碼中的標(biāo)記位置，如上面所示匯編代碼的.L1和.L3。X可以是NE（是否不等于），E（是否等于）或者可以填表示無條件跳轉(zhuǎn)。
將探索更優(yōu)排序算法表示為強化學(xué)習(xí)問題
AlphaDev 將 CPU 匯編指令層面的算法優(yōu)化過程轉(zhuǎn)化為一個單玩家的游戲。
游戲每一步的狀態(tài)定義為 : St = 。
其中， Pt 表示游戲到至今為止所生成的算法，Zt 則表示在給定輸入的前提下執(zhí)行完 Pt 里的指令之后，內(nèi)存和寄存器的狀態(tài)。
如上圖所示，在時間步 t ，AlphaDev 接受到當(dāng)前狀態(tài) St 和 所要執(zhí)行的動作 at （比如 mov ），也就是往當(dāng)前生成的算法 Pt 中添加的合法匯編指令。
在添加完指令之后，就是計算獎勵分?jǐn)?shù) rt （包括評估算法的正確性和延遲）。
算法正確性評估
正確性評估就是將 N 組測試序列輸入到算法 Pt 中，得到N 組輸出，和正確的排序結(jié)果最比較來計算獎勵分?jǐn)?shù)。
論文中給出了3種正確性評估函數(shù)，首先定義 P 為輸入序列長度， PCt 為在時間步 t 序列中，位置正確的值的個數(shù)，這里我理解應(yīng)該是和正確的排序結(jié)果逐個位置對比，統(tǒng)計相等的個數(shù)。
三個函數(shù)分別定義如下：
func1 = (P - PCt) / Pfunc2 = sqrt(func1)func3 = sqrt(PCt)
論文中提到采用第三個函數(shù)效果最好。
延遲評估
延遲分?jǐn)?shù)的計算可以是：
對系統(tǒng)增加代碼長度計算懲罰，因為代碼的長度一般都是和耗時高度相關(guān) 直接計算算法的真實耗時
整個強化學(xué)習(xí)的游戲在執(zhí)行有限步驟之后就會被終止。只有生成正確而又低延遲的匯編代碼才算贏得游戲。而不管是生成了錯誤的代碼還是正確但低效的實現(xiàn)都視為游戲輸了。
AlphaDev 采用的強化學(xué)習(xí)算法是對 AlphqaZero 算法的擴(kuò)展，也是采用深度神經(jīng)網(wǎng)絡(luò)來引導(dǎo)蒙特卡洛樹搜索（MCTS）的規(guī)劃過程。網(wǎng)絡(luò)模型的輸入是 St ，輸出是對動作策略和獎勵的預(yù)測。
整個游戲過程簡單來說就是，用一個固定參數(shù)的網(wǎng)絡(luò)模型，通過給定的當(dāng)前狀態(tài)執(zhí)行一個蒙特卡洛樹搜索過程，然后采取下一步動作。然后可以用生成的游戲過程（包含每一步的狀態(tài)和獎勵）去訓(xùn)練和更新網(wǎng)絡(luò)的參數(shù)。
網(wǎng)絡(luò)模型結(jié)構(gòu)
模型包含兩部分：
一個 Transformer 編碼器模塊，用于建模算法，輸入是至今為止生成的匯編指令序列 一個 CPU 狀態(tài)編碼器 MLP 模塊，輸入當(dāng)前寄存器和內(nèi)存的狀態(tài)
兩個網(wǎng)絡(luò)的輸出 embedding 會合并在一起來表示當(dāng)前的狀態(tài)。
網(wǎng)絡(luò)模型整體的結(jié)構(gòu)如下：
Transformer 編碼器模塊具體圖示
如上圖所示，把當(dāng)前生成的匯編代碼序列的每一條指令的操作碼和操作數(shù)都轉(zhuǎn)換為 one-hot 編碼序列，然后輸入到網(wǎng)絡(luò)中。
但是具體的 one-hot 編碼規(guī)則、詞表怎么設(shè)置、還有對于 CPU 狀態(tài)編碼網(wǎng)絡(luò)寄存器和內(nèi)存的狀態(tài)是怎么表示為網(wǎng)絡(luò)的輸入的等等，這些細(xì)節(jié)我在論文里沒找到。
然后兩個網(wǎng)絡(luò)的輸出 embedding 會合并到一起接著輸入到幾個函數(shù)頭里計算，分別是預(yù)測下一步策略的函數(shù)頭，預(yù)測算法正確性的函數(shù)頭和預(yù)測算法真實延遲的函數(shù)頭。
網(wǎng)絡(luò)參數(shù)超參設(shè)置
論文的補充資料中提供了網(wǎng)絡(luò)的參數(shù)和三個函數(shù)頭的具體配置。
而對于策略的預(yù)測，論文中提到為了簡化問題和提高收斂性，而對動作空間做了一些限制，規(guī)則如下：
必須按照升序方式讀取內(nèi)存 寄存器按照升序分配 cmp和cmovX指令的操作數(shù)不能出現(xiàn)內(nèi)存地址對每個內(nèi)存位置，只能讀取和寫入一次 每個寄存器在使用之前，必須初始化 不能連續(xù)調(diào)用 cmp指令
訓(xùn)練細(xì)節(jié)
AlphaDev 的訓(xùn)練采用了 TPU v3，每個 TPU 核的 batch size 是 1024 ，總共用了 16 個 TPU 核，總共訓(xùn)練了 100 萬次迭代。而在對于玩游戲積累訓(xùn)練數(shù)據(jù)來說，則是在 TPU v4 上進(jìn)行，總共用了 512 個 TPU 核。
實驗結(jié)果表明，最多只需2天模型就能訓(xùn)收斂。
實驗結(jié)果
生成的算法和人工調(diào)優(yōu)對比
從實驗結(jié)果表格可以看到，對于短序列排序算法 AlphaDev 生成的代碼長度更短，而且平均耗時也更低。
對生成算法延遲的評估方式，比如對于 sort3 則是在 100 臺機器上做評估，每臺機器隨機生成 1000 條 3個數(shù)的序列，然后每條序列輸入到算法中，對這 1000 次評估取第5百分位數(shù)作為最終的評估結(jié)果（排除 cache miss 和 任務(wù)搶占 等因素）。
耗時采用的是 CPU_CLK_UNHALTED.CORE 這個計數(shù)器結(jié)果， 其計數(shù)值表示在一個特定時間段內(nèi)，處理器內(nèi)核的時鐘周期數(shù)。這個值越高，意味著處理器內(nèi)核在該時間段內(nèi)執(zhí)行了更多的指令。
AlphaDev 發(fā)現(xiàn)新的算法
對于定長序列排序，當(dāng)應(yīng)用到排序網(wǎng)絡(luò)算法[6]（sorting network algorithm）的時候 AlphaDev 生成的代碼中包含了一些有趣指令序列，相對于原始指令序列可以減少一條匯編指令，論文中稱之為：
AlphaDev swap move AlphaDev copy move
啥是排序網(wǎng)絡(luò)算法？
排序網(wǎng)絡(luò)算法（Sorting Network Algorithm）是一種能夠?qū)σ唤M輸入數(shù)據(jù)進(jìn)行排序的并行算法，其具有較好的并行性能適用于多處理器或多核心系統(tǒng)。
該算法的特點是，它將所有的比較和交換操作預(yù)先規(guī)劃好形成一個固定的結(jié)構(gòu)，然后將輸入數(shù)據(jù)按照這個結(jié)構(gòu)進(jìn)行排序。
排序網(wǎng)絡(luò)由比較器（comparator）和線（wire）組成，如下圖所示：
水平線表示 wire，每條水平線持有一個待排序的值。兩條 wire 之間的垂直線段就表示一個比較器，比較器對比兩條水平線的值，如果比較器下方的值小于上方的值則交換兩條橫線的值，否則則不交換。
一個優(yōu)化過的排序網(wǎng)絡(luò)可以以最少的比較器，并將這些比較器放置在特定位置上，來實現(xiàn)對任意序列進(jìn)行排序。
下圖是對一個構(gòu)造好的排序網(wǎng)絡(luò)，輸入真實待排序序列的例子：
可見初始輸入是 [2, 3, 1, 4]，這些隨機數(shù)從左到右按順序經(jīng)過這些比較器之后，就得到了排序好的序列 [1, 2, 3, 4]。
AlphaDev swap move
先來看這個排序網(wǎng)絡(luò)，只看紅圈部分的功能就是對給定的輸入 [A, B, C] 將其轉(zhuǎn)換為 [min(A,B,C), max(min(A,C),B), max(A,C)]。
然后經(jīng)過 AlphaDev 優(yōu)化之后，可以將第一個輸出的 min(A,B,C) 改為只計算 min(A,B)，原因是因為前面的 B 和 C橫線之間經(jīng)過比較器之后已經(jīng)有了前置條件 B <= C。
而通過這個優(yōu)化就能省去一條匯編指令，下圖是紅圈部分的偽代碼實現(xiàn)：
左邊是原始偽代碼實現(xiàn)，右邊是經(jīng)過 AlphaDev 優(yōu)化之后的實現(xiàn)，可以看到少了一條匯編指令 mov S P。
AlphaDev copy move
接下來看對4個元素進(jìn)行排序的排序網(wǎng)絡(luò)，是在對 sort8 這個算法優(yōu)化過程中發(fā)現(xiàn)的。該排序網(wǎng)絡(luò)對于輸入序列 [A, B, C, D] 轉(zhuǎn)換為 [min(A, B, C, D), max(B, min(A, C, D), max(C, min(A, D)), max(A, D) ]。
該排序網(wǎng)絡(luò)是 sort8 的一個子排序網(wǎng)絡(luò)，而根據(jù)比較器的放置位置來看，A 和 D 比較之后后續(xù)就不再和其他元素比較了，所以D出來的結(jié)果就是四個元素中最大的，所以隱含了一個條件就是 D >= min(A, C)。
因此對第二個輸出元素的計算可以從 max(B, min(A, C, D)) 改為 max(B, min(A, C))，就可以節(jié)省一條匯編指令。
偽代碼如下:
左邊是原始偽代碼實現(xiàn)，右邊是經(jīng)過 AlphaDev 優(yōu)化之后的實現(xiàn)，可以看到少了一條匯編指令 mov P T。
總結(jié)
這篇文章只是對 AlphaDev 論文中的主要內(nèi)容作解讀，對于更多的內(nèi)容和細(xì)節(jié)感興趣的讀者可以查閱原論文和論文的補充資料 [2,3]，DeepMind 也也開源了一份偽代碼實現(xiàn) [7]。
參考資料
[1] https://ee.usc.edu/~redekopp/cs356/slides/CS356Unit5_x86_Control [2] https://www.nature.com/articles/s41586-023-06004-9#MOESM1 [3] https://static-content.springer.com/esm/art%3A10.1038%2Fs41586-023-06004-9/MediaObjects/41586_2023_6004_MOESM1_ESM.pdf [4] ? D118029 Introduce branchless sorting functions for sort3, sort4 and sort5. (llvm.org) [5] 小信豬的原始部落: PC Assembly Language 學(xué)習(xí)筆記(5) - Control Structures (godleon.blogspot.com) [6] https://en.wikipedia.org/wiki/Sorting_network#:~:text=as%20the%20contrapositive.-,Constructing%20sorting%20networks,are%20often%20used%20in%20practice. [7] https://github.com/deepmind/alphadev