Code Repair with LLMs gives an Exploration-Exploitation Tradeoff

Authors: hao  tang
keya  hu
jin peng zhou
sicheng  zhong
wei-long  zheng
xujie  si
kevin  ellis
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 20:01:51

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.

Main paper: 
 introduction: 
 An emerging paradigm for problem-solving with large language models (LLMs) is to have the language model correct, repair, or debug its initial outputs [1, 2, 3, 4, 5] , which we refer to here as refinement. For example, when generating the source code of a program, refinement prompts the LLM with a buggy program it previously generated, potentially with diagnostic information such as a stack trace, then asks it to fix its code. This strategy alleviates the need for the LLM to predict perfect code on the first try, and approximates the iterative code-writing process of software development.Complex programming problems often require several rounds of refinement, with each round requiring more LLM calls, each of which is stochastic, and so yields multiple possible outcomes. Therefore, this process generates a tree of possible programs (Fig. 1 left ). This tree is infinitely deep, because every refinement can itself be further refined. It is also infinitely wide, because the LLM can return infinitely many possible refinements. Success should, in principle, depend on exactly what policy is used to expand this tree. Recent work [1, 6, 3] adopts simple expansion policies (e.g., breadth-first),• • • • • • • • • • • • • • • initial prompt initial samples refined once refined twice initial prompt 80% correct exploit 60% correct explore 0% correctFigure 1 : Left: The tree of possible refinements is infinitely deep and has infinite branching factor. Each node is a program and each edge is an LLM sample. Right: Explore-Exploit tradeoff for a search state after performing 3 node expansions. Exploit by sampling another child of a program that is nearly correct, or Explore by sampling a child of a program that has been expanded fewer times.with mixed success: On close examination, the gains from refinement, at least with these basic policies, are marginal compared to repeatedly sampling i.i.d. programs from the initial prompt [3] .Our work here reanalyzes refinement in terms of a tradeoff between exploiting the refinement of programs that are closer to being correct (i.e., pass more testcases), and exploring programs that have been refined fewer times. Fig. 1 (right) diagrams this tradeoff for a tree after the first few node expansions. This exploration-exploitation tradeoff is made more challenging by the fact that every time we refine a program, we create another brand-new program, giving an ever-increasing set of possible actions (possible next refinements). Our problem is however not solvable by standard approaches to Monte Carlo Tree Search (MCTS)-a bandit-based node expansion policy-because our branching factor is infinite, our transition function is stochastic, and "rollouts" would demand a prohibitively expensive series of LLM calls to refine down to some maximum depth.Our primary contribution is an algorithm for efficiently performing refinement, which we call REx (REfine, Explore, Exploit). REx should be seen as a strategy for conducting LLM-guided search that constructs and navigates a tree of refinements. We derive REx via a multiarmed bandit framing: different actions correspond to refining different programs; reward corresponds to the quality of a newly generated program; and maximizing discounted future reward corresponds to solving the programming problem in the minimum number of LLM calls.The resulting algorithm is broadly applicable to LLM-based code generation tasks. We describe applications to competition programming problems, challenging software verification problems involving generating nonlinear loop invariants, and visual reasoning puzzles from the Abstraction and Reasoning corpus (ARC: [7] ). Across every domain, REx solves more problems in fewer LLM calls, typically reducing the amount of API calls by a factor of 1.5x-5x. REx is also able to consistently solve a modest number of difficult problems that prove out-of-reach for other approaches.Thompson Sampling does the following at timestep t:Sample θ a ∼ P (θ a | {r t : a t = a}), for each a ∈ A Do the above using Bayes Rule: P (θ a | {r t : a t = a}) ∝ P (θ a )t:at=a P (r t |θ a )Pull arm a t = arg max a∈AE [r | θ a ]Get reward r tBy choosing different priors and likelihoods, a variety of Thompson sampling strategies can be designed. For example, prior knowledge of the value of each arm can be embedded in P (θ a ), and constraints on the range of possible rewards can be embedded in P (r|θ a ). 
