【翻译】使用深度强化学习发现更快的排序算法

发布时间 2023-06-09 16:27:23作者: Larcvz

Faster sorting algorithms discovered using deep reinforcement learning | Nature
因为是机翻感觉阅读起来有问题的地方建议阅读上面链接的原文。
文中提到的附录部分本文并未收录,请使用原文链接查看。


Faster sorting algorithms discovered using deep reinforcement learning

基础算法,如排序或哈希,在任何一天都会被使用数万亿次。随着计算需求的增长,这些算法的性能变得至关重要。尽管过去取得了显著的进展,但进一步提高这些程序的效率对于人类科学家和计算方法来说都是具有挑战性的。在这里,我们展示了人工智能如何通过发现迄今未知的程序超越当前的技术水平。为了实现这一目标,我们将寻找更好的排序程序的任务定义为单人游戏。然后,我们训练了一个新的深度强化学习代理,AlphaDev,来玩这个游戏。AlphaDev从头开始发现了小型排序算法,超越了以前已知的人类基准。这些算法已经被集成到LLVM标准C++排序库中。这种对排序库部分的更改代表着用一种使用强化学习自动发现的算法替换了一个组件。我们还在额外领域展示了结果,展示了该方法的通用性。


人类的直觉和知识在改进算法方面起到了至关重要的作用。然而,许多算法已经达到了一个阶段,即人类专家无法进一步优化它们,导致计算瓶颈不断增长。在经典程序合成文献中,跨越了许多十年,旨在生成正确的程序和/或使用延迟代理来优化程序。这些包括枚举搜索技术和随机搜索,以及最近使用深度学习进行程序合成以生成正确程序的趋势。使用深度强化学习(DRL),我们可以更进一步,通过优化CPU指令级别的实际测量延迟,更有效地搜索和考虑正确和快速程序的空间,与以前的工作相比,生成正确和高性能的算法。

计算机科学中的一个基础问题是如何对一个序列进行排序。这在世界各地的初级计算机科学课程中都有教授,并被广泛应用于各种应用程序。几十年来,计算机科学研究一直致力于发现和优化排序算法。实用解决方案的一个关键组成部分是对短序列元素进行小型排序;当对使用分治方法的大型数组进行排序时,这种算法会被反复调用。在这项工作中,我们关注两种类型的小型排序算法:(1)固定排序和(2)可变排序。固定排序算法对固定长度的序列进行排序(例如,sort 3只能对长度为3的序列进行排序),而可变排序算法可以对不同大小的序列进行排序(例如,可变排序5可以对从1到5个元素的序列进行排序)。

我们将发现新的、高效的排序算法的问题定义为一个单人游戏,称为AssemblyGame。在这个游戏中,玩家选择一系列低级CPU指令,我们称之为汇编指令,将它们组合起来产生一个新的、高效的排序算法。这是具有挑战性的,因为玩家需要考虑汇编指令的组合空间,以产生一个既正确又快速的算法。AssemblyGame的难度不仅来自搜索空间的大小,它类似于极具挑战性的游戏,如国际象棋(\(10^{120}\)局)和围棋(\(10^{700}\)局),而且还来自奖励函数的性质。在AssemblyGame中,一条错误的指令可能会使整个算法无效,使得在这个游戏空间中探索变得极其困难。

为了玩这个游戏,我们引入了AlphaDev,一个学习代理,它被训练来搜索正确和高效的算法。这个代理由两个核心组件组成,分别是(1)一个学习算法和(2)一个表示函数。AlphaDev学习算法可以结合DRL和随机搜索优化算法来玩AssemblyGame。AlphaDev中的主要学习算法是AlphaZero的扩展,这是一个著名的DRL算法,其中一个神经网络被训练来指导搜索以解决AssemblyGame。表示函数是可互换的,并且捕捉汇编程序的底层结构。AlphaDev的主要表示基于Transformers。

使用AlphaDev,我们从头开始发现了固定和可变的排序算法,它们都是新的,并且比现有的人类基准更有效。由AlphaDev发现的sort 3、sort 4和sort 5的固定排序解决方案已经被集成到LLVM标准C++库中的标准排序函数中。这个库被数百万用户使用,包括大学和众多国际公司。此外,我们分析了新的算法发现,将AlphaDev与随机搜索优化方法进行比较,并将AlphaDev应用到更多领域,以展示该方法的通用性。

image

图1 | C++和汇编程序之间的关系。a,一个C++实现的可变排序2函数,可以对任何长度为2的输入序列进行排序。b,a中的C++实现被编译成这个等效的低级汇编表示。

