用强化学习让大语言模型为汇编代码提速：斯坦福团队的优化探索

在2025年5月，来自斯坦福大学、伊利诺伊大学香槟分校、卡内基梅隆大学和Visa研究院的研究团队发表了一篇引人注目的论文。这篇由Anjiang Wei、Tarun Suresh、Huanmi Tan、Yinglun Xu、Gagandeep Singh、Ke Wang和Alex Aiken共同完成的研究，探讨了如何利用大语言模型(LLM)优化汇编代码性能。该论文题为《通过强化学习改进大语言模型的汇编代码性能》(Improving Assembly Code Performance with Large Language Models via Reinforcement Learning)，发表于arXiv预印本平台(arXiv:2505.11480v1)。

近几年来，大语言模型在编程领域展现出了惊人的能力，从代码生成到程序修复，表现都相当出色。然而，在代码优化特别是底层代码优化方面的潜力却尚未被充分挖掘。想象一下，如果你拥有一个智能助手，它不仅能帮你写代码，还能自动让这些代码运行得更快、更高效，那将会是多么便利的体验。这正是这项研究想要实现的目标。

传统上，优化汇编代码的工作主要由编译器完成。就像厨师按照固定的菜谱制作料理一样，现代编译器如GCC会按照预设的规则对代码进行一系列优化转换。但这种方法面临一个经典问题：优化步骤的顺序会大大影响最终生成代码的性能，就像烹饪步骤的顺序会影响菜肴的口感一样。由于优化空间异常庞大，编译器经常无法找到最优解，留下了很大的性能提升空间。

另一种方法是所谓的"超级优化"(superoptimization)，它会在所有可能的等价程序中搜索性能最佳的版本。这有点像不按食谱烹饪，而是尝试所有可能的烹饪方法，直到找到最美味的那一个。理论上，这种方法可以产生最优的代码。但实际上，随着程序大小的增加，搜索空间呈指数级增长，使得穷举搜索在实际应用中变得不可行。而且，以往的超级优化工作主要针对没有循环的直线代码，因为这类代码更容易进行形式化验证，确保优化后的代码与原始代码在功能上完全等价。

面对这些挑战，研究团队提出了一个创新性的思路：利用大语言模型的强大能力来优化汇编代码。与高级语言相比，汇编代码更接近硬件，提供了对执行过程的精细控制，能够实现高级语言难以表达的优化。然而，这种设置也带来了新的挑战。汇编代码在预训练语料库中相对稀少，可能导致大语言模型难以有效推理其行为。此外，GCC等工业级编译器经过了数十年的精心调校，要在其最高优化级别(-O3)的基础上进一步提升性能，无疑是一项技术挑战。

研究团队巧妙地将强化学习技术应用到大语言模型的训练中，以解决低级代码优化的挑战。如图1所示，他们使用了近端策略优化(Proximal Policy Optimization, PPO)来训练大语言模型，并设计了同时考虑正确性和性能的奖励函数。正确性通过程序特定的测试用例来评估，而性能则通过与gcc -O3基线相比的加速比来衡量。

为了支持这个设定，研究团队构建了一个包含8,072个真实世界汇编程序的新数据集，这些程序源自竞争性编程提交。每个实例包括输入-输出测试用例和由编译器在最高优化级别生成的基准汇编代码，作为进一步优化的起点。

在这项研究中，他们评估了包括开源和商业模型在内的21个最先进的语言模型。从基础模型Qwen2.5-Coder-7B-Instruct开始，经过PPO训练后，他们的模型Qwen2.5-Coder-7B-PPO在编译通过率和测试通过率上都达到了96.0%，平均加速比达到1.47倍，远超所有其他评估模型，包括Claude-3.7-sonnet。

下面，让我们一起深入了解这项研究的细节，探索它如何突破传统编译优化的限制，为代码性能提升开辟新的可能性。

一、研究背景与问题定义

想象一下你正在开发一个需要处理海量数据的应用程序。即使是程序执行速度提升1%，也能在长时间运行中节省大量计算资源和等待时间。在这个场景下，代码优化变得尤为重要。

传统上，编译器负责将我们写的高级语言代码（如C语言）转换成计算机能直接执行的汇编代码。在这个过程中，编译器会应用各种优化技术来提高代码效率。就像一位经验丰富的厨师知道如何在不改变菜肴口味的前提下，通过调整烹饪步骤来缩短烹饪时间一样，编译器会在保持程序功能不变的前提下，使程序运行得更快。

