[Paper Review] 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 Sciences3 citations

TL;DR

This paper introduces REx, a novel LLM-based program synthesis method that frames code repair as an explore-exploit tradeoff using Thompson Sampling in a multi-armed bandit setting. By balancing refinement of high-performing programs (exploitation) and exploring less-visited programs (exploration), REx reduces LLM calls by 1.5x–5x while solving more problems across code generation, loop invariant synthesis, and visual reasoning tasks.

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.

Motivation & Objective

• Address the inefficiency of existing code refinement strategies that use simple greedy or breadth-first policies, which yield marginal gains.
• Reframe iterative code repair with LLMs as an explore-exploit tradeoff: choosing between refining programs that pass more test cases (exploit) or those refined fewer times (explore).
• Develop a scalable, bandit-based strategy for LLM-guided program synthesis that avoids the prohibitive cost of Monte Carlo Tree Search in infinite-branching, stochastic trees.
• Enable efficient and effective program synthesis across diverse domains, including competition programming, software verification, and visual reasoning puzzles.
• Minimize the number of LLM calls required to solve complex programming problems by optimizing the exploration-exploitation balance in refinement.

Proposed method

• Model code refinement as a multi-armed bandit problem where each program is an 'arm' and the reward is the number of passing test cases.
• Use Thompson Sampling to balance exploration and exploitation by sampling from posterior distributions over the expected reward of each program.
• Maintain probabilistic beliefs over the reward distribution of each program using Bayesian updating after each LLM call.
• At each step, select the program to refine by sampling its expected reward from the posterior and choosing the one with the highest sampled value.
• Dynamically expand the search tree by refining the selected program, creating new nodes (refinements) with new LLM calls.
• Apply the bandit strategy recursively to navigate the infinite tree of possible refinements efficiently, avoiding deep rollouts via discounted future reward maximization.

Experimental results

Research questions

• RQ1 Can a bandit-based strategy that balances exploration and exploitation outperform simple greedy or breadth-first refinement policies in LLM-based code generation?
• RQ2 How does the explore-exploit tradeoff in iterative code repair affect the number of LLM calls required to solve complex programming problems?
• RQ3 To what extent can Thompson Sampling be adapted to handle infinite branching and stochastic transitions in LLM-guided program synthesis?
• RQ4 Does the proposed REx method generalize across diverse programming tasks such as loop invariant synthesis, competition programming, and visual reasoning puzzles?
• RQ5 Can the method consistently solve difficult problems that are out of reach for standard refinement baselines?

Key findings

• REx reduces the number of LLM calls by a factor of 1.5x to 5x compared to baseline methods across all evaluated domains.
• REx solves more problems than greedy and breadth-first strategies, particularly on challenging problems that are otherwise intractable.
• In loop invariant synthesis, REx successfully generates correct invariants for problems where prior methods fail.
• For visual reasoning puzzles from the ARC benchmark, REx achieves higher success rates with fewer LLM calls than standard refinement baselines.
• The method consistently outperforms i.i.d. sampling from the initial prompt, demonstrating that structured refinement via explore-exploit is more effective than random re-sampling.
• Thompson Sampling effectively balances exploration and exploitation in the infinite, stochastic tree of code refinements, enabling efficient search without expensive rollouts.