将算法表示为低级CPU指令

当从高级语言(如C++)将算法编译为机器码时(例如,图1a中的排序函数),算法首先被编译成汇编(图1b)。然后,汇编程序将汇编程序转换为可执行的机器码。在这项工作中,我们在汇编级别优化算法。在典型的汇编程序中,值从内存复制到寄存器,然后在寄存器之间进行操作,最后写回内存。支持的汇编指令集取决于处理器架构。为了这项工作,我们专注于x86处理器架构使用AT&T语法支持的一组汇编指令。每条指令都是Opcode⟨OperandA, OperandB⟩的格式。一个示例指令是mov<A,B>,它被定义为从源(A)移动一个值到目标(B)。更多指令定义,如比较(cmp<A,B>)、条件移动(cmovX<A,B>)和跳转(jX<A>)可以在扩展数据表1中找到。在图1b中的示例中,%eax%ecx%edx%edi对应四个不同的寄存器位置,而(%rsi)、4(%rsi)对应两个不同的内存位置。符号$2是一个占位符,代表一个常数值,它对应于本例中向量的长度。我们在这项工作中交替使用汇编程序和汇编算法这两个术语。这是因为AlphaDev每次玩AssemblyGame时都会从头开始构建一个汇编程序,从一开始无序的指令集合中定义一个新的、高效的算法。

DRL for discovering faster algorithms

在这一节中,我们将优化算法在CPU指令级别的问题形式化为强化学习(RL)问题,在其中环境被建模为一个我们称之为AssemblyGame的单人游戏。这个游戏中的每个状态都被定义为一个向量\(\mathbf{S_t = ⟨P_t, Z_t⟩}\),其中\(\bf P_t\)是到目前为止在游戏中生成的算法的表示,\(\bf Z_t\)表示在执行当前算法对一组预定义输入后内存和寄存器的状态。如图2a所示,在时间步\(t\),玩家接收当前状态\(\bf S_t\)并执行动作\(\bf a_t\)。这涉及将一个合法的汇编指令(例如,mov<A,B>)附加到到目前为止生成的当前算法中。收到一个奖励\(r_{\bf t}\),它包括算法正确性和延迟的度量。算法正确性(图2b)涉及将一组\(N\)个测试序列输入到当前算法\(\bf P_t\)中以生成\(N\)个输出。然后将这些输出与预期输出进行比较,并计算出一个正确性奖励\(r_{\bf t}\)。延迟奖励可以通过(1)惩罚代理增加算法长度(当长度和延迟高度相关时),我们称之为算法长度奖励,或(2)测量算法的实际延迟来生成。游戏执行有限数量的步骤,之后游戏终止。赢得游戏对应于使用汇编指令生成一个正确、低延迟的算法。输掉游戏对应于生成一个不正确的算法或一个正确但效率低下的算法。

image

图2 | AssemblyGame和算法正确性计算。a. AssemblyGame由AlphaDev玩,它接收到目前为止生成的当前汇编算法\(\bf S_t\)作为输入,并通过选择要执行的动作来玩游戏。在这个例子中,动作是一个mov<Register0,Memory1>汇编指令,它被附加到当前算法中。代理收到一个奖励,它是算法正确性(在b中讨论)和算法延迟的函数。玩家发现一个低延迟、正确的算法就赢得了游戏。b. 程序正确性和延迟计算用于计算奖励\(r_{\bf t}\)。在这个例子中,测试序列被输入到算法中;例如,在排序三个元素的情况下,测试输入包括所有未排序元素的长度为3的序列。对于每个序列,将算法输出与预期输出进行比较(在排序的情况下,预期输出是排序后的元素)。在这个例子中,输出D′与预期输出B′不匹配,因此算法是不正确的。

image

a,与人类基准相比,AlphaDev在优化算法长度时的表现。AlphaDev从零开始发现算法,在每种情况下都能匹配或改进人类基准。b,与人类基准相比,AlphaDev在直接优化延迟时的表现。在这种设置下,AlphaDev发现的算法延迟明显低于人类基准。置信区间表示为延迟±(下限,上限),其中延迟对应于100台不同机器上延迟测量的第五百分位数。下限和上限分别指这个百分位数的95%置信区间的边界。


