Planning In Natural Language Improves LLM Search For Code Generation
Evan Wang Federico Cassano Catherine Wu Yunfeng Bai Will Song Vaskar Nath Ziwen Han Sean Hendryx Summer Yue Hugh Zhang



: Assessing Contextual Integrity Norms in Language Models
    
Yan Shvartzshnaider

York University


Vasisht Duddu

University of Waterloo


John Lacalamita

York University



    
=======================================================================================================================








§ ABSTRACT 
While scaling training compute has led to remarkable improvements in large language models (LLMs), scaling inference compute has not yet yielded analogous gains.
We hypothesize that a core missing component is a lack of diverse LLM outputs, leading to inefficient search due to models repeatedly sampling highly similar, yet incorrect generations.
We empirically demonstrate that this lack of diversity can be mitigated by searching over candidate plans for solving a problem in natural language.
Based on this insight, we propose PlanSearch,
a novel search algorithm which shows strong results across HumanEval+, MBPP+, and LiveCodeBench (a contamination-free benchmark for competitive coding).
PlanSearch generates a diverse set of observations about the problem and then uses these observations to construct plans for solving the problem. By searching over plans in natural language rather than directly over code solutions, PlanSearch explores a significantly more diverse range of potential solutions compared to baseline search methods.
Using PlanSearch on top of Claude 3.5 Sonnet achieves a state-of-the-art pass@200 of 77.0% on LiveCodeBench, outperforming both the best score achieved without search (pass@1 = 41.4%) and using standard repeated sampling (pass@200 = 60.6%).
Finally, we show that, across all models, search algorithms, and benchmarks analyzed, we can accurately predict performance gains due to search as a direct function of the diversity over generated ideas.





“If you fail to plan, you plan to fail.” 

— Mastermind, Taylor Swift









§ INTRODUCTION



The bitter lesson <cit.> famously posits that two forms of scaling trump everything else: learning and search. 
While recent advances in large language models have removed all doubt on the effectiveness of learning, search has not yet proven its value for large language models, despite its success with classical machine learning techniques <cit.>. 

Here, we refer to search as any method of spending compute at inference time to improve overall performance <cit.>. In this work, we focus our efforts on improving LLM search for code generation, one of the most important current applications of LLMs.
We hypothesize the major bottleneck preventing widespread use of search at inference time for code is a lack of high-level diversity in model outputs. This lack of diversity is likely in part due to specific post-training objectives commonly used to train LLMs as chatbots, in which models are oftentimes optimized to produce a single correct answer <cit.>.
We empirically demonstrate that this is the case for many open-source language models which have undergone significant post-training. Specifically, we show that in many cases, despite instruction tuned models outperforming base models by large margins on a single sample regime (pass@1), this trend disappears—sometimes even reversing—when generating many samples. We refer to Figure <ref> as an example of this phenomenon.

Furthermore, the lack of diversity is particularly harmful for search algorithms. In the most egregious of cases with little to no diversity, such as greedy decoding, repeated sampling from the model returns highly similar programs, resulting in minimal gain from additional inference-time compute. 
This diversity problem is also not reflected in many public leaderboards (e.g. LMSYS Chatbot Arena <cit.>, LiveCodeBench <cit.>, OpenLLMLeaderboard <cit.>), which often report only the pass rate from a single sample of the model, ignoring an entire dimension along which to compare models.
While the performance of one sample is the primary metric of relevance for applications such as chatbots, as users typically are sensitive to latency, this single scalar is insufficient to fully capture the quality of a model when it is allowed to use more inference time compute.



In this paper, we explore several directions for improving the diversity of LLMs at inference time. We hypothesize that the right axis of diversity to search over is the natural language conceptual/idea space, and we validate our hypothesis across several experiments. First, we show that models can produce the correct final program when fed the correct solution sketches, where these sketches have been “backtranslated” from passing solution code into sketches in idea space (Section <ref>).
Second, we show that when models are asked to generate their own ideas before implementing them on LiveCodeBench (IdeaSearch), their accuracy conditioned on a particular sketch trends towards either 0% or 100%, suggesting that most of the variance in passing a particular problem is captured by whether the sketch is correct rather than any other factor.
These two experiments suggest a natural method to improving LLM search for code generation: by searching for the correct idea to implement.


Guided by this principle of maximizing exploration of ideas, we propose . In contrast to many existing search methods that search over individual tokens <cit.>, lines of code <cit.>, or even entire programs <cit.>,  searches over possible plans for solving the problem at hand, where a plan is defined as a collection of high level observations and sketches helpful to solve a particular problem (Figure <ref>). To generate novel plans,  generates a number of observations about the problem, before combining these observations into a candidate plan for solving the problem. This is done for every possible subset of the generated observations to maximally encourage exploration in idea space, before the codes are eventually all translated into a final code solution (Section <ref>). We find that searching over plans outperforms both standard repeated sampling and directly searching over ideas (IdeaSearch, introduced in Section <ref>) in terms of effectively using compute at inference time.