hyperparameter analysis: 
 Table 1 : Analyze the hyperparameter sensitivity of methods by their performance rankings on different benchmarks. The rankings are determined by (AUC + final_success_rate) / 2. The hyperparameter, empty value for Greedy, denotes the heuristic value assigned for the empty solution. The hyperparameter choices for ARC are different from the others because of its domain convention (maximum # LLM requests=64 instead of 300). REx consistently outperforms or competes with the best baselines when C = 20 for all difficult benchmarks.LoopInv Solving for loop invariants is a core problem in the program verification domain. For each program, we can write down its required specifications as logical constrains on its inputs and outputs. One can then statically step through the program execution to check if the minimal preconditions on the inputs can ensure the post conditions on the outputs.However, loops in the program have dynamic repetitions, which cannot be directly described under Satisfiability Modulo Theories (SMT) checkers. This raises the question to find strong invariant conditions about the loop, which should summarize the loop's behaviour and allow the SMT checker to reason on the preserved properties of the variables. pre ⇒ inv, (inv ∧ lc ∧ trans) ⇒ inv', (inv ∧ ¬lc) ⇒ post   true ⇒ I[-50/x] (pre) I ∧ x < 0 ⇒ I[(y + 1)/y, (x + y)/x] (inv) I ∧ x ≥ 0 ⇒ y > 0 (post) (c) The desired loop invariant is (x < 0 ∨ y > 0).Figure 9 : A concrete example of a loop invariant taken verbatim from Si et al. [40] . Loop invariant connects the pre-conditions with the post-conditions to prove, using logical properties that always hold during the iterations. "/" denotes the value assignment. I[-50/x] means substituting x in the loop invariant, I, with -50. More details regarding loop invariant are available in Si et al. [40] .The specifications and the loop behaviour together describe the verification conditions (VC) for the program. Finding an invariant statement that satisfies these requirements proves the correctness of the loop's functionality. Specifically, the invariant is established if it can be deduced from the preconditions; it is preserved if it maintains itself after each loop iteration; and, it is strong if leading to the post conditions.Existing methods often fit for invariants by running program traces or by guessing heuristics based on the feedback from the checker. In this paper, we incorporate the large language model's code understanding and reasoning ability to analyze a program's behavior and to produce sufficient loop invariants.The verification of some programs require a particular subset of loop invariants that contain non-linear polynomial arithmetic (NLA) terms with regards to the program's variables. Finding NLA invariants is not well addressed by state-of-the-art formal solvers. Nevertheless, LLMs can produce arbitrary formula and are not constrained by the forms of the invariants. Thus, utilizing these models is a promising approach to address the performance gap and to discover novel invariants [41] .Here, we present a set of programs that require NLA invariants to prove their functional correctness [42] . All of these programs are meaningful algorithms on arithmetics that have real-world applications and have (multiple) non-linear loops. Examples include sums of powers, extended Euclidean algorithm to calculate greatest common divisors, and binary division algorithm to speedup calculation in hardware.Each instance in the benchmark contains the source code in C and the verification conditions for each of its loops. For programs with multiple loops, at least one loop requires NLA invariants. We manually verify the VCs and provide a ground truth reference invariant for each loop. We write the VCs and the reference invariants in the Z3Py format. We then use the Z3 theorem prover to check each invariant against the VCs. This way, we can statically verify whether the program's behaviour satisfy the specifications.Program Loop # InvType G-CLN [21] GSpacer [26] Loopy-GPT4 [41] RExEQUATIONTable 2 : Benchmark evaluation on the NLA programs. For programs with multiple loops, each loop is labeled with a number, top to down, from outer to inner.We generate potential invariants with GPT-4 by feeding in the C program and asking for Z3Py invariants. To identify the targeted loop, we mark the start of the loop with a comment, where the LLM should analyze the semantics on.We use Z3's simplification function to split the invariants into conjunctive normal form and check each item against the VCs. When an item from the CNF invariants does not satisfy all steps in the checking process, we filter out the candidate using the Houdini pruning algorithm, which is commonly used in invariant inference [41] . The final result is evaluated after performing Houdini.In the refinement setting, we use feedback from the Z3 checker to tell the LLM which step of the check failed, so that it knows whether a stronger or a more general invariant is needed. More detailed examples are shown in the prompts in Appendix A.6.1.G-CLN is a learning-based method that fits for invariants by running program traces [21] and achieved the state-of-the-art result on the NLA benchmark. Although it is designed to perform polynomial fitting on the program variables, it relies on manually crafted features that are instance-specific heuristics (i.e., each instance requires a different set of features). In addition, this dynamic approach requires sampling different program runs, which can omit corner cases, and is unable to handle programs with non-deterministic termination condition. Our method and many others avoid this problem using the static analysis paradigm.Examples where G-CLN fails on the NLA benchmark can be summarized in several reasons:• The loop requires high order invariants beyond the specified heuristic, as in knuth• The method fits on and generates excessive unnecessary weak invariants that inhibit the checker, as is the case in ps5• The method finds an invariant based on the global traces but is too strong on the current loop The Abstraction and Reasoning Corpus(ARC) is a visual reasoning task proposed by Chollet [7] . It contains 400 training and 400 evaluating problems. Each problem will provide the task taker with several input-output images expressed by 2D metrics pairs that follow the same transformation rule, such as rotating, expanding, and color-changing. The task taker is then asked to summarize the rule and give the correct output metric of the test input. Figure 10 shows two examples of the problem of the ARC dataset. The size of the metrics ranges from 1 × 1 to 30 × 30, and there are only 10 colors for each pixel represented by number 0 to number 9. In our paper, we use the subset of ARC proposed by Johnson [18] , which contains 40 ARC problems with human experiments. We select one hypothesis written by human beings for each problem as a hint for LLMs. successfully finds the candidate code that summarizes the transformation rule that can solve the problem, then the search process will stop.We run our experiment with a maximum iteration of 64 for a fair comparison with Hypothesis Search [23] . We set the hyperparameters of each method accordingly as follows:1. Greedy: empty value= 0, 0.5 For each method and each hyperparameter, we run them five times with five random seeds and report the average performance.We also scale up the maximum iteration to 200 and demonstrate the benefits of REx in this setting as well. The hyperparameters are as follows:1. Greedy: empty value= 0, 0.5 We use the exact same subset of APPS as Olausson et al. [3] for fair comparison. We refer the reader to its Appendix H for details. We use all 60 problems for the competition-level and the introductory-level, and randomly choose 60 of 120 problems for the interview-level.Most experiments are conducted with GPT-4 and temperature=1.0 by default. The choices of hyperparameters for each method are listed in Table 1 . We also experiment with other LLMs including GPT-3.5-turbo, an older model to alleviate concerns regarding data contamination.Figure 13 shows the results on APPS-Interview. As expected, REx performs competitively with the best baseline, Greedy, and outperforms the other two baselines. It consistently outperforms the state-of-the-art method, Olausson et al. [3] , which uses the fixed-width method with width=25/50. It also aligns with our hypothesis that greedy performs competitively with REx when the dataset is relatively easy.We list all the prompts that are used in our experiments in this section. The functionality for each prompt is stated in its subsection name. We highlight the dynamic information as yellow and the main instruction as blue. The dynamic information includes the codes synthesized so far by previous LLM calls.Initializing loop invariant. It asks LLMs to generate loop invariant to help understand the code and prove its correctness.-----Role: system --------------------You are an expert software engineer and you are working on a project that requires you to write loop invariants for a loop in the code. The loop invariant is a condition that must be true before and after each iteration of the loop. The loop invariant helps you understand the behavior of the loop and ensures that the loop is working correctly. The loop invariant should be strong enough to prove the correctness of the loop as well as the functionality of the whole code snippet. Please analyze the loop variant in your own language. The loop invariant should be a condition that must be true before and after each iteration of the loop. The loop invariant should help you understand the behavior of the loop and ensure that the loop is working correctly. The loop invariant should also be strong enough to prove the correctness of the loop as well as the functionality of the whole code snippet.What are the loop invariants in code? The loop invariant should be written in the format as Python-Z3 such as follows: ''' And(x >= 0, x <= 10, Implies(x >= 0, a == 0), y == x * x) ''' Please surround the loop invariant with triple backticks and write it in the format as Python-Z3.Refining loop invariant. It asks LLMs to refine loop invariant given the error messages. In detail, the loop invariant should satisfy three conditions in order to be valid and sufficient:• Established: meaning that loop invariant should be true given the preconditions, i.e., it should be true at the start of the loop.• Preserved: meaning that loop invariant should also be true after each iteration of the loop.• Sufficient: meaning that the loop invariant should be strong enough to help us prove the post conditions, i.e., the specifications/properties assigned by the users.The symbolic checker utilizes SMT solvers to evaluate these three conditions individually and returns the error messages. If any of the conditions is not satisfied, it will return the error type, the specific assertion, and its model to help debug. LLMs then refine the loop invariant given that information.Note that, we deliberately omit the detailed information about post conditions to avoid label leakage.Our loop invariant synthesis system can then also be applied in scenarios where user specifications are not explicitly available in logical forms.-----Role: system --------------------You are an expert software engineer and you are working on a project that requires you to write loop invariants for a loop in the code. The loop invariant is a condition that must be true before and after each iteration of the loop. The loop invariant helps you understand the behavior of the loop and ensures that the loop is working correctly. The loop invariant should be strong enough to prove the correctness of the loop as well as the functionality of the whole code snippet.-----Role: user - ------------------ The previous loop invariant is wrong because:-'Or(Not(x >= 0), u == a*x + b)' is neither established nor preserved, meaning it is not even true in the beginning of the loop and is neither true after each iteration of the loop. For example, we can set [p = 1, q = 0, y = 7, r = 0, v = 0, s = 1, x = 15, b = 7, a = 15, u = 7] to find a counterexample for establishing 'Or(Not(x >= 0), u == a*x + b)', since it conflicts with the assertion 'Implies(And(a > 0, b > 0, x == a, y == b, u == b, v == 0, p == 1, q == 0, r == 0, s == 1), Or(Not(x >= 0), u == a*x + b)) '.-'Or(Not(y >= 0), v == a + b*y)' is neither established nor preserved, meaning it is not even true in the beginning of the loop and is neither true after each iteration of the loop. For example, we can set [v = 0, s = 1, p = 1, x = 16, q = 0, u = 13, r = 0, y = 13, a = 16, b = 13] to find a counterexample for establishing 'Or(Not(y >= 0), v == a + b*y)', since it conflicts with the assertion 'Implies(And(a > 0, b > 0, x == a, y == b, u == b, v == 0, p == 1, q == 0, r == 0, s == 1), Or(Not(y >= 0), v == a + b*y)) '.-The metrics are neither enough to imply the postconditions.Please correct the previous loop invariants and provide the correct loop invariants for the loop in the code snippet. The loop invariant should be written in the format as Python-Z3 as before. Please surround the loop invariant with triple backticks and write it in the format as Python-Z3.Initializing code. It asks LLMs to generate code to solve python programming problems.-----Role: system --------------------You are a helpful assistant that can solve python programming problems and debug incorrect programs -----Role: user --------------------You are given a hard python programming contest problem:Question:You are holding a party. In preparation, you are making a drink by mixing together three different types of fruit juice: Apple, Banana, and Carrot. Let's name the juices A, B and C.You want to decide what fraction of the drink should be made from each type of juice, in such a way that the maximum possible number of people attending the party like it.Each person has a minimum fraction of each of the 3 juices they would like to have in the drink. They will only like the drink if the fraction of each of the 3 juices in the drink is greater or equal to their minimum fraction for that juice.Determine the maximum number of people that you can satisfy.-----Input------One line containing an integer $T$, the number of test cases in the input file.For each test case, there will be:-One line containing the integer $N$, the number of people going to the party.-$N$ lines, one for each person, each containing three space-separated numbers ''$A$ $B$ $C$'', indicating the minimum fraction of each juice that would like in the drink. $A$, $B$ and $C$ are integers between $0$ and $10000$ inclusive, indicating the fraction in parts-per-ten-thousand. $A + B + C \leq 10000$.You may assume that $1 \leq T \leq 2$ and $1 \leq N \leq 5000$.-----Output------$T$ lines, one for each test case in the order they occur in the input file, each containing the string ''Case #$X$: $Y$'' where $X$ is the number of the test case, starting from 1, and $Y$ is the maximum number of people who will like your drink.-----Examples-----Sample Input: 2 3 10000 0 0 0 10000 0 0 0 10000 3 5000 0 0 0 2000 0 0 0 4000 Sample Output: Case #1: 1 Case #2: 2Can you write a correct solution and make sure it passes the example test cases? (Please start the solution segment with '''python as usual and use Standard Input format)Refining code. It asks LLMs to refine code given the test case it failed on and the error message.-----Role: system --------------------You are a helpful assistant that can solve python programming problems and debug incorrect programs -----Role: user --------------------You are given a hard python programming contest problem:Question:You are holding a party. In preparation, you are making a drink by mixing together three different types of fruit juice: Apple, Banana, and Carrot. Let's name the juices A, B and C.You want to decide what fraction of the drink should be made from each type of juice, in such a way that the maximum possible number of people attending the party like it.Each person has a minimum fraction of each of the 3 juices they would like to have in the drink. They will only like the drink if the fraction of each of the 3 juices in the drink is greater or equal to their minimum fraction for that juice.Determine the maximum number of people that you can satisfy.-----Input------One line containing an integer $T$, the number of test cases in the input file.For each test case, there will be:-One line containing the integer $N$, the number of people going to the party.-$N$ lines, one for each person, each containing three space-separated numbers ''$A$ $B$ $C$'', indicating the minimum fraction of each juice that would like in the drink. $A$, $B$ and $C$ are integers between $0$ andThe bottom-most row that contains colored squares remains the same from test input to test output. Meanwhile, the other rows all start one column more to the right in the test output compared to the test input.The number in the input grid can be mapped to the following colors:0:black; 1:blue; 2:red; 3:green; 4:yellow; 5:grey; 6:fuschia; 7:orange; 8:teal; 9: brown Just reply with the implementation of transform_grid(input_grid: np.ndarray[ int]) in Python and nothing else, each cell in the output should only be numbers from 0 to 9. Please contains the necessary import statements.Refining image-manipulating program. We refine the codes that fail to generate correct answers by providing LLMs with previously failed codes and their output results.This is the refinement prompt for codes that raise runtime error.-----Role: user --------------------Example 0: Input: [[0 0 0 0 0 0 0 0 0] [0 6 6 6 0 0 0 0 0] [0 6 0 0 6 0 0 0 0] [0 0 6 0 0 6 0 0 0] [0 0 0 6 0 0 6 0 0] [0 0 0 0 6 6 6 0 0] [0 0 0 0 0 0 0 0 0] [0 0 2 2 2 0 0 0 0] [0 0 2 0 0 2 0 0 0] [0 0 0 2 2 2 0 0 0] [0 0 0 0 0 0 0 0 0] [0 0 0 0 0 0 0 0 0] [0 0 0 0 0 0 0 0 0] [0 0 0 0 0 0 0 0 0]] Output:[[0 0 0 0 0 0 0 0 0] [0 0 6 6 6 0 0 0 0] [0 0 6 0 0 6 0 0 0] [0 0 0 6 0 0 6 0 0] [0 0 0 0 6 0 6 0 0] [0 0 0 0 6 6 6 0 0] [0 0 0 0 0 0 0 The bottom-most row that contains colored squares remains the same from test input to test output. Meanwhile, the other rows all start one column more to the right in the test output compared to the test input.EQUATIONThe number in the input grid can be mapped to the following colors:0:black; This is the refinement prompt for codes that produce wrong answers.-----Role: user -------------------- [[0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 6, 6, 6, 0, 0, 0, 0, 0], [0, 6, 0, 0, 6, 0, 0, 0, 0], [0, 0, 6, 0, 0, 6, 0, 0, 0], [0, 0, 0, 6, 0, 0, 6, 0, 0], [0, 0, 0, 0, 6, 6, 6, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 2, 2, 2, 0, 0, 0, 0], [0, 0, 2, 0, 0, 2, 0, 0, 0], [0, 0, 0, 2, 2, 2, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0]] Correct Output: [[0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 6, 6, 6, 0, 0, 0, 0], [0, 0, 6, 0, 0, 6, 0, 0, 0], [0, 0, 0, 6, 0, 0, 6, 0, 0], [0, 0, 0, 0, 6, 0, 6, 0, 0], [0, 0, 0, 0, 6, 6, 6, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 2, 2, 2, 0, 0, 0], [0, 0, 0, 2, 0, 2, 0, 0, 0], [0, 0, 0, 2, 2, 2, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0]] Code Output: [[0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 6, 6, 6, 0, 0, 0, 0], [0, 0, 6, 0, 0, 6, 0, 0, 0], [0, 0, 0, 6, 0, 0, 6, 0, 0], [0, 0, 0, 0, 6, 0, 0, 6, 0], [0, 0, 0, 0, 0, 6, 6, 6, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 2, 2, 2, 0, 0, 0], [0, 0, 0, 2, 0, 0, 2, 0, 0], [0, 0, 0, 2, 2, 2, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0]] Task 2: Input: [[0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 8, 8, 8, 8, 8, 0, 0, 0], [0, 8, 0, 0, 0, 0, 8, 0, 0], [0, 0, 8, 0, 0, 0, 0, 8, 0], [0, 0, 0, 8, 0, 0, 0, 0, 8], [0, Justification: We use public benchmarks and follow the experimental settings of the stateof-the-art methods. For example, we set the maximum number of LLM requests to 64 to be consistent with Hypothesis Search [23] . We also attach the code which includes all the details. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
problem statement and assumptions: 
 Definitions: Specification, (⊢). We are given a programming problem to solve which comes equipped with an efficiently-checkable specification Φ. Our objective is to construct a program, written ρ, which satisfies the specification, written ρ ⊢ Φ. Generally, the specification decomposes into a conjunction of more basic constraints, notated {ϕ k } K k=1 , whereΦ = ϕ 1 ∧ ϕ 2 ∧ ϕ 3 ∧ • • • ∧ ϕ K (2)For example, a common form for Φ is a collection of input-output examples, with each ϕ k corresponding to a specific input-output pair. As another example, when synthesizing loop invariants for software verification problems, each ϕ i corresponds to a different verification condition.Definition: Counterexamples. Given a program ρ and a specification Φ of the form ϕ 1 ∧ϕ 2 ∧• • • ϕ K , a counterexample is a conjunct ϕ k where ρ ̸ ⊢ ϕ k . We define the set of all counterexamples as COUNTEREXAMPLES ρ,EQUATIONRefinement. Refining a program prompts an LLM to improve it in ways that are likely to cause it to satisfy the specification. We write P refine (•|ρ, Φ) to mean the distribution of possible refinements. Although the exact design of this distribution depends on the program synthesis domain, the refinement operator is generally implemented by prompting an LLM with a random counterexample:EQUATIONHeuristic measures of progress. Our method will organize its search strategy so as to prioritize programs that appear to be on the right track toward satisfying the specification. In general we assume access to a blackbox heuristic estimator of program goodness, h(ρ), with values ranging from zero to one. Within this paper we work with a very basic heuristic that simply reports the fraction of specifications that are satisfied:EQUATIONMore complex heuristics are possible, such as those which weigh certain parts of the specification more heavily, or which apply more fine-grained notions of approximate correctness, such as having low edit distance to a target output.4 REx: Refine, Explore, ExploitAt a high level, we frame this problem as an arm-acquiring bandit, meaning a bandit problem where new arms arrive over time [8] . Here each arm is a program, and "pulling" an arm corresponds to refining a program. Because refinement is stochastic, we receive a random reward after each refinement, depending on whether the newly generated program satisfies the specification. New programs arrive over time, because with each refinement, we generate a new program, hence the framing as "arm-acquiring." Our approach derives from Thompson Sampling, a stochastic algorithm that pulls an arm with probability equal to the probability that it is the best arm. Therefore we need to calculate probabilistic beliefs of the optimality of a given arm, which will be biased by the heuristic function h to prioritize programs with high heuristic value. The derivation of our Thompson Sampler will also automatically penalize programs that have been refined many times.Concretely, we receive a reward of 1 if pulling an arm (refining a program) yields a new program that satisfies Φ, and otherwise receive reward zero. Therefore the reward r follows a Bernoulli distribution, whose parameter we will write θ ρ :EQUATIONAs we solve this bandit problem using Thompson Sampling [9, 10] , we need to maintain probabilistic beliefs about each arm's expected reward, θ ρ , which get updated with each refinement. This means we need a prior over θ ρ , for which we choose the Beta distribution, which is conjugate to the Bernoulli, yielding simple Bayesian updates [11] . We choose a prior that places more mass on θ ρ whenever h(ρ) is also large, effectively injecting heuristic knowledge into the prior:EQUATIONEQUATIONEQUATIONwhere C is a hyperparameter. Large C gives greedier behavior by increasing the importance of h.The posterior beliefs over ρ get updated whenever ρ is refined. Assume the program (arm) has been refined (pulled) N ρ times, all with no success (no reward). Then the posterior is EQUATION*p.heuristic_value, # 1 + C × h(ρ) in paper 1 + C*(1-p.heuristic_value)+failed_cnt[p] # 1 + C × (1 -h(ρ)) + Nρ )) new_program = program.refine(problem) # ρ ′ ∼ P refine (•|ρ, Φ) if is_solved(problem, new_program): # ρ ′ ⊢ Φ in paper return new_program failed_cnt[program] += 1 programs.add(new_program)Understanding the behavior of REx. Intuitively, we begin with optimistic beliefs about the utility of refining a program. The strength of that optimism is controlled by h, with higher heuristic-value programs having higher initial estimates of θ ρ . With each unsuccessful refinement, we update our beliefs to be less optimistic, with the expected reward decaying toward zero with increasing number of failures: N , the number of times it was previously refined, and (2) h, the heuristic estimate of how close we are to satisfying the specification (larger h is better). Left: Expected benefit of refining decreases the more we refine, and asymptotically decays to zero (Eq. 11). Middle/Right: Posterior beliefs initially center around h and shift toward zero with each additional refinement. Same colored curves show same values of h for different values of N . The hyperparameter C modulates the rate of decay with each additional refinement, and also affects the initial concentration of the density around h. However, the system actually works with the full posterior distribution over θ ρ , not merely its expectation. Fig. 2 (middle-right) illustrate how the posterior evolves depending on the number of times that a program has been refined, showing that it initially concentrates its mass around h(ρ), and shifts downward with each refinement, making it progressively less likely to be refined further, but maintaining the property that every program always has a nonzero chance of being the next refinement. The variance also decays following VarEQUATION[θ ρ | N ρ ] = O N -3 ρ, meaning that a program that has already been heavily refined is not only lower in expectation, but also less likely to be randomly 'bumped up' and promoted as the next arm to pull by Thompson Sampling. 
experimental results: 
 We study three different domains that each involve code generation. See Fig. 3 . (Followup work [12] also applies REx to learning world models, where its scales to synthesizing 200+ line programs):1. Competition Programming. We benchmark on APPS, one of the most challenging LLM programming problem benchmarks [13] . APPS problems involved generating self-contained Python code that solves an algorithmic puzzle, given a natural-language problem description. APPS is split into three difficulty levels: Competition, Interview, and Introductory. Competitionlevel problems are very challenging, with landmark LLMs such as AlphaCode only solving 8% of APPS Competition [14] , making APPS substantially more difficult than basic benchmarks such as HumanEval [15] and MBPP [16] . 2. Visual Reasoning. We take visual reasoning puzzles from the Abstraction and Reasoning Corpus (ARC [7, 17] ). In ARC the goal is to infer an image-manipulating program from input-output examples, effectively combining inductive reasoning and visual/geometric puzzle solving. We work with a 40-problem subset of ARC introduced in [18] that has been annotated with naturallanguage descriptions of the target function. 3. Loop Invariant Synthesis for Software Verification. Discovering loop invariants is an important formal verification task; see Appendix A.3 for a primer on this problem statement. We collect 38 non-linear loop invariant synthesis tasks [19] from [20, 21] . These previous works rely on manually supplied templates/features and on dynamic program analyses (i.e., running the code). In contrast, we solely use source code without any dynamic execution traces or feature engineering.We check all three criteria of being a sufficiently strong inductive invariant (i.e., precondition, induction, and postcondition) with the Z3 solver [22] . The Z3 solver timeout is set to 2 minutes.We use these datasets to study several research questions: (1) For a given compute budget, which approaches to refinement solve the most problems? ( For visual reasoning, the goal is to synthesize an image-manipulating program that translates input images to the correct outputs. For software verification, the goal is to synthesize logical conditions as a valid loop invariant, in order to formally verify the functionality of the code. For competition programming, the goal is to generate an algorithm in Python.1. Greedy baseline refines the program with the highest heuristic value. It has a hyperparameter corresponding to the heuristic value of the initial (blank) program.2. Breadth-First Search expands the refinement tree breadth-first. Its single hyperparameter is the width (branching factor), searching deeper and wider the larger the compute budget.3. Fixed-Width generates a fixed number of initial programs, and then refines them roundrobin, while never refining the same program more than once [23, 24] . This generates a refinement tree of fixed width, searching deeper (but not wider) the larger the compute budget. There is a single hyperparameter corresponding to the initial width.Problems solved as a function of compute budget. Fig. 4 and Fig. 5 show programming problems solved as the number of refinements increases, varying method and hyperparameters. We consider both the number of problems solved at each level of compute budget, as well as the total Area Under Curve (AUC). While the exact rank order of methods depends on the dataset and the LLM, we see consistently that REx solves the most problems. However, the final number of problems solved is typically only modestly larger for REx. What method counts as second-best varies by domain.For example, Fixed-Width is nearly as good as REx on APPS-Competition and ARC, but on Loop Invariants and APPS-Introductory, Fixed-Width is the worst-performing method.Figure 4 : Comparing REx with alternatives using GPT-4 (temp=1). BFS and FW are Breadth First Search and Fixed Width, respectively. AUC denotes Area Under the Curve, and Final denotes the success rate at the maximum # LLM calls (64 for ARC and 300 for others due to domain conventions). Dark lines show performance with the best hyper-parameter setting for each method. Light lines show each run on each hyperparameter. The inset box plots show the distribution while varying the hyper-parameters. APPS baselines: Parsel [25] , AlphaCode [14] , and Olausson et al. [3] . Nonlinear Loop Invariant baselines: Z3/GSpacer [26] and Yao et al. [21] . ARC baseline: Hypo. Search [23] .More results on APPS Interview-Level and ARC in Figure 11 and Figure 13 We hypothesize that our robustness across datasets may come from the ability of REx to adaptively choose whether to search wide or deep. The other methods commit to a particular depth/width tradeoff, and how to optimally balance depth/width can vary by dataset. Indeed, examining the trees produced by REx shows that the algorithm searches quite broadly but refines deeply along more promising branches (Fig. 7 ), in ways that are reminiscent of Monte Carlo Tree Search. Speedups. Given the high monetary and environmental cost of frontier LLMs, it is valuable to understand whether different methods demand different levels of compute when solving the same problems. Fig. 6 analyzes these speedups, finding that REx requires less compute than alternatives: roughly 2× less on APPS-Competition, 2-5× less on loop invariants, and 1.1-1.6× less on ARC. We also see that there is little advantage to our method on easier benchmarks: for example, there is no speedup on APPS Introductory. By definition, harder problems take more time, therefore results on the hardest problems are the most relevant for saving time and compute. Box plot shows range of results across hyperparameters. "x" denotes the factor by which our method is faster than the compared baseline on this benchmark (ratio of median #LLM calls).Solving hard problems. In the large compute limit REx asymptotically solves modestly more problems than the baselines, particularly on harder datasets (Fig. 4 , APPS Competition vs Intro). How do these performance levels compare to other recent works on the same problem domains? On ARC, the state-of-the-art is Hypothesis Search [23] , which we replicate, finding REx solves more problems. 1 On Loop Invariants, we solve more problems than general purpose solvers such as Z3-GSpacer [26, 22] as well as the state-of-the-art solver that is specially designed for this dataset [21] .The specially designed solver individually tuned the feature space and the hyper-parameters for each benchmark problem. It also assumes access to runtime traces of the code, and the ability to run the code thousands of times. Yet REx outperforms it by a large margin (73.7% v.s. 60.5%2 ), with no special tuning. REx is also, to the best of our knowledge, state-of-the-art on APPS-Competition with GPT4 (Fig. 4 ). REx therefore not only solves problems that our baselines cannot, but can set new state of the arts on the specific datasets we apply it to. For the other methods, the optimal hyperparameter varies more widely: For example, even though Olausson et al. [3] suggests larger width for the fixed-width approach, we find smaller width can be better, depending on the dataset: Fixed-Width benefits from larger widths on ARC and Nonlinear Loop Inv, but performs better on APPS at smaller widths. Appendix Tbl. 1 shows these and other results in more detail, finding that our method is significantly more robust to variation in hyperparameters, while the other methods degrade more when their hyperparameters are not tuned to each specific dataset. We speculate that this may be due to the same reason that REx is more robust across domains: that it can adaptively decide when and where to search wide or deep. 
related work: 
 Code refinement. Prompting an LLM to fix problems with code that it generates has been explored in recent work [1, 4, 3] . To the best of our knowledge, our work is the first to systematically explore different policies for guiding refinement, and to introduce new bandit-based policies. Other works have explored fine-tuning specialized models to perform refinement, which allows using smaller models [6, 27] , and which could be synergistic with our work here. More broadly within computer science, there is a long history of research on automated program repair using non-neural methods such as genetic search and constraint solving [28, 29, 30, 31, i. a.], which generally considers a different problem setting: assisting the fixing of human-written code, often within the context of a large code base. It is also possible to combine these classical methods with contemporary machine learning [32, 33] .Bandits and Tree Search. That tree search algorithms introduce an explore-exploit tradeoff has been long appreciated. Monte Carlo Tree Search (MCTS) is the canonical tree search algorithm that relies on this insight [34] . Although REx derives from a similar insight, its structure is very different from MCTS. Unlike MCTS, we do not perform rollouts, backups, or tree traversals. Instead, REx exploits unique properties of the refinement process: Any node can be further expanded (infinitely), every node can be treated as a terminal node, and node expansions are very expensive (so rollouts would not be practical). To the best of our knowledge it is not possible to apply MCTS or any of its standard variants [35, 36, 37] to our problem statement out-of-the-box because of these unique properties. Bandits-based algorithms have also been applied for fuzz testing [38, 39] where code coverage is more important and code mutation is performed instead of code refinement. 
limitations and future directions: 
 We see that our method is only modestly more effective at actually solving more problems overall in the large-compute limit: Its advantages are largest when viewed through the lens of minimizing cost, and of being robust across hyperparameters and datasets. Progress on cracking the hardest problems, such as solving all of ARC or all of APPS, is more likely to come from improvements to the base model, and less likely to come from improved ways of organizing the refinement search process.We merely use the pass-rates of the programs to guide the search. There are richer information and more advanced bandits-based algorithms to integrate for better performance. A natural extension of our REx is to use contextualized bandits algorithms conditioned on the problem statement, the program, and the refinement history.A. θ ρ ∼ Beta(1 + C × h(ρ), 1 + C × (1 -h(ρ)) + N ρ ) ▷ sample ρ ← arg max ρ∈progs θ ρ ▷ select arm ρ ′ ∼ P refine (•|ρ, Φ) ▷ call LLM to refine N ρ ← N ρ + 1 ▷ increment count of number of refinements ▷ Arm arrival progs← progs ∪ {ρ ′ } ▷ add to set of arms N ρ ′ ← 0 ▷ initialize counter to zero until ρ ′ ⊢ Φ return ρ ′ A.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: We visualize and analyze the distribution using box plots for each method in different metrics varying the hyper-parameters in Figure 4 and Figure 6 . We also plot the learning curves for each experiment in Figure 4 . Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: Code is implemented by ourselves. The benchmarks and data splits are publicly available.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes]Justification: MIT license.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: The paper does not involve crowdsourcing nor research with human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.Our replication gets higher numbers than[23], potentially due to using different subsets of ARC.Our evaluation also checks if the invariant is valid (Fig.8), so the number is lower than reported in[21]. 
 
 Appendix:  


Decision: This paper proposes tackling the search process of (multi-turn) code repair as an armed bandit. In an evaluation of multiple tasks and a few LLMs this method shows to perform mildly better than baselines reducing the cost for achieving similar results to a baseline.
This method seems useful to the community, seems relatively simple to implement, and achieves tangible improvements. While additional information can be used by the banding (_e.g._ depth) and further improvements can be made, this technique is technically sound and hence all the reviewers and myself believe that there is no reason to reject this work. 

I would kindly ask the authors to strongly consider incorporating some of the additional ablations suggested by reviewer `ikA5` in their camera-ready.

Reviews: review: paper_summary: paper_summary: The paper identifies that LLM refinement process can be formulated as (arm-acquiring) non-contextual bandit problem which can be solved optimally (in the limit) using principled bandit-algorithms like Thompson Sampling against heuristic based solutions. It applies this idea to three code refinement tasks and demonstrates empirical improvements from their approach.
strength_weakness: strength_weakness: 1. **Non-contextual Bandit formulation**. While the bandit formulation is clean and simple, it leaves more to be desired and makes simplifying assumptions. Particularly, since the bandit formulation neither contextualizes on the problem statement nor considers "depth" of the arm as a special attribute. This seems different from expected behavior

    1.1 **Problem Statement**. Problem statement links is directly linked with the complexity of task. For very challenging tasks, optimal strategies would differ from easy tasks, where for challenging tasks more exploration across "width" might be expected.
    1.2 **Depth of Arm**. It might be reasonable to have different priors for "arms" at different "depths". For example, pulling the "empty program arm" seems special compared against other arms.

At the same time, without explicitly modeling these contextual components the proposed approach seems strong -- perhaps accounted for by more "diverse" hyper-parameter search space (discussed further in Questions section point 1.). 

2. **Choice of LLM**. As authors mentioned, choice of the LLM might play a role in the performance of this work. It might be useful to study this axis further and see if the findings change by taking a weaker model like `gpt-3.5-turbo`. 

3. **Invariant Task Grading**. The invariant task is a core task studied in the paper and the authors point out that they check 
> being a sufficiently strong inductive invariant (i.e., precondition, induction, and postcondition). 

The statement is not precise enough and needs to be further clarified. Additionally, is it possible for models to generate trivial invariants (e.g. - `True is True`) and fool the grader?

4. **Contamination**. APPS benchmark is considerably older and released before the cutoff date for newer GPT-4 models (authors do not specify the GPT-4 version used here).

5. **Missing details**. There are some missing/hard-to-find experimental details in the Appendix. Particularly, what feedback is used for refinement of invariant task based on the types of mistakes it makes -- weak invariant, failing invariant, solver timeout.
questions: questions: 1. Hyper-parameter variance. It seems that optimal hyper-parameters vary considerably across tasks (Table 1 in Appendix). Perhaps, it can be viewed that the hyper-parameter choice C allows better control over the exploration-exploitation search space allowing it to exhibit widely different behaviors -- accounting for the lack of contextual features discussion in weaknesses point 1?

2. **Connecting the Appendix**. There are many interesting discussions and details in the Appendix (e.g. A.2) etc which are not connected in the main paper. I would recommend the authors to add them (and at a minimum link) to the main paper. Similarly, lot of experimental settings are scattered throughout the Appendix which should be linked from the main paper via forward references making them more accessible.

3. **Figure 3**. Figure three while highlights the tasks considered is perhaps considerably larger and authors can truncate the examples and provide more experimental details and results in the main paper.
limitations: limitations: Thoughtful discussion on limitations of the approach in a.) allowing solving more problems and b.) differences arising from the choice of LLM is already provided.
paper_summary: paper_summary: This paper proposes to improve the iterative code refinement process by prioritizing the “good” programs, where the goodness is defined by a heuristic estimator -- the program that passes the more test cases is better. To balance exploration (explore a lesser refined program) and exploitation (refine the best program), the paper formulates this problem as a bandit problem, and solves with Thompson Sampling. The proposed method achieves modest but consistent improvements across various program synthesis benchmarks (loop invariant, visual transformation program, competition problems).
strength_weakness: strength_weakness: ### **Technique**

