Code Repair with LLMs gives an Exploration-Exploitation Tradeoff
Authors: Hao Tang
Keya Hu
Jin Peng Zhou
Si Cheng Zhong
Wei-Long Zheng
Xujie Si
Kevin Ellis
Venue: NeurIPS 2024 poster (2024)
Keywords: Program Synthesis
LLM
Primary area: natural_language_processing

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.

Code Repair with LLMs gives an Exploration-Expooittiio Tradeoff
Hao Tang Cornell University
Keya Hu
Shanghai Jiao Tong University hu hu_keya@sj tu. edu. cn
haotang@cs. cwraggesscconnlll... www.ell..eeeeoooooooooooeooeeeoe edu
Jin Peng Zhou
Sicheng Zhong
Cornell University pzzhou@cs cornell.
University of Toronto sicheng. zhong@mai 1 .utooonto.caa ca
Wei-Long Zheng Shanghai Jiao Tong University we Long@sj tu. edu.cn cn
Xujie Si University of Toronto CIFAR Al Chair, Mila sixêcs. tor
Kevin Ellis Cornell University
ke kellis@cornell.edd
Abstract
Iteratively improving and repairing source code with large language models (LLMs), known as refinement, has emerged as popular way of generating pro- grams that would be too complex to construct in one shot. Given bank of test cases, together with candidate program, an LLM can improve that program by being prompted with failed test cases. But it remains open question how to best iteratively refine code, with prior work employing simple greedy or breadth-firs. strategies. We show here that refinement exposes an explore-exploit tradeoff:: ex- ploit by refining the program that passes the most test cases, or explore by refining lesser considered program. We frame this as arm ren-acqurinng 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.
Introduction
An emerging paradigm for problem-solving with large language models (LLMs) is to have the language model correct, repair, debug its initial outputs [1] 2 which we refer here as refinement. For example, when generating the source code of a program, refinement prompts the I.LM with a buggy program previously generated, potentially with diagnostic information such as stack trace, then asks it to fix its code. This strategy alleviates the need for the L.LM 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. [eftt. This tree is infinitely deep, because every refinement can itself be further refined. It also infinitely wide, because the LLM can return infinitely many possible refinements. Success should, in principle, depend on exactly what policy is used expand this tree. Recent work 6 adopts simple expansion policies (c.g., breadth-fiin first),,
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

initial prompt
initial prompt
initial samples
exploit
explore
80% correct
60% correct
refined once
0% correct
refined twice
Figure 1: Left: The tree of possible refinements is infinitely deep and has infinite branching factor. Each node is program and each edge is L.LM sample. Right: Explore-Explos tradeoff for search state after performing node expansions. Exploit by sampling another child of a program that is nearly correct, or Explore by sampling child of 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 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. [ (right) diagrams this tradeoff for tree after the first few tradeoff is made more challenging by the fact that
node expansions. This 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) bandit- based node expansion our branching factor is infinite, our transition function is stochastic, and "rollouts" would demand 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 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 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 A Abstraction and Reasoning corpus (ARC: (7). Across every domain, REx solves more problems 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.
Background: Bandits and Thompson Sampling
Bandit problems concern maximizing total reward given set of actions, A. each with an unknown reward distribution, P(rla), for each action a E A. At each time step I,, an action a, is chosen and a reward is received. Bandit problems are challenging because they involve balancing exploration and exploitation: On the one hand, each action must be tried enough times to estimate its average reward (cxploration), but asymptotically, only the actions with the highest expected reward should be chosen (exploitation). Actions are called arms, and taking an action is called pulling that arm.
Thompson Sampling is framework for designing bandit strategies that performs probability matching: pulling an arm with probability equal to the odds that that arm is optimal (has the highest expected reward). It maintains probabilistic beliefs about the reward distribution of each arm, and updates these beliefs following Bayes rule. Writing 0,, for the parameters of the reward distribution for arm this means that P(rla) = P(r(0.) (by definition), and that Og is treated as random variable with a prior P(O.). Beliefs are updated after each reward, and the next arm can be selected by sampling 6, from the posteriors, and pulling the arm with the highest (expected) reward. Concretely,

---

Thompson Sampling does the following at timestep t:
Sample 0a ~ P(Oal (r, at = a)), for each a E .A
Do the above using Bayes Rule: P(Oal frr ar = a}} x P(O,) II P(re(0a) ta,=0
Pull arm de = arg max E [r 0a]
aeA
Get reward rz.
By choosing different priors and likelihoods, variety of Thompson sampling strategies can be designed. For example, prior knowledge of the value of each arm can be embedded in P(Oa). and constraints on the range of possible rewards can embedded in P(r(0.).
Problem Statement and Assumptions
Definitions: Specification, (H). We are given programming problem solve which comes equipped with an setccinnll-checkabble specification ④. Our objective is to construct a program, written p. which satisfies the specification, written p - $. Generally, the specification decomposes into a conjunction of more basic constraints, notated {Øx} 1 where
+ Φ1 02 O3 A A...A. ΔΦ
For example, common form for & is collection of onput-output examples, with each Φ corre- sponding specific input-output pair. As another example, when synthesizing loop invariants for software verification problems, each o; corresponds to a different verification condition.
Definition: Counterexamples. Given a program p and a specification & of the form A ΦK. a counterexample is a conjunct Ø% where of Φk. We define the set of all counterexamples as
COUNTEREXAMPLES (p, (0 to
dx
Refinement. Refining program prompts LLM improve in ways that are likely to cause to satisfy the specification. We write ( b) to mean the distribution of possible refinements. Al- though the exact design of this distribution depends on the program synthesis domain, the refinement operator is generally implemented by prompting an LLM with random counterexample:
lp.
E Prefine (p' |p, BO-E([COUNTEEXXAMPEES....))
PPIM (p' prompt(p, o; +))]
Heuristic 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 assume access to a blackbox heuristic estimator of program goodness, h(p), with values ranging from zero one. Within this paper we work with very basic heuristic that simply reports the fraction of specifications that are satisfied:
h(p) Σx 1 Φx
(5)
More 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 a target output.
4 REx: Refine, Explore, Exploit
At high level, we frame this problem where new arms arrive over time g. Here each arm is program, and "pulling" an arm corresponds to refining a program. Because refinement is stochastic, we receive random reward after each refinement, depending on whether the newly generated program satisfies the specification. New programs arrive over time, because with each refinement, generate new program, hence the
arm-acquiring bandit, meaning bandit problem

---

framing as "arm-acquie acquiring, Our approach derives from Thompson Sampling, stochastic algorithm that pulls an arm with probability equal to the probability that it the best arn. Therefore we need calculate probabilistic beliefs of the optimality of given arm, which will be biased by the heuristic function to prioritize programs with high heuristic value. The derivation our Thompson Sampler will also automatically penalize programs that have been refined many times.
Concretely, receive reward of pulling an arm (refining program) yields new program that satisfies 4, and otherwise receive reward zero. Therefore the reward r follows Bernoulli distribution, whose parameter we will write pp:
P = P(r = lp, 1)
E
[p J +]]
p' P-
we need to maintain probabilistic belief' about each arm's expected reward, 0p. which get updated with each refinement. This means need a prior over pp. for which we choose the Beta distribution, which is conjugate to the Bernoulli, yielding simple Bayesian updates [T]]. We choose a prior that places more mass 8p whenever h())) is also large, effectively injecting heuristic knowledge into the prior:
As we solve this bandit problem using Thompson Sampling (9] 10].
P(e,) Beta prior 3ptior
(ap
(7)
a = + C x hp))
(8)
|ptior C x (1 h(p)
where a hyperparameter. Large C' gives greedier behavior by increasing the importance of
The posterior beliefs over P 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 P)),IN)) P(N,)Op)
(1 (1) 0p) Beta( ggriior
oc Beta(app prior ,prio + Np)
BBtea
+
(10)
The above equation essentially defines REx. For each program we track its heuristic value and how many times has been refined, and to select the next program to refine, we sample from the Beta distributions in Eq. 10 and refine the program with highest sampled 0p. An implementation is about ten lines of Python (see below and Appendix Alg D. Refining an empty program means generating program from scratch. The heuristic value of empty program is unknown set C' hnulll)
x h(null) = NULL == O.
REx
def REx (Exobbbemm C)
programs (problem.empty sopty_soluttion)) fai led_eedcm cnt Deeaultadictllamdds: 0)
# No in paper
while True:
program = max key=lambda p: beta Exp.hhuristiic_wale,
1 1++Cx C h(p) paper
1 + Cx (c(_p.hhurristic__allee) failee
11 h(p)) + No
))
new_program program.
# p² Prefinoo-
is (prbble, new_program)
p' + in paper
return new_program
cetted_ent(program +=
programs wdw
Understanding the behavior of REx. Intuitively, we begin with optimistic belief's about the utility of refining a program. The strength of that optimism is controlled by ht with higher heuristic-value programs having higher initial estimates of 0p- With each unsuccessful refinement, update beliefs to be less optimistic, with the expected reward decaying toward zero with increasing number of failures:
1+C h(p))
E E N₂ =
(11)
2+ 2C'

---