我们将玩这个单人游戏的代理称为AlphaDev。该代理的主要学习算法是AlphaZero代理的扩展,它使用深度神经网络指导蒙特卡洛树搜索(MCTS)规划过程。神经网络的输入是状态\(\bf S_t\),输出是策略和价值预测。策略预测是一个关于动作的分布,价值函数是一个关于累积回报\(R\)的预测,代理应该期望从当前状态\(\bf S_t\)接收到这些回报。在游戏中,代理接收到当前状态\(\bf S_t\)作为输入。然后,代理执行一个MCTS过程并使用它来选择下一个要采取的动作。生成的游戏然后用于更新网络的参数,使代理能够学习。

对于AlphaDev来说,至关重要的是具有能够表示复杂算法结构以有效地探索指令空间的表示。为了实现这一点,我们引入了AlphaDev表示网络(representation network)(扩展数据图1a)。该网络包括两个组成部分,即(1)一个a transformer encoder网络,为代理提供算法结构的表示,以及(2)CPU状态编码器网络,帮助代理预测算法如何影响内存和寄存器的动态。CPU状态编码器网络包括一个多层感知器,它接收给定输入集合中每个寄存器和内存位置的状态作为输入。这些网络各自输出嵌入,这些嵌入被组合起来产生AlphaDev状态表示。

Transformer encoder

Transformer是自然文本编码器,在最近的语言模型中取得了巨大的成功。因此,这促使我们将标准Transformer适应到模拟汇编指令。我们开发并整合一个transformer编码器,我们对MultiQuery transformer编码器的改编,纳入到AlphaDev表示网络中来表示汇编指令。每个汇编指令的操作码和相应的操作数被转换为独热编码(one-hot encodings)并连接起来形成原始输入序列。这通过一个多层Transformer编码器进行馈送,将其映射到相应的嵌入向量(见扩展数据图1b以获得说明)。

提醒:因为机翻的关系虽然稍微修正了一些但难免有漏掉的,如果看到变换器指的就是Transformer

Latency value functions

延迟是一个重要的奖励信号,用于指导代理发现高性能的算法。为了更好地估计延迟,我们实现了一个双值函数设置,其中AlphaDev具有两个值函数头:一个预测算法正确性,另一个预测算法延迟。延迟头用于直接预测给定程序的延迟,通过在训练期间使用程序的实际计算延迟作为AlphaDev的蒙特卡洛目标。这种双头方法在优化真实延迟时比原始的(vanilla)、单头值函数设置取得了更好的结果。

Results

Discovering faster sort algorithms

我们从零开始训练AlphaDev代理,生成一系列固定排序和可变排序算法,这些算法都是正确的,并且比最先进的人类基准具有更低的延迟。

Fixed sorting algorithms

我们考虑了三个基本算法:排序3、排序4和排序5。这些算法的最先进人类基准是排序网络,因为它们生成高效的、无条件分支的汇编代码。这意味着所有指令都按顺序执行,没有分支。改进这些算法是具有挑战性的,因为它们已经高度优化。如表1a所示,AlphaDev能够找到比人类基准少指令的算法,对于排序3和排序5,它能够匹配最先进的性能。这些较短的算法确实导致了更低的延迟,因为算法长度和延迟在无条件分支情况下是相关的;详见附录B。我们还探索了使用AlphaDev的一个变体来扩展到稍大一点的排序。我们设法在排序6上节省了三条指令,在排序7上节省了两条指令,在排序8上节省了一条指令,这为未来的工作提供了一个有希望的基础。详见附录C以获得概述。

Variable sorting algorithms

我们考虑了三种可变排序算法:VarSort3,VarSort4和VarSort5。在每种情况下,人类基准被定义为一种算法,对于给定的输入长度,调用相应的排序网络。在这种情况下,需要分支,这大大增加了问题的复杂性,因为代理需要(1)确定它需要构造多少个子算法,并且(2)并行构建主算法的主体。代理也可能需要从其他子算法调用子算法。在这种情况下,优化长度会导致比人类基准明显更短的算法,如表1a所示。然而,由于分支引入的复杂性,延迟和长度并不总是相关的;详见补充信息。因此,我们实现了一种测量程序实际延迟的程序,通过在100台不同机器上取延迟测量的第五百分位数,并计算置信区间,并优化这个指标。详见方法以获得完整的基准测试设置。当优化延迟时,代理在每种情况下都显著改善了人类基准,如表1b所示。

New algorithm discoveries

AlphaDev发现的解决方案包括新颖和令人兴奋的算法发现,这些发现导致了更高效的性能。在固定排序设置中,我们发现AlphaDev发现了两个有趣的指令序列,当应用于排序网络算法时,每次都可以减少一条汇编指令。我们将每个指令序列分别称为(1)AlphaDev交换移动和(2)AlphaDev复制移动。

