DeepMind新作AlphaDev----強化學習探索更優排序算法

前言

DeepMind 最近在 Nature 發表了一篇論文 AlphaDev[2, 3]，一個利用強化學習來探索更優排序算法的AI系統。

AlphaDev 系統直接從 CPU 匯編指令的層面入手去探索更優的排序算法，因為相對于高級編程語言來說，在匯編指令層級對存儲和寄存器的操作可以更加的靈活，所以能發現更多潛在的調優策略。

在 AlphaDev 的論文中，只關注探索短序列排序：

• 定長序列排序（比如 sort3 算法只能對長度為3的序列進行排序）
• 變長序列排序（比如 variable sort5 算法可以對長度為1~5的變長序列進行排序）

而對于長序列的排序，可以被分解為短序列的排序。

DeepMind 通過 AlphaDev 發現了比目前人工調優算法更優的定長短序列排序算法 sort3，sort4 和 sort5 ，并且已經將代碼提交到了 LLVM 標準 C++ 庫[4] 。

簡單來說，AlphaDev 將探索更高效排序算法的過程，建模為一個單玩家的匯編游戲（single-player game， AssemblyGame）。

游戲的過程就是玩家從 CPU 匯編指令集合中，選取一系列的指令組合得到一個新的排序算法。不過這個過程是非常有挑戰的，玩家需要考慮，匯編指令的組合空間并最終得得到一個正確和高效的算法。

該游戲主要包括以下難點：

• 匯編游戲的搜索空間和圍棋類似（10^700）
• 只要有一條指令沒弄對，可能就會導致整個算法錯誤

AlphaDev 系統詳解

將排序算法表示為 CPU 匯編指令

首先來看一個簡單的變長（variable sort2）短排序函數的 C 代碼實現，排序結果從小到大：

voidvariable_sort_2(intlength,int*a){
switch(length){
case0:
case1:
return;
case2:
inttmp=a[0];
//a[0]保存兩者之間的最小值
a[0]=(a[1]0])?a[1]:a[0];
//a[1]保存兩者之間的最大值
a[1]=(a[1]1];
return;
}
}

通過 gcc 生成對應的匯編代碼，我用的 gcc 版本是 11.3.0，命令 gcc -S -O1 -o sort2.s sort2.c

匯編代碼只保留了核心部分，生成的結果和論文中的示例有些許不同但是原理是一致的：

variable_sort_2:  
.LFB0:
; %edi 寄存器保存參數 length 的值
; cmpl 指令對比 %edi 和 常量 2
cmpl$2, %edi 
; 相等就跳轉到 .L3 標簽處，
        ; 對應 C 代碼的 case 2
je.L3
.L1:
; 不等于 2 就直接返回，
        ; 對應 C 代碼 case 0 和 1
ret 
.L3:
; 將 a[0] 賦值給寄存器 %edx 
movl(%rsi), %edx
; 將 a[1] 賦值給寄存器 %eax 
movl4(%rsi), %eax
; 對比 %edx 和 %eax
cmpl%edx, %eax
; 將 %edx 賦值給 %ecx
movl%edx, %ecx
; cmov 是條件移動指令根據 cmpl 
; 指令的結果判斷是否執行
; 如果 %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

一般來說匯編程序所做的事情基本都是，將內存的值復制到寄存器，然后對寄存器的值作修改，再將寄存器的值寫回到內存中。

而 AlphaDev 系統只關注 x86 處理器架構所支持的匯編指令集合的一個子集。

每條匯編指令的格式均為：操作碼<操作數A, 操作數B> 比如：

• mov 移動指令，表示將 A 的值賦值給 B

• cmp 比較指令，相當于 執行 A - B 操作，但是不會對 A 和 B 做修改，而是根據相減的結果設置特殊的 flag 寄存器，更多內容可以參考[5]

• cmovX 條件移動指令，根據 X 和 flag 寄存器的值判斷是否執行將 A 賦值給 B 的操作，一般都是出現在 cmp 指令之后。X 可以是 L （是否滿足小于條件）， G （是否滿足大于條件），LE （是否滿足小于或等于條件），GE （是否滿足大于等于條件）。