expected reward, C= 10
density of mean reward, C - 10
density of mean reward, C 20
0.8
0.0.0
N-0
0.3
5
N=
N=1
0.6
0,6
N
E(a(M)
p(e(N)
h= 0.9
0.2
0.0
9 12 15
0.0
0.2 0.4 0.6 0.8 1.0
0.0
0.2
0.4
0.6 0.8
1.0
num refines, N
mean reward, e
mean reward, 8
Figure 2: How the model's beliefs about the benefit of refining a program, 0, change we vary (1) N, the number of times it was previously refined, and (2) h. the heuristic estimate of how close we are satisfying the specification (larger is better). Left: Expected benefit of refining decreases the more we refine, and asymptotically decays to zero (Eq. T 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 0p. not merely its expectation. Fig. riiddle-righh) 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(p), and shifts downward with each refinement, making progressively less likely to be refined further, but maintaining the property that every program always has nonzero chance of being the next refinement. The variance also decays following Var (Op 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. (Followup work [2] also applies REx to learning world models, where its scales to synthesizing 200+ line programs):
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-lanuuage problem description. APPS is split into three difficulty levels: Competition, Interview, and Introductory. Competition- level 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 [15 and MBPP [16].
HumanEval
Visual Reasoning, We take visual reasoning puzzles from the Abstraction and Reasoning Corpus (ARC [7). In ARC the goal is to infer an maage-manipulating program from input-output examples, effectively combining inductive reasoning and visual/geometric puzzle solving. We work with a 40- -0-problem subset of ARC introduced in [18 that has been annotated with natural- language descriptions of the target function.
3. Loop Invariant Synthesis for Software Verification. Discovering loop invariants is 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/fcatures and on dynamic program analyses (i.c., 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 sufficiently strong inductive invariant (i.e., precondition, induction, and postcondition) with the Z3 solver (22]. The Z3 solver timeout is set minutes.
We use these datasets to study several research questions: (1) For a given compute budget, which approaches to refinement solve the most problems? (2) Can REx solve hard problems other methods cannot? (3) How much of a speedup/cost-rrddettio methods to hyperparameter choices, and what are reasonable default hyperparameter settings for this new method? To investigate these questions, study range of baselines:
does REx grant? (4) How sensitive are different
5

---

Visual Reasoning Abstract and Reasoning Corpus Visual Reasoning Puzzle
Software Verification Nonlinear Loop Invariants Code to Verify
int mainQ x, yX
assert(x 1);
assert(y 1);
int 8F%.b=yp==1..00.000-==1 whilee1XX
//Loop Invariant Z3:
if(a==b) break;
(a>b)
a-b;
p=p-q:
r= P-;;
else b
9 q-P:
??
= s-r:)
/assent(x%a==0 && y%a==0)
//asser([vv > a,19 0 y% 0)
Human Hypothesis: connect same colored blocks diagonally.
return
Image-Manipulating Program
Loop Invariant
(Logical properties always hold during all iterations)
def transform_ inputt
color_list np.unique (input_grid)
for cotor in color_ Sitor_list::
indices = np. mp.arguhere m (input_grid m color)
if ten(indices] > indices(a]
p p*x+r*y==
y0,
rangel1. (on(indices));
y1, raxlabs(y2 ye),, abs(xl x8))
&& q x = b
steps eaxlsteps, 1)
[yl yet steps
&& (p * - x = 1
dx (x1 x01 steps
for 5 in rangeli, steps):
p * - q * = -1
y roundty@ dy
griddly]
&& a 2 0 && 0
return input_grid
Competition Programming, APPS Python Code
Programming Problem
Theta has been learning about powers of 2 school. She notices that some numbers when written out contain powers of their digit representation: for instance, 12560 contains 256 which is power of 2. She has been wondering how many such numbers there are.
power2 striz p ten(power2)
10" (p 1) do 1110 for rangela + 201 (=100"0000=-
rarge(21) for rangellen(n)
start potential match
part pover
fel
for 1 i3 rangeil, pl:
3 nect_pnerr?
pover2[s] pover2(i]:
next_ power2ts -1]
sower211) pover2151:
rest_power2[s]]
ranpel1en(n1): In for in rangelp else int(n[i)) 1)
Can you write a program that counts how many numbers contain given power of 2?
r_range
rangeld_rangel:
d int(nisI)
while and strid) power2[ak]:
ak - poxt_power2Ins striel pover2lnkl
dpli []]]] + dp(i](]](4) [k]
[ientn))] @b[lenatlit11p1
Figure 3: Evaluation domains. For visual reasoning, the goal is to synthesize an mange-manipulatinn 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 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 hyperparameter corresponding the heuristic value of the initial (blank) program.
2. Breadth-First Search expands the refinement tree breadth-first first. Its single hyperparameter is the width (branching factor), searching deeper and wider the larger the compute budget.
3. Winth generates fixed number of initial programs, and then refines them round- robin, while never refining the same program more than once [23] 24 This generates refinement tree of fixed width, searching deeper (but not wider) the larger the compute budget. There is single hyperparameter corresponding to the initial width.
Problems solved function of compute budget. Fig- and Fig. 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, well as the total Area Under Curve (AUC). While the exact rank order of methods depends on the dataset and the L.M,, we see consistently that REx solves the most problems. However, the final number of problems solved typically only modestly larger for REx. What method counts as second-best varies by domain. For example, Fixed-Width is nearly good REx APPS-Competinn and ARC, but Loop Invariants and IPPS-Introductory. Fixed-Wii Width is the worst-performing method.
Nonlinear Loop Invariant
ARC
80
60
Our replicaloon
Hypo. Search
70
Success Rate (%)
60
50
Yao et al. 2020
50
Final (%)
84
40
40
70
30
(%)
60
30
AUC (96)
58
20
70
Z3/GSpacer
60
55
10
REx
BFS Greedy
FW
20
REx
BFS Gready FW
0
50
100
150
200
250
300
10
20
30
40
50
60
APPS Competition-Level
APPS Introductory-Level
100
60
Olausson et al. 2023
90
Success Rate (%)
50
laassson el al. 2023
80
40
65
Final (%)
Final
60
70
30
Parsel
85
95
60
20
85
AUC (%)
AUC
10
AlpheCode
50
94
REx BFS Greedy FW
Greedy
100
150 200
250
300
Sample Budget LLM Requests)
Sample Budget Requests)
Figure 4: Comparing REx with alternatives using GPT-4 (temp=l). BFS and FW 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 syper-parammters. APPS baselines: Parsel 255 AlphaCode [14), and Olausson et al. [. Nonlinear
Loop Invariant baselines: Z3/GSpacer (26) and Yao al. (2]. ARC baseline: Hypo. Search 23
More results on APPS Interview-Lent and ARC in Figure 11 and Figure

---

We hypothesize that our robustness across datasets may come from the ability of REx to adaptively choose whether search wide or deep. The other methods commit to a particular depth/ 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.
GPT-3.5-turbo
Claude-3.5 Srtude-3.5-oomee 5-Sonnet
Llama-3 3. 1--005B
50
50
40
Supcess Rate (%)
20
15
20
15
AUC
10
10
35
30
5
Greedy
REX
Groody
BP8 Groody
0
10
20
30
40
50
0
10
20
30
40
50
0
10
20
30
40
50
Sample Budget
LLM Requests)
Sample Budget LLM Requests)
Sample Budget
LLM Requests)
Figure 5: Comparing REx with alternatives with other LLMs on competition programming (APPS Competition-Levee More results on ARC are available in Appendix in Figure 12
Speedups. Given the high monetary and environmental cost of frontier LLMs, valuable to understand whether different methods demand different levels of compute when solving the same problems. Fig. 6 amalyzes these speedups, finding that REx requires less compute than alternatives: roughly 2x less on APPS-Competition, 2-5x less on loop invariants, and 1.1-1.6x less on ARC. We also see that there is little advantage to our method on easier benchmarks: for example, there is no speedup APPS Introductory. By definition, harder problems take more time, therefore results on the hardest problems are the most relevant for saving time and compute.
LLM Roquests needed
Nonlinear Loop Invariant get best base ine
ARC
APPS Compettton-Lewe
APPS Introductory-LLev eve get ot 2023
to got Hypo. Search
get Olausson et al. 2023
300
300
60
200
50
200
200
40
100
100
30
100
REX
BFS
Greedy
FW
REx
BFS
Greedy
FW 1.12x
REx
BFS
Greedy
FW
REx 1.00x
BFS Greedy FW
1.00x
2.13x
2.42x 4.96x
1.00x
1.2 1.23x
59x
1 00x
2.88x
1.8 1.86x
2.46x
1.46x
0.48x 1.51x
Figure 6: Number of LLM Requests needed to get the best LLM-based state-of-the-ar art methods (for ARC, APPS) and the best LLM-basee -bsee baseline (for Nonlinear Loop Invariant). The lower the better. 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 of-te-of-the-art is Hypothesis Search (33. which we replicate, finding REx solves more problems On Loop Invariants,
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 [2I]. The specially designed solver individually tuned the feature space and the syper-parameters for each benchmark problem. also assumes access to runtime traces of the code, and the ability to run the code thousands of times. Yet REx outperforms it by large margin (73.7% v.s. 60.5%)) with no special tuning. REx is also, to the best of our knowledge, state-of-the-art on APPS-Competition with GPT4 (Fig. REx therefore not only solves problems that our baselines cannot, but can new state of the arts on the specific datasets we apply it to.
I Our replication gets higher numbers than 2 potentially due to using different subsets of ARC.
OOur evaluation also checks if the invariant is valid (Fig. the number lower than reported in [2].
8

---

10.0
80
20.3
20.0
so.3
200.
Figure 7: Example search tree generated by REx. Blue->Yellow gradient indicates heuristic value (blue: low, yellow: high). The order of children from left right indicates the sequence of generations (left: older, right: newer). See also appendix Fig. 14 17
Hyperparameter sensitivity. We consider range of hyperparameters for each method. Box-and- whisker plots in Fig. 4 5 and Fig. illustrraee the distribution of results for different methods, while varying hyperparameters. For REx, large C' values work well on all datasets (C = 20). For the other methods, the optimal hyperparameter varies more widely: For example, even though Olausson et al. 3 suggests larger width for the fixed-wiie width approach, find smaller width be better, depending on the dataset: Fixed-Wine benefits from larger widths on ARC and Nonlinear Loop Inv, but performs better APPS smaller widths. Appendix Tbl. shows these and other results in more detail, finding that our method significantly more robust to variation in hyperparameters, while the other methods degrade more when their hyperparameters are not tuned 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 search wide or deep.
6 Related Work
Code refinement. Prompting an LLM to fix problems with code that it generates has been explored in recent work 3. To the best of our knowledge, our work is the first systematically explore different policies for guiding refinement, and to introduce new bandit- have explored fine-tuning specialized models perform refinement, which allows using smaller models 27). and which could be synergistic with our work here. More broadly within computer science, there is long history of research on automated program repair using non-neural methods such as genetic search and constraint solving 28 29
policies. Other works
30 31 i.a.), which generally considers
different problem setting: assisting the fixing human- code, often within the context of
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-explae 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 similar insight, its structure is very different from MCTS. Unlike MCTS, we not perform rollouts, backups, or tree traversals. Instead, REx ex- ploits 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 is not possible 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 [8] 39] where code coverage is more important and code mutation is performed instead of code refinement.
7 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 solving all of ARC or all of APPS, more likely come from improvements 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 richer information and more advanced bandits-basee based algorithms to integrate for better performance. A natural extension of our REx is contextualized bandits algorithms conditioned on the problem statement, the program, and the refinement history.
Acknowledgements. This work was supported by NSF grant #2310350, gift from Cisco, and in part, by Individual Discovery Grants from the Natural Sciences and Engineering Research Council of Canada, and the Canada CIFAR AI Chair Program. We thank the Microsoft Accelerating Foundation Models Research (AFMR) program for providing generous support of Azure credits.
References
[1] Xinyun Chen, Maxwell Lin, Nathanael Schärli, and Denny Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2344.05128 2023.
[2] Sean Welleck, Ximing Lu, Peter West, Faeze Brahman, Tianxiao Shen, Daniel Khashabi, and Yejin Choi. Generating sequences by learning to self-correct. arXiv preprint arXiv:2211.0005:: 2022.
[3] Theo X. Olausson, Jeevana Priya Inala, Chenglong Wang, Jianfeng Gao, and Armando Solar- Lezama. Is self-repair silver bullet for code generation?, 2023.
[4] Noah Shinn, Beck Labash, and Ashwin Gopinath. Reflexion: an autonomous agent with dynamic memory and self-reflection. arXiv preprint 1xXXiv:333.11300 2023.
[5] Aman Madaan, Niket Tandon, Prakhar Gupta, Skyler Hallinan, Luyu Gao, Sarah Wiegreffe, Uri Alon, Nouha Dziri, Shrimai Prabhumoye, Yiming Yang, et al. Self-refine: Iterative refinement with self-feedback. Advances Neural Information Processing Systems, 36, 2024.
[6] Hung Le, Yue Wang, Akhilesh Deepak Gotmare, Silvio Savarese, and Steven Hoi. CodeRL: Mastering code generation through pretrained models and deep reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL http:::///ppereview net/
forum??dd-Wawbbo0zyy
[7] François Chollet. On the measure of intelligence, 2019.
[8] Peter Whittle. Arm-acquiring bandits. The Annals of Probability, 9()): 1981.
[9] Daniel J Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, Zheng Wen, al. A tutorial on thompson sampling, Foundations and Trends@ in Machine Learning, 11(1):1-96, 2018.
[10] Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. In
J. Shawe- Taylor, Zemel, P. Bartlett, Pereira, and K.Q Weinberger, editors, Ad- vances in Neural Information Processing Systems, volume 24. Curran Associates, Inc., 2011 /ptps///e//o/coooooccooooooo mewpeeedddinnnnnenineeeeeeeeeeee nwwrinuurp...e cc cttpsper_ppppppppeeeeeeee e5328222766228774445555533326663
URL https: /
10