研究团队关注的问题是：尽管现代编译器如GCC已经相当强大，但它们仍然依赖于预设的规则和转换序列，这使得它们在面对复杂程序时往往无法找到最优解。这就像一位只按照固定菜谱烹饪的厨师，虽然能做出美味佳肴，但可能错过了更高效的烹饪方法。

研究团队提出的任务可以定义为：给定一个用高级语言（如C）编写的程序C和由gcc -O3编译生成的汇编程序P，目标是生成一个新的汇编程序P'，使得P'在保持与P功能等价的前提下，运行速度更快。

在理想情况下，我们希望新生成的程序P'在整个输入空间中都与原程序P功能等价。但这个属性的验证在一般情况下是不可判定的。因此，研究团队采用了一种实用的近似方法：使用有限的测试集T = {(xi, yi)}来验证程序的行为，其中每个输入-输出对(xi, yi)都捕捉了程序C预期的行为。

一个汇编程序P'被认为是有效的，如果它能成功地组装和链接成可执行二进制文件。研究团队用valid(P') ∈ {True, False}来表示这个属性。而正确程序的集合则被定义为：

S(P) = {P' | valid(P') ∧ ?(xi, yi) ∈ T, P'(xi) = yi}

这表示所有既能成功编译又能通过所有测试用例的程序集合。

对于性能评估，研究团队测量了程序P和P'在测试集T上的执行时间，分别记为t(P)和t(P')。P'相对于P的加速比被定义为：

Speedup(P') = t(P) / t(P')，如果P'∈S(P)且t(P') < t(P) 1，其他情况

这意味着只有当新程序正确且比原程序快时，才会计算实际的加速比；否则，加速比默认为1。

优化目标是生成一个候选程序P'，使其加速比Speedup(P')最大化。只有S(P)中的程序才有资格获得加速；任何无法编译成二进制文件或产生错误输出的候选程序都被赋予默认加速比1。这反映了一种实用的回退策略：当生成的程序无效时，系统可以退回到基准P（由gcc -O3编译），从而确保至少获得1倍的参考性能。

二、数据集构建与研究方法

要训练一个能够优化汇编代码的模型，首先需要高质量的数据。研究团队利用CodeNet，这个包含大量竞争性编程提交的庞大语料库，构建了自己的数据集。就像建造一座房子需要优质的建材一样，构建良好的机器学习模型同样需要精心准备的数据。

研究团队的数据集由三部分组成：原始C源代码(C)、使用gcc -O3优化级别编译生成的对应x86-64汇编代码(P)，以及测试集(T)。这就像是给模型提供了一个"菜谱"(C源码)、一个"初始成品"(gcc生成的汇编)，以及一套"品鉴标准"(测试集)，让模型学习如何在保持食物原味(功能正确性)的前提下，使烹饪过程更加高效(代码更快)。

由于CodeNet包含超过800万个C和C++提交，研究团队无法全部使用。他们采用了一种聪明的抽样策略：选择那些从gcc -O0(无优化)到gcc -O3(最高优化)能获得最高相对加速比的程序。这种策略有两个好处：一是倾向于选择含有复杂逻辑的程序，这些程序在-O0下表现不佳，但能被-O3有效优化；二是创造了更具挑战性的环境，因为基准代码已经经过编译器的积极优化。

最终的数据集包含7,872个训练程序和200个保留用于评估的程序。这些程序平均包含约22行C代码和130行汇编代码，每个程序配有约9个测试用例。

接下来，研究团队采用强化学习方法来优化大语言模型。他们将任务概念化为一个标准的上下文多臂赌博机问题，由上下文空间S、动作空间A和奖励函数r : S × A → R组成。

想象你在玩一个赌博机游戏，但每次拉动赌博机的手柄前，都会看到一些背景信息(上下文)。在这里，上下文s ∈ S代表一个问题实例，包括源程序C、基线汇编P和相关测试用例T。而动作a ∈ A对应于生成候选汇编程序P~。奖励函数r(s, a)根据生成程序的正确性和性能评估其质量。