Applying  on top of Claude 3.5 Sonnet achieves a state-of-the-art pass@200 of 77.0% on LiveCodeBench, outperforming both the best score achieved without search (pass@1 = 41.4%) and the standard best-of-n sampling score (pass@200 = 60.6%). Furthermore, consistent with recent findings on the effectiveness of search on top of small models <cit.>, allowing  based on a small model (GPT-4o-mini) outperforms larger models not augmented with search after merely 4 attempts. Evaluations of  across two other coding benchmarks, HumanEval+ and MBPP+ <cit.>, suggest similar improvements.


Finally, we measure the diversity of output code over the idea space of all search methods via an LLM-as-a-judge procedure (Section <ref>) and show that the resulting diversity score is highly correlated with the performance gains generated by that search method. This provides further support for our hypothesis that the effective exploration of plans in idea space is key to LLM search for code generation (Figure <ref>).

























§ RELATED WORK

We reiterate that search as defined in the context of our paper refers to any method which expends inference time compute to improve performance. We further specify planning as any form of high level observation or abstract thought that assists a model in generating a final solution.
Our work builds off a long history of work in scaling search and planning.


 §.§ Search in Classical AI


Classical search algorithms like breadth-first search, depth-first search, and A* search have been widely used for pathfinding, planning, and optimization <cit.>. More advanced search techniques like Monte Carlo tree search (MCTS) have achieved remarkable success in domains like game playing, enabling superhuman performance in Go <cit.>, poker <cit.> and Diplomacy <cit.>.
More recently,
<cit.> find scaling laws for the performance of AI systems in board games, where ELO improves logarithmically with the amount of compute spent at inference.



 §.§ Search with Language Models


Applying search on top of LLMs has been a topic of much interest, especially with an eye towards code generation <cit.>. Historically, methods such as beam search significantly improved performance for translation systems <cit.>. Closer to the present day, several recent works have explored repeated sampling <cit.> as a search method for improving performance. Repeated sampling is a method which directly generates candidate code solutions from the model many times at moderate to high temperatures in hopes that one of the resulting generations will be correct. However, although these works address the roughly linear increase in pass@k with respect to log k, they only focus on the most basic version of repeated sampling, without searching in idea space.

When combined with a verifier, reward model, or other filtering algorithm to select the best generation (in cases where pass@k is not a viable metric due to lack of test cases), it is also known under the name of best-of-n sampling <cit.>. Many works show somewhat good results under intelligent selection of such a filtering algorithm <cit.>.
Recently, several approaches have demonstrated the power of repeated sampling. For example, repeated sampling from a small model can sometimes outperform taking a single sample from a large model on an equalized compute bases <cit.>.
Unlike algorithms such as repeated sampling, which search over the output space, the key insight of PlanSearch is that it is far more effective to instead search plans over the latent idea space. By explicitly searching over different natural language plans before generating the code, we significantly increase the diversity of the final code outputs and thus, the resulting pass@k scores for sufficiently large k.

