[论文解读] Code Repair with LLMs gives an Exploration-Exploitation Tradeoff

Hao Tang, Keya Hu+5 more|arXiv (Cornell University)|May 26, 2024

Scientific Computing and Data ManagementDecision Sciences被引用 3

一句话总结

本文提出 REx，一种基于大语言模型（LLM）的程序合成新方法，将代码修复建模为多臂赌博机设置下的探索-利用权衡问题，采用汤普森采样（Thompson Sampling）进行优化。通过在高性能程序的精炼（利用）与较少被访问的程序探索（探索）之间取得平衡，REx 在代码生成、循环不变量合成和视觉推理任务中，将 LLM 调用次数减少了 1.5 倍至 5 倍，同时解决了更多问题。

ABSTRACT

Iteratively improving and repairing source code with large language models (LLMs), known as refinement, has emerged as a popular way of generating programs that would be too complex to construct in one shot. Given a bank of test cases, together with a candidate program, an LLM can improve that program by being prompted with failed test cases. But it remains an open question how to best iteratively refine code, with prior work employing simple greedy or breadth-first strategies. We show here that refinement exposes an explore-exploit tradeoff: exploit by refining the program that passes the most test cases, or explore by refining a lesser considered program. We frame this as an arm-acquiring bandit problem, which we solve with Thompson Sampling. The resulting LLM-based program synthesis algorithm is broadly applicable: Across loop invariant synthesis, visual reasoning puzzles, and competition programming problems, we find that our new method can solve more problems using fewer language model calls.

研究动机与目标

• 解决现有代码精炼策略效率低下的问题，这些策略通常采用简单的贪婪或广度优先策略，带来的收益微乎其微。
• 将基于 LLM 的迭代式代码修复重新建模为探索-利用权衡问题：在精炼通过更多测试用例的程序（利用）与精炼次数较少的程序（探索）之间进行选择。
• 开发一种可扩展的、基于赌博机的策略，用于 LLM 引导的程序合成，避免在无限分支、随机性树结构中使用蒙特卡洛树搜索所带来的高昂计算成本。
• 在多种领域实现高效且有效的程序合成，包括竞赛编程、软件验证和视觉推理谜题。
• 通过优化精炼过程中的探索-利用平衡，最小化解决复杂编程问题所需的 LLM 调用次数。

提出的方法

• 将代码精炼建模为多臂赌博机问题，其中每个程序视为一个“臂”，奖励定义为通过的测试用例数量。
• 使用汤普森采样通过从每个程序预期奖励的后验分布中采样，来平衡探索与利用。
• 在每次 LLM 调用后，使用贝叶斯更新维护对每个程序奖励分布的概率信念。
• 在每一步中，通过从后验分布中采样每个程序的预期奖励，并选择采样值最高的程序进行精炼。
• 通过精炼选定程序动态扩展搜索树，使用新的 LLM 调用创建新节点（精炼结果）。
• 递归应用该赌博机策略，高效导航可能的精炼无限树结构，通过折扣未来奖励最大化避免深度回溯。

实验结果

研究问题

• RQ1 基于赌博机的策略是否能在基于 LLM 的代码生成中，通过平衡探索与利用，优于简单的贪婪策略或广度优先策略？
• RQ2 在迭代式代码修复中，探索-利用权衡如何影响解决复杂编程问题所需的 LLM 调用次数？
• RQ3 汤普森采样在多大程度上可被适配以处理 LLM 引导程序合成中的无限分支和随机转移？
• RQ4 所提出的 REx 方法是否能在包括循环不变量合成、竞赛编程和视觉推理谜题在内的多样化编程任务中实现泛化？
• RQ5 该方法是否能持续解决标准精炼基线无法触及的困难问题？

主要发现

• 在所有评估领域中，REx 相较于基线方法将 LLM 调用次数减少了 1.5 倍至 5 倍。
• REx 解决的问题数量多于贪婪策略和广度优先策略，尤其在原本难以处理的难题上表现更优。
• 在循环不变量合成任务中，REx 能够成功生成先前方法失败的正确不变量。
• 在 ARC 基准的视觉推理谜题中，REx 以更少的 LLM 调用次数实现了更高的成功率，优于标准精炼基线。
• 该方法始终优于从初始提示中独立同分布（i.i.d.）采样，表明通过探索-利用结构化精炼比随机重采样更有效。
• 汤普森采样在代码精炼的无限、随机树结构中有效平衡了探索与利用，实现了高效搜索，且无需昂贵的回溯操作。