---

[11] Christopher M. Bishop. Pattern Recognition and Machine Learning. Springer- Verlag New York, Inc., 2006.
[12] Hao Tang, Darren Key, and Kevin Ellis. Worldcoder, model-based Il agent: Building world models by writing code and interacting with the environment. NeurIPS, 2024.
[13] Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, et al. Measuring coding challenge competence with apps. arXiv preprint arXiv:2105.099.0 2021.
[14] Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition-leve code generation with alphacode. Nature, 2022.
[15] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107. (0t.:22007.03374,, 2021.
[16] Jacob Austin, Augustus Odena, Maxwell Nye. Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoe Le, and Charles Sutton. Program synthesis with large language models, 2021.
[17] Francois Chollet. Abstraction and reasoning challenge. https://ww..r wwwp:///wwkkkggggt Le. com/ c aatttaaactioonnnnd--asaaoiiiiiii Leeneeeennnnnccallleeeeeeeeeee Lenge, 2019.
[18] Aysja Johnson, Wai Keen Vong. Brenden M Lake, and Todd M Gureckis. Fast and flexible:
Human program induction in abstract reasoning tasks. arXiv preprint arXiv:2104 05823, 2021.
[19] Enric Rodríguez-Carboned Carbonell and Deepak Kapur. Generating all polynomial invariants in simple loops. J. Symb. Comput., 42(4):443 476, 2007. doi: 10. 10.10166..SC..00000111
URL
https://don org/ 1016 jsc. 2007. @1 002
[20] Thanh Vu Nguyen, Deepak Kapur, Westley Weimer, and Stephanie Forrest. Using dynamic analysis to discover polynomial and array invariants. In Martin Glinz, Gail C. Murphy, and Mauro Pezzè, editors, 34th International Conference on Software Engineering. ICSE 2012, June 2-9, 2012, Zurich, Switzerland, pages 683-693. IEEE Computer Society, 2012. doi: 10. 109/ ICSF.2012.62271 149. URL pottps://dos org/ 10. 1109 ICSE 2012. 20111.222112.
[21] Jianan Yao, Gabriel Ryan, Justin Wong, Suman Jana, and Ronghui Gu. Learning nonlinear loop invariants with gated continuous logic networks. In Proceedings of the 41 ACM SIGPLAN Conference on Programming Language Design and Implementation, PLDI 2020, page 106-120, New York, NY, USA, 2020. Association for Computing Machinery, ISBN 9781 77781450376136 36. doi: 10 / 3388599122.38888
10.1455/885412.3385588
URL (otthhtts::////
[22] Leonardo De Moura and Nikolaj Bjørner. Z3: An efficient smt solver. In International conference on Tools and Algorithms for the Construction and Analysis of Systems, pages 337-340. Springer, 2008
[23] Ruocheng Wang, Eric Zelikman, Gabriel Poesia, Yewen Pu, Nick Haber, and Noah D. Goodman. Hypothesis search: Inductive reasoning with language models. ICLR, 2024.
[24] Linlu Qiu, Liwei Jiang, Ximing Lu, Melanie Selar, Valentina Pyatkin, Chandra Bhagavatula, Bailin Wang, Yoon Kim, Yejin Choi, Nouha Dziri, and Xiang Ren. Phenomenal yet puzzling: Testing inductive reasoning capabilities of language models with hypothesis refinement. ICLR, 2024.
[25] Eric Zelikman, Qian Huang. Gabriel Poesia, Noah D Goodman, and Nick Haber. Parsel: Algorithmic reasoning with language models by composing decompositions. NeurIPS 2023, 2023.
[26) Hari Govind Vediramana Krishnan, YuTing Chen, Sharon Shoham, and Arie Gurfinkel. Global guidance for local generalization in model checking. In Shuvendu K. Lahiri and Chao Wang, editors, Computer Aided Verification, pages 101-125, Cham, 2020. Springer International Publishing. ISBN 978-3-030-53291-8

---

[27] Yangruibo Ding, Marcus Min, Gail Kaiser, and Baishakhi Ray. Cycle: the code generation. Proceedings of the ACM on Programming Languages, 8(OPPLAA): 392-418, 2024.
Learning to self-refine
[28] Westley Weimer, Thanh Nguyen, Claire Le Goues, and Stephanie Forrest. Automatically finding patches using genetic programming. In 2009 IEEE 31st International Conference on Software Engineering, pages 364-374. IEEE, 2009.
[29] Hoang Duong Thien Nguyen, Dawei Qi, Abhik Roychoudhury, and Satish Chandra. Semfix: Program repair via semantic analysis. In 2013 35th International Conference on Software Engineering (ICSE), pages 772-781, IEEE, 2013.
[30] Claire Le Goues, Michael Pradel, and Abhik Roychoudhury. Automated program repair. Communications of the ACM, 62(12):56-65, 2019.
[31] Rishabh Singh, Sumit Gulwani, and Armando Solar-Lezama. Automated feedback genera- tion for introductory programming assignments. In Proceedings of the 34th ACM SIGPLAN conference on Programming language design and implementation, pages 15-26, 2013.
[32] Zhiyu Fan, Xiang Gao, Martin Mirchev, Abhik Roychoudhury, and Shin Hwei Tan. Auto- mated repair of programs from large language models. In 2023 1EEE/ACC 45th International Conference on Software Engineering (ICSE), pages 1469-1481. IEFE, 2023.
[33] Chunqiu Steven Xia, Yuxiang Wei, and Lingming Zhang. Automated program repair in the of large pre-trained language models. In Proceedings of the 45th International Conference on Software Engineering, ICSE '23, page 1482-1494. IEEE Press, 2023. ISBN 978 665457019. doi: 10. 10991155481199200225 200.00000.....00009999 URL htttts::///ottt org/ 10. 1109 ICSD///CSE4869
[34] Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pages 72-83. Springer, 2006.
[35] Timothy Yee, Viliam Lisy, Michael H Bowling, and S Kambhampati. Monte carlo tree search in continuous action spaces with execution uncertainty. In IJCAI, pages 690-697, 2016.
[36] Guillaume M Jb Chaslot, Mark HM Winands, H Jaap van den Herik, Jos WHM Uiterwijk, and Bruno Bouzy. Progressive strategies for monte-carlo tree search. New Mathematics and Natural Computation, 4(03))343-351) 2008.
[37] David Silver and Joel Veness. Monte-carlo planning in large pomdps. Advances neural information processing systems, 23, 2010.
[38] Dezhi Ran, Hao Wang, Wenyu Wang, and Tao Xie. Badge: prioritizing ui events with hierar- chical multi-armed bandits for automated ui testing. In 2023 IEEF/ACM 45th International Conference on Software Engineering (ICSE), pages 894-905. IEFE, 2023.
[39] Chenyuan Yang. Yinlin Deng. Runyu Lu, Jiayi Yao, Jiawei Liu, Reyhaneh Jabbarvand, and Lingming Zhang. Whitefox: White-box compiler fuzzing empowered by large language models. arXiv preprint urXiv:2310. 15991, 2023.
[40] Xujie Si, Hanjun Dai, Mukund Raghothaman, Mayur Naik, and Le Song. Learning loop invariants for program verification. In S. Bengio. H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors. Advances in Neural In- formation Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceeddnngs www.:ss.eenuuuss.ccoco cc/pape files /F1es//aaper/201//iie 5bbbes2ccBfddc221ss1333668822223
[41] Adharsh Kamath, Aditya Senthilnathan, Saikat Chakraborty, Pantazis Deligiannis, Shuvendu K. Lahiri, Akash Lal, Aseem Rastogi, Subhajit Roy, and Rahul Sharma. Finding inductive loop invariants using large language models, 2023.
[42] E. Rooddguue--arrbae Carbonell and D. Kapur. Generating all polynomial invariants in simple loops. Journal of Symbolic Computation, 42(4))443 2007. ISSN 0747-7171. doi: https://doi. org/10.10166.jss.2000... http:///wp/////....ccccceeeerree anticcle/pii/577777777000002
12