The scope of this paper is limited to solving small, isolated programming challenges given a relatively large compute limit. It is unclear whether such techniques can be applied to more realistic settings (where we cannot afford large numbers of full program samples). For example, to repair a bug in a repository, it might be more important to improve the code refinement step itself, either with better prompting or training.

It is also unclear whether similar results can be achieved using more advanced prompting techniques (e.g., few-shot or CoT).

Formulating the problem as a bandit problem seems a little unnatural, as the algorithm terminates as soon as a positive reward is observed. In other words, the accumulated reward is always 0 before termination. It’d be great if the paper can include more theoretical discussion regarding how this special assumption would affect the theoretical guarantee of the algorithm. 

### **Evaluation**
The improvement of REx seems marginal.
- Compared with alternative search algorithms, REx often brings only marginal improvements given a large enough budget (Figure 4). 
- Also on the some dataset (loop invariant), BFS (with the optimal hyperparameter) is even better than REx with smaller sampling budget.

The authors should list the cost of baselines for fair comparison.
- In Figure 4, the existing LLM-based methods should be marked as data points in the plot, with their corresponding sample budget (if they use the same base LM).


According to Table 1 in appendix, the Greedy baseline seems to be very strong, even though the paper only considered two values for its single hyperparameter. This leads me to wonder whether such a simple baseline has more potential.
- The hyperparameter, specifically the heuristic value of the initial program, controls when to resample a new program instead of refining existing ones. This seems to be highly problem-specific. Instead of having a fixed value for all problems, can we maintain a (moving) average of all sampled solutions for one particular problem?