• jX 條件跳轉指令，根據 X 和 flag 寄存器的值判斷是否執行跳轉到指定標記位置操作，A 可以是匯編程序代碼中的標記位置，如上面所示匯編代碼的 .L1 和 .L3。X 可以是 NE （是否不等于），E （是否等于）或者可以填表示無條件跳轉。

將探索更優排序算法表示為強化學習問題

AlphaDev 將 CPU 匯編指令層面的算法優化過程轉化為一個單玩家的游戲。

游戲每一步的狀態定義為 : St = 。

其中， Pt 表示游戲到至今為止所生成的算法，Zt 則表示在給定輸入的前提下執行完 Pt 里的指令之后，內存和寄存器的狀態。

如上圖所示，在時間步 t ，AlphaDev 接受到當前狀態 St 和 所要執行的動作 at （比如 mov ），也就是往當前生成的算法 Pt 中添加的合法匯編指令。

在添加完指令之后，就是計算獎勵分數 rt （包括評估算法的正確性和延遲）。

算法正確性評估

正確性評估就是將 N 組測試序列輸入到算法 Pt 中，得到N 組輸出，和正確的排序結果最比較來計算獎勵分數。

論文中給出了3種正確性評估函數，首先定義 P 為輸入序列長度， PCt 為在時間步 t 序列中，位置正確的值的個數，這里我理解應該是和正確的排序結果逐個位置對比，統計相等的個數。

三個函數分別定義如下：

• func1 = (P - PCt) / P
• func2 = sqrt(func1)
• func3 = sqrt(PCt)

論文中提到采用第三個函數效果最好。

延遲評估

延遲分數的計算可以是：

• 對系統增加代碼長度計算懲罰，因為代碼的長度一般都是和耗時高度相關
• 直接計算算法的真實耗時

整個強化學習的游戲在執行有限步驟之后就會被終止。只有生成正確而又低延遲的匯編代碼才算贏得游戲。而不管是生成了錯誤的代碼還是正確但低效的實現都視為游戲輸了。

AlphaDev 采用的強化學習算法是對 AlphqaZero 算法的擴展，也是采用深度神經網絡來引導蒙特卡洛樹搜索（MCTS）的規劃過程。網絡模型的輸入是 St ，輸出是對動作策略和獎勵的預測。

整個游戲過程簡單來說就是，用一個固定參數的網絡模型，通過給定的當前狀態執行一個蒙特卡洛樹搜索過程，然后采取下一步動作。然后可以用生成的游戲過程（包含每一步的狀態和獎勵）去訓練和更新網絡的參數。

網絡模型結構

模型包含兩部分：

• 一個 Transformer 編碼器模塊，用于建模算法，輸入是至今為止生成的匯編指令序列
• 一個 CPU 狀態編碼器 MLP 模塊，輸入當前寄存器和內存的狀態

兩個網絡的輸出 embedding 會合并在一起來表示當前的狀態。

網絡模型整體的結構如下：

Transformer 編碼器模塊具體圖示

如上圖所示，把當前生成的匯編代碼序列的每一條指令的操作碼和操作數都轉換為 one-hot 編碼序列，然后輸入到網絡中。

但是具體的 one-hot 編碼規則、詞表怎么設置、還有對于 CPU 狀態編碼網絡寄存器和內存的狀態是怎么表示為網絡的輸入的等等，這些細節我在論文里沒找到。

然后兩個網絡的輸出 embedding 會合并到一起接著輸入到幾個函數頭里計算，分別是預測下一步策略的函數頭，預測算法正確性的函數頭和預測算法真實延遲的函數頭。

網絡參數超參設置

論文的補充資料中提供了網絡的參數和三個函數頭的具體配置。

而對于策略的預測，論文中提到為了簡化問題和提高收斂性，而對動作空間做了一些限制，規則如下：

• 必須按照升序方式讀取內存
• 寄存器按照升序分配
• cmp 和 cmovX 指令的操作數不能出現內存地址
• 對每個內存位置，只能讀取和寫入一次
• 每個寄存器在使用之前，必須初始化
• 不能連續調用 cmp 指令

訓練細節