---

A Appendix
A.1 Pseudocode
Algorithm
Bandit formulation of program synthesis
Input:
logical constraint @ heuristic h(.), seed program p (such as an empty blank function)
Hyperparameter: C'
progs < {p}
initialize with one arm (one program) counter for each program tracking how many times it has been refined b initially, no refinements
N + dict()
No <
repeat
Thompson Sampling
E progs: 0p ~ BetaB1 + C (Xx h(p)) + N2)
sample select arm > call LLM to refine increment count of number of refinements
f arg maxxpppooee 0,
p' ~ Pretine (, +)
Np + N₂ +1
Arm arrival
progs progs U {/}
add to set of arms initialize counter to zero
N₂ < 0
until p
return p
A.2 Hyperparameter analysis
Table I:
Analyze the hyperparameter sensitivity of methods by their performance rankings on different benchmarks. The rankings determined by (AUC final success rate) / 2. The hyperparameter, empty value for Greedy, denotes the heuristic value assigned for the empty solution. The hyper- parameter choices for ARC are different from the others because of its domain convention (maximum L.M requests=64 instead of 300). REx consistently outperforms or competes with the best baselines when C 20 for all difficult benchmarks.
Looplnv
ARC
APPS-Comp.
APPS-Inter.
APPS-Intro.
Greedy (empty value=0)
9
7
1
Greedy (empty value=0.5)
13
10
3
BFS (branching factor=2)
12
BFS (branching factor=3)
10
9
8
BFS (branching factor=4)
7
BFS (branching factor=5)
12
12
Fixed- Width (width=2)
Width (width=4)
8
Tixed-Width (width=8)
4
Fixed-Widit (width=25)
12
6
3
7
Fixed-Widee Width (width=50)
10
8
5
10
Fixed-Width (width=100)
11
12
10
9
REx (C=5)
11
5
5
REx (C=10)
5
4
7
2
REx (C=15)
3
REx (C=20)
2
1
4
11
REx (C=25)
REx (C=30)
6
REx (C=50)
2
4
REx (C=100)
3
2
6
13

---

A.3 Loop Invariants
Solving for loop invariants is 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 check 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 A le A trans) + inv', (inv A -lc) - post
Figure 8: Verification conditions of loop: pre preconditions, Ic loop condition, trans state transition function, llc loop termination condition. post post conditions
-50: while (a < 0) x = x x
(b) A desirable loop invariant I is predicate over x. y such that:
Va, { true = [[-0/z]
(pre)
TAXSO = T(y + 1)/v (o + /)/>]
(inv)
= 1+
I IAx20
y>0
(post)
assert (y > 0)
(a) An example program.
(c) The desired loop invariant is (rr
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. L[-50/ x> means substituting x in the loop invariant, I, with 50. More details regarding loop invariant 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 program's 's behavior and produce sufficient loop invariants.
A.3.1 NLA Benchmark Dataset
The verification of some programs require particular subset of loop invariants that contain non-linear polynomial arithmetic (NLA) terms with regards to the program's variables. Finding NLA invariants not well addressed by state-of thet---rr formal solvers. Nevertheless, LLMs can produce arbitrary formula and not constrained by the forms of the invariants. Thus, utilizing these models is promising approach to address the performance gap and to discover novel invariants [41.
Here, we present set of programs that require NLA invariants prove their functional correctness [4). All 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 7.3 theorem prover to check each invariant against the VCs. This way, we can statically verify whether the program's 's behaviour satisfy the specifications.
14

---

Index
Program
Loop
InvType
G-CLN 21
GSpacer 26
Coopy-GPT 4
REx
cohencu
NL
2
NL
cohendiv
3
NL
4
Linear
dijkstra
5
NL
/
6
Linear
divbin
7
NL
/
8
egcd
NL
9
NL
egcd2
10
NL
11
NL
12
eged3
NL
13
NL
14
NI.
15
fermat
Linear
16
Linear
17
fermat2
NL
18
freirel
NL
19
freire2
NL
20
geol
NL
21
geo2
NL.
22
geo3
NL
23
NL
hard
24
NL
25
knuth
NL
26
NL.
27
lem1
NI.
28
NL.
29
lem2
NL
30
mannadiv
NL
31
prod4br
NL
32
prodbin
NL
33
ps2
NL
34
ps3
NL
35
ps4
NL
36
ps5
NL
37
ps6
NL
38
sqrt
NL
Total correct
24
5
28
Table Benchmark evaluation on the NLA programs. For programs with multiple loops, each loop is labeled with number, top to down, from outer to inner.
A.3.2 LLM Invariant Generation
We generate potential invariants with GPT-4 by feeding in the C program and asking for 73Py 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, 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 733 checker to tell the LLM which step of the check failed, so that knows whether stronger or more general invariant is needed. More detailed examples are shown in the prompts in Appendix A.
15

---

A.3.3 Baseline Methods
G-CLN is learning-based method that fits for invariants by running program traces 211 and achieved the state-of-the-arl 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-specifie 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 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 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 too strong on the current loop
Unable to handle edge logic involving floating point calculation
Unable to prove the correctness of the Greatest-Common--Die algorithms when having access numpy-ged.
Spacer is Z3's builtin fixedpoint engine. It can look for inductive invariants for a loop after giving it the VC definitions, formulated in Constrained Horn Clauses (CHCs). We on the newest version called GSpacer that incorporates global guidance during solving [26) which is the state-of-the-art CHC solver. Although GSpacer can effectively solve linear invariants using global insights, where previous myopic heuristics were to cause conterminntiig runs, we observe that it is only able to solve 5 problems from the NLA benchmark, with 3 of them requiring actual non-linear polynomial invariants.
A.3.4 Summary of Program Types
Problem Type
Programs
Integer Division Algorithms Integer Power Algorithms
divbin, mannadiv, cohendiv, hard cohencu eged, eged2, egcd3, lem1, lem2 fermatl, fermar2, knuth
GCD & L.M
Divisor for Factorisation
Geometric Sequences
geol, geo2, geo3
Integer Multiplication
prodbin, prod4br sqrt, dijkstra, freirel, freire2 ps2, ps3, ps4, ps5, ps6
Square Root Calculation
Power of Sums
Table 3: Categorized NLA problems
A.4
The Abstraction and Reasoning Corpus
A.4.1 Dataset
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 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 the metrics ranges from x to x 30, and there are only 10 colors for each pixel represented by number ( to number 9. In our paper, we use the subset of ARC proposed by Johnson[18] which contains ARC problems with human experiments. We select one hypothesis written by human beings for each problem as hint for LLMs.
16

---

Problem Example
Training tasks
Testing task
?
Human hypothesis: All colors fall to the x-axis and stack on top of one another
Code generated by REx
def transforn Corstorsssssdnnneoo sepptiinptt_gridd ap. ndarray) ap-nnarcayy grid
roMa, cole input input_geid.ent ahide
Creste new grid with the waze diransions initielized with sezoe outgut_ggrid
Iterate through mach cntumn the for col cangejepua);
det all, non-zezo valuea in the current colums
a00 aarooooooooss input_gzid colj| Liput col)
outget grid outpot trova len(scn sero, values)):, coll non, sezo valuea
respective
reture outgut_grid
Problem Example
Training tasks
Testing task
?
Human hypothesis: change the orange to gray
Code generated by REx
avansfore, gardd Engot__gridt Ldt ap. sg.ndarray) np-ndarray:
ourpet_arid sopot_srid.cog711
priprpt_drdooouppuut 71 ahange orasge (7) to gray (5)
return
Figure 10: Problem examples of the ARC dataset.
A.4.2 Baseline
We use the Human-Wrine Hypotheses experimental set in the paper hypothesis search 23 which gains the current state-of-art results of the ARC dataset, as our baseline. The paper uses a Fixed-Width Width searching policy to get the candidate programs, and it sets the width into 8 and the depth into 3. We run the baseline with the same searching policy and hyperparameter fixed width of and its depth with 3, which accounts for maximum #LLM requests= 32. The prompts in our experiments are also consistent with the paper.
A.4.3 Experimental setting
We use the subset of ARC, which contains 40 problems, our dataset. Each problem contains several training tasks as examples. We utilize GPT-4 Turbo to generate code that summarizes the transformation rules and refine code that fails to pass the training examples. For each problem, we set REx's heuristic reward each code as the pass rate of each problem, which means the percentage of correctly solved training tasks. If the code passes all the training examples, which means it
17