Also, each arm is not independent, as the refined program $\rho^\prime$ is correlated with the original program $\rho$.
- I am curious about whether it is possible to better incorporate the observations of the refined program(s) to update the posterior belief of the original program. Although the current reward (refined program $\rho^\prime$ satisfies all input-output pairs) is very sparse, we may use the heuristic metric $h(\rho^\prime)$ as an alternative.
- In fact, MAB has been applied for fuzz testing [a,b], where code mutation is performed instead of code refinement. [a] formulates the problem as hierarchical multi-armed bandits, and [b] also formulates the problem as arm-acquiring bandits. It could be an interesting line of related work to check.

[a] Ran, Dezhi, et al. "Badge: prioritizing UI events with hierarchical multi-armed bandits for automated UI testing." 2023 IEEE/ACM 45th International Conference on Software Engineering (ICSE). IEEE, 2023.

[b] Yang, Chenyuan, et al. "White-box compiler fuzzing empowered by large language models." arXiv 2023



### **Implementation**

Some implementation details are unclear. For example, what is the temperature for ARC and loop invariant experiments? Appendix A5 lists temperature=1.0 for APPS, is it the same for all considered benchmarks?

### **Minor**

Formula (11):  $2+C+N_\rho$ should be $2+2C+N_\rho$

Figure-6: The x-axis is somewhat unclear. Does it represent the number of requests needed to achieve performance comparable to the best baselines? Please clarify.
questions: questions: Could you please explain and address the weaknesses?
limitations: limitations: The limitations are well addressed.
paper_summary: paper_summary: This paper explores a variety of code generation tasks, taking the approach of using a LLM to generate solutions, and conditioning the generation of each solution on the repair of a previously generated solution. It frames this as an arm-acquiring bandit problem, where each solution generated is an arm, and pulling an arm means refining a previous solution with an LLM (specifically, with GPT-4). The paper applies Thompson Sampling to this problem; specifically the paper's method is to use a heuristic-driven prior for which solution to refine at each step (i.e. the fraction of test cases passed by that solution), and then to penalize a solution each time its refinement does not lead to a perfect solution. The method is conceptually simple, and results in stronger performance (i.e. more problems solved; fewer model calls used) than a few baseline approaches.
strength_weakness: strength_weakness: I noted as the first strength listed that the method can readily be put into practice by practitioners, but I must couch this assessment: the problem statement and method assume access to an efficient checker for whether the task is solved, which restricts the space of tasks to which the method can be applied.