image

图3 | AlphaDev发现的排序网络和算法改进。a,三个输入的最优经典排序网络。被圈起来的比较器已经被AlphaDev改进。详见AlphaDev交换移动以获得更多细节。b,c,在应用AlphaDev交换移动之前(b)和应用AlphaDev交换移动之后(c)的汇编伪代码,导致删除了一条指令。d,一个最优的经典排序网络比较器配置,已经被AlphaDev改进。详见AlphaDev复制移动以获得更多细节。e,f,在应用AlphaDev复制移动之前(e)和应用AlphaDev复制移动之后(f)的汇编伪代码,导致删除了一条指令。

image

a,左边显示了在应用图3a中圈起来的操作符时,在经典排序网络中将输入A,B和C转换为输出的过程。右边显示了代替圈起来的操作符应用的AlphaDev交换移动转换。请注意,每次应用时都会节省一条指令的新转换以粗体显示。b,左边显示了根据图3d中的排序网络配置将输入A,B,C和D转换为输出的过程。右边显示了应用于此排序网络配置的AlphaDev复制移动转换。粗体显示的转换表示复制移动所做的更改,每次应用时都会节省一条指令。

AlphaDev swap move

图3a展示了三个元素的最优排序网络(请参阅方法以获得排序网络的概述)。我们将解释AlphaDev如何改进圈起来的网络部分。在各种大小的排序网络中都可以找到这种结构的许多变体,并且每种情况都适用相同的论点。网络的圈起来的部分(最后两个比较器)可以看作是一系列指令,它接受输入序列⟨A,B,C⟩并按照表2a(左)所示的方式转换每个输入。然而,在这个操作符之前有一个在B和C线上的比较器,因此保证了B≤C的输入序列。这意味着计算min(A,B)作为第一个输出就足够了,而不是如表2a(右)所示的min(A,B,C)。图3b,c之间的伪代码差异演示了AlphaDev交换移动每次应用时如何节省一条指令。

AlphaDev copy move

图3d展示了一个由三个比较器组成的排序网络配置,应用于四根导线上。这种配置在sort 8排序网络中可以找到,并且对应于一个操作符,它接受四个输入⟨A,B,C,D⟩并将它们转换为四个输出,如表2b(左)所示。可以证明,作为sort 8的一部分,流入操作符的输入满足以下不等式:D≥min(A,C)。这意味着可以通过应用表2b(右)中定义的AlphaDev复制移动来改进操作符,从而比原始操作符少一条指令。在图3e,f中分别可视化了原始操作符和应用AlphaDev复制移动后的代码之间的差异。

image

图4 | AlphaDev发现的根本不同的算法。a, 可变排序4(VarSort4)人类基准算法的流程图。在这个算法中,一系列未排序的数字被输入到算法中。如果序列长度为四个,三个或两个数字,则调用相应的sort 4,sort 3或sort 2排序网络来对结果序列进行排序。然后返回结果并由函数输出。b,AlphaDev发现的VarSort4算法。这个算法也接收长度为四个,三个或两个数字的序列作为输入。在这种情况下,如果长度为2,则调用sort 2排序网络并返回。如果长度为3,则调用sort 3对前三个数字进行排序并返回。然而,如果长度大于3,则调用sort 3,然后调用一个简化的sort 4例程来对剩余的未排序数字进行排序。正是这部分例程导致了显著的延迟节省。

New variable sort algorithms

AlphaDev发现的VarSort4算法特别有趣。人类基准算法和AlphaDev的流程图分别可以在图4a,b中看到。人类基准算法确定输入向量的长度,然后调用相应的排序网络对元素进行排序。AlphaDev的解决方案有一个完全不同的方法,如图4b所示。如果输入向量的长度严格大于2,则立即调用sort 3,从而对前三个元素进行排序。如果向量大于三个元素,则调用一个简化的sort 4算法,对输入向量中剩余的未排序元素进行排序。正是这个简化部分使得算法长度和延迟方面都取得了显著的收益。

Stochastic search optimization approaches随机搜索优化方法