---

Table 4: Analyze the hyperparameter sensitivity of methods by their performance rankings on different benchmarks on the ARC domain with maximum #LLM requests=200. The ranking principle
the same as Table
Methods
REx (C)
Greedy
BFS
Fixed-Wii Width
hyperparameter type
C
empty value
branching factor
width
hyperparameter value
10
25
50
100
0
0.5
4
16
10
25
50
ranking
5
2
1
9
6
3
8
10
12
11
ARC
70
70
Final (%)
68
Hypo. Search Our replication
Success Rate (%)
66
REX BFS Greedy FW
AUC (%)
62
61
50
100
150
200
REx
BFS Greedy FW
Sample Budget (# LLM Requests)
Figure 11: Results on the ARC domain with maximum #LLM requests=200. Denotations are the same as Figure
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 223 We set the hyperparameters of each method accordingly follows:
1. Greedy: empty value= 0, 0.5
2. BFS: branching factor= 2, 3,
3. Fixed-Width: width= 2, 2,4, 8
4. REx: C=5, 10, 15, 20, 25, 30
For each method and each hyperparameter, we run them five times with five random seeds and report the average performance.
A.4.4 Scaled-up results with 200 LLM requests
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.5
2. BFS: branching factor= 2, 16
3. Fixed-Width: width= 10, 25, 50
4. REx: C=10, 25, 50, 100
Results are shown in Figure IT and Table REx overall achieves better performance than the other methods on the scale of 200 as well. uses fewer LLM requests to reach certain passing rate and gains overall higher final success rate. REx achieves relatively better performance when C. > 20.
18

---

GPT-3.5-turke
Mraude-3.5-oomeee
Llama-3. 3..--05B
Success Rate (%)
16
40
14
Firal
45
30
12
40
10
AJC 1S0
35
(S)
8
REx
BF8 Growdy PW
REX
BPS Crewdy F4
REX
BP8 Crewdy Fi
30
40
60
40
60
40
50
Sample Budget Requests)
Sample Budget (# Requests)
Sample Budget Requests)
Figure 12: Comparing REx with alternatives with other LLMs on ARC: GPT-3.5-turbo, Claude-3.5. Sonnet, and Llama-3.1-4053
A.5
Competition Programming, APPS
A.5.1 Dataset
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. use all 60 problems for the competition-level and the introductory-level, and randomly choose 60 of 120 problems for the interview-level.
A.5.2 Experimental setting
Most experiments are conducted with GPT-4 and temperature=1.0 by default. The choices of hyperparameters for each method are listed in Table We also experiment with other LLMs including GPT-3.5-turbo, an older model to alleviate concerns regarding data contamination.
A.5.3 Results on APPS-Interview
Figure 13 shows the results on APPS-Intervient As expected, REx performs competitively with the best baseline, Greedy, and outperforms the other baselines. consistently outperforms the state-of- ofte-offthhe-rt method, Olausson et al. (3), which uses the fixed-width method with width=25/50. also aligns with our hypothesis that greedy performs competitively with REx when the dataset is relatively easy.
A.6 Prompts
We list all the prompts that used in our experiments in this section. The functionality for each prompt is stated its subsection name. We highlight the dynamic information as yellow and the main instruction blue. The dynamic information includes the codes synthesized so far by previous LLM calls.
A.6.1 Loop Invariant Synthesis
Initializing loop invariant. It asks LLMs to generate loop invariant to help understand the code and prove its correctness.
Role: system
You are expert software engineer and you working project that requires you write loop invariants for loop in the code. The loop invariant
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 invar should be strong enough to prove the correctness of the loop well the functionality the whole code snippet.
19

---

APPS Interview-Level
Final (%)
85
90
80
80
Success Rate (%)
70
Olausson et al. 2023
AUC (%)
80
78
60
76
50
#LLM Requests needed
APPS Intrerview-
to get Olausson et al. 2023
60
40
50
40
30
30
20
0
50
100
150
200
250
300
REX
BFS 1.50x
Greedy
FW 2.39x
Sample Budget LLM Requests)
1.00x
0.93x
Figure 13: Results on APPS-Interview. Denotations are the same as Figure
Role: user
You are given the following code snippet in C:
...
#include <stdio.h>
int main(int a, int b
;a>=)(
assert ;b>=);
int u,,,
x=a;
y=b;
u=b;
v=0;
p=1;
q=0;
r=0;
s=1;
while(1) {
// loop invariant in Z3:
if (x!x=))) break;
while (1)
if(1))))) break;
x=x-y;
vev+u;
p-q;
r-s;
20

---

while (1){
(x(y)) break;
y-y-x;
u=uty;
q-P;
s-r;
}
int r utv;
return r; //lcm
Do you understand the code snippet? What function does the code snippet perform? Can you explain the code snippet part part? What is the functionality of each part of the code snippet? What is the function of the loop in the code snippet?
The loop that needs loop invariant has been identified in the code snippet with line starting the function of loop the code snippet? Can you explain function of the loop in the code snippet? What is the purpose of the loop in the code snippet?
/// loop invariant in
you understand
snippet? How are the variables used and updated in the loop in the code snippet? what is the relationship among the variables in the loop in the code snippet?
What variables are there the loop
Please analyze the loop variant your own language. The loop invariant should be 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 correctness the loop as well as
should also strong enough to prove the functionality of the whole code snippet.
what are the loop invartants in code? The loop invariant should be written in the format Py such follows:
And(x 0, 10, Implies(x 0, 0), x
Please surround the loop invariant with triple backticks and write it in the format Python-Z3.
Refining loop invariant. It asks L.LMs to refine loop invariant given the error messages. detail, the loop invariant should satisfy three conditions in order to be valid and sufficient:
21

---

Established:
meaning that loop invariant should be true given the preconditions, i.e., should be true at the start of the loop.
Preserved: meaning that loop invariant should also be after each iteration of the loop.
Sufficient: meaning that the loop invariant should be strong enough to help prove the post conditions, i.., the specifications/proerrtic assigned by the users.
The symbolic checker utilizes SMT solvers to evaluate these three conditions individually and returns the error messages. any of the conditions not satisfied, will return the error type, the specific assertion, and its model help debug. L.MMs 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 you are working on project that requires you to write loop invariants for a loop in the code. The loop invariant is 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 invar iant should be strong enough to prove the correctness of the loop as well as the functionality of the whole code snippet.
Role: user
You are given the following code snippet in C:
#include <stdio.h>
int main(int a, int {
assert(a>=1))
assert (b>=1);
int v,,
x=a;
y=b;
u=b;
v=0;
p=1;
q=0;
r=o;
s=1;
while(1)
// loop invarfant in Z3:
(x!x)))) break;
while (1)
ix(1))))) break;
x=x-y;
v=v+u;
p p-q;
r-s;
while (1)
22

---

ix(y)) break;
y-y=x;
u=utv;
q-P;
s-r;
}
int utv; return r; //lcm
}
Do you understand the code snippet? What function does the code snippet perform? Can you explain code snippet part by part? What is the functionality each part of the code snippet? What is the function the loop in the code snippet?
The loop that needs a loop invariant has been identified in code snippet with a line starting with *// loop invariant in Z3: '. Do you understand the function of the loop in the code snippet? Can you explain the function of the loop in the code snippet? what is the purpose of the loop in the code snippet?
What variables are there in the loop in the code snippet? How are the variables used and updated in the loop in the code snippet? What is the relationship among the variables in loop the code snippet?
According those analysis, could you refine the following loop invariants that you generated before?
And
>= 1,
1,
X >= y, a*b == x*y
The previous loop invariant wrong because:
''OrNNot 0), meaning is even true b)* is neither established beginning the loop and neither after each iteration of the loop. For example, we can set [p = 1, q 0, 7, = 0, v 0, counterexample for establishing 'Or 'Or (x >= 0), u == a*x + b)s, since it conflicts with the assertion 'Indlies(And 0, 5 1), Or (Not 0),
preserved,
1, 15, b
15, u u = to find a
0, b 0,
a, y ==
b, 0, p 1, q == 0,
aex b))
'or 'Or (y >= 0), == b*y) is neither established nor preserved, meaning it is not true in the beginning of the loop and is neither true after each iteration of the loop. For example,
set 0, 1, p
23