研究团队使用近端策略优化(PPO)来训练模型。PPO是一种一阶策略梯度算法，通过约束每次策略更新与前一个保持接近，从而稳定训练过程。具体来说，PPO最大化形式为Es,a[min(ρ(θ)A, clip(ρ(θ), 1 - ε, 1 + ε)A)]的裁剪替代目标，其中ρ(θ) = πθ(a | s)/πθold(a | s)是当前策略与先前策略之间的概率比率，A是动作a在状态s中的估计优势，ε是限制策略更新到小信任区域的裁剪系数。

研究团队设计了两种不同的奖励函数变体来指导模型学习。可以将奖励函数想象成"教练"的反馈，告诉模型它的生成效果如何。

第一种是"正确性引导的加速比"(CGS)，其定义为：

r(s, a) = -1，如果P~编译失败 pass(s, a)，如果某些测试失败 1 + α · speedup(s, a)，如果所有测试通过

其中，pass(s, a)表示通过的测试用例比例，speedup(s, a)表示相对于基线P的加速比，α是一个控制全部正确情况下加速比相对重要性的常数。

第二种是"仅加速比"(SO)，其定义为：

r(s, a) = 0，如果P~编译失败或任何测试失败 speedup(s, a)，其他情况

CGS奖励函数提供了密集的信号，即使在最终目标未达成时也能指导策略。它就像是在学习烹饪过程中，教练不仅关注最终菜肴的口感，还会对准备食材、火候控制等中间步骤给予反馈。相比之下，SO定义了一个更直接、更稀疏的目标，只在程序既正确又高效时才给予非零奖励，类似于教练只关注最终菜肴的质量，不给中间过程任何反馈。

三、实验设置与评估方法

为了全面评估大语言模型优化汇编代码的能力，研究团队设计了一套严谨的实验方案。就像科学家测试新药效果需要精确的实验设计一样，评估AI模型的性能同样需要周密的规划。

首先，研究团队为每个实例构建了一个包含原始C程序和使用gcc -O3生成的汇编代码的提示词(prompt)。所有测试用例都对模型保密，这就像是给学生出题，但不提前告诉他们答案。模型被指示只生成优化后的x86-64汇编代码，不包含任何其他文本。

评估指标包括正确性和性能两个方面。编译通过率(Compile pass)是生成的汇编能成功编译为二进制可执行文件的问题百分比。测试通过率(Test pass)是编译后的代码通过所有测试用例的问题百分比。对于性能，研究团队测量了相对于gcc -O3基线的加速比。如前所述，任何无法编译、测试失败或比基线慢的候选程序都被赋予默认加速比1倍。这反映了实际应用中的情况：当新方法不起作用时，系统可以回退到gcc -O3输出，不会造成性能损失。为了捕捉分布行为，研究团队报告了加速比的25、50(中位数)和75百分位数，以及整个评估集的平均加速比。

研究团队评估了21个最先进的语言模型，涵盖了多种架构。他们的评估包括顶尖的专有模型，如gpt-4o、o4-mini、gemini-2.0-flash-001和claude-3.7-sonnet，以及开源模型家族，如Llama、DeepSeek和Qwen。此外，他们还评估了基于DeepSeek-R1蒸馏的Qwen和Llama模型，以及基于Code Llama并专为编译器相关任务设计的最新编译器基础模型。所有开源模型都经过指令微调。

为了确保性能评估的准确性，研究团队使用了hyperfine，这是一个基准测试工具，通过执行预热运行和多次计时执行来减少测量噪声。对于每个程序的执行，他们丢弃前三次运行，并报告接下来十次运行的平均运行时间。

实验是在配备四块A100 GPU的单个节点上进行的。在所有评估模型中，他们选择Qwen2.5-Coder-7B-Instruct进行训练，原因是它表现出最强的正确性结果，并且性能还有很大的提升空间。他们有意避免使用特定于编译器的基础模型，以保持通用性。

四、实验结果与分析

经过严谨的实验设计，研究团队得到了一系列令人振奋的结果。就像美食评论家会从多个维度评价一道菜品一样，研究团队也从多个角度分析了各个模型的表现。

从表2可以看出，大多数模型在生成高性能汇编代码方面面临挑战：大部分只能达到1.00倍的加速比，编译通过率和测试通过率较低。在基线模型中，claude-3.7-sonnet和DeepSeek-V3表现最佳，测试通过率超过40%，平均加速比分别达到1.22倍和1.21倍。值得注意的是，一些模型如DeepSeek-R1完全无法生成有效的汇编，而o4-mini的测试通过率仅为4.5%。这些结果凸显了任务的难度，也说明了需要针对性方法的必要性。

