Planning with Large Language Models for Code Generation
Authors: Shun Zhang
Zhenfang Chen
Yikang Shen
Mingyu Ding
Joshua B. Tenenbaum
Chuang Gan
Venue: ICLR 2023 poster (2023)
Keywords: Large Language Model
Code Generation
Planning

Abstract: Existing large language model-based code generation pipelines typically use beam search or sampling algorithms during the decoding process. Although the programs they generate achieve high token-matching-based scores, they often fail to compile or generate incorrect outputs. The main reason is that conventional Transformer decoding algorithms may not be the best choice for code generation. In this work, we propose a novel Transformer decoding algorithm, Planning-Guided Transformer Decoding (PG-TD), that uses a planning algorithm to do lookahead search and guide the Transformer to generate better programs. Specifically, instead of simply optimizing the likelihood of the generated sequences, the Transformer makes use of a planner that generates candidate programs and tests them on public test cases. The Transformer can therefore make more informed decisions and generate tokens that will eventually lead to higher-quality programs. We also design a mechanism that shares information between the Transformer and the planner to make our algorithm computationally efficient. We empirically evaluate our framework with several large language models as backbones on public coding challenge benchmarks, showing that 1) it can generate programs that consistently achieve higher performance compared with competing baseline methods; 2) it enables controllable code generation, such as concise codes and highly-commented codes by optimizing modified objective.

Published as a conference paper at ICLR 2023
PLANNING WITH LARGE LANGUAGE MODELS FOR CODE GENERATION
Shun Zhang, Zhenfang Chen, Yikang Shen MIT-IBM Watson Al Lab
Mingyu Ding
The University of Hong Kong
Joshua B. Tenenbaum MIT BCS, CBMM, CSAIL
Chuang Gan
UMass Amherst, MIT-IBM Watson Al Lab
ABSTRACT
Existing large language model-based code generation pipelines typically use beam search or sampling algorithms during the decoding process. Although the pro- grams they generate achieve high loken-matching-baee based scores, they often fail to compile or generate incorrect outputs. The main reason is that conventional Trans- former decoding algorithms may not be the best choice for code generation. In this work, we propose a novel Transformer decoding algorithm, Transformer Decoding (PG-TD), that uses a planning algorithm to do lookahead search and guide the Transformer to generate better programs. Specifically, in- stead of simply optimizing the likelihood of the generated sequences, the Trans- former makes use of a planner to generate candidate programs and test them on public test cases. The Transformer can therefore make more informed decisions and generate tokens that will eventually lead to higher-quality programs. We also design mechanism that shares information between the Transformer and the planner to make our algorithm computationally efficient. We empirically evaluate our framework with several large language models as backbones on public coding challenge benchmarks, showing that I) it can generate programs that consistently achieve higher performance compared with competing baseline methods; 2) it en- ables controllable code generation, such as concise codes and highly-commented codes by optimizing modified objective'.
Saaning-Guidded
INTRODUCTION
Large language models like Transformer (Vaswani et al., 2017a) have shown successes in natural language processing, computer vision, and various other domains. Thanks to Transformer's power on sequence modeling, it has been adopted for code generation (Wang et al., 2021; Ahmad et al., 2021) 1) by treating programs as text sequences. Transformer has achieved significant improvements on the benchmarking tasks of code translation (Roziere et al., 2022), code completion (Chen et al., 2021a), and solving coding challenge problems (Hendrycks et al., 2021). Recently, AlphaCode (Li et al., 2022) even achieved competitive-level performance in programming competitions with the help of large Transformer models pre-trained on a large programming corpus.
Transformer- based pipelines like AlphaCode follow the tradition of natural language processing and use sampling methods (Fan et al., 2018; Dabre & Fujita, 2020) during the generation process. Specifically, they sample a large number of complete programs using pre-trained code generation Transformer, evaluate these programs using the public test cases provided in the dataset, and output the program that passes the most number of test cases. Compared with beam search-based methods, these sampling followed by filtering algorithms (which we will refer to as sampling + filtering) can take advantage of test cases and indeed improve the quality of the generated programs. However, during the Transformer generation process, they do not consider the test cases at all. Instead, they only use the test cases to evaluate the programs after all the candidate programs are generated. This can make their algorithms sample inefficient. Different from natural languages, programs may fail
Project page: https: / cowea mot.s, .i tht.... 1o. Correspondence to: shun. zhang@ bbm com.

---

Published as a conference paper at ICLR 2023
completely with even a single incorrect generated token. So these algorithms need to exhaustively sample a large number of programs to find a correct solution.
The main reason behind the sample efficiincy issue of these algorithms is that the Transformer beam search algorithm and the sampling algorithm (Vaswani et al., 2017b) may not be the best choices for code generation. An ideal code generation algorithm should stop early in the generation process when it knows the program it currently generates would certainly fail, and bias the generation pro- cess towards generating successful programs that pass more test cases. To achieve such a goal, we contribute to applying a planning algorithm in the Transformer generation process. Since a planning algorithm can use the pass rates of the generated programs as its objective, we use it to determine the quality of the generated codes and make the Transformer model make more informed decisions.
In this paper, we investigate the following research question: Can we integrate a planning algo- rithm with pre-trained code generation Transformer, achieving an algorithm that generates better programs than the conventional Transformer generation algorithms and the well-accepted sam- pling + filtering scheme in the literature? To answer this question, we propose novel algorithm, SoaminngGGuided Transformer Decoding (PG-TD). During the code generation process, a planner does lookahead search and finds tokens that will lead to high-quality codes. The planner alone may not efficiently find high-quality codes due to the large search space of codes, and that is where a pre- trained code generation Transformer comes into play. Specifically, the Transformer beam search algorithm and the next-token probabilities are used inside the planner to provide useful heuristics. We find that a straightforward integration between the planner and the Transformer be compu- tationally inefficient. So we design mechanisms that allow the Transformer and the planner to share their information make the overall algorithm more efficient.
We emphasize that our algorithm is model-agnostic,, that is, any standard code generation Trans- former model can be used as the backbone Transformer. Importantly, our algorithm does not require acquiring more sample solutions or finetuning the Transformer model to improve its performance. We empirically find that our proposed algorithm generates higher-quality programs under multiple accepted metrics compared with competing baseline methods. Additionally, we also empirically show that our algorithm has the following advantages. 1) By changing the reward function of the planner, our algorithm becomes versatile and can optimize different objective functions without the necessity of finetuning the Transformer model. 2) Our algorithm can generate solutions that are used to finetune onde-generation Transformer model to improve the Transformer's 's performance. More precisely, we have the following contributions in this paper.
First, we propose novel algorithm, Planning-Gin Guideed Transformer Decoding (PG-TD), that uses a planning algorithm for lookahead search and guide the Transformer to generate better codes. Our algorithm is model-agnostic, which can work with any standard Transformer model, and does not require knowledge of the grammar of the generated programs.
Second, a direct integration of the planning algorithm with the Transformer decoding process can cause redundant uses of the Transformer beam search algorithm. We contribute to design- ing mechanisms that significantly improve the computational efficiency of the algorithm. Third, we evaluate our algorithm on competitive programming benchmarks and empirically show that our algorithm can consistently generate hetter programs in terms of the pass rate and other metrics compared with the baseline methods. We also show that our algorithm is versatile and can optimize objectives other than the pass rate for controllable code generation, such as generating concise codes and codes with more comments.
2 RELATED WORK
Transformers for program synthesis. Our work based on Transformer for program synthe- sis (Roziere et al., 2020; Austin et al., 2021). Inspired by their capacities on a range of natural language tasks, modern transformer-based language models (Devlin et al., 2019; Radford et al., 2019; Raffel et al., 2020) have been adopted for program synthesis by treating programming lan- guages in the same way as natural languages. A family of BERT-based Transformers are developed for code syntax (Kanade et al., 2020; Feng et al., 2020; Devlin et al., 2019; Guo et al., 2020). Later, CodeX (Chen et al., 2021a) and CodeT5 (Wang et al., 2021) adopted GPT2 (Radford et al., 2019) and T5 (Raffel et al., 2020), respectively, as backbones for both code understanding and generation. Different learning methods including learning from examples (Ellis et al., 2021) and nyural-symbolic methods (Nye et al., 2020) were explored. Recently, AlphaCode (Li et al., 2022) combined large
2

---

Published as a conference paper at ICLR 2023
transformer models pre-trained massive program data with large-scale sampling, showing com- petitive performance in programming competitions. All these works mainly focused on training more powerful code-generation models and still used beam search (Graves, 2012) or sampling (Fan et al., 2018) during the Transformers' generation process.
Test cases for program synthesis. Our work is also related to using unit tests (Myers, 1979) for program synthesis. Tufano et al. (2020a;b) propose to generate test cases and corresponding accurate assert statements with Transformer models (Vaswani et al., 2017a). Recently, Roziere et al. (2022) leverage automatically generated unit tests to construct parallel training data for unsupervised code translation. Chen et al. (2018); Gupta et al. (2020); Chen et al. (2021b) directly synthesize domain- specific programs from input-output pairs without problem description. Ellis et al. (2019) use test cases to train a reinforcement learning agent and use a sequential Monte-Carie sampling method for code generation. Unlike the prior work, we use unit-testing results as reward signals for a tree-search-basee based planning algorithm, which is further integrated with a Transformer-basee large language model to generate better codes.
Planning and reinforcement learning for code generation. The code generation problem has been formulated as sequential decision-making problem (Bunel et al., 2018; Chen et al., 2018), which enables designing and applying reinforcement learning (RL) and planning algorithms for code generation. RL has been used for both the training phase and the decoding phase of the Transformer for code generation. In the training phase, Le et al. (2022); Xu et al. (2019b) use an RL objective that optimizes the correctness of the generated programs instead of optimizing their similarity the reference solutions. In the decoding phase, the Monte-Carlo tree search (MCTS) algorithm has been applied to search for high-quality codes. However, MCTS itself is unable to scale larger domains. It is only used to generate domain-specific languages or for restricted programming synthesis tasks like assembly code generation (Xu et al., 2019a), Java bytecode translation (Lim & Yoo, 2016), and robot planning (Matulewicz, 2022). These tasks have a much smaller scale than generating Python codes in our work. Therefore, their frameworks do not require a large language model and do not need to address the challenge of integrating a planning algorithm with a large language model.
MCTS is more efficient and applicable to large domains when combined with deep learning or with default policy prior knowledge (Gelly & Silver, 2011; Silver et al., 2016; Simmons-Edler et al., 2018). We consider the same overall recipe in our work. Specifically, we contributed to integrating the large language model with a tree search algorithm to design novel algorithm that is capable of solving competitive programming problems, and designed mechanisms to improve its efficiency.
Planning in natural language generation. Planning algorithms like MCTS have also been used to find the optimal text outputs for different natural language processing (NLP) tasks. For example, Scialom et al. (2021); Leblond al. (2021); Chaffin et al. (2022) pre-trained discriminators or pre-defined metrics as reward functions. We want to emphasize that we are the first to combine a tree search algorithm with large language models for general-purpose programming language gener- ation. We design the interfaces between these two components and deal with the unique challenges of making the framework computationally efficient. Concurrently, other search algorithms like A* algorithm (Hart et al., 1968) are also applied in the Transformer decoding process. Lu et al. (2021) consider the constrained text generation problem and integrate A* with the beam search algorithm. However, their constraints are expressed as a formal logic expression, which is different from maxi- mizing the pass rate in our problem.
3 METHOD
3.1
OVERVIEW
We consider the code generation problem for competitive programming, illustrated in Fig. 1. An agent given the natural language description of a coding problem. It requires the agent to under- stand the problem description and generate a program that solves the problem. Similar to Li et al. (2022); Chen et al. (2018); Ellis et al. (2019), we assume that the agent has access to a set of test cases, where a test case is pair of input, output strings. Given the input string, the agent is expected to generate a program that produces an output that exactly matches the test case's 's output string. The objective is to generate program that passes the most number of test cases. To determine if the agent is able to generate programs that generalize to unseen test cases, we divide the test cases into public test cases and private test cases, following the terms in Li et al. (2022). The agent can only
3