Regarding searching over plans in natural language, several approaches have also proposed generalizing chain-of-thought <cit.> reasoning into a search-like process, such as Tree of Thoughts <cit.> and Reasoning via Planning <cit.>.
However, prior methods have largely demonstrated effectiveness on somewhat contrived problems designed to highlight the power of search, such as the game of 24, or classic planning benchmarks such as Blocksworld <cit.>, where both benchmarks are easier to solve by explicitly considering many options, and where the `steps' over which to search over are fairly obvious.
By contrast, most real-world planning is used to assist in domains that are complex enough to benefit from, but not require, the additional exploration of plans. We demonstrate that , which plans in natural language, outperforms baseline search methods in one such domain: code generation.
Moreover, our analysis reveals the underlying reason that such search is effective: it increases the diversity of the generated ideas, allowing more efficient search relative to other methods which repeatedly submit highly similar, incorrect solutions.


diversity in LLM paper, cite all of natural llm history

bruh filtering section, code approaches section? maybe like reflexion, allphacode, codet, parsel





§ MOTIVATION



Coding is a powerful area in which search should excel. While search in other domains requires both generating many solutions and selecting the correct solution amongst all the resulting generations, coding often only requires the former, as any valid piece of code can be tested via code execution against given test cases. This allows code search algorithms to sidestep many of the issues that plague search algorithms for more open-ended domains (e.g. generating poetry) due to difficulty in selecting correct solutions out of all the generated solutions.








 §.§ Defining the Search Space

Perhaps the most important question for eliciting strong search capacities is determining which space to search over, as finding the proper layer of abstraction is critical to progress in the field. Prior approaches have varied, with many people searching over individual tokens <cit.>, lines of code <cit.>, or even entire programs <cit.>.
We hypothesize that the key factor is obtaining the correct solution sketch, which we define as a description of the correct program in natural language space. Intuitively, conducting the reasoning process in natural language space allows us to effectively harness the training process of LLMs, which have observed many human reasoning traces in both pre- and post-training. Prior work <cit.> citehas observed strong positive effects from being allowed to conduct such reasoning in natural language, making it a natural place to search over. We describe two experiments providing evidence for this hypothesis by testing on the LiveCodeBench benchmark using GPT-4o-mini as our model. 



 §.§ Backtranslation


To investigate the hypothesis whether the idea space, instantiated as solution sketches, is the right area of exploration, a natural question is whether LLMs can correctly implement a correct code solution given a correct sketch. Inspired by approaches to backtranslation in machine learning <cit.>, we experiment with “backtranslating” passing code solutions back into idea space.
First, we generate code solutions using GPT-4o to generate 1000 attempts to solve the problem and filter out problems without any passing solutions.
As we also do not have a dataset of correct solution sketches associated with each solution, we generate a candidate correct idea via backtranslation. 
We do this by feeding an LLM both the problem and code solution and asking the LLM to convert said solution into a natural language description of the solution. Additionally, we vary the detail of the backtranslated idea via instructions to the LLM in the prompt (e.g. `in w words'). A full description of the prompts can be found in Appendix <ref>, alongside several example backtranslated solutions of various lengths.

We observe that prompting a model with a backtranslated idea significantly improves accuracy, increasing with the length of the translated idea (Figure <ref>), which suggests that having a correct sketch is sufficient to produce the correct final solution with relatively high accuracy, even only after 10 tokens of backtranslated solution. This suggests that the correct direction of search is to explore through idea space to maximize the chance of arriving at a correct idea.








 §.§ Conditioning on Idea Quality

In a follow-up experiment, we prompt an LLM to generate its own sketches to solve LiveCodeBench problems instead of providing it with golden ones via backtranslation. First, we generate 5 ideas per problem using IdeaSearch, defined in Section <ref>. For each idea, we then sample 25 candidate solutions and measure their pass rate.
For this experiment, we filter out any problem that GPT-4o-mini solves with either a 100% or a 0% solve rate, since such problems are either too easy or too hard for the model and would not be informative for this experiment. We end with 75 problems and 375 sketches.

To test our hypothesis that generating a correct sketch is a critical factor for solving problems, we compare the distribution of solve rates for generating correct code solutions conditioned on a given sketch to the distribution over solve rates given a sketch drawn at random, i.e., just the distribution over solve rates. 

Formally, for any problem P_i, we sample some sketch I from some conditional distribution with probability mass P(I| P_i). The probability of solving P_i is then P(solve| P_i, I). We compare the solve-rate distribution, P(solve| P_i, I) over all problems and all sketches versus the solve-rate distribution of ∑_I P(solve| P_i, I) · P(I | P_i) = P(solve| P_i) over all problems.

While verifying whether a sketch is correct or incorrect is difficult without access to external labels, a key insight is that if generating the correct idea is a critical factor in solving the problem, then conditioning on a particular sketch should polarize the distribution of solve rates towards {0, 1}. If the model is given a correct sketch, it should consistently generate correct solutions, while if given a bad sketch, it should consistently generate incorrect solutions.

Our results confirm this to be the case. Figure <ref> shows the distribution of solve rates across problems, both unconditionally (in red) and conditioned on each sketch (in blue). We notice that when grouping by sketches, the solve rates indeed become polarized towards {0, 1}.
This result has important implications for improving code generation, suggesting that a large portion of variance in performance may come from whether the model is able to generate a correct idea or not. Therefore, a natural path for improvement is to focus on the sketch generation step and search for correct sketches and observations in idea space before generating solution code.







§ METHODS


We provide a description of the various methods of search we explore in our work. If additional background on competitive programming and related notation is desired, we provide more (optional) information in Appendix <ref>.


 §.§ Baselines



  §.§.§ Repeated Sampling

We consider the basic prompting approach as a baseline, in which we use few-shot prompting by providing the LLM with a number of problem-solution pairs before asking it to solve the desired question <cit.>. A full example of the prompt is given in Appendix <ref>. In code generation, the most common variant of search utilized is repeated sampling, where models are repeatedly sampled from until they generate an output that passes the test or the maximum number of samples is reached. Refer to the Related Work for more information (Section <ref>).


  §.§.§ IdeaSearch

A natural extension of the Repeated Sampling approach discussed in Section <ref> is to avoid prompting the LLM for the solution code immediately.
This can be viewed as an application of the commonly used “chain-of-thought” prompting to programming problems <cit.>, although we find that IdeaSearch shows non-negligible performance boosts over standard “chain-of-thought” prompting (see Appendix <ref>).

In IdeaSearch, the LLM is given the problem P and is asked to output a natural language solution S of the problem. Then, a separate instance of the LLM is given P and S, and tasked to follow the proposed solution S to solve the problem P. The purpose of IdeaSearch is to isolate the effectiveness of having the correct “idea/sketch” for solving the problem. Empirically, we find that explicitly forcing the search algorithm to articulate an idea for solving the problem increases diversity.
See Appendix <ref> for detailed prompts.



 §.§ PlanSearch


While both Repeated Sampling and IdeaSearch are successful and lead to improvement in the results on benchmark results, we observe that in many of the cases,
prompting multiple times (pass@k) (even at high temperatures) will only lead to small, narrow changes in the output code that change minor aspects but fail to improve upon pitfalls in idea. fix





  §.§.§ Prompting for Observations


Starting from the problem statement P, we prompt an LLM for “observations”/hints to the problem.
We give examples of observations generated in Appendix <ref>.


We denote these observations as O^1_i, where, i ∈{1, …, n_1} due to the fact that they are first-order observations. Typically, n_1 is on the order of 3 to 6. The exact number depends on the LLM output.
To use these observations to inspire future idea generation, we create all subsets with size at most 2 of S^1={O^1_1, …, O^1_n_1}. Each of these subsets is a combination of observations, and for clarity we denote each subset as C^1_i, i ∈{1, …, l_1}, where l_1 = 1 + n_1 + n_12.



  §.§.§ Deriving New Observations

The set of all observations can be thus defined as a directed tree with depth 1, where the root node is P, and an edge exists for each C^1_i pointing from P to C^1_i. We then repeat this procedure from Section <ref> on each leaf node C^1_i to generate a set of second order observations, S^2_i={O^2_i,1, …, O^2_i,n_i, 2}. To obtain second order observations, we prompt the model with both the original problem P and all observations contained in C^1_i, framed as primitive observations that are necessary in order to solve P. The LLM is then prompted to use/merge the observations found in C^1_i in order to derive new ones. 

The same procedure as Section <ref> is used to create all subsets C^2_i, j, for all i ∈{1, …, l_1}. This process may be arbitrarily repeated, but we truncate the tree at depth 2 for computational constraints.


Note that there is no assumption any of the observations generated are correct. In fact, it is critical to note that many of them may be incorrect. , and we give examples in Appendix <ref>.
The observations merely serve to elicit the model to search over a more diverse set of ideas.



  §.§.§ Observations to Code

After the observations have been made, they must be implemented as ideas before being translated into code. For each leaf node, we prompt the model with all observations, along with the original problem P, in order to generate a natural language solution to the problem P. 
To add more diversity, for each generated idea, we generate an additional idea by supposing the idea is wrong, and asking an LLM to give criticisms/feedback, thus increasing our proposed ideas by a factor of 2.

These natural language solutions are then translated into pseudocode, which are subsequently translated into actual Python code. We take a more granular approach to reduce the translation error (which may cause the model to revert to its original mode, disregarding the reasoned-through observations).
We provide all prompts for all sections in Appendix <ref>.









§ EXPERIMENTAL RESULTS







 §.§ Datasets


We evaluate our search methods on three benchmarks: MBPP+, HumanEval+  <cit.>, and LiveCodeBench <cit.>. MBPP <cit.> and HumanEval <cit.> are some of the most widely used code benchmarks in the field. However, since both benchmarks provide only a few test cases, <cit.> updates both benchmarks with additional test cases that increase the benchmarks' robustness to reward hacking. LiveCodeBench is a benchmark for coding that consists of competitive programming problems which typically require advanced reasoning capabilities. Given the reality that coding data is often highly upsampled during pre-training <cit.>, LiveCodeBench differentiates itself from other benchmarks by taking care to segregate problems by date to avoid data contamination concerns. For this paper, we use only the subset of problems between May 2024 and September 2024 to avoid possibilities of contamination.
We choose May 2024 as the cutoff date to ensure that our results with our best performing model (Claude 3.5 Sonnet) are not due to contamination, because Claude 3.5 Sonnet has a knowledge cutoff of April 2024.
To ensure fair comparison, we use the same cutoff for all models evaluated, even though the precise cutoff dates for other models may vary slightly from May 2024.


 §.§ Experiment Details

For all search algorithms, we require that all output code be in the correct format specified, and we mark a solution as incorrect if it does not follow the intended formatting. The extracted code is then run through all tests of the program and marked as correct if and only if it passes all tests.
 
All models are run with temperature 0.9 and top-p of 0.95. Temperature was determined through a coarse hyperparameter sweep on Repeated Sampling and IdeaSearch from T∈{0.0, 0.1, 0.2, …, 1.2}, which we describe in Appendix <ref>.

Both Repeated Sampling and IdeaSearch generate exactly n codes, whereas PlanSearch generates a variable number of codes, usually ranging on the order of 300 to 400. To compute pass@k, we use the unbiased estimator in Equation <ref> <cit.>[Note that the estimator in Equation <ref> theoretically requires that the number of successes follows a binomial distribution. Repeated Sampling and IdeaSearch obey this, but PlanSearch generations may not be independent. See Appendix <ref> for more discussion.].
If k > n, we assume the remaining generations did not pass.
To compute pass@k for filtering, we limit the pool of codes to those that are filtered, meaning that both n and c may shrink in size. This can be thought of as a conditional probability, where the condition is that the code passes public tests.



 §.§ Results


Our summarized results for Repeated Sampling, IdeaSearch, and PlanSearch can be found in Table <ref>, Figure <ref>, and Figure <ref>.
Additionally, we plot our full pass@k curves for all methods, models, and datasets in Appendix <ref>. For sake of easy comparison, we also plot all relative gains compared to Repeated Sampling@1 averaged over all models in Appendix <ref>. For a compute-normalized comparison between Repeated Sampling and PlanSearch, see Figure <ref>. 


 §.§ Public Test Filtering


Public test filtering is a method which only chooses samples out of the original pool n which pass the public tests.
This is particularly useful in settings such as code deployment where executing the full suite of tests may be computationally costly or otherwise undesirable (e.g. in a coding contest where every incorrect submission is penalized).
Thus, instead of submitting all n codes, after public test filtering, only codes c_i would be submitted such that c_i(x_j) = y_j for all j ∈{1, …, u}, where c_i(x) refers to the output from running the code on some input x. The primary effect of public test filtering is to shift the pass@k curve leftward, since public test filtering will discard low quality candidate solutions that either fail to compile or fail elementary test cases for the problem.

All problems in MBPP+, HumanEval+, and LiveCodeBench come with a few public tests which are usually used to sanity check any submissions. We can further improve performance by filtering on these public tests before a final submission, as described.
Applying public test filtering reduces the number of samples to achieve the same accuracy by tenfold:  to achieve a 77.1% accuracy on LiveCodeBench after just 20 submissions (pass@20) compared to a pass@200 of 77.0% without using public filtering (see Figure <ref>). We provide full results for the other datasets in Appendix <ref>.





§ ANALYSIS


Our results suggest that both  and IdeaSearch outperform basic sampling by a wide margin (Figures <ref>, <ref>, <ref>), with  achieving the best score across all methods and models considered.
We show the detailed pass@k results for each dataset in Figures <ref>, <ref> and <ref>.
We also compare with Chain-of-Thought <cit.> in Appendix <ref>. Interestingly, we find that IdeaSearch performs somewhat better, which we speculate comes from differences in splitting solution sketch into two model responses, instead of doing both chain-of-thought and code solution in one model response.

Investigating the differences in specific models, we notice that trends exhibited by the pass@k curves are not uniform across all models; in fact, each curve seems unique.
We hypothesize that these differences are in part due to changes in idea diversity, as investigated in Figures <ref>, <ref>, <ref>. From the figures, we can see that our approximate diversity score accounts for much of the variance we see in the relative improvement that arrives from scaling-up inference-time compute. This correlation holds across all methods and models on the same dataset, thus suggesting that diversity score can be used as a proxy to predict for relative pass@k improvement. For further discussion on the specifics of the diversity score, see Section <ref>.

One interesting point of observation is that  often hurts pass@1 for several models, including most notably Sonnet 3.5 on LiveCodeBench, our best performing combination. 
Intuitively, this is because increasing the diversity across ideas likely dilutes the probability that any particular idea is generated, while simultaneously increasing the chance of having at least one correct idea within said pool. Therefore, pass@1 may be slightly lower than usual, yet pass@k will likely surpass “pools” of ideas lacking diversity for this reason. See Figure <ref> for a graphical intuition.

Finally, in Table <ref> and Figure <ref>, we present our main results normalized across attempts/completion, where each search method is allowed k attempts to solve each problem. An alternative method of normalizing across methods is to equalize the amount of compute spent on each method. Since  and IdeaSearch first plan out an idea before implementing the final solution, they both spend more compute at inference time per solution generated. In Appendix <ref>, we report the equivalent plots normalized across compute. Our findings are highly similar and suggest that  outperforms all other methods if sufficient compute is expended at inference time.



 §.§ Measuring Diversity