A weakness of the paper is that only GPT-4 is tested. As the paper acknowledges, GPT-4 is among the strongest models for editing code. It is also among the more expensive models to run. By only evaluating the method with GPT-4, we are left to wonder about the cost tradeoffs of using smaller cheaper models with more samples vs this large model with fewer samples. Since one of the stated goals of the paper is to reduce cost (the paper states this as the goal of reducing the number of model calls), this would be valuable to understand. Since the results are high variance and of modest effect size, understanding how the method's performance varies across models would be valuable.

It is not clear that single-sample-refinement is the best framing of the problem; there are lots of ways to prompt a model beyond refinement from a single prior solution, and the paper does nothing to suggest that single-solution-refinement is a promising strategy. I think this is the biggest weakness of the problem statement; it assumes a rather rigid set of ways to apply LLMs to solving these programming tasks, and then optimizes within that rigid constraint, without justifying adequately that it is a valuable constraint.

---

Overall, given the constraints of the problem statement, the method presented is a step forward in terms of performance, formulating an idea from bandits for this program synthesis setting and achieving improved results. A key aspect of the paper's contribution is that the paper frames the search tree as infinite width and infinite depth; I cannot meaningfully comment of the novelty of this framing, though given the limited search budgets of a few hundred LLM calls the claim of infinite width and infinite depth doesn't seem critical. The paper is written clearly and the method presented is simple to follow and implement. The main restriction on the significance of the work is that the constraints made by the problem statement -- that the method works by iteratively refining a single preexisting solution at a time -- are overly constraining and not representative of the range of ways in which people apply LLMs today.
questions: questions: In the related work section you state that it is not possible to apply MCTS or any of its standard variants to the problem statement out-of-the-box. Could you elaborate on (or refine) this position, and comment on whether small modifications would be sufficient to apply MCTS to this problem?

