[論文レビュー] 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

ひとこと要約

本稿では、多腕バンディット設定におけるトレードオフとしてコード修復を定式化し、Thompson Sampling を用いて探索と活用のバランスを取ることで、LLMベースのプログラム合成手法 REx を提案する。高いパフォーマンスを示すプログラムの改善（活用）と、訪問回数が少ないプログラムの探索（探索）をバランスさせることで、コード生成、ループ不変式の合成、視覚的推論タスクの全分野において、LLMの呼び出し回数を 1.5x–5x 減少させながら、より多くの問題を解決することができる。

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.

研究の動機と目的

• 単純なグリーディ戦略や幅優先探索（BFS）ポリシーに依存する従来のコード精錬戦略の非効率性に対処する。これらはわずかな向上しか得られない。
• LLMを用いた反復的コード修復を、探索と活用のトレードオフとして再定式化する。具体的には、より多くのテストケースに合格するプログラムの改善（活用）と、まだあまり修正されていないプログラムの探索（探索）のどちらを選ぶかを決定する。
• 無限分岐・確率的木構造においてモンテカルロ木探索（MCTS）の高コストを回避する、スケーラブルなバンディットベースの戦略を、LLMガイドドのプログラム合成に適用する。
• コンpetitionプログラミング、ソフトウェア検証、視覚的推論パズルなど、多様な分野において、効率的かつ効果的なプログラム合成を可能にする。
• 精錬における探索と活用のバランス最適化により、複雑なプログラミング問題を解くために必要なLLM呼び出し回数を最小限に抑える。

提案手法

• 各プログラムを「腕」と見なし、報酬を通過するテストケースの数とする多腕バンディット問題としてコード精錬を定式化する。
• 各プログラムの期待報酬の事後分布からのサンプリングを通じて、探索と活用のバランスをThompson Samplingで実現する。
• 各LLM呼び出し後にベイズ更新を用いて、各プログラムの報酬分布に関する確率的信念を維持する。
• 各ステップで、各プログラムの期待報酬を事後分布からサンプリングし、最も高いサンプル値を持つプログラムを選択して精錬する。
• 選択されたプログラムを精錬することで、新たなノード（精錬）を生成し、新たなLLM呼び出しを伴って探索木を動的に拡張する。
• 無限木の精錬可能性を効率的に探索するため、報酬の割引を用いた将来報酬最大化により、深すぎるロールアウトを回避するように、バンディット戦略を再帰的に適用する。

実験結果

リサーチクエスチョン

• RQ1 探索と活用のバランスを取るバンディットベースの戦略は、LLMベースのコード生成において、単純なグリーディ戦略や幅優先探索（BFS）戦略を上回る性能を発揮できるか？
• RQ2 反復的コード修復における探索と活用のトレードオフは、複雑なプログラミング問題を解くために必要なLLM呼び出し回数にどのように影響を与えるか？
• RQ3 Thompson Sampling は、LLMガイドドのプログラム合成における無限分岐および確率的遷移を効果的に扱えるか？
• RQ4 提案されたREx手法は、ループ不変式の合成、コンテストプログラミング、視覚的推論パズルなど、多様なプログラミングタスクに一般化可能か？
• RQ5 標準的な精錬ベースラインでは到達できない難易度の高い問題を、この手法は一貫して解くことができるか？

主な発見

• REx は、全評価分野において、ベースライン手法と比較してLLM呼び出し回数を 1.5x から 5x 減少させた。
• グリーディ戦略や幅優先探索（BFS）戦略よりも多くの問題を解けるが、特に従来では到達不能とされていた難易度の高い問題で顕著である。
• ループ不変式の合成タスクでは、先行手法が失敗する問題に対しても、REx は正しく不変式を生成できた。
• ARCベンチマークの視覚的推論パズルにおいて、REx は標準的な精錬ベースラインよりも高い成功確率を達成し、LLM呼び出し回数も少ない。
• 初期プロンプトからのi.i.d.サンプリングよりも一貫して優れた性能を示しており、探索活用に基づく構造的な精錬が、単なる再サンプリングよりも効果的であることを示している。
• Thompson Sampling は、コード精錬の無限かつ確率的な木構造において、探索と活用のバランスを効果的に実現し、高価なロールアウトを回避した効率的な探索を可能にした。