---

1, x = 16, = 8, u = 13, r= = 0, y 13, a = 16, b = 13] to find a counterexample for establishing 'Or (Not >= 0), v = b,y)', since it conflicts with the assertion 'Implies (a 0, b > 0, b, o, 1, == 0, 8, s 1), Or (Not 0), == mat bxy))
a, ==
The metrics are nei ther 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 the format as Python-Z3 as before. Please surround the loop invariant with triple backticks and write it the format as Python-Z3.
A.6.2 Code Generation
Initializing code.
It asks LLMs to generate code to solve python programming problems.
Role: system
You are 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. 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.
to decide what fraction of drink should be made from each type of juice, such a way that the maximum possible number of people attending the party like it.
Each person has minimum fraction of each of the juices they would like to have the drink. They will only like drink 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 integer STS, the number of test cases in the input file.
For each test case, there will be:
One line containing the integer $NS, the number of people going to the party.
-
$NS lines, numbers "$SA $BS $c$**, indicating the minimum fraction of each juice that would like in the drink. $AS, $B$ and SC$ are integers between $0$ and $10000$ inclusive, indicating the fraction in C lee 18000s.
for each person, each containing three space-separated
$A + B
You may assume that $1 \leq \leq 2$ and $1 \leq \leq 5000s.
-.--Ooutut-----
$TS lines, for each test case in the order they occur in the input file, each containing the string 'Case #$x$: $YS* where $XS is the number
24

---