Since the LLM is shown the full specification (Page 3, eq (4)), there is risk of it generating solutions that satisfy the specification literally without generalizing to satisfy the intent of the specification. A simple example of this would be if the model writes special case handling for each of the constraints in the specification, rather than capturing the underlying problem to be solved. Do you observe this in practice at all?

How do methods that fit the problem statement, i.e. methods that operate by only refining one solution at a time to produce new solutions, compare with other approaches to using LLMs to solve these tasks. Showcasing what other approaches have been applied and how they perform, or alternatively making a strong case that this particular formulation is a good choice, would improve the paper.

At Line 92 you consider an alternative heuristic of edit-distance to a target output. Am I correct in understanding that having a target output available for the heuristic would obviate the need for applying the method?

One limitation of the approach is that as information comes in about proposed solutions, the system does not use that to learn anything about other previous solutions, even if they are heavily related. Similar to weakness 3 and question 3, approaches that leverage this source of information sit outside the problem statement. Could you comment on the relevance of such approaches, either comparing with them or justifying their absence?

nit: Figure 1, Right: I would encourage you to label the outgoing edges "exploit" and "explore" rather than labeling the nodes or incoming edges as is currently done. This is because the current figure makes it look as if generating labelled nodes was done as exploitation vs exploration, where in fact it is refining those nodes that would be exploitation or exploration.
limitations: limitations: The authors address the limitations of the work in Section 7, Limitations. They explain the effect of only testing their approach on GPT-4, and discuss the limited effect size of their approach at solving more problems overall.
paper_summary: paper_summary: This paper studies the code refinement problem: given a candidate program and the reason why the program fails to satisfy the user specification, an LLM is called to generate an improved program. The paper frames this as an arm-acquiring bandit problem and solves this using Thompson sampling. The evaluation is performed on three code generation tasks: competitive programming, visual program reasoning, and loop invariant generation. The results show that while the proposed approach solves only moderately more tasks than baseline given a large-compute limit, it is significantly faster to achieve a certain performance.
strength_weakness: strength_weakness: REx seems like a straightforward application of Thompson sampling. Moreover, it only gives moderately improvement over baseline search methods in a large-compute setting. I do appreciate simplicity when it is combined with effectiveness, but this is not the case for REx. Therefore, I consider the technical contribution of this paper rather thin.

