天天快资讯：DeepMind 新作 AlphaDev ---- 强化学习探索更优排序算法

AlphaDev 将 CPU 汇编指令层面的算法优化过程转化为一个单玩家的游戏。

游戏每一步的状态定义为 : St = 。

其中， Pt表示游戏到至今为止所生成的算法，Zt则表示在给定输入的前提下执行完 Pt里的指令之后，内存和寄存器的状态。

如上图所示，在时间步 t，AlphaDev 接受到当前状态 St和 所要执行的动作 at（比如 mov），也就是往当前生成的算法 Pt中添加的合法汇编指令。

在添加完指令之后，就是计算奖励分数 rt（包括评估算法的正确性和延迟）。

正确性评估就是将 N组测试序列输入到算法 Pt中，得到N组输出，和正确的排序结果最比较来计算奖励分数。

论文中给出了3种正确性评估函数，首先定义 P为输入序列长度， PCt为在时间步 t序列中，位置正确的值的个数，这里我理解应该是和正确的排序结果逐个位置对比，统计相等的个数。

三个函数分别定义如下：

论文中提到采用第三个函数效果最好。

延迟分数的计算可以是：

整个强化学习的游戏在执行有限步骤之后就会被终止。只有生成正确而又低延迟的汇编代码才算赢得游戏。而不管是生成了错误的代码还是正确但低效的实现都视为游戏输了。

AlphaDev 采用的强化学习算法是对 AlphqaZero 算法的扩展，也是采用深度神经网络来引导蒙特卡洛树搜索（MCTS）的规划过程。网络模型的输入是 St，输出是对动作策略和奖励的预测。

整个游戏过程简单来说就是，用一个固定参数的网络模型，通过给定的当前状态执行一个蒙特卡洛树搜索过程，然后采取下一步动作。然后可以用生成的游戏过程（包含每一步的状态和奖励）去训练和更新网络的参数。

模型包含两部分：

两个网络的输出 embedding 会合并在一起来表示当前的状态。

网络模型整体的结构如下：

Transformer 编码器模块具体图示

如上图所示，把当前生成的汇编代码序列的每一条指令的操作码和操作数都转换为 one-hot 编码序列，然后输入到网络中。

但是具体的 one-hot 编码规则、词表怎么设置、还有对于 CPU 状态编码网络寄存器和内存的状态是怎么表示为网络的输入的等等，这些细节我在论文里没找到。

然后两个网络的输出 embedding 会合并到一起接着输入到几个函数头里计算，分别是预测下一步策略的函数头，预测算法正确性的函数头和预测算法真实延迟的函数头。

网络参数超参设置

论文的补充资料中提供了网络的参数和三个函数头的具体配置。

而对于策略的预测，论文中提到为了简化问题和提高收敛性，而对动作空间做了一些限制，规则如下：

训练细节

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这个计数器结果， 其计数值表示在一个特定时间段内，处理器内核的时钟周期数。这个值越高，意味着处理器内核在该时间段内执行了更多的指令。

对于定长序列排序，当应用到排序网络算法[6]（sorting network algorithm）的时候 AlphaDev 生成的代码中包含了一些有趣指令序列，相对于原始指令序列可以减少一条汇编指令，论文中称之为：

啥是排序网络算法？

排序网络算法（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]。

标签：