of the test case, starting from 1, and $Yş is the maximum number of people le who will like your drink.
Examples-
Sample Input:
2
10000
10000
9 10000
5000 0
2000 0
4000
Sample Output:
Case #1: 1
Case #2: 2
Can you write correct solution and make it passes the example test cases? (Please start the solution segment "'pthon as usual and use Standard Input format)
Refining code. It asks LLMs refine code given the test case it failed on and the message.
Role: system
You are a helpful assistant that can solve python programming problems and debug incorrect programs
Role: user
given hard python programming contest problem:
Question:
You are holding party. In preparation, you making 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, such way that the maximum possible number of people attending party like it.
Each person has a minimum fraction of each of the juices they would like to have in the drink. They will only like the drink if the fraction of each of 3 juices in the drink is greater equal their minimum fraction for that juice.
Determine the maximum number of people that you can satisfy.
-..--IIptt--.........
One line containing an integer ST$, the number of test cases in the input file.
For each case, there will be:
One line containing the integer $NS, the number of people going to the party.
for each person, each containing three space-separated numbers " ''SA$ $BS $ indicating the minimum fraction of each juice that would like in the drink. $AS, $B$ SC$ are integers between $0$
$NS lines,
25

---

$10000$ inclusive, indicating the fraction in parts-per-ten-tioosand $A + B c lleq 10000S.
You may assume that $1 Vleq \leq 2$ $1 Vleq \leq 5008s.
-----ouptt---
$TS Lines, for each test case in the order they occur in the input file, each containing the string 'Case #$x$: $Vş** where $X5 is the number of the test case, starting from 1, and $Vş the maximum number people who will like your drink.
Examples-m----
Sample Input:
3
10000
10000
10000 3
5000 0 0
2000
4000
Sample Output:
Case #1: 1
Case #2: 2
Below is failed attempt to solve the problem:
int(input())
for case in range(1, 1):
= (Inpuut)))
total_fractioos {'A': 8, 'B': 'C': 8} for person in range(n):
a, b, c map (int, (mput().sllite. total_fractionss ['A]] += a
['B'] += b
total_fracttoon("CC) +=
scale min(1, sum(value value total_ sol_ffacction.valuess 10800) staled_fraction (key: value scale for key, value total_ fractifn ttmes )))
count =
for person in range(n):
a, b, c map (int, spllitt))
a <= scaled fractioos(''''" and b <=
and c
<= scaled fracctiottins''
count += print('Case #(): O' (crsa count))
The failure is due to runtime error at test
case 0 for input 2\n3 0 olne 10980 0'n0 0 (8060(n3\)55000 0 0' 2000 0' 0 4000\n and expected output Case #1: 1\nCase #2: 2\n. Error type: ValueError, detailed error message: got 1) at line 14, a, b, map(int, input().split())
enough values unpack (expected 3,
Overall evaluation: 0 out of 2 test cases passed
Can you
1. Analyze why the runtime error occurs?
2. Correct the code and make sure it passes the example test
cases? (Please start the solution segment with "'pyhhon as usual and use Standard Input format)
26

---

A.6.3 ARC
Initializing mangglaanipulating program. We ask LLMs to generate code that transforms the input matrices into the output matrices correctly, following the same transformation rule. The prompts are derived from Hypothesis Search 23]
Role: user
Examplee
0:
Input:
[[O
[e]
01
[e]
6
[0]
0
6
0
0
[0
61
00
01
E0
[
[0 0
2
0
0
2
81
[0
8
2
2
61
0
0
[0
Output:
[[O
[0
[0 [0
[0
[0
0
0
[
9
2 2
2
8]
[6
2
61
[0
9
6
[0
6]]
Example 1:
Input:
[[O
[e
8
8
8
8
6]
[0
61
[0
0
8
0
81
00
8
8]
[0
61
e]l
Output:
[[o
[0
8
8
6]
8
[0
0
0
8
81
[0
0
8
0
8
81
10
01
0]]
Now, please write a python program transform_pn (iwpgrrdddddddrrrrere np. ndarrryy int
1) -> np. www.ddaar(nnne that transforms the input grid to the corresponding output grid.
Hint: You may want use the following guidance implement function:
27

---

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 the test input.
The number in the input grid can be mapped the following colors:8: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 seansforn.grid(tnputgrridd np. adarray int]) in Python and nothing else, each cell in the output should only be numbers from 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
Input:
[[O
[0
6
6
6
0]
[0
6
0
[0
8]
[0
0
6
6]
[0
8
0
0
[0
01
[
0
2
2
2
6]
[
0
9
0
0
8]
[6
[0
6]]
Output:
[[o
[0 [0
[0
[ 0
9
2
0
2
61 6]
8
2
2
[
[0
[0
0
0
0
6]]
Example
1:
Input:
[[o
[0
8
8
8
8
0]
[0
8
8
8]
[0
81
81
[0
0
0
0
9
8
[0
0
0
8]
Output:
[[0
28

---

[
8
8
6]
8
6]
r [0
81
[0
0]
[0
0]]
Now, please write a python program transform_gre (iippurrdddinnggge. np. np.ndarray (int
1) -> [inn that transforms the input grid to the corresponding output grid.
the following guidance implement the function: The bottom-most row that contains colored squares remains 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.
Hint: You may want
same from test
The number in the input grid be mapped the following colors:0:black;
1:blue; 2:red; 3:green; 4:yellow; 5:grey; fusuchi; 7:orange; 8:teal; brown
Just reply with the implementation of prnesfrrmgrrddinnuutt
np. ndarrayy int]) in Python and nothing else, each cell in the output should only be numbers from to
This is the code you wrote last time: se.
import numpy as
def transform lransform_gridtiiiee nrows, ncols input_grid, shape output_grid = np.zeros_ like (iinut__grrd
no.ndarray(int]) -> (int)):
Initialize variables to trace contiguous block of colored in coloree_block False top_color_inee index
cottom_color_iee index
Identify the contiguous main block of colored rows from bottom top for idx range(nrows 1, -1, -1):
if any if not ciclored_blacke
:= 0):
nottom_color__e in jn_colored_ble block True
idx
cop_color_indes index idx
else:
if colorrd_bblock: break
Apply the transformation to the identified colored block shift =
for in range<bootto__color index, top_color_inee index 1, -1): 0):
if output_gridie shift: (ncols (ncols ncols shift))] input_grid(r, 0: (ncols shift)] shift +=
return output_grid
It generates error:
wrong output format: Error: could broadcast input array from shape (8,)
into shape (7,) The correct format should be np.array
Please correct the error and generate the Just reply with the implementation of transform.grrid(tpput, np.ndarray[int]) in Python and
29

---

the and nothing else, each cell in the output should only be numbers from to 9. Please contains the necessary import statements.
This is the refinement prompt for codes that produce wrong
Role: user
Example o:
Input:
[[O
[e
6
[0]
6
9
0
6
[0
0
6
6
[0
01
8
[
[0 0
2
2
2
9
81
[
2
0
61
0
8
0
9
0
[0
Output:
[[O
[0
[0
[0
[0
[
0
0
[0
2
2
61
8
9
9
0
0
8
8
6
80
6.
8]]
Example 1:
Input:
[[O
[e
81
[0
0
8
0
0
61
[0
8
81
[0
81
[0
0]]
Output:
[[o
[0
0]
8
81
[0
0
8
0
81
0
8
0
8
81
10
81
0]
[0
0]]
Now, plaase write python program cinnsforn_grid(npputtt np. ray int 1) np.ndrraay(int) that transforms the input grid to the corresponding output grid.
30

---

Hint: You may want to use the following guidance to implement the function: The bottom-most row that contains colored squares remains the same from test all start column
input to test output. Meanwhile, the other to the right in test output compared to the test input.
The number in the input grid can be mapped to the following colors:0:blacks 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 (nn.frrrnggdddiitt___rr.... np.ndarray[ int]) Python and nothing else, each cell in the output should only be numbers from 0 to 9.
This is the code you wrote last time:
import numpy as np
def transform_gr
np.ndarray(int]) -> np.ndarray[int]: nrows, ncols input_grid.shap output_grid np.zeros_ like wwwp:llleeeieennnneereee (input_grid)
Finding the bottom-most row with non-zero entries
for row in range (nrow(ns -1, -1):
if np. (onntttutggridrroo != 0): last_coloreddrow row break
Copy the last colored row is: wrtpput.grddllsst..ooore.....
input_grid(last.ccooree.
Shift each row above the bottom-most colored row one column to the right
for row in range(tast_coloreerrvww... -1, -1):
mp..ny(tnput_griddr output_grid[roww, 1:] input_grid(row, :-1]
0):
return output_grid
'..
These are the failed examples of the code:
Task 1:
Input:
[[o, 0, 0, 0,
0,
[0] ,
6, 6,
6
0
[8,
6,
8,
0,
6,
9, 0 0 0], [0,
0
6,
0
0
6,
0
0
0
0
6,
0,
0,
6,
0,
8]
[0]
8, 0, 0, 6, 6, 6
0
0
[0
0
0
0
0
8
0
0
[0]
0.
2
2
2
0
0
8, 0], [0,
2
8
0
2
0
0
0
00
8
0
2,
2
2
0
0
81
[0
0
0, 0, 0].
10
0
0
8
81
0
8
0],
[0 0, O,
0,
0]]
Correct Output:
[[O, 0, 0, 0
0
0,
0
0]
[0,
0 6,
6,
6,
0
0
0,
8].
[8,
0,
6, 0, 0,
6,
0].
6
0
0
6
0]
[0]
0
0
6,
0
6,
0
8]
[0,
0 6,
6,
0
00
0
81
0
2,
2,
8, 0] [0.
0
2
0
01
0
0
2,
2,
2
0
81
[8
0
0 0, 01,
8, 0 0 0
0
0
0
0
0
0
0
8]
[0,
0
0
0
0
0
8
[0, 0 O
0
0]]
Code Output:
[[o, 0, 0,
0,
0]
[0]
6,
6,
6,
0
8,
0,
[8,
0,
6,
0, 0,
6, 0 0, 0]
,
[0.
0
6,
0
0
6,
01
[0]
0
0
0
6,
0,
0
6,
8]
[0,
8, 0 0
6,
6
0
10
0
0
0
8
8
0
81
[0
0
2
2
2,
0,
8, 0] 00
0
2
0
0
0
8
0
8
0
2
2,
2
8
0
81
[0
0
0 0],
10
0
8
8]
[0,
8
0,
[0, 0 0
0
8
0
]]]
Task 2:
Input:
[[O, 0,
0
9
01
[0
8,
8
8
8
8
[8.
0,
8, 8, 0, 0],
0
0
[0
0
[0.
31

---

8, 0, , 8, 8, 8, 8, 8J,
[0,
8
0
9,
0
8]
[6
0
8
0
8
0,
0,
8, o]]
Correct Output:
[[o, ,, 9 8
0
8
8
[0
8
8
81
[8
8
8
8, , 8, 0] 00,
8
0 8
0
0
81
00
0
0
8
0
8
0
81
[0
8, 0, , 8, 8, 8 o]] Code Output:
81
80
0
0
0
0
88
[0
0
0
8
0,
[[o, 8, 0, 8 0, 8 8, 0,
0
[0.
8,
8,
8:
8,
8,
8
0
81
[8.
0
8,
0
8
8, 0 8, 0] 00
8
0 8,
0
0
0
8]
00
0
0
8
0
0
81
[8
8, 0 , 8, 8
8,
8
81
00
0
0
0
8
0
81
[0
0
0
0
0,
0
8, 0]]
Please correct the error and generate the code. Just reply with the implementation of cransforn_grid(innut.. nothing else, each cell in the output should only be numbers from to Please contains necessary import statements.
np. ndddrrayyiitt] in Python and
A.7 Example trees generated by REx
32

---

Figure 14: Search Tree
33

---

Figure 15: Search Tree
34

---

100.0
3.6
Figure 16: Search Tree
35

---

Figure 17: Search Tree
36

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the
paper's contributions and scope?
Answer: [Yes]
Justification: We show that REx, which solves the ppploration-cxploittaiin tradeoff with Thompson Sampling, solve more problems using fewer language model calls in three domains, as shown in Figure 4 The exploration and exploitation tradeoff in refinement is illustrated in Figure 1 and Figure 7
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Please check Section for details.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
37

---

Answer: [NA]
This paper does not include theoretical results, even though standard theoretical analysis of arm-acquiring bandits is directly applicable.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: All prompts are provided in Appendix. The full algorithm easy to implement and has been shown in [
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
38

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We public benchmarks and attach the code in the supplementary material for reproduction.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurlPS code and data submission guidelines (https://nips.ca
/ C/ce//oddeessssseeeeeeeeee wrrrrissioPollcyyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to reproduce the results. See the NeurIPS code and data submission guidelines (https)
//nips. public hwwplllccguiess/eeee guw Seaeiidee//odeuumiiii
onoolnoolooccoooooooooooo Picyl for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We use public benchmarks and follow the experimental settings of the state- of-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 the core the paper to level 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, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We visualize and analyze the distribution using box plots for each method in different metries varying the learning curves for each experiment in Figure
in Figure 4 and Figure 6
We also plot the
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
39

---

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.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We mainly use GPT-4 and the maximum number of LLM requests for each experiment is 300.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics htttpp:://eerrrpssccc//ubb/////p
Answer: [Yes]
Justification:
The research conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work makes code generationnrefinemeett cheaper and more accessible. This may potentially have positive and negative social impacts most LLMs.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
40

---

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), faimess 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 tied to particular applications, let alone deployments. However, if there is 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 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 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 system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 challenging, and many papers do not require this, but encourage authors to take this into account and make 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 publicly available.
Guidelines:
The answer NA means that the paper does not 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 used and, 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.
41

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, -aper thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. 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/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent 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.
14. Crowdsoureing 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 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 the supplemental material is fine, but if the main contribu- tion 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.
15. 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 equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending 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.
42

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere 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.
43


Reviews:
[{"review_id": "UEthfdl3Qz", "reviewer": "Reviewer_mETF", "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.", "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?\n\n2. **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.\n\n3. **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": "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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. **Novelty**. The proposed formulation although simple is novel in the context of multi-turn LLM applications and provides insights discovered beyond existing works. The proposed solution using Thompson-Sampling with Beta priors is straightforward but effective at improving efficiency.\n\n2. **Writing and Clarity**. The paper was pleasant to read and provides appropriate context connecting prior works. There is some missing information detailed in weaknesses.\n\n3. **Results**. Comprehensive experiments across three (widely varying tasks) demonstrating empirical effectiveness over simpler heuristic based exploration-exploitation baselines.", "weaknesses": "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\n\n    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.\n    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.\n\nAt 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.). \n\n2. **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`. \n\n3. **Invariant Task Grading**. The invariant task is a core task studied in the paper and the authors point out that they check \n> being a sufficiently strong inductive invariant (i.e., precondition, induction, and postcondition). \n\nThe 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?\n\n4. **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).\n\n5. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and the supportive review. Below we answer your main questions.\n\n> 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\n\nGreat idea! In the global response we have attached a PDF showing results for GPT3.5, Llama, and Claude. Although the absolute numbers are different, the qualitative outcomes are highly similar, and in fact, the advantage of REx seems larger for gpt3.5.\n\n> 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\n\nThis would be interesting to explore. Conditioning on depth or program/problem embeddings could be very promising, but would introduce more free parameters. We’ll add it to the discussion section as a future direction.\n\n> 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?\n\nIt is not possible to \"fool\" the verifier with trivial invariants, which we will clarify by revising the main text to explain precondition/induction/postcondition. Briefly the *precondition* means that the invariant holds at the beginning of the loop; *induction* means when the invariant is true it stays true after another loop iteration; and *postcondition* means if the invariant is true and the loop terminates then the assertions after the loop are satisfied. `True is True` would not satisfy the postcondition constraint.\n\n> Contamination. APPS benchmark is considerably older and released before the cutoff date for newer GPT-4 model\n\nWe'll revise to mention that this as an important concern for APPS. Originally we (incorrectly) assumed that, given the popularity of benchmarking GPT4 on APPS, there must be good reason to assume no contamination, but upon digging through the GPT4 tech report this week, we could find no such justification. Good catch.\n\nIn the global response we show new results on GPT3.5, which deriving from GPT3 is a lot less likely to have seen APPS in pretraining.\n \n> 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?\n\nYes, it's possible that with a richer set of features for the context, there could be a universal optimal hyperparameter setting. We wanted to explore many hyper parameters to understand sensitivity of different methods, and REx was almost always the best both in terms of median and max hyperparameter settings.\n\n> 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... 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.\n\nThank you for suggesting these improvements. We are revising to incorporate this feedback, especially adding links to the appendix from the main text.\n\nThanks again for your helpful feedback and please let us know if there are any further questions we can answer.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ze6Ce98DL4", "reviewer": "Reviewer_rHDY", "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).", "questions": "Could you please explain and address the weaknesses?", "limitations": "The limitations are well addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper’s formulation of the iterative program refinement problem is insightful.\n- The exploration-exploitation tradeoff is an important yet surprisingly overlooked aspect for effective code repair. This paper systematically analyzes different search policies.\n- The proposed solution with bandit algorithms is very simple yet effective, and seems to be robust across various repair benchmarks.", "weaknesses": "### **Technique**\n\nThe 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.\n\nIt is also unclear whether similar results can be achieved using more advanced prompting techniques (e.g., few-shot or CoT).\n\nFormulating 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. \n\n### **Evaluation**\nThe improvement of REx seems marginal.\n- Compared with alternative search algorithms, REx often brings only marginal improvements given a large enough budget (Figure 4). \n- Also on the some dataset (loop invariant), BFS (with the optimal hyperparameter) is even better than REx with smaller sampling budget.\n\nThe authors should list the cost of baselines for fair comparison.\n- 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).\n\n\nAccording 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.\n- 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?\n\nAlso, each arm is not independent, as the refined program $\\rho^\\prime$ is correlated with the original program $\\rho$.\n- 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.\n- 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.\n\n[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.\n\n[b] Yang, Chenyuan, et al. \"White-box compiler fuzzing empowered by large language models.\" arXiv 2023\n\n\n\n### **Implementation**\n\nSome 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?\n\n### **Minor**\n\nFormula (11):  $2+C+N_\\rho$ should be $2+2C+N_\\rho$\n\nFigure-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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful input and for your support. Please see below our responses.\n\n> REx often brings only marginal improvements given a large enough budget (Figure 4)\n\nWe agree! REx isn't magic: It is simply a more hyperparameter-robust, cost-saving refinement policy that also modestly improves the number of solved problems overall. We're excited about REx because it seems to help generically across the board, so we're optimistic that it could have an impact for many researchers, and why we've emphasized just how easy it is to implement in Python. (But please see below our response for the next comment on where we've taken REx recently, and the global response for some cases where REx makes a big difference when using cheap models)\n\n> The scope of this paper is limited to solving small, isolated programming challenges given a relatively large compute limit\n\nWe have since applied REx to much harder problems involving programs with hundreds of lines of code. We will reference these results in the camera ready. In fact, the original reason we considered a relatively large compute limit was because we believe that these harder problems require many more rounds of refinement.\n\n> the Greedy baseline seems to be very strong, even though the paper only considered two values for its single hyperparameter\n\nWe originally did this because greedy's hyperparameter is unique in that it has very few reasonable settings, as it corresponds to the heuristic value of an empty program. Setting it to e.g. 1 would mean that the search policy would never even try doing refinement; setting it close to zero means that it would only ever refine the initial program once. To double-check these intuitions we will rerun the greedy experiments with more hyperparameters.\n\n> Also on the some dataset (loop invariant), BFS (with the optimal hyperparameter) is even better than REx with smaller sampling budget.\n\nThe important point here is that you can't count on BFS always being the best: on a different dataset (APPS), it's actually the worst performing method! We want methods that are robust across hyperparameters and robust across datasets, which means acknowledging that sometimes special hyperparameters can make a method *seem* superior on a particular datasets for a particular sampling budgets. For loop invariants in particular, the box and whisker plots show performance across a range of hyperparameters, for which BFS actually tends to be worse than REx (Fig 4/6). \n\n> It is also unclear whether similar results can be achieved using more advanced prompting techniques (e.g., few-shot or CoT).\n\nREx is orthogonal to the prompting technique: Given a refinement prompt (e.g., few-shot or CoT), REx then works through an outer loop that repeatedly uses that prompt.\n\n> 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?\n\nYes, temperature=1 for all considered benchmarks.\n\n> In fact, MAB has been applied for fuzz testing [a,b]...\n\nThat work is indeed very related! We will cite and discuss the papers you mention.\n\nThank you for your review and for your support. Please let us know if we can answer further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kHXdYWKq60", "reviewer": "Reviewer_ikA5", "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.", "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?\n\nSince 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?\n\nHow 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.\n\nAt 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?\n\nOne 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?\n\nnit: 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": "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.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The method presented is conceptually simple, and is easy to implement and adapt to new tasks. Therefore this method can readily be put into practice by practitioners. (However, see weakness no. 1).\n\nThe connection to bandit problems is sound, and potentially a valuable framing of the problem. (However, see weakness no. 3.) With regard to significance, making this connection could suggest a wide range of approaches from bandit literature. That said the paper does not explore this potential benefit of the framing beyond its application of Thompson Sampling.\n\nThe range of tasks considered is another strength of the paper, and a key strength of the paper is that the method demonstrates performance improvements on three of the four tasks (the exception being the easiest APPS problems). The variance of the results on NLI and ARC is high though, so while the method is an improvement, the improvement is not substantial or guaranteed.\n\nThe modesty/honesty of the limitations section is refreshing too, providing clear commentary on the biggest limitations of the work.", "weaknesses": "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.\n\nA 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.\n\nIt 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.\n\n---\n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful review! Please see below for our responses, and see the global response PDF for new experimental results motivated by your suggestions.\n\n> only GPT-4 is tested\n\nThanks or the suggestion of testing other models. Please see the global response PDF for new results on other LLMs. Although the results are still coming in, and the absolute numbers are somewhat different, the overall qualitative conclusion is the same, with one exception: We've discovered that our method seems to help a lot with cheaper models, which means that it might be more broadly applicable then we originally sold it as.\n\n> you state that it is not possible to apply MCTS... Could... small modifications would be sufficient to apply MCTS to this problem?\n\nWhile such modifications could, in principle, likely be invented, they would probably be more complicated than REx, which is considerably simpler than MCTS, despite capturing MCTS-esque dynamics. However, we'd be happy to try any specific modifications you suggest next week during the discussion period.\n\n> Since the LLM is shown the full specification... there is risk of it generating solutions that satisfy the specification literally without generalizing to satisfy the intent of the specification\n\nWe evaluate on holdout inputs for ARC. For loop invariants we use a formal verifier to check that the invariant actually holds for all possible inputs. APPS, unfortunately, does not come with designated holdout tests, and conventionally has been evaluated without them [e.g. Olausson ICLR '24; Zelikman NeurIPS '23]. However, APPS averages 27.6 tests/problem, so it would be very hard to \"overfit\" without recovering the true algorithm, and in manually inspecting 15 program solutions we observe no memorization of isolated test-cases. The revision will mention all these issues.\n\n> other [non-refinement] approaches to using LLMs to solve these tasks\n\nWe are happy to try any particular baseline that you suggest next week during the discussion period. We do think that refinement is an especially popular and simple approach, though, so we think focusing on refinement is a reasonable research strategy for this paper.\n\n> Am I correct in understanding that having a target output available for the heuristic would obviate the need for applying the method?\n\nWe'd still need to apply our method because we would still need to generate a program that maximizes the heuristic value.\n\n> 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?\n\nThis would be a fascinating direction to explore, for example by incorporating a distance metric between programs and using probabilistic kernel methods for the prior (such as a Gaussian process). But this requires a metric between program source code in order to determine what programs are similar, which introduces its own complexities and would likely be more brittle and sensitive to hyperparameters and/or require in-domain training data. Still, it could be an interesting direction to explore, and we will add it to the future work.\n\nThanks again for the review and for the suggestions of future work. Please let us know if you have any further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t4w5NFKp64", "reviewer": "Reviewer_UuYA", "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.", "questions": "Please consider addressing the points raised in the “Weakness” section.", "limitations": "The paper has sufficiently addressed the points concerning limitations and societal impact.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is very well-written and easy-to-follow. I thank the authors for providing the details regarding the evaluation tasks and the prompts. Moreover, I appreciate the paper’s message on emphasizing the importance of developing better base models over constructing better refinement strategies. While it sounds like a negative result for the paper, it is definitely useful for the community.", "weaknesses": "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.\n\nMoreover, 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.\n\nAnother 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and for the suggestions. We have run new experiments that we really believe can address your concerns. Please see below.\n\n> 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?\n\nThank you for this idea. We ran last week other (weaker) language models, and discovered that the advantage of REx is sometimes much more pronounced: For example, on GPT3.5/APPS-competition, our method solves about ~40% more problems (relative improvement) over the second best approach. \n\nThank you again for suggesting this, which we think significantly strengthens the paper by showing more models as well as incidentally discovering cases where REx makes a much bigger difference than we showed in the original submission.\n\n> 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\n\nWe hope that you can reconsider this in light of the above result from the experiment that you suggested. We also hope one can see technical contribution in cost reduction and increased hyperparameter robustness.\n\n> 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.\n\nWe'll include videos we made recently of REx searching through refinement trees, which clearly show a balancing of exploring and exploiting. We also gave a theoretical (mathematical) analysis in L120-L132/Fig 2. \n\nPlease let us know if you have further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Multiple reviewers raised the concern that we only evaluated our method on GPT4. 
To address this concern we are in the process of running our experiments on GPT3.5, Llama3, and Claude3. The attached PDF shows preliminary in-progress results.

Although the results are still preliminary, the advantage of REx seems to like it might be more pronounced for some weaker models such as gpt3.5: Speculatively, this may be due to a "saturating" effect when the model is powerful enough. We think this makes the work stronger because it reveals a regime that we have not discovered previously where our method especially shines. Thank you to the reviewers who suggested these experiments.

Meta-review: 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.
Decision: Accept (poster)