编译器基础模型(以llm-compiler-为前缀)是在汇编代码和编译器中间表示上预训练的。其中，llm-compiler-13b在正确性和加速比方面都表现出较强的性能。相比之下，经过微调的变体(-ftd)表现不佳，可能是因为它们适应了将x86-64和ARM汇编反汇编为LLVM-IR等任务，而非优化汇编代码以提高执行性能。

研究团队选择Qwen2.5-Coder-7B-Instruct作为强化学习训练的起点，因为它具有较强的编译通过率(79.0%)和最高的测试通过率(61.0%)。经过PPO微调后，该模型在两项指标上都达到了96.0%，并将平均加速比从1.10倍提高到1.47倍。值得注意的是，它是唯一一个在25百分位上展示出明显加速的模型，并且在所有评估指标上都优于其他所有模型，包括正确性、平均加速比和加速比百分位数。

研究团队还对奖励函数设计进行了深入分析。他们评估了两种奖励设计：正确性引导的加速比(CGS)和仅加速比(SO)。如表3所示，两种变体都实现了较高的编译通过率和测试通过率，但SO产生了更好的性能。去除中间的塑形信号似乎有助于模型专注于终端目标。这就像是学习一项复杂技能时，有时直接针对最终目标练习，比分解成多个简单步骤学习更有效。研究团队还尝试了不同的CGS缩放因子α(5或10)，发现它对结果影响不大。

这些结果表明，稀疏的、终端的奖励(SO)在这种情况下更有效。由于基础模型已经达到61.0%的测试通过率，正确性不是瓶颈；直接优化加速比提供了更强的训练信号。

研究团队还进行了一项关键的消融研究：大语言模型是否能直接从C代码编译出高性能汇编，而不依赖编译器生成的基线汇编？为了调查这个问题，他们比较了两种评估设置：(1)提供C源代码和gcc -O3汇编(默认，表4中的"w/ O3")，(2)只提供C源代码("w/o O3")。

表4显示，移除基线汇编导致性能严重下降。例如，Qwen2.5-Coder-7B-PPO从96.0%的正确性和1.47倍的加速比降至0.0%和1.00倍。即使是强大的模型如Claude-3.7-sonnet也出现了显著下降。

这些结果表明，当前大语言模型直接从C编译到优化汇编的能力仍然有限。编译器输出为大语言模型提供了可靠的参考。这支持了研究框架的设计：使用gcc -O3作为强化学习的有效起点。虽然未来的工作可能会探索直接从C生成，但由于缺乏编译器指导，这可能会更加具有挑战性。

五、案例分析与实际应用

为了更直观地理解大语言模型如何优化汇编代码，研究团队展示了一个代表性例子，如图2所示。这个例子清晰地展示了语言模型如何发现比先进编译器更高效的优化方案。

原始C函数计算"人口计数"(也称为位计数，即计算一个数二进制表示中有多少个1)。这个函数通过重复移位输入并累加其最低有效位来计算结果。想象一下，这就像是你手动数一串灯泡中有多少个是亮着的，一个一个地检查。

gcc -O3生成的汇编代码保留了这个循环结构，依赖于显式的位操作和条件分支来计算结果。这就像是编译器按部就班地遵循了C代码的逻辑，虽然做了一些基本优化，但基本思路没变。

而大语言模型生成的代码则显著不同：它用一条popcnt指令替换了整个循环。这条指令是现代x86-64架构支持的，能在一次操作中执行相同的计算，从而显著减少了指令数量和运行时开销。这就像是发现了一个专用的计数器，可以立即告诉你有多少灯泡是亮的，而不必一个一个地检查。

这种转换超出了gcc -O3的能力范围，因为编译器应用的是预设的基于规则的优化传递序列，而不进行这种语义级别的重写。大语言模型能够合成功能等价的代码，利用编译器未使用的硬件级指令。这个例子展示了语言模型通过探索更广泛的语义保留程序转换空间来优化汇编的潜力。

这种优化能力在实际应用中有着广泛的潜在价值。在性能敏感的领域，如高频交易、实时图像处理或游戏引擎，即使是微小的性能提升也能带来显著的实际好处。例如，在处理大量数据的服务器应用中，1.47倍的平均加速可以直接转化为能源成本的节约和服务容量的提升。

