基于检索增强搜索的 LLM 程序优化 LLM Program Optimization via Retrieval Augmented Search
提出RAS和AEGIS方法，通过上下文检索和束搜索优化LLM代码性能。
前置知识
黑盒适应
黑盒适应是指不调整模型权重，仅通过提示策略改进LLM性能的技术，包括指令提示、上下文学习、思维链提示等方法。这类方法利用预训练模型的现有能力，通过精心设计的提示模板、少样本示例或多轮交互来引导模型生成更优质的输出，具有低计算成本和快速部署的优势。
本文的核心方法RAS和AEGIS都是黑盒适应策略，理解这个概念有助于明白为什么作者选择不微调模型，而是通过检索和搜索来改进性能，以及如何设计有效的提示模板。
检索增强生成
检索增强生成（RAG）是一种结合信息检索和生成式AI的技术，在生成响应前先从知识库中检索相关文档作为上下文，然后基于检索到的内容生成答案。在代码领域，这通常意味着从训练集中检索相关的代码示例，然后将这些示例作为少样本提示提供给LLM，从而提高代码生成和优化的准确性。
本文的核心创新是将动态检索与束搜索结合，特别是在代码优化中首次系统性比较了基于代码检索和基于上下文检索的差异，理解RAG有助于理解论文的动机和方法设计。
束搜索
束搜索是一种启发式搜索算法，在每一步保留多个最优候选（称为束宽），而不是只保留单一最优路径。在代码优化中，这意味着在每次迭代中生成多个优化版本，选择性能最好且正确的版本作为下一步的基础，通过多轮迭代逐步改进程序性能，类似于编译器通过多遍扫描逐步优化代码。
RAS的核心就是束搜索策略，理解这个概念有助于明白为什么作者能够通过多轮迭代显著提升性能，以及如何控制搜索空间和计算成本。
编译器Pass
编译器Pass是编译器优化过程中的一个独立阶段，每个Pass专注于执行特定类型的优化，如常量传播、循环展开、内联函数等。现代编译器通过一系列精心设计的Pass将源代码转换为高效的目标代码，每个Pass接收中间代码并输出改进后的代码，Pass之间相互配合形成完整的优化流水线。
AEGIS的设计灵感直接来源于现代编译器的多Pass架构，理解这个概念有助于理解为什么作者要将大范围的优化分解为原子编辑序列，以及这种方法如何提高可解释性。
研究动机
现有LLM在程序优化任务上表现不佳，因为传统的训练数据集缺乏程序性能相关信息。具体而言，在PIE基准测试中，未经适应的LLM优化C++程序的平均加速比仅为1.73到2.89倍，远低于人类专家的3.63倍水平。虽然之前的PIE工作提出了动态检索方法，能够在测试时从训练集中检索相关的慢快程序对作为上下文示例，实现了4.23倍的平均加速比，但这种方法仍然存在显著局限。它基于源代码本身的嵌入进行检索，忽略了算法和数据结构的本质相似性。它只执行单次优化，没有像现代编译器那样通过多轮逐步改进程序。
本文的目标是本文的目标是提出更有效的黑盒适应策略，显著提升LLM在程序优化任务上的性能，使其能够接近甚至超越人类专家的优化水平。具体而言，作者希望在PIE C++基准上超越动态检索的4.23倍加速比，实现更接近人类专家3.63倍水平的性能。在Mercury Python基准上，目标是缩小较小模型与更大推理模型之间的性能差距，使得7B参数规模的模型也能达到接近更大模型的性能。同时，作者还希望通过将大范围优化分解为增量步骤，提高优化过程的可解释性，使每个修改更容易理解和验证，从而提升用户对优化结果的信任度。
与已有工作不同的是，本文的独特切入角度是将编译器的多Pass设计理念与检索增强生成相结合，同时创新性地提出基于自然语言描述的上下文检索。具体而言，作者观察到现有动态检索方法基于代码本身进行检索，但竞技程序员在不同问题中常使用相似的变量名和操作，导致代码嵌入相似但算法本质不同。而LLM生成的代码也不太可能像人类程序员那样编写优化版本，加剧了检索难度。作者提出先用LLM生成程序的自然语言描述，提取核心算法和数据结构，再基于描述进行检索，从而抽象掉实现层面的差异。此外，作者还提出将端到端的优化分解为一系列原子编辑，每个编辑对应一个具体的、可解释的代码改进，从而模仿编译器的多Pass优化流程。
核心方法
RAS（Retrieval Augmented Search）的整体思路是结合束搜索和动态上下文检索，通过多轮迭代逐步优化程序。给定待优化程序 $p_0$ 和训练集 $\Pi_{train}$，RAS在每一步迭代中首先使用LLM生成当前程序 $p_{i-1}$ 的自然语言描述，然后基于描述的嵌入向量从训练集中检索最相似的 $k$ 个慢快程序对，接着将检索到的示例和当前程序一起提供给LLM生成 $k$ 个优化版本，最后选择运行最快且通过所有测试用例的版本作为下一步的基础。重复这个过程 $m$ 次后返回最终优化版本，类似于现代编译器的多遍扫描，每次迭代都基于更相关的上下文示例进行优化从而持续改进性能。
RAS的核心创新点有两个：第一是基于LLM生成的自然语言描述进行上下文检索而不是基于源代码本身的嵌入。这种方法的本质区别在于自然语言描述能够抽象掉实现细节如变量名和具体代码结构，聚焦于算法和数据结构的本质相似性。例如两个使用不同变量名但都实现动态规划算法的程序，其代码嵌入可能相距甚远但自然语言描述都会强调动态规划这一核心算法，从而在嵌入空间中相近。第二是将束搜索与检索结合，每一步都重新检索最相关的示例而不是在整个过程中使用固定的检索集，这使得后续迭代能够根据已经应用的优化进行调整检索到更适合当前状态的示例从而实现更好的增量优化效果。
方法步骤详情
RAS方法的步骤如下：预处理阶段对训练集中每个程序对 $(p, p')$ 使用LLM生成慢程序 $p$ 的描述 $s = F_{context}(p)$ 再用嵌入模型 $\Psi$ 计算嵌入向量 $\Psi(s)$ 构建索引。优化阶段给定程序 $p_0$ 对于迭代步 $i$ 从1到 $m$：步骤1用 $F_{context}$ 生成当前程序 $p_{i-1}$ 的描述并计算其嵌入；步骤2基于L2距离 $d_\phi(p_{i-1}, p)$ 检索训练集中最相似的 $k$ 个程序对 $\Pi_i$；步骤3对每个检索到的程序对 $\pi_i^j \in \Pi_i$ 用LLM生成优化版本 $p_i^j \sim F_{opt}(\pi_i^j, p_{i-1})$；步骤4选择最快且正确的版本 $p_i = \text{arg max}_{j \in [k]} R(p_i^j)$。如果所有版本都未通过测试则保持 $p_i = p_{i-1}$。最终返回 $p_m$。关键超参数为 $k=8$ 和 $m=4$。
技术新颖性
RAS的技术新颖性体现在三个方面：首次在程序优化任务中系统性比较了上下文检索与代码检索的效果，证明了基于自然语言描述的检索显著优于基于代码嵌入的检索；首次将束搜索与动态检索结合，每步都重新检索最相关的示例而不是使用固定检索集；首次通过上下文检索的描述更新机制，使得后续迭代能够根据已应用的优化进行调整检索到更适合当前状态的示例。与之前的动态检索方法相比，RAS在PIE基准上实现了2.06倍的性能提升；与纯指令提示相比，实现了高达5.1倍的提升。这种性能提升的核心在于自然语言描述能够捕获算法和数据结构的本质相似性，使得LLM能够学习到真正可迁移的优化策略而不仅仅是表面模式的匹配。
实验结果
在PIE C++基准测试中RAS显著超越了所有基线方法。使用Qwen-3-Coder时RAS实现了8.70倍的平均加速比而之前的最佳方法动态检索仅达到4.23倍提升了2.06倍。使用GPT-4o时RAS达到8.03倍动态检索为4.43倍Instruct Only为2.31倍。使用DeepSeek 3.2时RAS更是达到了9.18倍的加速比动态检索为7.03倍。AEGIS方法虽然性能略低于RAS但仍然显著优于基线：使用GPT-4o时达到6.08倍加速比比动态检索提升37%同时将平均编辑规模减少了17%比首次迭代减少30%证明了其增量编辑的可解释性优势。在Mercury Python基准测试中RAS将Qwen2.5-7B-Instruct的平均运行时百分位数提升了10.27从77.58提升到87.85显著缩小了其与更大推理模型的性能差距。在准确性方面RAS在973个测试程序中只有5个选择了不正确的程序准确率为99.5%AEGIS则达到100%的准确率。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| C++程序优化(PIE基准) | 平均最佳加速比 | 8.70倍(Qwen-3-Coder) | 4.23倍(Dynamic Retrieval) | 提升2.06倍 |
| C++程序优化(PIE基准) | 平均最佳加速比 | 8.03倍(GPT-4o) | 4.43倍(Dynamic Retrieval) | 提升1.81倍 |
| C++程序优化(PIE基准) | 平均最佳加速比 | 6.08倍(AEGIS, GPT-4o) | 4.43倍(Dynamic Retrieval) | 提升1.37倍 |
| Python程序优化(Mercury基准) | 运行时百分位数(Beyond@1) | 87.85(RAS, Qwen2.5-7B-Instruct) | 77.58(Instruct Only) | 提升10.27 |
| C++程序优化(PIE基准) | 优化比例 | 98.56%(RAS, Qwen-3-Coder) | 77.49%(Dynamic Retrieval) | 提升27.2% |
局限与改进
作者明确承认的主要局限性包括计算开销较大：由于使用束搜索和动态检索，RAS和AEGIS比单次检索方法需要更多的LLM调用和计算资源。AEGIS还需要额外的训练时计算，因为它需要使用LLM生成的代码来构建原子编辑数据集。此外，作者推测这些方法可能难以扩展到更复杂、面向对象的大型代码库，因为那可能需要中间步骤来识别需要优化的特定代码片段。从实验结果中也可以观察到，RAS和AEGIS在某些特定类型的程序上表现不佳：AEGIS在涉及动态规划和二叉搜索的程序上失败比例较高，在未优化程序集合中，动态规划问题占45.05%，二叉搜索问题占12.09%，显著高于整体测试集中的比例。RAS则在涉及Kruskal算法的程序上失败率较高，在未优化程序集合中占20%，而Kruskal算法问题仅占整体测试集的1.03%。这些局限性表明，当前的检索策略可能无法很好地处理某些特定算法类别，需要更有针对性的检索或示例增强。
独立分析的弱点
从独立分析的角度看RAS的第一个主要弱点是在特定算法类别上的检索失败率较高特别是Kruskal算法相关的程序。这表明当前的上下文检索可能无法有效捕获某些算法的优化模式改进方向可以包括：为特定算法类别构建专门的检索索引或者在训练集中增加这些算法的高质量示例对。第二个弱点是束搜索的计算成本随迭代次数线性增长在实际部署中可能成为瓶颈。改进方向可以是引入早期停止机制当连续两轮的性能提升低于阈值时终止搜索或者使用更高效的搜索策略如启发式剪枝。第三个弱点是AEGIS的LLM分解过程不保证准确如果分解出的自然语言编辑不准确会导致原子编辑数据集质量下降。改进方向可以包括引入人工验证或自动化检查机制过滤掉质量较低的原子编辑。第四个弱点是当前方法主要针对竞技编程风格的短程序对大型软件项目的适用性未经验证。改进方向可以是扩展方法以支持模块化优化先识别关键模块或热点函数再针对性地应用RAS或AEGIS。
未来方向
作者提出的未来研究方向包括：针对特定算法和数据结构增强检索策略通过分析失败模式来识别需要更多训练示例的类别；探索更高效的搜索算法减少计算开销；将方法扩展到更复杂、面向对象的代码库；研究如何将AEGIS的原子编辑与传统的编译器优化Pass结合形成混合优化框架。基于论文成果的可延伸方向包括：将上下文检索的思想应用于其他代码相关任务如bug修复、代码重构、文档生成等；研究如何自动生成高质量的原子编辑描述减少人工成本；探索跨语言的迁移能力例如用C++的优化示例来指导Python代码的优化；研究如何与形式化验证工具集成确保优化后的代码语义等价性；探索在联邦学习或隐私保护设置下的应用通过检索而非共享代码来提升性能。
复现评估
论文提供了较好的复现支持：代码和数据已在GitHub开源作者承诺遵循开源许可。使用PIE基准的Apache 2.0许可和Mercury基准的Creative Commons Attribution Non Commercial 4.0许可。实验细节较为详细：PIE实验使用4080个高质量训练对和973个测试对在gem5模拟器上使用5个随机测试用例评估性能与完整测试集的Pearson相关系数为0.89Spearman相关系数为0.86证明了评估方法的可靠性。超参数设置明确：RAS使用k等于8个检索和m等于4个搜索步每步采样h等于1个版本；基线方法使用k等于4个检索和h等于32个采样。模型使用明确：GPT-4o作为上下文生成函数、分解函数、编辑函数和泛化函数Qwen-3-Coder和DeepSeek V3.2作为优化函数嵌入模型使用OpenAI的text-embedding-3-large。计算资源要求中等：需要在服务器上运行gem5模拟器以及AWS实例用于测量Python运行时百分位数。总体而言复现难度中等主要挑战在于计算资源需求特别是gem5模拟器的运行成本。
论文图表
