LLM Program Optimization via Retrieval Augmented Search

📄 arXiv: 2501.18916v1 📥 PDF

作者: Sagnik Anupam, Alexander Shypula, Osbert Bastani

分类: cs.LG

发布日期: 2025-01-31

💡 一句话要点

提出检索增强搜索（RAS）优化LLM程序，并用AEGIS提升可解释性。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 程序优化 大型语言模型 检索增强学习 黑盒优化 代码生成 可解释性 原子编辑

📋 核心要点

  1. 程序优化是编程语言研究的关键挑战，现有方法难以有效利用大型语言模型（LLM）的潜力。
  2. 论文提出检索增强搜索（RAS），通过检索上下文示例指导LLM进行程序优化，提升优化效果。
  3. 实验表明，RAS优于现有黑盒自适应策略1.8倍，AEGIS在提升可解释性的同时，性能提升1.37倍。

📝 摘要（中文）

随着大型语言模型（LLM）的出现，人们对其在解决复杂编程任务中的应用产生了浓厚的兴趣。最近的研究表明，LLM在程序优化方面具有潜力，而程序优化是编程语言研究中的一个关键挑战。我们提出了一种名为检索增强搜索（RAS）的黑盒自适应方法，该方法在候选优化方案上执行束搜索；在每个步骤中，它从给定的慢速-快速程序对训练数据集中检索上下文示例，以指导LLM。至关重要的是，我们发现基于LLM生成的自然语言描述执行上下文检索明显优于基于源代码的检索。此外，我们提出了一种名为AEGIS的方法，通过将训练示例分解为“原子编辑”来提高可解释性，这些原子编辑本质上更具增量性。我们表明，RAS的性能比先前的最先进的黑盒自适应策略高1.8倍，而AEGIS的性能高1.37倍，同时执行的编辑操作明显更小。

🔬 方法详解

问题定义：论文旨在解决如何利用大型语言模型（LLM）进行程序优化的问题。现有的黑盒优化方法通常难以有效利用LLM的上下文学习能力，导致优化效果不佳，且优化过程缺乏可解释性。

核心思路：论文的核心思路是利用检索增强搜索（RAS）来指导LLM进行程序优化。RAS通过检索与当前优化步骤相关的上下文示例，为LLM提供更丰富的背景信息，从而提高优化效果。此外，论文还提出AEGIS方法，将训练样本分解为更小的“原子编辑”，以提高优化过程的可解释性。

技术框架：RAS的技术框架主要包括以下几个步骤：1) 初始化：从原始程序开始，作为束搜索的初始状态。2) 检索：使用LLM生成的自然语言描述作为查询，从训练数据集中检索相关的慢速-快速程序对。3) 优化：利用检索到的上下文示例，指导LLM生成候选的优化程序。4) 评估：评估候选优化程序的性能。5) 选择：选择性能最佳的候选程序，加入束搜索的队列中。重复步骤2-5，直到达到预定的搜索深度或找到满足要求的优化程序。AEGIS方法则是在训练数据预处理阶段，将程序优化示例分解为更小的、更易于理解的原子编辑。

关键创新：论文的关键创新在于：1) 提出基于LLM生成自然语言描述的上下文检索方法，显著优于基于源代码的检索。2) 提出AEGIS方法，通过分解训练示例为原子编辑，提高了优化过程的可解释性，并能生成更细粒度的优化。

关键设计：RAS的关键设计包括：1) 使用LLM生成自然语言描述，作为检索的查询，更好地捕捉程序的语义信息。2) 使用束搜索算法，探索不同的优化路径。3) AEGIS的关键设计在于如何将复杂的程序优化分解为一系列更小的、独立的原子编辑，例如变量替换、循环展开等。具体的分解策略需要根据具体的编程语言和优化目标进行设计。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，RAS的性能比先前的最先进的黑盒自适应策略高1.8倍。AEGIS方法在提高可解释性的同时，性能也提升了1.37倍，并且生成的优化编辑更加细粒度。这些结果验证了RAS和AEGIS在程序优化方面的有效性和优越性。

🎯 应用场景

该研究成果可应用于各种软件开发场景，例如编译器优化、代码重构、性能调优等。通过自动化的程序优化，可以显著提高软件的运行效率和资源利用率，降低开发和维护成本。此外，AEGIS方法提高了程序优化的可解释性，有助于开发者理解和信任优化结果。

📄 摘要（原文）

With the advent of large language models (LLMs), there has been a great deal of interest in applying them to solve difficult programming tasks. Recent work has demonstrated their potential at program optimization, a key challenge in programming languages research. We propose a blackbox adaptation method called Retrieval Augmented Search (RAS) that performs beam search over candidate optimizations; at each step, it retrieves in-context examples from a given training dataset of slow-fast program pairs to guide the LLM. Critically, we find that performing contextual retrieval based on an LLM-generated natural language description significantly outperforms retrieval based on the source code. In addition, we propose a method called AEGIS for improving interpretability by decomposing training examples into "atomic edits" that are significantly more incremental in nature. We show that RAS performs 1.8$\times$ better than prior state-of-the-art blackbox adaptation strategies, and that AEGIS performs 1.37$\times$ better while performing significantly smaller edits.