此外，这种方法还可以应用于资源受限的环境，如嵌入式系统或移动设备，在这些环境中，计算资源和电池寿命都是宝贵的。通过生成更高效的代码，可以延长设备使用时间或实现更复杂的功能。

在教育方面，这项研究也提供了宝贵的见解。通过观察大语言模型如何优化代码，开发者可以学习到新的优化技术和模式，这些可能不在传统编译器教科书中。就像从大师的棋谱中学习一样，开发者可以从AI生成的优化中获得启发，提升自己的编程技能。

六、局限性与未来工作

尽管这项研究取得了令人瞩目的成果，研究团队也坦诚地讨论了当前方法的几点局限性。就像任何工具都有其适用范围一样，他们的方法也面临一些挑战。

一个关键的局限是缺乏形式化的正确性保证。虽然研究团队通过输入-输出测试用例验证生成的程序，但这种测试本质上是不完整的，可能会忽略边缘情况。就像用几个样本测试一道菜的口味，无法保证在所有情况下都美味一样，测试无法涵盖所有可能的输入。因此，与Stoke等完全验证的系统不同，大语言模型生成的汇编可能会在某些情况下产生错误结果或表现出未定义行为。这一局限反映了编程语言领域的一个更广泛的挑战：在一般情况下，验证两个任意程序之间的语义等价性是不可判定的问题。

另一个局限是在实际硬件上测量性能的固有随机性。尽管研究团队通过重复测量来减少噪声，但低级硬件波动仍可能导致加速比估计出现变异。这种不确定性在现实环境中很难完全消除。虽然前人工作采用了基于模拟器的评估，但模拟器可能无法忠实捕捉实际硬件性能。

最后，观察到的性能提升可能无法跨机器泛化。模型可能隐式学习并利用特定硬件特性，如缓存大小。因此，在一台机器上训练的模型在部署到不同机器时可能无法保持其有效性。这就像一道适合特定烹饪设备的食谱，在使用不同设备时可能需要调整。

展望未来，研究团队提出了几个有前途的研究方向。首先，虽然他们使用了近端策略优化(PPO)，未来可以探索替代的强化学习算法，如GRPO。其次，扩展到更大、更多样化的数据集，特别是涉及竞争性编程之外的性能关键代码，将使设置更加现实和适用。第三，结合强化学习和监督微调可能会带来额外收益，尽管目前尚不清楚在gcc -O3输出上进行训练是否会提供额外价值。

另一个有趣的方向是融入交互式改进循环，模型可以使用错误或性能测量的反馈迭代更新其输出。最后，将这种方法从x86-64扩展到其他架构，如MIPS、ARM或GPU编程，可以扩大其适用性和影响力。

七、结论与总结

本研究探索了使用大语言模型优化汇编代码的潜力，这是一个能够通过精细控制执行过程实现性能提升的领域，而这些优化在高级语言中难以表达。传统编译器依赖固定的基于规则的转换，但由于优化空间的复杂性，它们在生成最优代码方面面临根本性挑战。

为了解决这个问题，研究团队应用强化学习来微调大语言模型，使用近端策略优化(PPO)和基于正确性和性能的奖励函数。为了支持这项研究，他们引入了一个包含8,072个真实世界C程序的基准，这些程序配有编译器生成的基线汇编和测试用例。

经过训练的模型Qwen2.5-Coder-7B-PPO在编译通过率和测试通过率上都达到了96.0%，平均加速比达到1.47倍，优于所有其他评估模型。消融研究表明，强调最终加速比而非中间正确性信号的奖励函数能够更有效地指导训练。

这些结果表明，强化学习能够释放大语言模型作为汇编代码性能的有效优化器的潜力。通过学习识别和应用编译器可能错过的优化机会，经过训练的模型能够生成既正确又高效的代码，为性能敏感的应用提供了宝贵的工具。

虽然仍存在一些局限，如缺乏形式化正确性保证和性能测量的固有随机性，但这项研究为未来探索打开了新的可能性。随着这一领域的发展，我们可以期待大语言模型在代码优化方面发挥越来越重要的作用，与传统编译器形成互补，共同推动软件性能的边界。

在这个计算资源日益珍贵的时代，能够自动生成更高效的代码不仅具有学术价值，更有着广泛的实际应用前景。从高性能计算到资源受限的嵌入式系统，这种技术都可能带来显著的效益，使软件不仅能够完成任务，还能以最高效的方式完成。