P LANNING WITH L ARGE L ANGUAGE M ODELS FOR C ODE G ENERATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 complete 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 objectives 1 .
I NTRODUCTION
Large language models like Transformer (Vaswani et al., 2017a) have shown successes in natural language processing, computer vision, and 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) 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 coding challenge competence (Hendrycks et al., 2021).
Recently, AlphaCode (Li et al., 2022) even achieved a 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 a pretrained 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, their algorithms are not sample efficient.
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.
Different from natural languages, programs may fail completely with even a single incorrect generated token.
So these algorithms may need to exhaustively sample a large number of programs to find a correct one.
The main reason behind the sample efficiency 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 process 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 algorithm with a pretrained code generation Transformer, resulting in an algorithm that generates better programs than the conventional Transformer generation algorithms and the well-accepted sampling + filtering scheme in the literature?
To answer this question, we propose a novel algorithm,
PlanningGuided Transformer Decoding (PG-TD).
In this algorithm, we have a backbone Transformer that is pretrained on code generation.
In each step that the Transformer generates a token, the planner does lookahead search and suggests the token that passes the most number of test cases.
The planner alone may not find high-quality codes due to the large space of possible codes.
So the Transformer beam search algorithm and the next-token probabilities are used inside the planner to provide useful heuristics.
We find a straightforward integration between the planner and the Transformer can be computationally inefficient.
So we design mechanisms that allow the Transformer and the planner to share their information to make the overall algorithm more efficient.
We emphasize that our algorithm is model-agnostic , that is, any standard code generation Transformer 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.
Our algorithm can generate solutions that are used to finetune a code generation Transformer model to improve the Transformer’s performance.
More precisely, we have the following contributions in this paper.
• First, we propose a novel algorithm,
Planning-Guided 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 into the Transformer decoding function can cause redundant uses of the Transformer beam search algorithm.
We contribute to designing 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 better 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 concise codes and highly-commented codes.
R ELATED W ORK
Transformers for program synthesis.
Our work is based on Transformer for program synthesis (Roziere et al., 2020; Austin et al., 2021).
Inspired by their impressive 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 languages in the same way as natural languages.
A family of BERT-based Transformers are developed to 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 neural-symbolic methods Nye et al.
(2020) were explored.
Recently, AlphaCode (Li et al., 2022) combined large transformer models pre-trained on massive program data with large-scale model sampling, showing competitive performance in programming competitions.
All these models mainly focused on training more power code-generation models and simply use sampling (Fan et al., 2018) or beam-search (Graves, 2012) 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 automated unit tests to construct parallel training data for unsupervised code translation.
Chen et al.
; Gupta et al.
; 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 sequential Monte-Carlo sampling method for code generation.
Unlike the prior work, we use unit-testing results as reward signals for the Monte-Carlo tree search (MCTS) planning algorithm, which is further integrated with a Transformer-based large language model to generate better codes.
Planning and reinforcement learning for code generation.
The code generation problem has been formulated as a 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. ; Xu et al. (2019b) use an RL objective that optimizes the correctness of the generated programs instead of optimizing their similarity to reference solutions.
In the decoding phase, the Monte-Carlo tree search (MCTS) algorithm has been applied to search for high-quality codes.
However, MCTS is unable to scale to larger domains and 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 framework does not require a large language model and does not need to address the challenge of incorporating a planning algorithm with a large language model.
MCTS is more efficient and applicable to large domains when combined with deep learning or with a default policy as 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 MCTS to design a novel algorithm that is capable of solving competitive programming problems, and designed mechanisms to improve its efficiency.
Planning in natural language processing.
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.
; Leblond et al.
; Chaffin et al.
(2022) use pre-trained discriminators or pre-defined metrics as reward functions.
We want to emphasize that we are the first to combine MCTS with large language models for general-purpose programming language generation.
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 maximizing the pass rate value in our problem.
M ETHOD 3. O VERVIEW
We consider the code generation problem for competitive programming, illustrated in Fig.
An agent is given the problem description of a coding problem.
It requires the agent to understand the problem from the natural language description and generate a program that solves the problem.
Similar to Li et al.
; Chen et al.
; Ellis et al.
(2019), we assume that the agent has access to a set of test cases, where a test case is a 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 output string. The objective is to generate a 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.
The agent can only access the public test cases during the program generation process, while we use the private test cases to evaluate the programs it generates.
Problem Statement
Given is a string S .
Replace every character in S with x and print the result.
Constraints
S is a string consisting of lowercase English letters.
The length of S is between 1 and 100 (inclusive).
[Figure]
Input is given from Standard Input in the following format: S
[Figure]
Replace every character in S with x and print the result.
[Figure]
Transformer models (Li et al., 2022; Hendrycks et al., 2021) have been widely applied for code generation thanks to its capacity in sequence-to-sequence modeling.
In the Transformer’s generation process, beam search (Graves, 2012) and sampling (Fan et al., 2018; Dabre & Fujita, 2020) are adopted to generate code sequences.
However, these algorithm cannot easily optimize an objective different from what it is trained on (the similarity to the reference solutions).
So we cannot directly use these generation algorithms to generate programs aiming to pass the most test cases.
On the other hand, a planning algorithm can directly optimize the pass rate or any desirable programming objective.
To use a planning algorithm, we follow Bunel et al.
; Ellis et al.
(2019) to first formulate the code generation problem as a Markov decision process (MDP) (Sutton & Barto, 2018).
A state s is the concatenation of the problem description and a 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, which contains the termination action (the terminal token) that indicates that the agent believes the program is complete.
The transition function deterministically concatenates the state s with the 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 a complete program ( i.e., when the last token of s is the terminal token).
The reward of a partial program is always 0.
There is rich literature on finding the optimal policy for an MDP.
In this paper, we consider a tree search-based planning algorithm inspired by Monte-Carlo tree search (MCTS), illustrated in Fig.
The algorithm maintains a tree structure where nodes correspond to states and edges correspond to actions.
The algorithm starts from the 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) a value function for each edge s, a , which is the maximum value obtained by starting in s and taking action a .
The algorithm would visit nodes with either higher Q values (as they lead to higherquality programs) or with smaller visit numbers (as they are under-explored).
In the following part of this section, we describe how we integrate the MCTS algorithm in the generation process of a Transformer.
3. P LANNING -G UIDED T RANSFORMER D ECODING
Now we are ready to answer the question we asked in the introduction:
Can we integrate a planning algorithm with a pretrained code generation Transformer to generate better programs?
We design a a, , <PD> a x a= x, a x = , <PD> a x a= x, a x = , generates complete programs using beam search a,
b
a,\n a, , <PD> a x a= x, a x = , a, , <PD> a x a= x, a x = , suggests next tokens …
Selection
Expansion
Evaluation
Backpropagation
Transformer
Evaluate on public test cases
Transformer generation algorithm where MCTS is called iteratively to perform lookahead planning and return with the token that will lead to high-quality programs.
The MCTS algorithm alone may not be able to find high-quality codes due to the large search space.
So the conventional Transformer beam search algorithm and the next-token probabilities are used within the MCTS algorithm to guide the search process.
Algorithm
Require: root : the current state; rollouts : the number of
MCTS rollouts; k max : the maximum number of children of any node; b : the number of beams for Transformer beam search 1: for i ← 1 , 2 , . .
. , rollouts do 2: node ← root 3: # Selection 4: while | node.children |
= k max do 5: node ← UCB SELECT ( node.children ) 6: end while 7: # Expansion 8: children num ← | node.children | 9: act ← TOP K ( node, k = children num + 1) 10: new node ← the state reached by taking action act in node 11:
Add new node to node.children 12: # Evaluation 13: program list ← GENERATE SEQ ( new node, b ) 14: r ← max p ∈ program list GET REWARD ( p ) 15: # Backpropagation 16:
Update and the values of
The main generation process is described in Algorithm 2.
In each step, PG-TD is called to generate a token a t .
We provide the pseudocode of PG-TD in Algorithm 1 and illustrate the whole process in Figure 2.
The PG-TD algorithm follows the same framework as the standard
MCTS algorithm.
We specify k max to be the maximum number of children that any node may have.
As our problem has a huge action space (the set of all tokens), we specifically constrain the number of children so that the agent would not explore too many actions of the same state.
For completeness, we provide all the steps in the pseudocode and in Sec.
C.1.
Here, we focus on the expansion and the evaluation steps where the Transformer is used to provide heuristics for the search process.
The Transformer provides two functions for PGTD: GENERATE SEQ ( s, b ) generates a list of b programs using the Transformer beam search algorithm with the prefix s , where b is the beam size; TOP K ( s, k ) returns the k -th most likely next token starting from s .
In the expansion step, after a node in the tree is selected, we add a new node to its children which corresponds to a possible next state (for succinctness, a node also refers to the state that it represents).
Adding a random token as in standard
MCTS may very likely cause a syntax error.
So we call the TOP K function to get the most likely next tokens (or actions).
When the current node already has children num children, we get the ( children num +1)
-st most likely next token, and add the corresponding next state, next node , to the children of node (Line 9-11).
In the evaluation step, we need to evaluate the quality of new node .
Note that new node may still be a partial program.
We cannot directly evaluate the quality of a partial program as we do not know how it will be completed, and how many test cases it will pass.
Here, we use the Transformer again by calling the GENERATE SEQ function to generate complete programs from the current node using the beam search algorithm.
We set the size of beam search, b , to indicate how many complete programs that we want to generate from the current partial program.
We run these generated programs on the public test cases to get their rewards, and set the value of new node to be the highest reward of the complete programs (Line 13-14).
3. I NFORMATION S HARING B ETWEEN T RANSFORMER AND MCTS
a,b a,c a, , <PD> a x a= x, a x = , a,b, a,b= a,c= a,c, … … … … …
a, , <PD> a x a= x, a x = , a,b, a,b= a,b
Iteration t
Iteration t+
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 Transformer beam search algorithm also implicitly builds a tree structure , which can be used by the future iterations of MCTS.
In this section, we describe how we improve the algorithm’s efficiency by sharing information in the Transformer beam search algorithm with MCTS.
Consider the example in Fig.
3, which shows two iterations of PG-TD.
In the evaluation step of the t -th iteration, the Transformer beam search algorithm implicitly builds a tree to find the most likely sequences within a beam, visualized in red color in Fig.
3 (left).
Because we only keep b partial programs in the beam, it is a tree where only b nodes with the highest likelihood are expanded at each level (in the illustration, b = 2 ).
Other nodes are dropped and no longer considered by the beam search algorithm.
In the ( t + 1) -st iteration, if MCTS selects “ a, ” , such a state is already visited in the Transformer beam search in the t -th iteration.
When the algorithm needs to 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 .
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 in Sec.
C.2.)
To improve the efficiency of
GENERATE SEQ , 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 .
One may be concerned about the space efficiency of tree structure caching and sequence caching, as they keep saving trees and sequences for each node in the tree.
Fortunately, we only need to keep the caches that are consistent with the current state.
Once a new action is taken, we maintain the cached data that are useful for the new state.
We will provide more details on the caching methods in Sec.
C.2.
E MPIRICAL E VALUATION
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 [Table] tree search algorithm an effective planning algorithm in PG-TD?
Is it better than other planning or sampling-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 the computation time?
Q5:
Can we use the samples generated by PG-TD in its search process and finetune the Transformer to generate even better solutions?
Datasets and models.
Recently, several competitive programming datasets have been made available for benchmarking code generation algorithms.
For each programming problem, they include the program description in natural language, 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 in AlphaCode (Li et al., 2022).
APPS does not split public vs. private 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 is 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, GPT-Neo that are finetuned on the APPS training dataset (Hendrycks et al., 2021).
Algorithms.
We compare PG-TD with the following algorithms.
Beam search only uses the Transformer beam search algorithm to generate the whole program, without using the test cases.
This is the method used in Hendrycks et al.
; Chen et al.
(2021a).
We use the beam size of 5 in the
Transformer generation function, which is the same design choice as in Hendrycks et al.
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 is in the top-3 most-likely tokens in each step.
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 of programs than them.
To determine the effectiveness of the tree search algorithm, we considered another baseline,
Sequential MonteCarlo-Guided Transformer Decoding (SMCG-TD) , that replaces the the tree search algorithm component with a sequential Monte-Carlo algorithm (Ellis et al., 2019).
It is an iterative algorithm that maintains a population of partial programs.
It determines the fitness of a partial program using the Transformer in the same way as 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 C.3.
For PG-TD, we set the parameters k max = 3 , rollouts = 3 , b = 1 .
For the baseline methods, we sample 512 programs in S+F, and maintain a set of 10 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 APPS introductory dataset (1000 problems) (Hendrycks et al., 2021) and CodeContests (165 problems) (Li et al., 2022).
We use the same metrics as in Hendrycks et al.
(2021), pass rates and strict accuracies of the private test cases.
The private test cases are hidden from the agent during the code generation process.
The pass rate is the average percentage
[Figure]
[Figure] [Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
of the private test cases that the generated programs pass.
The strict accuracy is the percentage of the problems where the generated programs pass all the private test cases.
Effectiveness of PG-TD.
As shown in Table.
1, PG-TD consistently outperforms all other baselines in all the settings under the pass rate metric.
As we optimize the pass rate, we expect our algorithm to outperform other baselines under strict accuracy as well.
This is almost true except that using the GPT-2 (2.7B) model on the CodeContests dataset, our algorithm is slightly outperformed by
SMCG-TD.
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 search-based planning algorithm is 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) in Sec.
A.
Efficiency of PG-TD.
To show PG-TD’s efficiency, we report the pass rates vs. the number of generated programs by Transformer and the computation time on a smaller dataset in Fig.
We set k max = 3 , rollouts = 3 , and b = 1 as the default setting and only changes one parameter at a time.
We vary these parameters to be k max = 2 , 3 , 4 , rollouts = 6 , 9 , and b = 3 , 5 .
As we expect, increasing k max , rollouts , and b all helps PG-TD achieve higher pass rates, while increasing b helps improve the pass rate the most as the Transformer beam search algorithm generates more high-likelihood candidates.
We also consider different sample sizes for S+F.
We choose sample sizes of 128, 256, and 512.
With the same number of samples, S+F generates programs with much lower pass rates (answering Q3 ).
This confirms our observation that S+F, like AlphaCode (Li et al., 2022), generates samples only according to the Transformer’s generation probabilities, without taking their pass rates into consideration until the programs are all generated .
PG-TD actively considers the pass rates of the generated programs during the generation process, which achieves a higher efficiency.
For SMCG-TD, we set the size of population to be 10, 20, 50, and 100.
Since SMCG-TD also uses the Transformer and public test cases to guide the generation process, we expect that its performance is close to the default setting of PG-TD with a larger size of population.
However, due to its sampling nature, it cannot do multi-step lookahead search as in PG-TD.
So SMCG-TD’s performance is worse than the variations of PG-TD.
More details are discussed in Sec.
C.3.
Effectiveness of caching.
In terms of the design choice 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 to regenerate whole sequences, ending up with consuming
[Figure]
much
Finetuning transformer with PG-TD-generated samples.
Since we 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 and 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 set of APPS.
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
Transformoer model, we run beam search to generate solutions on the test set to see if it has a better performance.
We use the first 500 problems in the APPS interview-level test set for validation and the APPS introductory-level test set for testing.
The results are reported in Table 3.
After finetuning with
TD-generated solutions we see improvement in both pass rate and strict accuracy.
More details are in Sec
C
[Table]
rate
[Table]
Code length denotes the length of the generated code string and comment number denotes the number of code lines that contains comments.
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 encouragement .
As shown in Table 4, the code length penalty reward function makes the model generate more concise codes; the comment encouragement reward function encourages models to generate codes with more comments.
Both objectives still achieve reasonable pass rates on the public test cases.
We provide a qualitative example for comment encouragement in Figure 5 and more details of the reward functions in Sec.
B of the appendix.
D ISCUSSION AND C ONCLUSION
In summary, we proposed a novel algorithm that uses the powers of a pre-trained Transformer and the Monte-Carlo tree search algorithm.
We evaluate our algorithm and empirically show that our algorithm generates programs of higher quality compared with competing baseline algorithms in different settings.
We also design model structure-specific caching mechanisms which contribute to saving computational expenses.
We show that our algorithm is 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 integrating 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 a small number of test cases can help find better solutions (Sec.
A, Table 7).
Our algorithm is also more computationally expensive than the pure Transformer beam search algorithm since ours needs to run the beam search algorithm for multiple times within MCTS.
In the future, we consider improving the framework’s performance by using a value function to estimate the programs’ pass rates, similar to the evaluation function in Silver et al.
(2016) for mastering Go.
We can also learn a neural state representation for partial programs as in Chen et al.
We also consider using parallel MCTS (Chaslot et al., 2008) for the evaluation step to parallelize the computation for different states.
A PPENDIX
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 to help readers better understand our paper.
This appendix is organized as follows.
• In Sec.
A, we provide empirical results on more competitive coding datasets and 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 provide empirical evidence for our claims in the discussion section that our algorithm is versatile and can be used to optimize different code generation objectives other than optimizing the pass rates.
We consider the objectives of code length penalty and comment encouragement.
In Sec.
C, we provide more details on the components in PG-TD as well as the baseline algorithms.
• In Sec.
D, we illustrate more examples on the codes generated by our PG-TD algorithm and the baseline algorithms, and also provide a step-by-step visualization of execution of PG-TD.
• In Sec.
E, we discuss more on the advantages and the potential negative social impacts of our algorithm.
A E MPIRICAL E VALUATION
We report 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 dataset contains coding problems at three levels: introductory (1000 problems), interview (3000 problems), and competition (1000 problems).
Our experiments are run on machines with two Intel(R)
Xeon(R)
Gold 6258R
CPUs (@ 2.70GHz), and one V100-SXM2 GPU.
In Table 5, we can see that our algorithm consistently has the highest pass rate on all the datasets, which further supports Q1 in the main paper.
[Table]
In Fig.
6, we supplement Fig.
4 by reporting the strict accuracies of PG-TD and baseline algorithms.
They have the similar patterns as the pass rate results.
However, since we only optimize the pass rates, the performance of each algorithm does not monotonically increase by increasing the size of the tree in PG-TD or increasing the sample sizes for Sampling and SMCG-TD.
n@k and pass@k metrics.
We also report the performance of PG-TD and Sampling under the n@k and pass@k metrics (Li et al., 2022).
Given a 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 n submitted programs passes all the private test cases.
pass@k is the proportion of the problems where any of the k generated programs passes all the private test cases (effectively k@k).
For PG-TD, the k samples are the complete programs found in the first k rollouts.
The results are reported in Table 6, evaluated on the first 500 APPS introductory problems.
We observe that with the same number of generated samples ( k ), PG-TD finds programs that are more likely to pass all the private test cases.
[Figure] [Figure]
PG-TD 5.
S+F 5.
PG-TD 6.
S+F 5.
pass@ pass@ pass@
[Table]
[Table] [Table] [Table]
[Table] [Table]
Using different numbers of public test cases.
On the APPS dataset, we split the 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 a satisfactory performance, we consider using 1, 3, 5 public test cases and using the rest of the test cases as private test cases.
The results are reported in Table 7, evaluated on the first 500 APPS introductory problems.
We observe that even with a small number of public test cases,
PG-TD is able to find programs with higher a pass rate and strict accuracy than beam search.
This suggests that the Transformer can serve as a 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 on Codex.
We have run both beam search and PG-TD algorithms using Codex (Chen et al., 2021a) as the backbone. We follow the “best practices” provided on the Codex guide to put the problem description into a block comment in the beginning, and ask Codex to complete the rest. The prompt looks like the following. [Equation]
As we expect, PG-TD helps Codex generate better codes.
This further validates our claim that our method is model-agnostic and can help a pre-trained code generation Transformer generate better codes.
Our results are reported in the Table 8.
Due to the limits of the OpenAI APIs, we are only able to test the algorithms on a small set of data.
We also show a concrete example where PG-TD outperforms beam search in Fig. 15.
Using sampling-based MCTS.
We could use a sampling approach in the evaluation step of MCTS instead of using beam search, which would be a more faithful implementation of MCTS.
We experimented with the two versions of PG-TD and find that using beam search indeed has a better performance (Table 9).
The gap could be caused by that we only do sampling once to evaluate a new node.
We could get a more accurate by sampling multiple times.
However, that would make our algorithm more computationally expensive.
About failure modes.
In Table 10, we report the average percentages of test cases where compilation errors and runtime errors occur.
We use the error analysis utility provided in the original APPS codebase.
As we 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.
About termination conditions.
In Table 11, we find enabling early termination does (marginally) improve the algorithms’ efficiency.
We will include this implementation choice in the paper.
We expect early termination would be significantly more helpful when the code generation algorithms can solve most of the problems (achieving pass rates of 1 on most of them).
Using PG-TD for code translation.
A real-world use case of our algorithm is code translation.
The goal of code translation is to translate a script from one programming language to another.
In this problem, automatically generated test cases are used to determine if the translated codes are correct (Roziere et al., 2022).
There exist free or commercial tools like EvoSuite, Jtest, etc.
We conducted code translation experiments in the following table.
We use TransCoder-ST (Roziere et al., 2022) as the backbone Transformer, and compare beam search with PG-TD.
In PG-TD, it uses EvoSuite to generate public test cases, and use the pass rates on them as rewards.
The translated codes are evaluated on a set of private test cases.
We can see that in this problem,
PG-TD also outperforms the beam search algorithm (Table 12).
Assets and licenses.
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 et al., 2019) is released under Apache License 2.0 4 .
B P LANNING FOR O THER C ODE G ENERATION O BJECTIVES
Besides the default reward function that optimizes the pass rate, we show the versatility of the proposed algo- rithm by setting two new objectives, code length penalty and comment encouragement.
Code length penalty .
We make the planner to generate more concise codes by using the following reward function
[Equation]
where p is the average pass rate on the public test case set and l c is the length of the code string.
λ and t are hyperparameters that control the weight of code length penalty and are set to 0 .
1 and 20 , respectively.
As shown in the Figure 7, the generated solution becomes more concise and the code string length decreases from 187 to 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
[Equation]
where N cm is the number of “ # ” in the code string and λ 2 is set to 0 .
2 , controlling the weights of the comment lines.
If we simply add λ 2 × N cm as the comment encouragement term, the planner would generate code with repeated meaningless comment lines to get more rewards.
To handle this problem, we made the following two changes.
First, we add the code length penalty term, λ 1 × e − l c /t .
Second, we upper-bound the rewards for the number of comment lines by setting the comment encouragement term to be λ 2 × min(1 , N cm N max ) , where N max = [Table] [Table]
[Table] [Table] [Table]
Problem
Statement
Takahashi loves palindromes.
Non-palindromic strings are unacceptable to him.
Each time he hugs a string, he can change one of its characters to any character of his choice.
Given is a string S .
Find the minimum number of hugs needed to make S palindromic.
Constraints (1).
S is a string consisting of lowercase English letters.
The length of S is between 1 and 100 (inclusive).
Input
Input is given from Standard Input in the following format: S .
Output
Print the minimum number of hugs needed to make S palindromic.
Sample Test Inputredcoder
Sample Test Output
For example, we can change the fourth character to o and get a palindrome redooder.
1 # cook up a palindrome 2 s=input() 3 n=len(s) 4 c= 5 if n%2==0: for i in range(n//2): if s[i]!=s[n-i-1]: c+=  print(c) 10 else : for i in range(n//2): if s[i]!=s[n-i-1]: c+=  print(c)
Default reward function (Pass Rate: 1.00, code string length: 187 ).
1 s=input() 2 ans= 3 for i in range(len(s)//2): if s[i]!=s[˜i]: ans+= 5 print(ans)
Reward function with code length penalty (Pass Rate: 1.00, code string length: 78 ).
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 to get accepted.
Want to try?
Input
The input contains two integers a and b (0 ≤ a, b ≤ 10 3 ) , separated by a single space.
Output
Output the sum of the given integers.
Sample Test Input
Sample Test Output
Sample Test Input
Sample Test Output 1 a, b = map(int, input().split()) 2 print(a + b)
Default reward function.
1 #!/usrbin
/env python 2 # coding =utf-8 #py3 runs on py 3 a,b=map(int,input().split()) 4 print(a+b)
Reward function with comment encouragement (code with comments).
C A LGORITHM D ETAILS C. PG-TD
In this section, we describe our PG-TD algorithm with more details, following the Monte-Carlo Tree Search (MCTS) steps.
We show our framework in Fig.
10 again for convenience.
The algorithm maintains a search tree structure where nodes correspond to states and edges correspond to actions.
For each node s , it maintains the number of times that it has been visited, denoted by s.visits .
For a state, action pair, s, a (corresponding to an edge in the tree), the algorithm maintains a function Q ( s, a ) , which is the maximum value obtained by starting in s and taking action a .
Conventionally, Q ( s, a ) is the average return.
Since our the code generation problem is deterministic (that is, there are no stochastic transitions), we find the performance is better if we keep track of the best program generated from a node.
First, the algorithm selects a node for expansion.
It starts from the root node, selects substrees recursively, until it reaches a node which is not fully extended (whose children are fewer than k max ).
The common node selection criterion is upper confidence bound (UCB) (Kocsis & Szepesv´ari, 2006), which leverages between exploiting known-to-be-good states and exploring less-visited states,
[Equation]
where s ′ is the state reached by taking action a in s ; β is a hyperparameter.
We set β = 1 in our experiments.
The UCB SELECT function in Alg.
1 returns the next state that has the highest UCB value,
[Equation] Then the algorithm expands a node using GENERATE SEQ and evaluates the new node using GENERATE SEQ , as we described in the main paper.
Finally, the estimated value is backpropagated to its parents recursively until it reaches the root and update their Q values.
Suppose the estimated value is v .
Then for all state, action pairs, s ′ , a ′ , along the trajectory in the tree to reach the new state, Q ( s ′ , a ′ ) is updated accordingly using the new value: Q ( s ′ , a ′ )
← max( Q ( s ′ , a ′ ) , v ) .
Comparison with the standard MCTS algorithm.
It is worth noting that our implementation of
MCTS in PG-TD is 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.
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.
It is impractical to apply the standard MCTS algorithm to domains with large state spaces or large action spaces, as we 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, we use a pre-trained Transformer to resolve these limitations of the standard MCTS algorithm.
Transformer generation functions.
Our framework has a Transformer model that is pretrained on competitive coding datasets, and provides the GENERATE SEQ function that generates most-likely sequences from a partial program, and TOP K that generates the most-likely next tokens given a partial program.
We describe the implementations of these two functions in Alg.
4 and Alg.
They both make use of P Transformer ( a | s ) , the probability that token a is the next token given the partial program s , which is readily provided in a pretrained Transformer model.
Specifically, TOP K selects the top-k most-likely tokens according to P Transformer .
GENERATE SEQ follows the standard Transformer beam search algorithm to generate complete programs.
C. E FFICIENT I MPLEMENTATION BY I NFORMATION S HARING
As discussed in the main paper, there may be calls to the Transformer generation functions the perform redun- dant computations.
We have described tree structure caching in the main paper.
We provide more details on sequence caching in this section. Sequence caching.
The goal of sequence caching is to reduce the computational cost of the evaluation step of PG-TD.
In the evaluation step, the Transformer beam search algorithm is called to generate a complete program from the current node.
These complete programs and their rewards are both cached and used in a future iteration.
ab b <PD> a b aa ba a b a a <PD> a b aa ba a b a a
Selection
Expansion
Evaluation
Backpropagation ab b <PD> a b aa ba a b a a ab b <PD> a b aa ba a b a a
Select the next action in the action set
Generate a random policy and evaluate it a, , <PD> a x a= x, a x = , <PD> a x a= x, a x = , generates complete programs using beam search a,b
a,\n a, , <PD> a x a= x, a x = , a, , <PD> a x a= x, a x = , suggests next tokens …
Selection
Expansion
Evaluation
Backpropagation
Transformer
Evaluate on public test cases
However, we can only use sequence caching for beam size b = 1 .
When b > 1 , using cached sequences in the previous iterations may be suboptimal.
Let’s consider the example in Figure 12, where b = 2 .
Starting from “ a, ” (the left figure), suppose the beam search algorithm returns a 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 in the previous iteration, which only searches the subtree of “ a,b ” with one beam .
In other words, we may find a sequence with a higher likelihood by re-running the beam search algorithm with b = 2 starting from “ a,b ”.
In Alg. 4 and Alg. 5, we described how GENERATE SEQ and TOP K take advantage of the caching mechanism, highlighted in blue colors.
In Alg. 6, we described how we clean the cache so that the memory it takes does not increase without bound as we execute PG-TD.
Alg. 6 is called at the beginning of each function call of PG-TD.
We also cache the search tree between different function calls of PG-TD, which is the same design choice used in Silver et al.
The tree structure built in one iteration will be kept and used for following iterations.
So the following iterations can search even deeper in the tree.
C. B ASELINE A LGORITHMS
Sampling + Filtering.
We have described the sampling algorithm in the main paper.
The pseudocode is shown in Alg.
Sequential Monte-Carlo-Guided Transformer Decoding.
Sequential Monte-Carlo-Guided 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. 8.
One can think of it from a view of genetic algorithms.
SMCG-TD maintains a group of partial programs, called a generation , and updates them iteratively.
For each partial program in the generation, we sample a next token using the Transformer and append it to the a,b a,c a, , <PD> a x a= x, a x = , a,b, a,b= Transformer Beam Search ( b = 1) [Figure] a, , <PD> a x a= x, a x = , a,b a,b, a,b= … a,b=0, a,b=0,
[Figure]
,c [Figure] [Figure] [Figure]
[Figure]
,c
[Figure]
,c= a,c,
[Figure]
a,b a, a,b, a,b= [Figure] partial program, and put it in a new generation. We then evaluate each partial program in the new generation in the same way as PG-TD: We use the Transformer to generate a 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, we 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 it finds that has the largest reward.
One may notice that SMCG-TD resembles PG-TD on a higher level.
It also uses the Transformer and the public test cases to decide which next token to generate in the code generation process.
The key difference is that SMCG-TD uses a sampling-based approach and PG-TD uses a tree search algorithm. We are able to increase the number of rollouts in PG-TD to make it do multi-step look-ahead search, which is what SMCG-TD incapable of. So in our experiments, the performance of SMCG-TD only approximates the default [Table]
[Table] It is worth noting that sequence caching is also implemented for SMCG-TD when it calls GENERATE SEQ . So even if multiple partial programs in a generation can be the same, the algorithm would not call the Transformer to generate a 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. Require: root : the current state; rollouts : the number of MCTS rollouts; k max : the maximum number of children of any node; b : the number of beams for Transformer beam search 1: UPDATE CACHE ( root )
▷ Alg.2: for i ← 1 , 2 , . . . , rollouts do 3: node ← root 4: # Selection 5: while | node.children |
= k max do 6: node ← UCB SELECT ( node.children ) 7: end while 8: # Expansion 9: children num ← | node.children | 10: act ← TOP K ( node, k = children num + 1)
▷ Alg.11: new node ← the state reached by taking action act in node 12:
Add new node to node.children 13: # Evaluation 14: program list ← GENERATE SEQ ( new node, b ) ▷ Alg.15: r ← max p ∈ program list GET REWARD ( p ) 16:
# Backpropagation 17:
Update the values of new node its ancestors in the r 18: end for 19: return the action with the highest value starting from root
C. F INETUNING
In our experiments, we use the samples generated by the PG-TD to further finetune the Transformer models.
We then generate solutions using the finetuned Transformer model by running beam search, and observe improvement on both pass rate and strict accuracy compared with running beam search using the original model.
One challenge in our design is that solutions that PG-TD generates have different pass rates.
We could use only solutions with 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 pass rates of training samples, we use the loss function similar to Le et al.
[Equation]
where W denotes a complete program; r ( W ) is the pass rate of the program; PD is the program description.
In future work, we will investigate different loss functions and compare their effectiveness in improving the Transformer model’s performance.
Require: s : the current partial program. 1: for cached state ∈ sequence cache.keys () do 2: if s does not match the prefix of cached state then 3: delete sequence cache [ cached state ] 4: end if 5: end for 6: for cached state ∈ top k cache.keys () do 7: if s does not match the prefix of cached state then 8: delete top k cache [ cached state ] 9: end if 10: end Require sample samples
Transformer sample do p programs ← LIST 5: while | population | > 0 and t max steps do 6: new generation ← LIST () 7: new fitness ← NEW LIST () 8: for i ← 1 , . .
. , LEN ( generation ) do 9: s ← sampled from generation , where generation
[ j ] is selected with the probability of fitness [ j ] 10: a ← sampled from P Transformer ( ·| s ) 11: s ′ ← CONCATENATE ( s, a ) 12: if s = <
|endoftext| then 13: # sample finishes generation, save to output 14: add s to complete programs 15: else 16: add s ′ to new generation 17: p ← GENERATE SEQ ( s ′ , b = 1) 18: r ← GET REWARD ( p ) 19: add r to new fitness 20: end if 21: end for 22: normalize new fitness 23: generation ← new generation 24: fitness ← new fitness 25: t ← t +26: end while 27: return the program complete programs largest reward
D I LLUSTRATIVE E XAMPLES D. C ODE G ENERATION E XAMPLES
In this section, we provide more qualitative examples generated by our PG-TD and the baseline methods, shown in Figure 13 and Figure 14.
Our PG-TD algorithm generates solutions that have higher pass rates than the baseline methods.
Problem Statement
Takahashi is meeting up with Aoki.
They have planned to meet at a place that is D meters away from Takahashi’s house in T minutes from now.
Takahashi will leave his house now and go straight to the place at a speed of S meters per minute.
Will he arrive in time?
Constraints
1 ≤ D ≤(2).
1 ≤ T ≤
≤ S ≤(4).
All values in input 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 Y es ; otherwise, print No .
Sample Test Input[-1000 
-]
[Figure]
Sample Test Output
Yes
It takes 12.5 minutes to go 1000 meters to the place at a speed of 80 meters per minute.
They have planned to meet in 15 minutes so he will arrive in time.
1 # cook your dish here 2 try : a=list(map(int,input().split( )))
(cid:44) → (cid:44) → if (a[2]==1): print("Yes") else : print("No") 8 except : pass
Beam Search (Pass Rate: 0.125).
1 # cook your dish here 2 d,t,s=list(map(int,input().split ()))
(cid:44) → (cid:44) → 3 if (d==t): if (s*(s+1)/s==t): print("Yes") else : print("No") 8 else : if (s*(s+1)/s>t): print("Yes") else : print("No")
Sampling + Filtering (Pass Rate: 0.625).
1 n, t, s=map(int, input().split())
(cid:44) → 2 t1=n / s 3 if t1 <=
t: print("Yes") 5 else : print("No")
[Figure]
[Figure]
Rate: 1.00).
Problem Statement
Given is an integer r .
How many times is the area of a circle of radius r larger than the area of a circle of radius 1 ?
It can be proved that the answer is always an integer under the constraints given.
Constraints
1 ≤
r ≤
All values in input are
integers.
Input
Input is given from Standard Input in the following format: r .
Output
Print the area of a circle of radius r , divided by the area of a circle of radius 1 , as an integer..
Sample Test Input
Sample Test Output
The area of a circle of radius 2 is 4 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.
1 # cook your dish here 2 r=int(input())
if (r>=2*r): print(r*r*2) 5 else : print(0)
Beam Search (Pass Rate: 0.00).
1 # cook your dish here 2 r=int(input())
if (r>=2*r): print(r*r*2) 5 else : print(0)
Sampling + Filtering (Pass Rate: 0.00).
1 # cook your dish here 2 r=int( input() ) 3 print(2*(r**2)-r**2)
PG-TD (Pass Rate: 1.00).
Problem Statement
You are given an array a consisting of n integer numbers. Let instability of the array be the following value: n max i = 1 a i − n min i =1 a i . You have to remove exactly one element from this array to minimize instability of the resulting ( n − 1) -elements array.
Your task is to calculate the minimum possible instability.
Input
The first line of the input contains one integer n ( 2 ≤ n ≤ 10 5 ) — the number of elements in the array a .
The second line of the input contains n integers a 1 , a 2 , . .
. , a n ( 1 ≤ a i ≤ 10 5 ) — elements of the array a .
Output
Print one integer — the minimum possible instability of the array if you have to remove exactly one element from the array a .
Examples
Input
[Figure]
Note
In the first example you can remove 7 then instability of the remaining array will be 3 − 1 = 2 .
In the second example you can remove either 1 or 100000 then instability of the remaining array will be 100000 − 100000 = 0 and 1 − 1 = 0 correspondingly.
[Figure]
int(input())
[Figure]
a = list(map(int, input().split()))
[Figure]
print(max(a) - min(a) - 1)
[Figure]
[Figure]
Rate 0.00).
[Figure]
main()
[Figure]
n = int(input())
[Figure]
a = list(map(int, input().split()))
[Figure]
a.sort()
[Figure]
print(min(a[n-1] - a[1], a[n-2] - a[0]))
[Figure]
D. I LLUSTRATIONS OF S TEP - BY -S TEP E XECUTION OF PG-TD
In the example shown in Fig.
13, we see that PG-TD outputs the first tokens (“ n, t, . . . ”)
differently from the beam search (“ # cook your . . . ”).
It means that after evaluating the pass rates of the potentially generated programs, our algorithm disagrees with the Transformer beam search algorithm and outputs different tokens.
We use this problem as an example to visualize the execution of PG-TD.
We consider the setting where the maximum children of any node ( k max ) is 3, the number of rollouts is 6, and the number of beam size ( b ) is 1.
<PD> [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
n
[Equation]
[Figure]
(a) We start with decoding from the problem description,
<
PD > . Initially , there is only the root node
which is the problem description To expand this node, the algorithm calls TOP K to ask the Transformer to suggest the most-likely tokens.
In this case, Transformer suggests ‘#’, ‘n’, ‘import’ , in descending order of their likelihood.
(b) The algorithm adds the most-likely token, ‘#’ , to the tree.
To evaluate this node, the algorithm calls GENERATE SEQ to generate a complete program, and run it on test cases.
The pass rate of the generated program is 0.25 (it passes 25% of the public test cases).
(c)
In the next iteration, the algorithms adds the second most-likely next token, ‘n’ , to the tree.
Similarly, the algorithm generates a complete program from this new node and runs it on the test cases.
Its pass rate is 0.375.
[Figure]
[Figure]
[Figure] import [Figure] D
[Figure]
PD
[Figure]
# n import
[Figure]
m
[Figure]
[Table]
[Figure]
[Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table] [Table]
n,
[Table]
[Figure]
the
‘import’
The generated starting from it has 0 pass rate.
After the first three iterations, three nodes are added to the tree and their values are computed.
the selects ‘n’ for expansion, as it has the highest value.
The algorithm calls K to get the most-likely tokens.
In this case, they are ‘,
The algorithm adds ‘n,’ to the tree.
The program generated from it has the pass rate of 0.375.
<PD> # n import [Table]
0.375n,
[Table] n= int n=int(input())m= int(input())if n==1:if m%2==0:print("Yes") else:print("No")[-elifm==
-]
elifm==2:if n==1000:print("Yes")elifn==10000: print("Yes")else:print("No") else: if (n*(n+1)/2)<=m:print("Yes") else:print("No")
Pass rate: 0.
<PD> # n import [Table]
# cook
[Table]
[Table] ) (h) n, n=
(g)
The algorithm selects ‘n’ again and adds ‘n=’ to the tree.
The program generated from it has the pass rate of 0.
(h)
The algorithm selects ‘#’ for expansion and adds ‘# cook’ to the tree.
It no longer selects ‘n’ for expansion in this iteration as it believes the subtree of ‘n’ is well explored, and it instead explores the substree of ‘#’
Now the algorithm has finished six iterations of rollouts.
One can verify that six nodes have been added to the tree.
After the tree search, the algorithm will return the token ‘n’ as the action it believes to be optimal.
Our algorithm disagrees with the most-likely token suggested by the Transformer, as it finds that it can generate a program with a higher pass rate by outputting a different token (the second most-likely token suggested by the Transformer).
E M ORE D ISCUSSIONS
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 optimize different objectives.
These are done without fine-tuning the pre-trained Transformer models.
Our framework makes a pre-trained Transformer more versatile and adapts it to different tasks by designing different reward functions used by the planning algorithm.
If we re-train or fine-tune a Transformer 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 to train or fine-tune Transformers for different tasks.
Potential negative social impacts.
Automatic code generation makes it easier for anyone to generate programs that meets a specification.
Our hope is that this development will relieve the burden of software engineers and enable AI-assisted programming or even end-to-end automatic programming.
However, it may also make it easier to develop malwares.
Anyone that can specify a malicious goal in a 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.