We find that diversity as measured in idea space is highly predictive of search performance, as measured by the relative improvement between a model/method's pass@1 and its pass@200 (Figure <ref>).
While the most common measure of diversity is entropy <cit.>, entropy is insufficient for a number of reasons for the precise setting of LLMs <cit.>. As a simple example, consider two different language models, one of which generates minor variations of the same program while another generates a variety of programs with different underlying ideas. Even if both models have the same entropy, the latter model will be significantly better when augmented with search capabilities. 

In our setting, we measure diversity by grounding it in idea space using a simple pair-matching strategy across all generated programs. Formally, suppose we have a pool of n code generations, {c_1, …, c_n}.
We assume that each piece of code implements some sketch, which can be thought to exist in some latent `idea' space. We consider two sketches similar if they are within ϵ of each other in this latent space, for some choice of ϵ. As such, in this space, c_i having a similar idea to c_j and similarly for c_j and c_k does not imply c_i and c_k share a similar idea. 

To compute the diversity of such a given generation pool, we ask an LLM to judge the similarity of two ideas in the following manner.
First, we
construct each of the n 2 pairs. For each pair (c_i, c_j), we judge (using an LLM) whether both c_i and c_j implement the same idea. We define this as the function S(c_i, c_j) ∈{0, 1}, which evaluates to 1 if c_i and c_j implement the same idea and 0 otherwise.
Our overall diversity score for a particular problem is then defined as:

    D = 1 - ∑_i < jS(c_i, c_j)/ n 2

Models that output programs that all implement the same idea will have a score of D=0, while models that output completely unique programs will have a score of D=1. Overall, a score of D implies that if two codes are chosen at random, the probability that they are the same idea (as measured by the LLM) is D. In Appendix <ref>, we describe this measure in additional mathematical depth.

For a particular method, our reported diversity score is simply the diversity score over all problems in the considered dataset.
For computational feasibility, for large n, we instead sample a subset of 40 codes and test all pairs from that subset instead. In order to test code samples, we first backtranslate using an LLM to express the code in natural language before comparing each pair using both the code and the backtranslated idea.
We detail the prompts used in Appendix <ref> and use OpenAI's GPT-4o-mini as the supporting LLM. 



§ LIMITATIONS AND FUTURE WORK

While PlanSearch substantially improves diversity over idea space at inference-time, fundamentally, improvements in diversity should come at the post-training stage. This likely requires re-imagining the post-training pipeline for LLMs around search, instead of the current paradigm optimized for a single correct response. This may require both collecting high quality post-training data that is also sufficiently diverse, and new learning objectives that do not aim solely to maximize the expected reward of a given response.
We are optimistic around future work to design significantly improved post-training objectives that maximize both quality and diversity and which specifically optimized to use inference-time compute to maximum effectiveness.


In terms of methodological improvements to ,  currently searches all leaf nodes in the search tree uniformly. Because of this, it becomes quickly intractable to go further than a couple levels deep, and in our experiments, we are only able to go two levels down the tree. Several approaches based on Monte-Carlo Tree Search (MCTS), such as Tree of Thought <cit.> or Reasoning as Planning <cit.>, have suggested that some form of dynamic pruning and expansion of nodes can be very helpful. We are optimistic that  can be further improved by such methods.
Furthermore,  is a fairly elementary method taking advantage of the paradigm that searching over a conceptual or idea space is an effective method to improve diversity, and thus, downstream task performance. It is completely feasible to search at an even higher level of abstraction than observations, which may be used to inject even more diversity into the final generated outputs.

 and IdeaSearch tradeoff a slight deterioration of pass@1 performance for a large improvement in pass@k performance. However, in many such cases outside of code generation, it is infeasible to run an LLM-based model for more than a few attempts at most. For example, in Figure <ref>, PlanSearch does not significantly outperform Repeated Sampling until k ≥ 4.

Fortunately, many filtering algorithms exist, which implicitly bring pass@k (for high k) to pass@1 (or lower k), i.e. shifting the original pass@k curve leftward. A simple example of this is public test filtering. As seen in Figure <ref>, pass@1 of filtered PlanSearch significantly improves upon pass@1 of base Repeated Sampling, which gets even better as k increases. Moreover, most to almost all base models with public test filtering outperform their instruct model variants at pass@1, no matter the dataset (see Appendix <ref>), where clearly base models are known to be worse, yet trading off for somewhat higher diversity. Thus, we argue that there exists a potential for a new paradigm—developing search algorithms which tradeoff pass@1 performance for much stronger pass@k performance, then filtering the promising generated solutions to extract the pass@k back into the pass@1. 

Additionally, we focus on code generation in this paper and do not consider the applicability of  to a broader set of domains. One point of importance is that the pass@k metric heavily used throughout code generation may not be as applicable to other domains and a larger focus on selecting the correct solution out of all possible generated candidates may be required, instead of merely generating the correct solution. However, with good filtering methods, which we demonstrate can be simple in nature, pass@k, for medium k, can be effectively brought down to pass@1, emphasizing a similar paradigm of increasing diversity, then strengthening existing filtering methods.

Finally, a natural extension of this work is training the underlying model itself on successful plans and code solutions obtained from PlanSearch. This has the potential to distill the pass@k into the pass@1—without inference-time methods like filtering—by reducing the likelihood of the model going down branches of the search tree which do not lead to correct solutions. 
We believe that such training is likely to significantly improve the model and look forward to future work in this direction.







§ ACKNOWLEDGEMENTS

We would like to thank Jason Wei, Miles Turpin, Sail Wang, Horace He, Kenneth Li, Celia Chen, Rahul Chalamala, Alan Wu, and Kevin Chang for their helpful comments, suggestions and discussion over the course of this project.







iclr2024_conference






Appendix





§ FULL PASS@K CURVES FOR ALL MODELS AND ALL BENCHMARKS



See Figures <ref>, <ref>, <ref>. We plot all models and methods on HumanEval+, MBPP+ <cit.>, and LiveCodeBench <cit.>, respectively.









§ FULL PASS@K CURVES WITH PUBLIC FILTERING



See Figures <ref>, <ref>, <ref>. We plot all models and methods with public test filtering on HumanEval+, MBPP+ <cit.>, and LiveCodeBench <cit.>, respectively.








§ AVERAGE RELATIVE IMPROVEMENTS



See Figures <ref>, <ref>, <ref>. To create these graphs, the relative improvements of each point on all pass@k curves are computed and compared to the respective pass@1 of Repeated Sampling. Then these values are averaged over all models, so that there is one curve per method per dataset. The datasets are HumanEval+, MBPP+ <cit.>, and LiveCodeBench <cit.>, respectively. For the public test filtered versions, see Figures <ref>, <ref>, <ref>.

















§ COMPUTE NORMALIZED PASS@K GRAPHS



See Figure <ref>. For each run of a method in Appendix <ref>, we compute the number of generated tokens needed per completion, per problem, independently on each dataset. Then, we average across all datasets to obtain 244 generated tokens per completion per problem for Repeated Sampling, and 1,428 generated tokens per completion per problem for PlanSearch.







§ COMPARISON WITH CHAIN-OF-THOUGHT



See Figures <ref>, <ref>, <ref>, which are run on LiveCodeBench <cit.>, MBPP+, and HumanEval+ <cit.>, respectively. These are the same plots as Appendix <ref>, with CoT <cit.>. See Figures <ref>, <ref>, <ref> for the public test filtered versions.














§ ABLATION ON TEMPERATURE FOR REPEATED SAMPLING AND IDEASEARCH



See Figure <ref>. We sweep over temperature increments of 0.1 from 0.0 to 1.2, inclusive, with top-p of 0.95, on Repeated Sampling and IdeaSearch.








§ DIVERSITY SCORE VS SEARCH IMPROVEMENT PLOTS FOR MBPP+ AND HUMANEVAL+

See Figures <ref>, <ref>, <ref>. Each figure is made through running the diversity measure as described in Section <ref> on the generated codes of each run, then compared with the relative gain from pass@k compared to pass@1.









§ BASE MODELS VS. INSTRUCT MODELS FOR LARGE SAMPLES



We find that base models, despite performing poorly relative to their instruct counterparts for evaluated with pass@1, will frequently match or even exceed performance on pass@k for sufficiently high k. This is likely due to higher amounts of diversity in base models, which have not undergone post-training designed to elicit a single strong response from the model.

We see this effect across all models for HumanEval+ and MBPP+, but only the DeepSeek-Coder-V2 family for LiveCodeBench.

See Figures <ref>, <ref>, <ref> for Llama-3.1-8b pass@k comparisons.

See Figures <ref>, <ref>, <ref> for Llama-3.1-70b pass@k comparisons.

See Figures <ref>, <ref>, <ref> for DeepSeek-Coder-V2-Lite pass@k comparisons.

We also ran Llama-3.1-8b and DeepSeek-Coder-V2-Lite pass@k comparisons for k up to 10,000; see Figures <ref>, <ref>.



















§ BASE MODELS VS. INSTRUCT MODELS WITH PUBLIC TEST FILTERING



We repeat the graphs from Appendix <ref>, but with public test filtering. We find that base models with public test filtering almost always exceed the pass@1 of their instruct model variants.

See Figures <ref>, <ref>, <ref> for Llama-3.1-8b pass@k comparisons with public test filtering.

See Figures <ref>, <ref>, <ref> for Llama-3.1-70b pass@k comparisons with public test filtering.

See Figures <ref>, <ref>, <ref> for DeepSeek-Coder-V2-Lite pass@k comparisons with public test filtering.




















§ PROMPTS




 §.§ Backtranslation




  §.§.§ Backtranslate System Prompt





  §.§.§ Implement Backtranslation Idea






 §.§ Repeated Sampling






 §.§ Simple Idea






 §.§ PlanSearch




  §.§.§ Prompt for Observation Part 1





  §.§.§ Prompt for Observation Part 2






  §.§.§ Combining Observations







 §.§ Measuring Diversity






§ COMPETITIVE PROGRAMMING

Competitive programming is a popular subset of programming tasks that involve solving complex algorithmic reasoning. Typically, problems consist of a problem statement (written in natural language) P, with associated tests: (x_i, y_i), i ∈{1, …, m}, for which any solution must pass all of them. 

The number of tests m depends on the problem, but typically ranges on the order of 25 to 100. A small subset of the tests are typically given to the solver (we call these public tests) to use as validation that their program passes simple cases. The rest of the tests are hidden. Solutions to the problems must generally pass all the tests to be considered correct. Formally, we let f(x) denote the output of said code ran on input x. The solution code is considered correct (passing) if and only if f(x_i) = y_i for all i ∈{1, …, m}.

Each dataset consists of many (on the order of low-hundreds) independent problems, and models are evaluated on each of these problems independently.



§ A MODEL OF REPEATED SAMPLING: PASS@K

Consider a simplified model of repeated sampling for code generation.
Suppose we have a dataset D = {P_1, …, P_l} with l problems.
For some problem P_i, define the probability p_i as the probability that our code generation model solves the problem P_i in one submission.
The pass@k <cit.> metric (for problem P_i) is defined as the probability that our code generation model solves the problem P_i at least once out of k submissions. 
Thus, if we know the true p_i of our model, we may compute our pass@k simply:

    pass@k_i    = 1 - (1 - p_i)^k 
    pass@k   = ∑_i pass@k_i / l


However, it turns out that for k>1, the naïve estimator as seen in Equation <ref> is biased, if we sample n_i ≥ k from our code model to solve P_i, c_i ≤ n_i are correct, and compute p_i = c_i / n_i <cit.>.
Instead, pass@k_i is typically computed using the unbiased estimator:

    pass@k_i    = 1 - n-c k/ n k

Note that reporting pass@k on a dataset where l=1 is rather pointless, since pass@k can be derived using only pass@1_1 and n_1. Every curve, over a suitable range of k values, will look like the S-curve seen in Figure <ref> (as k is plotted on a log scale).



However, with datasets where l > 1, models are able to differentiate themselves through larger k, since the overall pass@k is an average of these l curves. For example, for l=3, it is less optimal to have solved probabilities of Set 1 = {0.001, 0.7, 0.9} versus Set 2 = {0.05, 0.1, 0.25}, in the regime of roughly k=20 to k=2,000 (in which both converge to 1), even though Set 1 has a pass@1 of 53% and Set 2 has a pass@1 of 13%. See Figure <ref>.

Although not shown in the graph, Set 2 converges close to 1 at roughly k=400, several orders of magnitude below Set 1. In addition, note that the slight notch seen in Set 1's curve at large k is due to the presence of low, but non-zero solve-rates, which can be seen in empirical pass@k curves later on. (These can be thought as the beginning of the `ramping-up' regime of the typical S-curves in Figure <ref>.)


























§ MATHEMATICS OF THE DIVERSITY MEASURE


While our choice of a diversity metric is intuitive, one should note that there are a number of intriguing details that result from our definition. In particular, it is not necessarily the case that a model that outputs k unique ideas out of n samples to achieve a diversity score of k/n. 
Consider an example of n=9 codes, separated into 3 cliques of 3, where each clique implements the same idea (and separate cliques implement separate ideas). In this setup, 1/3 of ideas are unique, but in our metric, there are 3 matching idea pairs (and 9 total matching idea pairs) out of 92 = 36, for a diversity score of 1 - 9/36 = 3/4.








§ BIASED ESTIMATOR FOR PASS@K DUE TO NON-INDEPENDENCE OF 



From a pure theoretical standpoint, the expression is biased (if using the same interpretation), but it still leads to a similar interpretation—computing the probability that a subset of size k drawn from the set of samples we already generated contains at least one success. (These given samples were generated by one run of PlanSearch.) As such, in theory, the estimator may be slightly biased in the PlanSearch case when computing its true pass@k. In practice, we do not believe this to be a large concern, especially as our primary results feature a relatively large k=200.


do all four below later