AlphaDev 的訓練采用了 TPU v3，每個 TPU 核的 batch size 是 1024 ，總共用了 16 個 TPU 核，總共訓練了 100 萬次迭代。而在對于玩游戲積累訓練數據來說，則是在 TPU v4 上進行，總共用了 512 個 TPU 核。

實驗結果表明，最多只需2天模型就能訓收斂。

實驗結果

生成的算法和人工調優對比

從實驗結果表格可以看到，對于短序列排序算法 AlphaDev 生成的代碼長度更短，而且平均耗時也更低。

對生成算法延遲的評估方式，比如對于 sort3 則是在 100 臺機器上做評估，每臺機器隨機生成 1000 條 3個數的序列，然后每條序列輸入到算法中，對這 1000 次評估取第5百分位數作為最終的評估結果（排除 cache miss 和 任務搶占 等因素）。

耗時采用的是 CPU_CLK_UNHALTED.CORE 這個計數器結果， 其計數值表示在一個特定時間段內，處理器內核的時鐘周期數。這個值越高，意味著處理器內核在該時間段內執行了更多的指令。

AlphaDev 發現新的算法

對于定長序列排序，當應用到排序網絡算法[6]（sorting network algorithm）的時候 AlphaDev 生成的代碼中包含了一些有趣指令序列，相對于原始指令序列可以減少一條匯編指令，論文中稱之為：

• AlphaDev swap move
• AlphaDev copy move

啥是排序網絡算法？

排序網絡算法（Sorting Network Algorithm）是一種能夠對一組輸入數據進行排序的并行算法，其具有較好的并行性能適用于多處理器或多核心系統。

該算法的特點是，它將所有的比較和交換操作預先規劃好形成一個固定的結構，然后將輸入數據按照這個結構進行排序。

排序網絡由比較器（comparator）和線（wire）組成，如下圖所示：

水平線表示 wire，每條水平線持有一個待排序的值。兩條 wire 之間的垂直線段就表示一個比較器，比較器對比兩條水平線的值，如果比較器下方的值小于上方的值則交換兩條橫線的值，否則則不交換。

一個優化過的排序網絡可以以最少的比較器，并將這些比較器放置在特定位置上，來實現對任意序列進行排序。

下圖是對一個構造好的排序網絡，輸入真實待排序序列的例子：

可見初始輸入是 [2, 3, 1, 4]，這些隨機數從左到右按順序經過這些比較器之后，就得到了排序好的序列 [1, 2, 3, 4]。

AlphaDev swap move

先來看這個排序網絡，只看紅圈部分的功能就是對給定的輸入 [A, B, C] 將其轉換為 [min(A,B,C), max(min(A,C),B), max(A,C)]。

然后經過 AlphaDev 優化之后，可以將第一個輸出的 min(A,B,C) 改為只計算 min(A,B)，原因是因為前面的 B 和 C橫線之間經過比較器之后已經有了前置條件 B <= C。

而通過這個優化就能省去一條匯編指令，下圖是紅圈部分的偽代碼實現：

左邊是原始偽代碼實現，右邊是經過 AlphaDev 優化之后的實現，可以看到少了一條匯編指令 mov S P。

AlphaDev copy move

接下來看對4個元素進行排序的排序網絡，是在對 sort8 這個算法優化過程中發現的。該排序網絡對于輸入序列 [A, B, C, D] 轉換為 [min(A, B, C, D), max(B, min(A, C, D), max(C, min(A, D)), max(A, D) ]。

該排序網絡是 sort8 的一個子排序網絡，而根據比較器的放置位置來看，A 和 D 比較之后后續就不再和其他元素比較了，所以D出來的結果就是四個元素中最大的，所以隱含了一個條件就是 D >= min(A, C)。

因此對第二個輸出元素的計算可以從 max(B, min(A, C, D)) 改為 max(B, min(A, C))，就可以節省一條匯編指令。

偽代碼如下:

左邊是原始偽代碼實現，右邊是經過 AlphaDev 優化之后的實現，可以看到少了一條匯編指令 mov P T。

總結

這篇文章只是對 AlphaDev 論文中的主要內容作解讀，對于更多的內容和細節感興趣的讀者可以查閱原論文和論文的補充資料 [2,3]，DeepMind 也也開源了一份偽代碼實現 [7]。