Moreover, a major selling point of the paper is to study an exploration-exploitation tradeoff. While the paper provides some preliminary study at Lines 179-182, I believe a more thorough analysis is needed for full understanding.

Another limitation, as also pointed out by the paper, is that only GPT-4 is studied. Would REx give better results in the large-compute setting when the self-repairing capability of the base model is weaker and the search strategy might play a more important role?
questions: questions: Please consider addressing the points raised in the “Weakness” section.
limitations: limitations: The paper has sufficiently addressed the points concerning limitations and societal impact.

Comments: title: Thanks for the rebuttal
Thanks for the rebuttal
comment: Thanks for your rebuttal. The inclusion of additional language models makes the results more compelling.

A key remaining concern is that it is not clear that single-sample-refinement is the best framing of the problem; there are lots of ways to prompt a model beyond refinement from a single prior solution, and the paper does nothing to suggest that single-solution-refinement is a promising strategy compared to these alternatives.

Here are a couple examples of categories of alternatives: single-sample prompting techniques like scratchpad/Chain-of-Thought and all its variants, LLM-chain approaches that generate a single result, like ReAct and its variants (one can view these as improvements on how to refine, rather than what to refine), and refinement approaches like EvoPrompting that consider more than one sample at a time.
Thanks for answering my questions. I will maintain my rating.
Thanks for answering all my questions! I'll be keeping my score of 6 and support the acceptance of the paper.
Thanks for running the experiments! The rebuttal has addressed my main concerns. Therefore, I am raising my score from 4 to 5. I hope the authors could consolidate the experiments on more LLMs, the application of REx to more challenging tasks, and the analysis of the exploration-exploitation tradeoff. Including these results would help strengthen the paper.