理解RL与其他程序优化方法的优缺点非常重要。因此,我们实现了一种最先进的随机超优化方法,并将其适应于排序设置,并将其用作AlphaDev中的学习算法。我们将这种变体称为AlphaDev-S(有关详细信息,请参阅方法)。我们运行这个算法,至少与AlphaDev具有相同的资源和墙上时钟时间。由于每次突变后都需要计算延迟,因此AlphaDev-S直接优化延迟所需的时间过多。因此,AlphaDev-S优化延迟代理,即算法长度,然后在训练结束时,我们搜索所有由AlphaDev-S生成的正确程序并对每个程序进行基准测试以找到最低延迟解决方案。总的来说,我们发现当从头开始学习而没有先前知识时,AlphaDev始终比AlphaDev-S表现更好。此外,随着程序大小的增加,与AlphaDev-S(最坏情况下为31万亿个程序)相比,AlphaDev探索的程序数量少几个数量级(最坏情况下为1200万个程序)。这可能是因为与更容易陷入局部最优解的广度优先随机搜索过程相比,AlphaDev能够更好地探索算法空间;有关此探索假设的概述,请参阅方法。此外,由于使用延迟值函数预测,因此AlphaDev在搜索期间从不评估延迟,并且因此只需要在不到0.002%的生成程序上计算实际测量的延迟。当将先前知识纳入AlphaDev-S中时,例如使用接近最优解决方案来启动学习算法,对于sort 3、sort 4和sort 5(无分支汇编算法),AlphaDev-S在计算上更有效率,并且在每种情况下都生成与AlphaDev竞争性低延迟算法。但是,对于需要分支(if-else语句)的算法,在其中算法长度和延迟不相关时,即使使用接近最优解决方案启动该算法,AlphaDev也会发现比AlphaDev-S更低延迟的解决方案。有关这些算法的深入分析,请参阅方法。

推广到其他领域

为了测试AlphaDev的普遍性,我们在一组附加领域上训练代理。这些包括下面介绍的协议缓冲区反序列化子例程VarInt和竞争性编码问题(有关详细信息,请参阅附录D)。竞争性编码领域的延迟性能在表1b中报告。 协议缓冲区是Google的开源数据格式,用于序列化结构化数据。此格式通常用于性能或网络负载是主要关注点的情况。VarInt算法是序列化和反序列化过程中的关键组件。我们像变量排序一样训练AlphaDev代理以优化VarInt反序列化函数,以正确性和测量延迟为准。对于正确性,我们奖励代理正确反序列化每个输入。我们使用一组80个输入和相应的输出,涵盖常见的protobuf用例。AlphaDev学习了一个优化的VarInt反序列化函数,并且能够显著优于人类基准的单值输入。我们的代理发现了一个无分支解决方案,它既短(表1a),又比人类基准快大约三倍(表1b)。在这样做时,代理还发现了一个新的VarInt赋值移动,在其中AlphaDev学会将两个操作组合成一个指令,从而节省延迟。有关此移动的完整概述,请参阅附录D.1。这是一个强烈的迹象,表明AlphaDev能够推广以优化非平凡的实际算法。

Libc++ sort patch

LLVM libc++标准排序库中的sort 3、sort 4和sort 5算法被更大的排序算法多次调用,因此是库的基本组件。我们将AlphaDev发现的sort 3、sort 4和sort 5的底层汇编排序算法逆向工程为C++,并发现我们的排序实现对于长度为5的序列提高了高达70%,对于超过250,000个元素的序列大约提高了1.7%。这些改进适用于ARMv8、Intel Skylake和AMD Zen 2 CPU架构的uint32、uint64和float数据类型;有关完整性能表,请参阅附录E。性能改进归功于AlphaDev生成的无分支条件装配以及新的AlphaDev交换移动。对于sort 5,我们使用了由AlphaDev发现的43长度算法,因为它导致了更高效的C++实现。这些算法已经发送进行审核,并已正式纳入libc++标准排序库中。这是这些子例程十多年来的第一次更改。这也是该排序库中任何组件被使用强化学习自动发现的算法替换的第一次。我们估计,每天有数万亿次调用这些例程。

Discussion

AlphaDev从头开始发现了新的、最先进的排序算法,这些算法已被纳入LLVM C++库中,被全球数百万开发人员和应用程序使用。AlphaDev和随机搜索都是强大的算法。未来研究的一个有趣方向是调查将这些算法结合在一起,以实现两种方法的互补优势。

重要的是要注意,理论上,AlphaDev可以推广到不需要对测试用例进行详尽验证的函数。例如,哈希函数以及加密哈希函数通过哈希冲突的数量定义函数正确性。因此,在这种情况下,AlphaDev可以优化以最小化冲突以及延迟。理论上,AlphaDev还可以优化大型、令人印象深刻的函数体内的复杂逻辑组件。我们希望AlphaDev能够提供有趣的见解并在人工智能和程序综合社区中激发新的方法。