---

Published as a conference paper at ICLR 2023
Problem Statement
Given a string S. Replace every character in S with a and print the result.
Constraints
(1). S a string consisting of lowercase English letters.
(2). The length of S is between 1 and 100 (inclusive).
Input
Input is given from Standard Input in the following format: S
Output
Replace every character in S with and print the result.
Sample Test Input
sardine
Sample Test Output
EIEETEE
s=input()
1 s=input()
s=str(input() for in range(len(s));
s=list(s)
2 s= (is
3 for 1 in range( (S))e(
3 for i in range(
if s !="x"
4
for j in range((en(s))
if [[i]=="x":
5
if s [11] =="x"
5
[[i]=" x
s[il=j print (s)
else:
print(s)
continue 8 print (""int("""" ("" join((s)
Beam Search (Pass Rate: 0.00).
Sampling + Filtering (Pass Rate:
PG-TD (Pass Rate: 1.00).
0.22).
Figure 1: A code generation example for competitive programming, with the problem description (top) and the programs generated by baseline algorithms and our PG-TD algorithm (bottom).
access the public test cases during the program generation process, while we use the private test cases to evaluate the programs it generates.
Transformer models (Li al., 2022; Hendrycks al., 2021) have been widely applied for code generation thanks to their capacity in to-suuene modeling. In the Transformer's genera- tion process, beam search (Graves, 2012) and sampling (Fan al., 2018; Dabre & Fujita, 2020) are adopted to generate code sequences. However, these algorithms cannot easily optimize an objective different from what trained on (usually the similarity the reference solutions). So we cannot directly use these generation algorithms to generate programs aiming to pass more test cases.
On the other hand, a planning algorithm can directly optimize the pass rate or any desirable pro- gramming objective. To use planning algorithm, we follow Bunel et al. (2018); Ellis et al. (2019) to formulate the code generation problem as a Markov decision process (MDP) (Sutton & Barto, 2018). In this formulation, state is the concatenation of the problem description and partial or complete program, where a complete program ends with a special terminal token. An action a is a token in the vocabulary set of the Transformer. There is special termination action (the terminal token) that indicates that the agent believes the program complete. The transition function de- terministically concatenates a state with token a, and an episode ends when the agent takes the termination action. The reward of state s is the pass rate of the program on the public test cases when s is complete program (i.e., when the last token of s the terminal token). The reward of partial program is always 0.
There is rich literature on finding the optimal policy in an MDP. In this paper, we consider a tree search-based planning algorithm inspired by Monte-Cario tree search (MCTS), illustrated in Fig. 2. Intuitively, the tree search algorithm maintains tree structure where nodes correspond to states and edges correspond to actions. The algorithm starts from the root node (the initial state) and searches the state space to find terminal states with high rewards. It maintains 1) the number of times each node is visited and 2) value function that maintains the maximum reward obtained by starting in node (or state) s and taking action a. The algorithm would visit and expand nodes with either higher values (as they lead to higher-quality programs) or with smaller visit numbers (as they under-explored) In the following part of this section, we describe how we integrate the tree search algorithm in the generation process of Transformer.
4

---

Published as a conference paper at ICLR 2023
Selection
Expansion
Evaluation
Backpropagation
<PD>
<PD>
<PD>
<PD>
8
a=
%
a=
a
suggests next tokens
8,b
a,
Evaluate
Transformer
public test cases
generates complete programs using beam search
Illustration of using the Monte Carlo tree search algorithm in the Transformer generation process for code generation. <PD> stands for problem description.
Figure 2:
3.2 PLANNING-GUIDER TRANSFORMER DECODING
Now we are ready to answer the question we asked in the introduction: Can we in- tegrate a planning algorithm with pre- trained code generation Transformer to generate better programs? We design a Transformer generation algorithm where a tree search algorithm is used to perform
Algorithm The PG-TD algorithm.
Require: roof: the current state; P-UCB exploration pa- rameter; A: the maximum number children any node; b: the number of beams for Transformer beam search. DICTIONARY()
1: arogram_dict
2: for 1,2, mar rollouts do
3:
node root
lookahead planning. The tree search al- gorithm alone may not be able to find high-quality codes due the large search space. So the conventional Transformer
4: Selection
5: while [node.childrent > 0 do node +- P UCB SELECT (RLECT(nndde.chlldee. end while
6:
7:
beam search algorithm and the next-token
8: Expansion
probabilities provided by the pre-trained
9: newL.okeens + TOP TOP_K(node, do
Transformer are used by the tree search algorithm to guide the search process.
10:
for next_loken E
11:
next_state CONCAT next token)
Create node new mode where new new mode neat_sstaee
We provide the pseudocode of our Soaning--Gudden Transformer Decoding algorithm (PG-TD) Algorithm and illustrate the whole process in Figure 2. The PG-TD algorithm follows the same framework as the standard MCTS algo-
13:
Add new_node to the children of node
14:
end for
15:
# Evaluation
p + BEAM SEARCH(node,
17:
r + GET GET GET.RREWRD(()
18:
program_dict(pp =
rithm, based on the implementation used
19:
# Backpropagation
in Silver al. (2017). Here, we focus how the Transformer is used in the tree search steps. We provide more details of our algorithm in Sec. D.I.
20:
Update and the values of node and its ancestors in the
tree with
21:
end for
22: return program in program dict with the highest reward
In the selection step, we follow Silver
et al. (2017) and use the P-UCB algorithm to select which branch of the tree we want to explore. In P-UCB, we weigh the exploration term by the probability of the next tokens determined by the Trans- former. So the tree search selects higher-probability tokens more often. The selection algorithm is parameterized by an exploration parameter, e, where higher value leads more exploration. We describe the details of P-UCB in Sec. D.I.
In the expansion step, after a node in the tree selected, we select the possible next tokens and add the corresponding next states new nodes to its children list (for succinctness, a node also refers to the state that it represents). Sampling a random token as the standard MCTS may very likely cause syntax error. So call TOP_K to get the most likely next tokens, where rOP_K(s. (s, k) returns the l most likely next tokens starting from s; le is the maximum number of children that any node may have. The corresponding lt next states the concatenations of the current state with each of

---

Published as a conference paper at ICLR 2023
the next tokens suggested by the Transformer. These next states added to the children list of the current node. (Line 9-14)
In the evaluation step, need evaluate the selected node. Note that node may still be partial program. We cannot directly evaluate the quality of a partial program as we do not know how will be completed and how many test cases it will pass. Here, use the Transformer again by calling the REAM_SEARCH function to generate a complete program from the current node, where BEAM SEARCI(S b) generates sequence using the Transformer beam search algorithm with the prefix s and beam size b. We run the generated program on the public test cases to get its reward, and set it be the value of node (Line 16-17). This value backpropagated up in the tree so that the values of its ancestors are updated (Line 20).
Information sharing between Transformer and tree search. A keen reader may notice that if we follow the algorithm described above, there may be a lot of repeated computations. The key observation is that the Trans- former beam search algorithm also implicitly builds a free structure, which can be used by future iterations of tree search. In the rest of this section, we describe how we im- prove the algorithm's 's efficiency by sharing information in the Transformer beam search algorithm with tree search.
Iteration t
Iteration NI
<PD>
a.b
a.c
a,b
Consider the example in Fig. 3, which shows two iter- ations of PG-TD. In the evaluation step of the t-th iter- ation, the Transformer beam search algorithm implicitly builds a tree to find the most likely sequences within beam (Fig. 3 (left)). Because we only keep b partial pro- grams in the beam, is tree where only b nodes with the highest likelihood are expanded at each level (in the illus- tration, b- 2). Other nodes are dropped and no longer considered by the beam search algorithm. In the (1-st iteration, if the tree search algorithm selects "a, ", such state is already visited in the Transformer beam search in the 1-th iteration. When the algorithm needs find the top-k most likely next tokens, such information is already obtained in the t-th iteration and can be reused without recomputation. In our implementation. we cache the tree structure generated by the Transformer beam search algorithm. We call this implementation tree structure caching.
a.b
a
0,0
a,c
0.b.
a,ba
Figure Illustration for caching in the PG-TD algorithm. The tree search part is visualized in black color and the Transformer beam search part is in red color.
We can also cache the complete programs generated during the PG-TD evaluation step. In the evaluation step of the t-th iteration (Fig. 3), suppose the greedily-optimal sequence is "a, b= In the (t + 1)-st iteration, to generate a sequence starting with "a, b", the Transformer beam search algorithm will generate the same sequence "a, b= as before. (This is not necessarily true when the beam size is larger than 1. We will clarify this Sec. D.2.) To improve the efficiency of we cache the sequences that have been generated during the evaluation step. In the evaluation step of future iterations, PG-TD will check if the current state matches the prefix of any sequence that has been generated before, and use the generated sequence directly without calling the Transformer beam search function. We call this implementation sequence caching. We will empirically confirm the effectiveness of using these caching methods in the next section.
4 EMPIRICAL EVALUATION
In this section, we empirically examine the effectiveness and efficiency of our PG-TD algorithm by answering the following questions. Q1: Does PG-TD generate better programs than using the Transformer beam search algorithm and other competing baseline methods? Is our algorithm model-agnostic, showing improvement when applied to different Transformer models? Q2: Is the tree search algorithm an effective planning algorithm in PG-TD? Is it better than other planning or sampling-bast based algorithms? Q3: Is PG-TD efficient in terms of the number of times it runs the Transformer beam search algorithm and also the computation time it consumes? Q4: Are the caching methods effective in saving computation time? QS: Can we use the samples generated by PG-TD in its search process and finetune the Transformer to generate even better solutions?

---

Published as a conference paper at ICLR 2023
Sinci Accuracy
Inser.
comp.
CodeConietts
comp.
APES
5.50
Sampling+Filerigg
13,80
2.30
SMCG-TD
24.10
21.98
10.37
17.47
11.20
5.50
2.10
4.24
PG-TD (e 4)
26.70
24.91
12.89
24,05
13:10
6.10
3.10
4.85
APPS Neo
Beam Search
14.32
9.30
6.39
5.73
6.70
2.00
2.10
0.00
Sampling4Fiee Filtering
27.71
24.85
12.55
25.26
15.50
530
3,00
4.24
SMCG-TD
25.09
20.3
9.16
15/44
13.80
5:10
1.80
3,03
29.27
25,69
13.55
26.07
15.50
Table I: Results of PG-TD and other algorithms. The maximum number of Transformer generations for all algorithms is 256.
28
28
SMCG-TD
Pass Rate 196
Pass Rate (%6
Sampling + Filtcring
26
26
PG-TD
PG-TD (c 4)
PG-TD (c 6)
200
400
600
1000 Computation Time (sec.)
2000
3000
# Transforme Generations
Figure 4: Pass rates of PG-TD and baseline algorithms vs. the number of Transformer generations (left) and the computation time (middle) on the introductory problems in the APPS test dataset (1000 problems), using the APPS GPT-2 Transformer model.
Datasets and models. Recently, several competitive programming datasets have been made avail- able for benchmarking code generation algorithms. For each programming problem, they include the natural language program descriptions, sample solutions, and test cases. We evaluated PG-TD and the baseline algorithms on some popular benchmark datasets: APPS (Hendrycks et al., 2021) and CodeContests AlphaCode (Li et al., 2022). The APPS dataset does not specify public vs. pri- vate test cases. So we split all the test cases of a program evenly into two sets, where the first set is used as the public test cases for the algorithms to optimize the pass rate, and the second set used as the private test cases for evaluating the generated programs. For CodeContests, we use their public and generated test cases as our public test cases, and their private test cases as our private test cases. To show that PG-TD is model-agnostic and can be applied to different pre-trained Transformers and achieve a better performance, we use two popular pre-trained code-generation Transformers in the literature: GPT-2 and GPT-Neo finetuned on the APPS training dataset (Hendrycks et al., 2021).
Algorithms. We compare PG-TD with the following algorithms. Beam search only uses the Trans- former beam search algorithm to generate the whole program, without using the test cases. This is the method used in Hendrycks et al. (2021); Chen et al. (2021a). We use the beam size of 5 in the Transformer generation function, which is the same choice as in Hendrycks al. (2021).
We also implemented two baselines that use Transformer to sample programs and filter them using the test cases. Sampling Filtering (S+F) generates a set of programs using the Transformer sampling algorithm. Once the programs are all generated, it computes the pass rates of all the programs on the public test cases and returns the program with the highest pass rate. To avoid generating low-probability tokens that may fail the program completely, we use top-3 sampling, that is, the Transformer only samples the token that sampling temperature is 1. This method is similar to the algorithm in AlphaCode (Li et al., 2022), except that we did not perform clustering as we generate a much smaller number programs.
the top-3 most-likely tokens in each step. The
To determine the effectiveness of the tree search algorithm, we considered another baseline, Se- quential Monte-Carlo-Gunn Guided Transformer Decoding (SMCG-TD), that replaces the tree search algorithm component with sequential Monte-Carlo algorithm (Ellis et al., 2019). It is iterative algorithm that maintains a population of partial programs. determines the fitness a partial pro- gram using the Transformer in the same way the evaluation step in PG-TD. Partial programs with higher fitness scores are more likely to be selected in the next iteration. We leave more details about the baseline algorithms in Sec D.3.
For PG-TD, we set the maximum number of children of any node (k) to be 3, and the beam size (b) to be 1 by default. For the baseline methods, we sample most 512 programs in Sampling

---

Published as a conference paper at ICLR 2023
Method
Time (scc.)
Method
Pass Rate (%)
Strict Acc. (%)
With both caching methods
1312.27
Original
14.32
6.70
W/ only sequence caching
1430.63
FT w/ S+F
15.24
7.90
W/ only tree structure caching W/ caching
1899.17
FT w/ PG-TD
16.54
7.90
2206.38
Table 3: Performance on APPS introduc-
Table 2: Affects of using caching for PG-TD on the first 100 problems on the APPS introductory dataset, using the APPS GPT-2 (1.5B).
tory dataset with finetuned (FT) APPS GPT-2 (2.7B), using Beam Search for code genera- tion.
Filtering, and maintain set of 200 partial programs in each iteration for SMCG-TD. We will use these parameters for the following experiments unless noted otherwise. The experiments are conducted on the test of the APPS dataset (5000 problems) (Hendrycks al., 2021) and the test set of CodeContests (165 problems) (Li et al., 2022). We use the same metries as in Hendrycks et al. (2021), which are pass rates and strict accuracies on the private test cases. Specifically, the pass rate is the average percentage of the private test cases that the generated programs pass over all the problems. The strict accuracy is the percentage the problems where the generated programs pass all the private test cases.
Effectiveness of PG-TD. To make fair comparison, we evaluate the best programs found by the algorithms when they use the same number of Transformer generations. Specifically, the num- ber of Transformer generations is the number of function calls of BEAM_SEARCH in PG-TD and SMCG-TD (when cached sequences are found and used in sequence caching), and the number of sampling function calls in S+F.
As shown in Table 1, PG-TD consistently outperforms all the other baselines on all the datasets under the pass rate metric. As we optimize the pass rate, we expect our algorithm to outperform other baselines under strict accuracy well. This is almost true except that our algorithm is matched or outperformed by S+F on the APPS introductory set, mostly due to the fact that this set of programs less challenging. The gaps between different algorithms are also small on this set. Overall, these results confirm that our algorithm indeed generates programs that pass more test cases than the baseline methods (answering Q1). Specifically, S+F and SMCG-TD both use test cases to filter programs generated by the Transformer, while their performance is overall outperformed by PG- TD. So the tree rearch-based planning algorithm indeed powerful enough and can be effectively integrated with Transformer (answering Q2). We also report the performance of PG-TD and S+F under the n@k and pass@k metrics (Li et al., 2022) Sec. A (Table 6).
Efficiency of PG-TD. To show PG-TD's efficiency, we report the pass rates of the best programs found by the algorithms using the same computation time (Fig. 4 (right)). The experiments are run on the introductory problems in the APPS test set. Since one may implement these algorithms dif- ferently (possibly using parallelization), we also report the results using the number of Transformer generations budget (Fig. 4 (left)).
For PG-TD, we set = 3, = and vary the P-UCB exploration parameter c to be 2, 4, and 6. We that PG-TD with 4 has the best performance, while setting c = 2 tends to under-explore and setting c = 6 over-explores. With the same computational budget, S+F generates programs with lower pass rates (answering Q3). This confirms our observation that S+F, like the algorithm used in AlphaCode (Li et al., 2022), generates solutions only according to the Transformer's s generation probabilities, without taking their pass rates into consideration until all the programs are generated. PG-TD actively considers the pass rates of the generated programs during the generation process, which achieves better efficiency. On the other hand, AMCG-TD also uses the Transformer and public test cases to guide the generation process. However, due to its sampling nature, it cannot do multi-step lookahead search as in PG-TD, which results worse performance. We leave additional results of varying k, 6 for PG-TD and varying temperatures for S+F in Sec. A.
Effectiveness of caching. In terms of the design choices of tree structure caching and sequence caching, we performed ablative studies to verify their efficiencies. In Table 2, we compare versions of PG-TD with and without tree structure caching and sequence caching. As we expect, without sequence caching, the algorithm needs regenerate whole sequences, ending up consuming much more time. Without tree structure caching, the algorithm is slightly slower as it needs to call Trans- former get the most likely next tokens (answering Q4).
8

---

test cases Published as a conference paper at ICLR 2023
Finetuning transformer with PG-TD PG-TD-gemerated samples. Since are able to generate solutions with high pass rates using our PG-TD algorithm, can we use these generated solutions to finetune the code generation Transformers to further improve their performance? This may effectively solve the problem that high-quality human-written programs that can be used to train the code generation Transformers are scarcely available. Concretely, we first run PG-TD) on the training We then add the generated solutions with pass rates larger than 80% to the APPS sample solution set, and use the augmented solution set to finetune the GPT-2 (2.7B) model. With the finetuned Transformer model, we run beam search to generate solutions on the test set to see if it has better performance. We use the first 500 problems in the interview-level problems in APPS test set for validation and the introductory-level problems in APPS test set for testing. The results are reported Table 3. After finetuning with PG-TTD-generated solutions, we see improvement both pass rate and strict accuracy. More details are in Sec. D.4.
of APPS.
Optimizing other code generation objectives. Beyond the pass rate, we can make the algorithm versatile by using different reward functions. We consider two new objectives, code length penalty and comment encour- agement. As shown in Table 4, the code length penalty reward function makes the model generate more con- cise codes; the comment encourage- ment reward function encourages mod- els to generate codes with more comments. Both objectives still achieve reasonable pass rates on the public test cases. We provide qualitative examples and more details of the reward functions in Sec. C of the appendix.
Code length V
Comment number
Pass rate
Methods
Default Length Penalty Comment Encouragement
248.42
0.68
23.144
190.73
22.82
3.11
21.65
Performance of controllable code generation. Code length denotes the length of the generated code string and comment number denotes the number of code lines that contains comments.
Table 4:
Using automatically- generaaed test cases. The datasets we use in this paper both contain test cases that can be used to compute rewards for PG-TD. However, in reality, human-specified test cases are not always available. Recently, Chen et al. (2022) observe that Transformer pre- rraitrined on code generation can also generate useful test cases by adding an assert keyword at the end of the prompt. We follow the prompt design in Chen et al. (2022) to automatically generate test cases test cases. Empirically, we confirm
and run our PG-TD algorithm using the automatically-generaat that compared with beam search, PG-TD still has a higher strict accuracy by using automatically- generated test cases to verify the generated programs. We provide the details of the method and the experiments are in Sec. B.
DISCUSSION AND CONCLUSION
In summary, we proposed novel algorithm that uses the power of a pre-trained Transformer and a tree search algorithm inspired by Monte-Carlo tree search. We evaluate our algorithm and em- pirically show that our algorithm generates programs of higher quality compared with competing baseline algorithms in different settings. We also design model structure-specific caching mecha- nisms which contribute to saving computational expenses. We show that our algorithm versatile and can generate codes under objectives other than the pass rate without finetuning the Transformer. We hope our work can inspire more ways of incorporating planning algorithms into the Transformer generation process for code generation or other problems.
One limitation of our algorithm is its reliance on test cases, although we find that even small number of test cases can help find better solutions (Sec. A, Table 7). To address this limitation, we provided results to show that our algorithm can take advantage of automatically-geeenaa (Sec. B). Our algorithm is also more computationally expensive than the pure Transformer beam search algorithm since ours needs to run the beam search algorithm multiple times within the tree search algorithm. In the future, we consider improving the framework's 's performance by using value function to estimate the programs' pass rates, similar to the evaluation function in Silver et al. (2016) for mastering the game of Go. We can learn neural state representation for partial programs as in Chen et al. (2018). We also consider using parallel tree search algorithms (Chaslot et al., 2008) for the evaluation step to parallelize the computation for different states.
9

---

Published as a conference paper at ICLR 2023
This work was supported by the MIT-IBM Watson AI Lab, DARPA MCS, and gift funding from MERL. Cisco, and Amazon. We would also like to thank the computation support from AiMOS, server cluster for the IBM Research Al Hardware Center.
Acknowledgements.
REFERENCES
Wasi Ahmad, Saikat Chakraborty, Baishakhi Ray, and Kai-Wei Chang. Unified pre-training for pro- gram understanding and generation. In Proceedings the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pp- 2655-2668, Online, June 2021. Association for Computational Linguistics.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski. David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton. Program Synthesis with Large Language Models. arXiv:2108. 07732 [cs], August 2021. URL htttp::/aaxiv.org/ arXiv: 2108.07732. 2
abs/2108 077
Rudy Bunel, Matthew Hausknecht, Jacob Devlin, Rishabh Singh, and Pushmeet Kohli. Leveraging Grammar and Reinforcement Learning for Neural Program Synthesis. ICLR, May 2018. URL http:://axiii www.orrg///o.... abs 8005 04276. 6. arXiv: 1805.04276. 3.4
Antoine Chaffin, Vincent Claveau, and Ewa Kijak. Ppl-mets: Constrained textual generation through discriminator-guideed mets decoding. In Proceedings of the 2022 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2953-2961 2022. 3
Guillaume MJ-B Chaslot, Mark HM Winands, and HJVD Herik. Parallel monte-carlo tree search. In International Conference on Computers and Games, pp. 60-71. Springer, 2008. 9
Bei Chen, Fengji Zhang, Anh Nguyen, Daoguang Zan, Zeqi Lin, Jian- Gunn-Gang Lou, and Weizhu Chen. CodeT: Code Generation with Generated Tests, July 2022. URL abb/2200 10397. arXiv:2207.10397 [cs]. 9, 17
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Nosss,t-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Eval- uating Large Language Models Trained on Code. 1.X0::2107.03374 [cs]. July 2021a. URL http:// rxiv. org/ abs/ /2107. 03374. arXiv: 2107.03374. 1, 2, 7, 15, 17
Xinyun Chen, Chang Liu, and Dawn Song. Execution-guider neural program synthesis. In Interna- tional Conference on Learning Representations, 2018. 3,
Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis beyond domain-specific languages. Advances in Neural Information Processing Systems, 2021b.
Raj Dabre and Atsushi Fujita. Softmax tempering for training neural machine translation models. arXiv, 2020. 1,
Jacob Devlin, Ming-We Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 10.18653/v//N19-1 1423. URL https:
//acl ho oog..org//i.. 9111199144 423.
10

---

Published as a conference paper at ICLR 2023
Kevin Ellis, Catherine Wong, Maxwell Nye, Mathias Saylé-Meyer, Lucas Morales, Luke Hewitt, Luc Cary, Armando Solar-Lezama, and Joshua B Tenenbaum. Dreamcoder: Bootstrapping in- ductive program synthesis with wake-sleep library learning. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation, pp- 835-850, 2021.
Kevin M. Ellis, Maxwell Nye, Yewen Pu, Felix Sosa, Joshua Tenenbaum, and Armando Solar- Lezama. Write, execute, assess: Program synthesis with a repl. Neural Information Processing Systems, 2019. 3, 7,7
Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 889-898, 2018. 3,4
Zhangyin Feng. Daya Guo, Duyu Tang, Nan Duan, Xiaocheng Feng, Ming Gong, Linjun Shou, Bing Qin, Ting Liu, Daxin Jiang, et al. Codebert: pre-trained model for programming and natural languages. In Findings of the Association for Computational Linguistics: EMNLP 2020, pp. 1536-1547, 2020. 2
Sylvain Gelly and David Silver. Monte-carlo tree search and rapid action value estimation in com- puter go. Artificial Intelligence, 175(11):1856-1877 2011.
Alex Graves. Sequence transduction with recurrent neural networks. In International Conference on Machine Learning: Representation Learning Workshop, 2012. 3,4
Daya Guo, Shuo Ren, Shuai Lu, Zhangyin Feng, Duyu Tang, LIU Shujie, Long Zhou, Nan Duan, Alexey Svyatkovskiy, Shengyu Fu, et al. Graphcodebert: Pre-training code representations with data flow. In International Conference on Learning Representations, 2020. 2
Kavi Gupta, Peter Ebert Christensen, Xinyun Chen, and Dawn Song. Synthesize, execute and debug: Learning to repair for neural program synthesis. Advances in Neural Information Processing Systems, 2020.
Peter E Hart, Nils J Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 4(2):100-107, 1968. 3
Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, and Jacob Steinhardt. Measuring coding challenge competence with APPS. NeurlPS, 2021. 4. 7, 8, 14, 17
Aditya Kanade, Petros Maniatis, Gogul Balakrishnan, and Kensen Shi. Learning and evaluating contextual embedding of source code. In International Conference on Machine Learning, pp. 5110-5121. PMLR, 2020.
Levente Kocsis and Csaba Szepesvári. Bandit based monte-carlo planning. In European conference machine learning, pp. 282-293. Springer, 2006. 21
Hung Le, Yue Wang. Akhilesh Deepak Gotmare, Silvio Savarese, and Steven C. H. Hoi. CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning. arXiv, July 2022. doi: 10. 488550/arxiv.227..000 01780. arXiv:2207.01 7087080 [cs]. 3, 24
URL http:: wrtp:///////////////////////////
Rémi Leblond, Jean- Jean-Baptist Baptiste Alayrac, Laurent Sifre, Miruna Pislar, Yean-Baptiste Lespiau, Ioannis Antonoglou, Karen Simonyan, and Oriol Vinyals. Machine translation decoding beyond heam search. arXiv preprint 2rXiv:2/04 104.3360 2021.
Yujia Li, David Choi, Junyoung Chung. Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, and Agustin Dal Lago. Competition-LLeve Code Genera- tion with AlphaCode. arXiv preprint arXiv:2203.078// 2022. 2, 3, 4, 7, 8, 14, 17
Jinsuk Lim and Shin Yoo. Field report: Applying monte carlo tree search for program synthesis. In International Symposium Search Based Software Engineering, pp. 304-310. Springer, 2016.

---

Published as a conference paper at ICLR 2023
Ximing Lu, Sean Welleck, Peter West, Liwei Jiang, Jungo Kasai, Daniel Khashabi, Ronan Le Bras, Lianhui Qin, Youngjae Yu, Rowan Zellers, Noah A. Smith, and Yejin Choi. NeuroLogic A A'esque Decoding: Constrained Text Generation with Lookahead Hcuristics, December 2021.
URL
org/ abs /tsss/2211. 08726. arXiv:21 12.08226 [cs].
Nadia Matulewicz. Inductive program synthesis through using monte carlo tree search guided by heuristic-based loss function. 2022.
Glenford Myers. The of software testing. 1979. 3
Maxwell Nye, Armando Solar-Lezama, Josh Tenenbaum, and Brenden M Lake. Learning compo- sitional rules via neural program synthesis. Advances in Neural Information Processing Systems, 33:10832-10842, 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019. 2, 28
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to- text transformer. Journal of Machine Learning Research, 2020. URL http:/ /jmmm.m.orrrr..
papers/v21/20- 0700-074.00 07..0m1l.. 2
Baptiste Roziere, Marie-Anne Lachaux, Lowik Chanussot, and Guillaume Lample. Unsupervised translation of programming languages. Advances in Neural Information Processing Systems, 33,
2020. 2
Baptiste Roziere, Jie M Zhang, Francois Charton, Mark Harman, Gabriel Synnaeve, and Guillaume Lample. Leveraging automated unit tests for unsupervised code translation. In International Conference on Learning Representations, 2022. 1,3
Thomas Scialom, Paul-Alexis Dray, Jacopo Staiano, Sylvain Lamprier, and Benjamin Piwowarski. To beam or not to beam: That is question of cooperation for language gans, Advances in neural information processing systems, 14:26585-266999 2021. 3
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schritwiesser, Ioannis Antonoglou, Veda Panneershelvam, and Mare Lanctot. Mas- tering the game of Go with deep neural networks and tree search. Nature, 529(7587):484, 2016. Publisher: Nature Publishing Group. 3, 9
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou. Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel. Timothy Lillicrap. Karen Si- monyan, and Demis Hassabis. Mastering Chess and Shogi by Self-Play with a General Reinforce- ment Learning Algorithm, December 2017. URL http: arxiv org /bss///b///t 1172..1115. arXiv:1712.0 1815 [cs]. 21
01815.
Riley Simmons-Fdlerr Anders Miltner, and Sebastian Seung. Program synthesis through reinforce-
ment learning guided tree search. arXiv preprint arXiv:1806.0200 2018. 3
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. MIT Press, October 2018. ISBN 178-0-262-35270-2 4
Michele Tufano, Dawn Drain, Alexey Svyatkovskiy, Shao Kun Deng, and Neel Sundaresan. Unit test case generation with transformers and focal context. arXiv, 2020a.
Michele Tufano, Dawn Drain, Alexey Svyatkovskiy, and Neel Sundaresan. Generating accurate assert statements for unit test cases using pretrained transformers. arXiv, 2020b. 3
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 2017a. 1.3
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones. Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. arXiv:1706.03762 [cs], De- cember 201 7b. URL htt: httpttttpp///......//// www.wiwww......a....o...ooooooow abs///////00 03762. arXiv: 1706.03762. 2
12

---

Published as a conference paper at ICLR 2023
Yue Wang, Weishi Wang, Shafiq Joty, and Steven C.H. Hoi. CodeTS: Identifier-aware Unified Pre-trained Encoder-Decoder Models for Code Understanding and Generation. In Proceedings of the 2021 Conference Empirical Methods in Natural Language Processing, pp. 8696-8708, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 2021.
wttp:////2221...pmmmmmmm... 685. 1,2
URL http:////aaaaaallloooooooooooooo ogy org/
Thomas Wolf, Lysandre Debut. Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierrie Cistac, Tim Rault, Rémi Louf, and Morgan Funtowicz. Huggingface" transformers: State- of-the-art natural language processing. arXiv preprint arXiv: /91000371, 2019. 17
Yifan Xu, Lu Dai, Udaikaran Singh, Kening Zhang, and Zhuowen Tu. Neural program synthesis by self-learning. arXiv preprint arXiv:1910.05860 2019a.
Yifan Xu, Lu Dai, Udaikaran Singh, Kening Zhang, and Zhuowen Tu. Neural Program Synthesis By Self-Learning. October 2019b. doi: 191000008....0.. URL https:///rrss
exg/abs/1es 1910. 05865v1. 3
13

---

Published as a conference paper at ICLR 2023
APPENDIX
In this appendix, we supplement the main paper by providing more thorough empirical evaluations to back up our claims and more detailed descriptions of the algorithms help readers better understand our paper.
This appendix organized as follows.
In Sec. A, we provide more comprehensive results of our algorithm and the baseline algorithms. We also include the license information of the datasets we use.
In Sec. B, we consider the scenario where test cases are not provided. We evaluate our PG-TD algorithm using autoomatically-generain test cases.
In Sec. C, we provide empirical evidence for our claims in the discussion section that our algorithm versatile and can used to optimize different code generation objectives other than the pass rate. We consider the objectives of generating concise codes and generating codes with more comments.
In Sec. D, we provide more details on the components in PG-TD as well as the baseline algo- rithms.
In Sec. E, we illustrate more examples of the codes generated by our PG-TD algorithm and the baseline algorithms.
In Sec. F, we discuss more on the advantages and the potential negative social impacts of our algorithm.
A EMPIRICAL EVALUATION
We reported the performance of our algorithm and other baselines on the whole APPS dataset (Hendrycks et al., 2021) and CodeContests (Li et al., 2022). The APPS test dataset contains coding problems three levels: introductory (1000 problems), interview (3000 problems), and competition (1000 problems).
In addition to results in Table 1 in the main paper, Table 5 shows the results where the budget of the number of Transformer generations is 512. As we expect, with a larger number of generated programs, the gap between PG-TD and the rest is smaller. Sampling + Filtering has a better strict accuracy than our algorithm on some subsets of datasets. We also experimented with other parameter settings PG-TD. In Fig. 5 and 6, we vary the beam size (b 1,3, 5) under 2, In Fig. vary the maximum number of children of any node (k = 2, 4).
Although one may expect expanding the width of the tree (increasing k) can help find better pro- grams, considering generating less-likely tokens suggested by the Transformer may not contribute to finding better programs. In fact, it wastes the budget of the number of Transformer generations and the computation time by considering less likely tokens. On the other hand, increasing the beam size (b) does help improve the pass rate when the number of rollouts (corresponding to the number of Transformer generations) is small. However, increasing the beam size under a larger number of rollouts does not always increase the pass rate (Fig. 7), which may be caused by an overfitting issue. Increasing also costs more computation time even though sequence caching is used.
In Fig. 8, we use different temperatures in the Sampling + Filtering algorithm (t = 0.6, 0.8, 1), confirming that our choice of t in the experiments in the main paper does have the best perfor- mance.
n@k and pass@k k metrics. We also report the performance of PG-TD and S+F under the n@k and passãk metrics (Li et al., 2022). Given problem, each algorithm is asked to generate k programs and submit n programs for evaluation (they submit the n programs with the highest reward on the public test cases). n@k is the proportion of the problems where any of the submitted programs passes all the private test cases. pass@k is the proportion of the problems where any of the generated programs passes all the private test cases (effectively k@k). For PG-TD, the k samples the complete programs found in the first k rollouts. The results are reported in Table 6, evaluated on the APPS introductory problems. We observe that for smaller k values, PG-TD finds programs
14

---

Published as a conference paper at ICLR 2023
Sinci Accuracy
comp.
comp.
Sumpling+Filerrng
25.75
SMCG-TD
24.40
22.18
10.60
17.72
11.80
5.87
2.50
4.24
PG-TD (e 4)
28.27
26.13
13.74
25.91
14.40
0.63
4,00
APPS Neo
Beam Search
14.32
9.30
6.39
5.73
6,70
2.00
2,10
0.00
Sampling4itherigg Filtering
90.23
26.58
13.12
26.57
16,70
7.40
3.10
4.24
SMCG-TD
24.77
20.50
9.19
15442
14.00
5.30
1,70
242
30.36
26.89
13.28
27.10
16.70
Table 5: Results of PG-TD and other algorithms. The maximum number of Transformer generations for all algorithms is 512.
1@
1@
PG-TD (c 4)
8.1
10.2
11.6
S+F
6.5
9.7
11.4
5@
5@50
5@ 100
PG-TD (c 4)
12.5
14.9
16.2
S+F
11.9
14.6
18.2
pass 10
pass @ 50
pass @ 100
PG-TD (c
14.7
23.7
27.3
S+F
14.7
23.2
28.3
Table 6: n@k and pass@k results of PG-TD and Sampling, evaluated on the APPS introductory problems.
Pass Rate
Strict Accuracy
O public test cases (the beam scarch baseline) public test public test cases 5 public test cases
13.61
4.2
23.27
7.4
33.02
13.4
32.99
14.2
Half of all test cases (10.44 public test cases average)
37.71
14.8
Table The performance PG-TD using different numbers of public cases, evaluated on the first 500 APPS introductory problems.
that are more likely to pass all the private test cases. When more programs are generated (k) or more programs submitted for evaluation (n), the performance of PG-TD is matched outperformed by S+F.
Using different numbers of public test cases. On the APPS dataset, we split all the test cases evenly into public test cases and private test cases. To investigate how many test cases are enough for PG-TD to achieve satisfactory performance, we consider using 3, and public test cases and using the rest of the test cases as private test cases. The results reported in Table 7, evaluated on the first 500 APPS introductory problems. We observe that even with small number of public test cases, PG-TD is able find programs with higher pass rate and strict accuracy than beam search. This also suggests that the Transformer can serve as regularizer in code generation. With a small number of public test cases, PG-TD still generates programs similar to human-written programs without overfitting the pubic test cases.
Results Codex. To evaluate the effectiveness of algorithms on more powerful code- generation Transformers, we run both beam search and PG-TD algorithms using Codex (Chen al., 2021a) as the backbone. We follow the "best practices" provided in the Codex guide to put the problem description into block comment in the beginning, and ask Codex to complete the rest. The prompt looks like the following.
Python
(probles em description)
15

---

St Published as a conference paper at ICLR 2023
28
28
9MC5-10
Pans (50
PG-TD Rotrnn
Pass Bute
24
24
100 200
300
400 500
1000
2000
3000
x Transformer Gererations
Computation Time (sec.)
Figure 5: Results of PG-TD (c 2) on the APPS introductory dataset, using the APPS GPT-2 Transformer model.
Ded Ratee 1991
Samping Ricating
26
New tatt
P6-TD 0 -3 FG-TD 55)
24
22
22
200
400
1000
3000 Compunelion Time [sec.)
2000
Thwnffommr Gererations
Figure 6: Results of PG-TD (c = 4) on the APPS introductory dataset, using the APPS GPT-2 Transformer model.
28
28
76-10 R x -
Pess Bale 00
Paus Bate
22
22
2000
Transformer
Com porpptaaiion Time
Figure 7: Results of PG-TD (c 4) with different &o the APPS introductory dataset, using the APPS GPT-2 Transformer model..
Pan Bane (N1
26
26
Paok Batte
F(tttll
24
24
22
100
200
300
400
500
500
2000 Computation Time (sec.)
1000
1500
TrTrnnotomeer Gonerations
Figure 8: Results of Sampling + Filtering with different temperatures on the APPS introductory dataset, using the APPS GPT-2 Transformer model.
As expect, PG-TD helps Codex generate better codes (Table 8). This further validates claim that our method is model-agnostic and can help a pre-trained code generation Transformer generate better codes. Due the limits of the OpenAl APIs, only able to the algorithms on subset of data. We also show concrete example where PG-TD outperforms beam search in the following sections in Fig. 16.
Using sampling-basee based tree search. We could use a sampling approach in evaluation step of PG-TD instead of using beam search, which would be more faithful implementation of MCTS. We experimented with the two versions PG-TD and find that overall using beam search indeed has better performance (Table 9). The gap could be caused by that only do sampling once to evaluate any node. We could estimate the value of nodes more accurately by sampling multiple times. However, it would make algorithm more computationally expensive.
16

---

Published as a conference paper at ICLR 2023
Model
Decoding Algorithm
Pass Rate (%)
Strict Accuracy (%)
APPS GPT-2 (1.5B)
Beam Search
11.65
APPS GPT-2 (1.5B)
PG-TD
36.42
Codex (code-davinci-000
Beam Search PG-TD
33 33,26
14
Codex code-davinci-002:
59.14
31
Comparison between using different Transformers (APPS and Codex) the backbone for code generation, evaluated on the first 100 APPS introductory problems.
Table 8:
Decoding Algorithm
Pass Rate (%)
Strict Accuracy (%)
# Transformer Generations
128
256
512
128
256
512
PG-TD using Beam Search
36.19
37.71
38.57
14.4
14.8
15.6
PG-TD using Sampling
33.95
35.67
37.87
13.8
15.4
17.2
Table 9: PG-TD using beam search vs, sampling in the evaluation step, evaluated on the first 500
APPS introductory problems.
Compilation Error (%)
Runtime Error (%)
APPS GPT-2 (1.5B), Beam Search
5.58
32.95
APPS GPT-2 (1.5B), Sampling Filtering
4.68
27.38
APPS GPT-2 (1.5B), PG-TD
1.93
19.5
Table 10: Failure case analysis. The percentages are averaged over the APPS introductory dataset.
Strict Accuracy (%)
Beam Search
26.82
PG-TD (using auto-generated test cases)
74.53
Table 11:
Comparison between beam search and TG-PD using automatically-geeeenei test cases, using the Codex model and evaluated on the HumanEval dataset.
Failure mode analysis. In Table 10, we report the average percentages of test cases where com- pilation errors and runtime errors occur. We use the error analysis utility provided in the original APPS codebase. As expect, PG-TD executes the generated programs and finds the ones with higher pass rates, so it dramatically reduces the percentage of both types of errors.
Assets and licenses. Our experiments are on machines with two Intel(R) Xeon(R) Gold 6258R CPUs (@ 2.70GHz), and one V100-SXM2 GPU. The APPS dataset (Hendrycks et al., 2021) is released under MIT License 2. The CodeContests dataset (Li et al., 2022) is released under Apache License 2.0 3, The Huggingface Transformer (Wolf al., 2019) is released under Apache License 2.0
B USING AUTOMATICALLL-GENER..N TEST CASES
When human-specified test generated test cases. We first follow Chen et al. (2022) to use the Codex model and evaluate the code generation algorithms on the HumanFyal dataset (Chen al., 2021a). We adopt the same procedure as Chen et al. (2022) to automatically generate test cases. Concretely, we use prompt that includes the natural language description of the problem and an assert statement at the end of the prompt. We also removed the example input, output pairs in the problem description to avoid
not available, our algorithm can still rely automatically-
http:////..thtttccoc////eeaaaaea hhttps://g thub. http://////.t....c..o////nnnnnno http://twwwt.c.ooooooooooeeeeeen /tt.
CENSS
CENSK
17

---

Published as a conference paper at ICLR 2023
the Transformer from directly copying the test cases in the problem description. The following is an example the prompt for test case generation.
from typing import List
def Mas_close_clement (numbers: List[floate threshold: float) -> bool: Check if in given 1ist of numbers, are any two numbers closer to each other than given thresheld. =**
pass
check the correctness of www.ccsss_eeeeemttee
def check (candidate): assert candidate
Does Codex generate correct test cases? In the HumanEval dataset, the generated solutions are evaluated by the test script by calling the check function. So we use the strict accuracy (pass@) metric that counts the percentage of the problems where the generated solutions pass the check function. To evaluate the quality of the autommatically-generatee test cases, we compute the strict accuracies of the sample solutions (correct solutions written by human programmers) on these test cases. Clearly, the strict accuracies of the sample solutions the ground-truth test cases should be 100%. We evaluate the sample solutions on the automatically -enerrted test cases and find the corresponding striet accuracy is 72.56%, which confirms that the automaticalll-generatt test cases are mostly correct.
Can PG-TD take advantage of the automatically-genene strict accuracy of the generated programs on subset of HumanEval problems in Table 11. We confirm that the strict accuracy of PG-TD is higher as uses high-quality automaticalll--geeraatdd test cases to verify the generated programs.
test cases? We report the average
C PLANNING FOR OTHER CODE GENERATION OBJECTIVES
Besides the default reward function that optimizes the pass rate, we show the versatility of the pro- posed algorithm by setting two new objectives, code length penalty and comment encouragement.
Code length penalty. We make the planner generate more concise codes by using the following reward function
Riength
A x -le/t
(1)
where p is the average pass rate on the public test case set and l. is the length of the code string. and are hyperparameters that control the weight of the code length penalty and are 0.1 and 20, respectively. As shown in Figure 9, the generated solution becomes more concise and the code string length decreases from 187 78 while still passing all the test cases.
Comment encouragement. We can also generate solutions with more comments. We achieve this goal by using the following reward function
N
Reomment p Å e -laft left + A2 X min(1,
(2)
N
where Nom is the number of "#" in the code string and X2 set to 0.2, controlling the weights of the comment lines. If simply add X₂ x Nom the comment encouragement term, the planner would generate code with repeated meaningless comment lines to get more rewards. To handle this problem, we add the code length penalty term, ê -let, We also upper-bound the rewards for the number of comment lines by setting the comment encouragement term be Xg X min(1, Nem.) ), where Nmax in the example shown in Figure 10. As shown in Figure 10, we can generate solutions with more comment lines with this designed reward function.
18

---

Published as a conference paper at ICLR 2023
Problem Statement
Takahashi loves palindromes. Non-palindromic strings are unacceptable to him. Each time he hugs string, he can change one its characters any character his choice. Given is string S. Find the minimum number of hugs needed to make palindromic. Constraints
(1). S is a string consisting of lowercase English letters.
(2). The length of is between and 100 (inclusive).
Input
Input given from Standard Input in the following Sormat:...
Output
Print the minimum number of hugs needed to make palindromic.
Sample Test Input 1
redcoder
Sample Test Output 1
For example, can change the fourth character to and get palindrome redooder.
cook up a palindrome s=input ()
s input () ans=0
n=len (s)
for in range (len (s) //2) if s[i] == i]: anst=]
c=0
if n#2==0:
print (ans)
for i in range (2):
if (-
c+-1
print (c)
10
else:
n
for i in range (n/ /2)
if [i] [eil!=sliin [n-i-1] c+=1
print (c)
Default reward function (Pass Rate: 1.00, code string length: 187).
Reward function with code length penalty (Pass Rate: 1.00, code string length: 78).
Figure 9: A code generation example with code length penalty. The problem description is shown on the top. The programs generated by default PG-TD and PG-TD with code length penalty shown at the bottom. We can generate more concise code by code length penalty reward.
19

---

Published as a conference paper at ICLR 2023
Problem Statement
A + B is often used as an example of the easiest problem possible to show some contest platform. However, some scientists have observed that sometimes this problem is not so easy get accepted. Want to try? Input
The input contains two integers a and b (0 S a, S 10²), separated by a single space.
Output
Output the sum of the given integers.
Sample Test Input 1
5 14
Sample Test Output 1
19
Sample Test Input 2
381 492
Sample Test Output 2 873
a, map (int, input () spplit ))) print (a + b)
# /usrbin/env python3
# coding =utf-8 #py3 runs on py2
a, b-mpap (int, input () sp)i(())
print (a+b)
Default reward function.
Reward function with comment encouragement
(code with comments).
Figure 10: A code generation example with code comment encouragement. We can generate codes with more comments by using the code comment encouragement reward.
20

---

Published as a conference paper at ICLR 2023
D ALGORITHM DETAILS
D.I PG-TD
In this section, we describe our PG-TD algorithm in more details. We show our framework in Fig. 12 again for convenience. The algorithm maintains a search tree structure where nodes correspond to states and edges correspond actions. For each node s, maintains the number of times that has been visited, denoted by s.visits. For state, action pair, s, (a stat, action pair corresponds to an edge in the tree), the algorithm maintains a function Q(s. a), which the maximum reward obtained by starting in and taking action a. Conventionally. (ss, a) is the average return. Since our code generation problem is deterministic (that is, there are no stochastic transitions), we find the performance is better if we keep track of the pass rate of the best program generated from node.
The algorithm first selects node for expansion. It starts from the root node and selects subtrees recursively until it reaches a node that has not been expanded before (with children). A com- mon node selection criterion is upper confidence bound (UCB) (Kocsis & Szepesvári, 2006), which leverages between exploiting known-to-be-ge good states and exploring less- states, We adapt the P-UCB heuristic used in Silver et al. (2017), which is a variation of UCB, as follows,
Vlog(s.visits)
P-UCB(s, Q(s, a) B(s) P Prrausformer (als)
+ 8 seisits
where s' is the state (deterministically) reached by taking action a in s; Prransformer (a)))) is the probability that token a is the next token given the partial program s, which is provided by a Trans- former model. 3 is the weight for exploration, which depends on the number of visits of s, defined as
s.visits Chase + 1)++ +
(ss) log(
Cbase
The P_UCB_SELECT function in Alg. returns the action that has the highest P-UCB value,
P_UCB SELECT c) = arg max P-UCB(ss a).
(5)
Note that parameterizes the exploration weight in Eq. 4. We omit it in Eq. 5 for notational con- ciseness.
Intuitively, (ICBBSELLCC((sc) would return action more often if 1) Q(s, a) is large, which means that has found high-quality programs starting from 8, at or 2) Prannsformer (a]s) is large, which means that the Transformer believes that is highly likely next token; 3) s.visits is large but s' s'visits is small, which means s' is under-explored. We did a hyperparameter search and set Cbase = 10, and used e 2, 4, 6 in our experiments.
In the following steps, the algorithm expands the node by adding its children to the tree and evaluates the node using BEAM SEARCH, ward of the completed program r is backpropagated to its parents recursively until reaches the root update their Q values. For all state, action pairs, a", along the trajectory in the tree from the root to the selected node, Q(s", a²) is updated accordingly using the new value: Qss", ad) t max(Q(s", "osss,a"))) r)).
described in the main paper. Finally, the
Comparison with the standard MCTS algorithm. It is worth noting that our implementation of MCTS PG-TD different from the standard MCTS algorithm. If we use the standard MCTS algorithm, it is computationally intractable to search the large space of possible programs to find high-quality ones.
For comparison, we visualize the standard MCTS algorithm in Fig. 11. In the expansion step, MCTS selects the next available action in the action set, and adds the state that can be reached by following the action. In the example, action "b" is taken, and the new state that is added to the tree is "ab". In the evaluation step, MCTS evaluates the new state by executing a random policy from the new state and computes the value of the policy.
is impractical to apply the standard MCTS algorithm to domains with large state spaces or large action spaces, cannot afford to try all possible actions in the expansion step. A random policy in the evaluation step also has a high variance in estimating the value of the new state. In PG-TD, use pre-trained Transformer resolve these limitations of the standard MCTS algorithm.
21

---

Published as a conference paper at ICLR 2023
Selection
Expansion
Evaluation
Backpropagation
<PD>
<PD>
<PD>
<PD>
b
b
b
b
a
a
a
88
aB
ba
aa
ba
aa
Select the next action in the action set
Generate random policy and evaluate
Figure 11: The standard Monte-Carlo tree search algorithm, without using a pre-trained Trans- former. <PD> stands for problem description.
Selection
Expansion
Evaluation
Backpropagation
<PD>
<PD>
<PD>
<PD>
X
X
X
X
X,
a,
X,
X
a
suggests next tokens
a/An
Evaluate
Transformer
public
generates complete programs using beam search
Figure 12: Our proposed framework, PG-TD, repeated here for convenience.
a,b a,c
a,b
a,b,
a,b=
a,c= a,c,
a,b,
a,b=
Beam
Beam 2
Beam
Beam
Figure 13: Illustration of example where sequence caching may return the correct sequence when b > 1.
D.2 EFFICIENT IMPLEMENTATION BY INFORMATION SHARING
As discussed in the main paper, there may be calls the Transformer generation functions the perform redundant computations. We have described tree structure caching the main paper. We provide details sequence caching in this section.
22

---

Published as a conference paper at ICLR 2023
Algorithm 2 Sampling + Filtering.
Require: sample_nuume nw:m: the number of samples to generate using Transformer.
1: for + 1, 1,2, 3.2,.... sample-nums do
2: Ps - FRANSFORMER.SAMPLEEI
5)
3: rewards[i] < GET REWRRDD()
4: end for
5: return p; with the largest rewards
Sequence caching. The goal of sequence caching is to reduce the computational cost of the eval- uation step of PG-TD. In the evaluation step, the Transformer beam search algorithm is called to generate complete program from the current node. These complete programs and their rewards both cached and used in a future iteration.
However, when b > using cached sequences in the previous iterations may be suboptimal. Let's consider the example in Figure 13, where = 2. Starting from "a, (the left figure), suppose the beam search algorithm returns sequence by searching two beams with the prefixes of "a, b² and "a, c", respectively (the bold lines). In the next iteration (the right figure), to evaluate "a, b" using the beam search algorithm, we may not use the sequence found the previous iteration, which only searches the subtree of aa, b" with one beam. In other words, we may find sequence with a higher likelihood by re-running the beam search algorithm with b starting from "a,b".
Although using sequence caching with b > 1 may underestimate the values of programs in the evaluation step in PG-TD, we still use sequence caching for all choices of b for the merits of com- putational saving.
D.3 BASELINE ALGORITIMS
Sampling + Filtering. We have described the sampling algorithm in the main paper. docode is shown in Alg. 2.
The pseu-
Sequential Monte-Carlo-Guider Transformer Decoding. Sequential Monte-Carlo- Transformer Decoding (SMCG-TD) is another algorithm that we compare PG-TD with, which also uses Transformer to evaluate partial programs. The pseudocode is shown in Alg. One can think of it from view of genetic algorithms. SMCG-TD maintains a group of partial programs, called generation, and updates them iteratively. For each partial program in the generation, we sample the next token using the Transformer and append it to the partial program, and put it in We then evaluate each partial program the new generation in the same way as PG-TD: We use the Transformer generate complete program from the partial program and compute its reward (pass rate). We then use these rewards as the partial programs' fitness scores. At the end of an iteration, re-sample from the partial programs according to their fitness. So partial programs with higher fitness scores are more likely to remain in the generation. In the end, the algorithm outputs the complete program finds that has the largest reward.
generation.
One may notice that SMCG-TD resembles PG-TD on higher level. It also uses the Transformer and the public test cases decide which next token to generate in the code generation process. The key difference is that SMCG-TD TD uses sampling-basee approach and PG-TD uses tree search algorithm. in our experiments, the performance of SMCG-TD worse than PG-TD.
It is worth noting that sequence caching is also implemented for SMCG-TD when it calls REAR_S So even if multiple partial programs in a generation can the same, the algo- rithm would not call the Transformer to generate complete program from the same prefix more than once, which effectively reduces the number of times that it needs to call the Transformer and saves the computation time.
D.4 FINETUNING
In our experiments, we use the samples generated by the PG-TD to further finetune the Transformer model. We then generate solutions using the finetuned Transformer model by running beam search,
23

---

Published as a conference paper at ICLR 2023
Algorithm 3 Sequential Monte Carlo-Guided Transformer Decoding (SMCG-TD).
Require: pop_size: the population size
1: generation <- [(PD)] pop_size e initialize the population with the problem description
2: filness + [1.0/ pop_size) pop_size # uniform distribution
3: ++0
4: complete programs NEW LIST()
5: while ipopulation > 0 and & < mar steps do
6: new.generation <- NEW_LISTE
7: new-fitness + NEW LIST()
8:
i LEN (genratioo) do 4 sampled from generation, where generationlj] is selected with the probability of fitness]j]
9:
10:
- sampled from formerformer-Is)
CONCAT(s,,))
12:
if 8 [-1] = < endoftext |> then
13:
# this sample finishes generation, save to output
14:
add s' to complete-prograns.
15:
else
16:
add s' to new-generation
17:
# BEAM SEARCH uses sequence cache as in PG-TD
18:
P BEAM SEARCH(s',
19:
r + GET _REWARD(()
add new-fitness
21:
end if end for
22:
23: normalize new ffenffstess
24:
generation + nEW sew_eeeeeriion
25:
filness +- new filness
26:
L ++1+1
27:
end while
28:
return the program in complete programs that has the largest reward
and observe improvement in both pass rate and strict accuracy compared with running beam search using the original model.
One challenge our design that solutions that PG-TD generates have different pass rates. We could use only solutions with a pass rate of 100%. However, that largely limits the number of samples we can use for finetuning. We instead collect all solutions that have pass rates larger than 80%. To make the Transformer aware of the different rates of training samples, we use the loss function similar to Le al. (2022),
((0) = Ew Erw>pulrWWWV Vr log po WWPD)),
where W denotes complete program: r(W) is the pass rate of the program; PD is the program de- scription. In future work, we will investigate different loss functions and compare their effectiveness in improving the Transformer model's performance.
E ILLUSTRATIVE EXAMPLES
In this section, provide methods, shown in Figure 14 and Figure 15. Our PG-TD algorithm generates solutions that have higher pass rates than the baseline methods. We also show concrete example where PG-TD out- performs beam search in Fig. 16 using Codex as the backbone Transformer.
qualitative examples generated by our PG-TD and the baseline
For step-by-step illustration of PG-TD, please visit our project website https:// codeaimets. gitggttuuu aithub... io.
24

---

Published as a conference paper at ICLR 2023
Problem Statement
Takahashi is meeting up with Aoki. They have planned meet place that is D meters away from Takahashi's ' house in I minutes from now. Takahashi will leave his house now and go straight the place speed meters per minute. Will he arrive in time?
Constraints
(1). D < 10000
(2). < T 10000
(3). 1 S 10000
(4). All values in inpul are integers.
Input
Input is given from Standard Input in the following format: D T S Output
If Takahashi will reach the place in time, print Yes; otherwise, print No.
Sample Test Input
1000 15 80
Sample Test Output Yes
takes 12.5 minutes to go 1000 meters to the place at speed of 80 meters per minute. They have planned to meet in 15 minutes he will arrive in time.
cook your dish here
cook your dish here s, (map (int,
s-map (int,
try:
input spli:
a=list (map (int,
4 input .piit ())
tl=n
+ input () spllpit
if t1 <= t: print ("Yes") else: print ("No")
11)
if(d--t):
if (a[2)-11)
if (s+ (s+1))/---))
print ("Yes")
print ("Yes")
else:
else:
print ("No")
print ("No")
except:
else:
pass
if (a* (s+1) /s>t)
10
print ("Yes")
II
else:
12
print ("No")
PG-TD (Pass Rate: 1.00).
Beam Search (Pass Rate: 0.125).
Sampling + Filtering (Pass Rate: 0.625).
Figure 14: A code generation example for competitive programming. The problem description is shown on the top. The programs generated by baseline algorithms and our PG-TD algorithm are shown the bottom.
25

---

Published as a conference paper at ICLR 2023
Problem Statement
Given is an integer How many times is the area of circle of radius larger than the area of circle of radius 1? It be proved that the answer is always an integer under the constraints given.
Constraints
(I). S S 100.
(2). All values in input are integers.
Input
Input given from Standard Input in the following format: r.
Output
Print the area of a circle of radius v, divided by the area of a circle of radius 1, as an integer.
Sample Test Input
Sample Test Output
The area of circle of radius times larger than the area of a circle of radius 1. Note that output must be an integer for example, 4.0 will not be accepted.
cook your dish here r=int (input )))
cook your dish here
# cook your dish here
r r-int (input ())
r=int { input ()
if (x>=2xx)
if (r>==2xz)
print (2+ (r* -
print (r*r*2)
print (r#r*2)
else:
else:
print (0)
print (0)
Beam Search (Pass Rate: 0.00).
Sampling Filtering (Pass Rate:
PG-TD (Pass Rate: 1.00).
0.00).
Figure 15: code generation example for competitive programming. The problem description is shown the top. The programs generated by bascline algorithms and our PG-TD algorithm are shown at the bottom.
26

---

Published as a conference paper at ICLR 2023
Problem Statement
You are given an array consisting of n. integer numbers.
Let instability of the array be the following value: max min i=1
You have to remove exactly one element from this array to minimize instability of the resulting
(n eleemmesnss array. Your task is to calculate the minimum possible instability.
Input
The first line of the input contains one integer (2 S n S 10") - the number of elements in the
array
The second line of the input contains n integers a1 02 On (1 S a; S 10 - elements of the
array
Output
Print one integer the minimum possible instability of the array you have remove exactly one element from the array a.
Examples
Input
4
1337
Output
2
Input
2
100000
Output
Note
In the first example you can remove 7 then instability of the remaining array will be 1 1 2. In the second example you can remove either or 100000 then instability of the remaining array will be 100000 100000 = 0 and - 1 = 0 correspondingly.
n int (input )
a 1ist (map (int, input () spplit ())
print (max (a) min (a) 1)
Beam search (Pass Rate 0.00).
def main():
int (input ())
1ist (map (int, input () spplit() )))
a.sort ()
print (a|nn] - a(1), a [ - a (01))
if
mair
main ()
PG-TD (Pass Rate 1.00).
Figure 16: A code generation example using Codex (code-davinci-002) Beam search finds an incorrect solution while PG-TD finds correct solution. This result further confirms that PG-TD is model-agnostic and could be effective with different backbone Transformers (GPT-2 (1.5B), GPT- Neo (2.7B) and Codex).
27

---

Published as a conference paper at ICLR 2023
F MORE DISCUSSIONS
Potential benefits. We have shown that our framework can not only help achieve higher pass rates compared with using only the Transformer generation process, but also generate codes that opti- mize different objectives. These are done without fine-tuning the pre-trained Transformer models. Our framework makes pre-trained Transformer more versatile and adapts designing different reward functions used by the planning algorithm. If we re-train or fine-tune Transformers like GPT-2 used in this paper, we will need to train billions of parameters (Radford et al., 2019). Our approach potentially saves the computational cost and energy consumption that is needed train or fine-tune Transformers for different tasks.
different tasks by
Potential negative social impacts. Automatic code generation makes it easier for anyone gen- erate programs that meet a specification. Our hope is that this development will relieve the burden of software engineers and enable Al-assisted programming or even end-to-end automatic program- ming. However, it may also make it easier to develop malware. Anyone that can specify malicious goal in natural language can use an automatic code generation tool to generate codes that meet the goal. When automatic code generation techniques become more mature, we may need a separate module that screens the natural language description and rejects the code generation requests that can lead to harmful codes.
28


Reviews:
[{"review_id": "oDh-L0Lo_LA", "reviewer": "Reviewer_fkta", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This work proposes PG-TD, which integrates Transformer with MCTS for code generation. PG-TD works by generating a tree of possible programs, first generating high likelihood sequences and then evaluating them on test problems. It then selects the best and searches for more sequences until it completes the program. The authors also propose to cache completed queries as many are repeated during search. The authors demonstrate results of pass rate and computation time on an established coding challenge, demonstrating PG-TD outperforms GPT-3 baselines.\n", "summary_of_the_review": "The authors propose a novel approach that is interesting and potentially useful, and the paper is clearly written. However, they only compare and leverage GPT-2 for coding problems, which leads to low performance across the board and making it unclear if they are solving a real problem -- they should run the examples with Codex and also with prompting. The algorithm is also quite slow, the code produced is overly complicated, and the metric of pass rate only is lacking (they should also report whether the code was correct). I do not believe caching to be a major contribution, but is presented as such and takes up a page. The paper may be improved by further demonstrating performance on other problems (perhaps not code) that have a reward structure.", "strength_and_weaknesses": "Strengths:\n* The paper is clearly written and described.\n* The website has a useful visualization and providing the code is appreciated (though there is a lot of code provided and maybe a subset would be preferred).\n* The paper approaches an important research area, improving generation for LLMs and proposes a reasonable approach for it that does not require any additional training and minimal code complexity.\n* The results demonstrate it outperforms baselines in performance.\n* The results showing varied rewards is interesting and useful and a nice benefit of the approach.\n\nWeaknesses\n* The results are all performed with GPT-2 and seem to be a relatively weak LLM for coding. For coding in particular, where any small mistake may be catastrophic (in that it doesn't compile or run), I am unable to fully contextualize the benefits with a weak model (and one not trained on code). Codex for example would be a much stronger approach a minimize issues like the second sentence in the abstract “ Although the programs they generate achieve high token-matching-based scores, they often fail to compile or generate incorrect outputs.” Indeed, I ran two of the proposed tasks (Figure 1) and the circle problem on the website and *zero shot* Codex solved both. Furthermore it solved both with much more reasonable code than the PG-TD code and did it in under 5 seconds. \n* The authors should report cost of running the algorithm.\n* The overall success rate, particularly based on the examples shown in the paper, seems quite low for all methods.\n* I believe all the results in this paper are zero-shot (or else I missed this), but presumably much of the issues would be solved with simple prompting with a few examples. The authors should compare to this.\n* The examples of code shown are odd, in that it tends to generate unnecessarily complicated code for quite simple problems, stating `print(2*(r**2)-r**2)` instead of `print(r**2)`.\n* Please discuss the failure modes, for both PG-TD and baselines, are they not compiling or failing test cases?\n* I find pass rate an odd metric. These problems are simple and the input is constrained, so if the code doesn't solve them 100% then is it not incorrect code? If nothing else, the authors should report the number of tasks solved 100%. \n* The computation times are quite slow and somewhat prohibitive, perhaps an improved LLM would help with this or a smaller one to search more quickly. Can you early terminate? What if a beam achieves 100% pass rate during search, do you stop?\n* The caching I do not see as a major contribution. Writing a dictionary for cache[query or partial] should suffice unless I misunderstood. \n* Though mostly clearly written, the statement of contributions, introduction, and abstract can be more clear what actually is done. What planning algorithm, what computationally efficient design mechanisms? These should be presented early and clearly for the reader.\n* The method itself is reasonable interesting, integrating MCTS to the search, and it might be useful to show the approach on more varied problems, perhaps other coding challenges or BigBench benchmarks.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "iGBnlGeI-z_", "reviewer": "Reviewer_kmKW", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes using a (MCTS-inspired) heuristic planning algorithm to guide the decoding process of a pre-trained code-generation transformer, for the specific task of test-case-guided competitive programming. At a high level, the proposed decoding strategy maintains a tree of prefixes, conducts beam-search-based \"rollouts\" and evaluates them on test cases to assess quality of these prefixes, then uses this information to determine which tree nodes to continue exploring. (I believe that the proposed method is deterministic, in contrast to MCTS which is stochastic by design.)\n\nSince beam search and greedy decoding already involve expanding a tree of tokens, the authors introduce caching mechanisms to make their proposed planner more efficient: they cache the tree structure generated during beam search and re-use it when expanding nodes for MCTS, and (when using greedy rollouts) they cache the greedy decoding output for any given prefix and re-use it when expanding the nodes along an already-explored trajectory.\n\nThe authors compare their approach to beam search, sample-and-filter approaches (such as the one used by AlphaCode), and the SMCG-TD decoding strategy proposed by Ellis et al. (2019). They argue that their decoding strategy outperforms these other strategies, although some of this may be due to the particular hyperparameters they used; for a fixed number of generated programs or fixed computation time their approach seems to perform fairly similarly to SMCG-TD.", "summary_of_the_review": "Although combining MCTS with a LLM for program generation is a good idea, there is not much technical novelty, and I have concerns about the validity of the empirical claims and whether it is appropriate to call the proposed approach MCTS given that it removes the stochastic part. I don't think this paper currently meets the bar for acceptance.\n\nI think this paper could be strengthened by\n- investigating the behavior of the planning algorithm when rollouts are actually drawn randomly rather than via greedy search or beam search (in other words, when the planning algorithm is actually MCTS)\n- being more careful about hyperparameter choices for the baselines to ensure any differences are due to the differences in action selection / exploration / planning, rather than being due to differences in search space size or generated program budgets.\n\n**Updated review:** The authors have addressed my concerns: they have clarified the relationship to MCTS, adjusted baselines to use more similar hyperparameters, and ran a comparison of beam search to sampling-based MCTS (for which their beam search method does better). In the discussion thread below, they have also included additional comparisons for Table 1 (labeled Table F in the comments below) that give a better sense of the performance of their approach; it consistently outperforms sampling-and-filtering, and given either a fixed computation time or a large enough budget, it also seems to outperform the SMCG-TD baseline (although only by a small margin). I have updated my score from 3 to 8, with the understanding that the updated results will be included in the final version of the paper if accepted.", "strength_and_weaknesses": "Strengths:\n\n- Integrating LM generation with MCTS is a good idea, allowing the search process to focus on prefixes with a higher chance of producing good programs.\n- The ability of their approach to cache and reuse outputs from the beam-search rollouts means the proposed strategy can be implemented fairly efficiently.\n- Their method seems to slightly outperform baseline methods given a fixed runtime budget, without needing to retrain the transformer model.\n\nWeaknesses:\n\n- **(mostly resolved in new revision)** Although the authors refer to their planning algorithm as MCTS, the \"Monte Carlo\" aspect appears to be missing. The proposed approach is deterministic: it selects actions based on the most-likely tokens following a prefix, and computes \"rollouts\" using either beam search or greedy decoding.\n  - The algorithm thus seems more like a heuristic best-first search than MCTS. Since a lot of the value of MCTS comes from aggregating over randomized rollouts, I would imagine randomization might lead to better performance. (On the other hand, the proposed \"sequence caching\" wouldn't work for random rollouts.)\n  - This choice is also confusing from a clarity standpoint; I expected the rollouts to be sample-based given that the method was based on MCTS.\n  - *The authors have rephrased this to not say their method is exactly MCTS, although I do have some remaining suggestions for how to make this clearer; see discussion.*\n- The main evaluation results are based on a seemingly arbitrary set of hyperparameters for each decoding strategy, and it's not clear to me that the comparisons (e.g. in Table 1) are fair.\n  - **(resolved in new revision)** As one example, their method explores only the top 3 tokens at each state, but their sample-and-filter baseline explores the top 5, and SMCG-TD doesn't have any top-k restriction (at least not stated in the paper).\n  - **(resolved in comments)** Additionally, the plots in Figure 4 show much smaller differences between their method and SMCG-TD than the differences in Table 1. This suggests that the differences in Table 1 are partially explained by generating different numbers of programs, or running for different amounts of time, due to the arbitrary hyperparameters selected for comparison.\n    - *The authors have included updated results in the comments below. Indeed their approach generates more programs by default. For a small budget, their approach is slightly worse than the SMCG-TD baseline, whereas for a large budget their approach is better. Their approach is also faster relative to the number of unique programs it generates.*\n- The novelty of the approach seems fairly small relative to ordinary MCTS and to existing applications of it to LLMs (e.g. the PPL-MCTS approach of Chaffin et al. (2022)).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "UiDwJykJiy", "reviewer": "Reviewer_FcyD", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper proposes an approach for improving pretrained models of source code for generating better code, if the input includes test cases. Given a pretrained model of code, a natural language (NL) description of the problem, and a few input-output examples, the authors use a Monte-Carlo tree search (MCTS)-based algorithm to generate code that satisfies the input-output examples, instead of just sampling / using beam search and \"hoping\" that the generated code will satisfy the input-output examples (=\"pass the tests\").\nThe idea is demonstrated on the APPS dataset, across GPT-2-1.5B and GPT-Neo-2.7B models, where the model passes more tests than the standard approaches of beam search or sampling + post-filtering.", "summary_of_the_review": "While the idea is interesting conceptually, I think that the settings and assumptions are overly unrealistic, and I do not understand how can the proposed approach be ever useful in real life.\n\nIt would have improved the paper if the authors could show that a model that was **only trained** with test cases - works better at test time **without test cases at test time**. Otherwise, if the user of the model needs to provide both a NL intent *and* test cases for every test-example, isn't it a bit too much of a burden for the user of such a model?\n\nI thus currently vote for rejection, but I am open to changing my mind if the authors convinced me otherwise.\n\n=========================\n## Post-discussion\n\nThe authors have resolved most of my concerns.\nI think that the main weakness of the paper remains its limited usability, i.e., the setting it addresses is overall unrealistic.\n\nI increased my score to 6, and I hope that the authors will complete the experiments where they generate the test cases as well, which will significantly strengthen the paper.", "strength_and_weaknesses": "### Strengths\n+ The problem is important and has a broad interest.\n+ The idea is interesting conceptually and shows improved results.\n+ The idea of treating code generation as an RL problem, where input-output pairs provide the reward, is novel and related to much work in program synthesis.\n\n### Weaknesses\n- **Usability** - I am finding it hard to see how the approach can be usable in realistic settings. This severe limitation is not explicitly mentioned, and is clear only to readers who are very familiar with the field and its datasets: while the paper places itself besides general models of code such as Codex,  CodeBERT, CodeT5 etc., the proposed approach is actually much more limited than them, and does not \"play the same game\". While Codex (for example) is only **tested** on executable test-cases to measure its execution accuracy, the proposed approach **requires training** on the input-output pairs. Further, not only that it requires input-output pairs at training time, at test time it can only be applied to examples that have input-output pairs as well.\n\nIn contrast, all other models (Codex,  CodeBERT, CodeT5 etc.) can be applied (at test time) to inputs that do not have input-output pairs, and this way they can serve as general code-completion / code-generation / bug-detection assistants.\nThe proposed approach kind of \"overfits\" on the existing benchmarks, while losing usability on the realistic tasks that these benchmarks are meant to be a proxy of.\nHow do the authors see the proposed model to be useful in real life? In which settings and tasks do we expect the user to provide both a natural language intent *and* test cases?\n\n- **Clarity** - discussed in the next section.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RQq88sPiP2n", "reviewer": "Reviewer_Xzif", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "", "summary_of_the_paper": "A new decoding strategy for language models for code is presented, based on the intuition of selecting tokens that lookahead search deems to be more likely to lead to correct sequence outputs. Additional ideas on managing the performance overhead of this strategy are presented. Finally, experimental results indicate that this new decoding strategy can outperform baselines by some margin.", "summary_of_the_review": "Novel decoding algorithm, good experimental results, no major concerns: this paper should be accepted, as it will be of substantial interest to the ICLR subcommunity of researchers interested in code generation.", "strength_and_weaknesses": "* (+) This is an interesting direction, focusing on a core part of neural code generation that often receives little attention.\n* (+) The empirical results indicate a significant improvement over baselines for the evaluated models.\n* (~) I found the discussion of related work unsatisfactory - it's presented as a list of related papers, with little discussion of commonalities and differences to the presented work", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster