CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning
Abstract
Program synthesis or code generation aims to generate a program that satisfies a problem specification.
Recent approaches using large-scale pretrained language models (LMs) have shown promising results, yet they have some critical limitations.
In particular, they often follow a standard supervised fine-tuning procedure to train a code generation model from natural language problem descriptions and groundtruth programs only.
Such paradigm largely ignores some important but potentially useful signals in the problem specification such as unit tests, which thus results in poor performance when solving complex unseen coding tasks.
To address the limitations, we propose “CodeRL”, a new framework for program synthesis tasks through pretrained LMs and deep reinforcement learning (RL).
Specifically, during training, we treat the code-generating LM as an actor network, and introduce a critic network that is trained to predict the functional correctness of generated programs and provide dense feedback signals to the actor.
During inference, we introduce a new generation procedure with a critical sampling strategy that allows a model to automatically regenerate programs based on feedback from example unit tests and critic scores.
For the model backbones, we extended the encoder-decoder architecture of CodeT5 with enhanced learning objectives, larger model sizes and better pretraining data.
Our method not only achieves new SOTA results on the challenging APPS benchmark, but also shows strong zero-shot transfer capability with new SOTA results on the simpler MBPP benchmark.
Introduction
Considering program synthesis as a sequence-to-sequence task, pretrained language models (LMs)
[Hendrycks et al., 2021, Chen et al., 2021a, Austin et al., 2021] can be adapted to receive input sequence as problem specification in natural language and generate a sequence of codes as the output program (see Figure 1, right, for an example).
While these models achieve promising results, especially in basic programming tasks [Chen et al., 2021a, Austin et al., 2021], we observe that they still fail to generate codes to solve complex problems [Hendrycks et al., 2021, Li et al., 2022].
There are two main limitations.
First, current models are trained using a conventional next-token prediction (NTP) objective which maximizes the next ground-truth token likelihood.
As noted in NLP domains [Bengio et al., 2015, Ranzato et al., 2016], training models only with next-token prediction objective in a "teacher-forcing" manner often leads to accumulating errors during test time when tokens are generated by conditioning on previously sampled tokens, not the ground-truth tokens.
This issue becomes more serious in the domain of program synthesis, where token-matching scores such as BLEU [Papineni et al., 2002, Ren et al., 2020] are more appropriate in partial program synthesis tasks (i.e. code completion)
[Husain et al., 2019] but have failed to measure the functional correctness
Actor-Critic RL Finetuning with
Returns
Problem Specification
A string is a palindrome if it reads the same from the left to the right and from the right to the left….If
there is such a substring in $s$ that is not a palindrome, print the maximum length of such a substring….
Example Input and Output:
Input: ‘hannah’
Output:
Solution Program
Unit Tests Input: wuffuw
Output: 5 Input: iiiiiii
Output: 0…
[Figure]
of complete programs
[Hendrycks et al., 2021, Chen et al., 2021a].
Training only with NTP objective is hence, not ideal to tackle full program generation to solve programming problems.
Secondly, current models fail to utilize the potential meaningful signals from unit tests, which directly determine the model performance by the functional correctness of programs.
Current approaches neglect this important signal during model optimization as well as generation procedure.
During optimization, unit tests could be factored into learning objectives to match the final goal of generating semantically correct programs.
During inference, since unit tests are often parts of problem description (i.e. example unit tests), they are potentially powerful to further improve output programs.
To address the above issues, we introduce “CodeRL”, a new framework to improve pretrained LMs for program synthesis tasks through reinforcement learning (see Figure 1, left).
Specifically, we propose a training strategy that optimizes pretrained LMs for program synthesis tasks in an actor-critic approach [Konda and Tsitsiklis, 1999, Sutton et al., 1999].
We treat the pretrained LM as an actor network and synthetically sample sequences from this actor, including both correct and incorrect programs.
These program samples are passed to a critic model which is trained as an error predictor to assess the functional correctness of these samples.
We use the token-level hidden states extracted from the learned critic model to estimate the values/scores of output tokens of these synthetic samples.
The actor network is then finetuned on these synthetic samples weighted by their critic scores.
During inference, we introduce a new generation procedure that involves example unit tests and a critic model to filter and select sub-sequences.
These sub-sequences are utilized as seeds that condition the model to resample new tokens and obtain new output programs.
This approach allows the model to automatically refine output programs based on their functional correctness during test time.
We extend CodeT5 with better pretraining strategies as the foundation model for CodeRL.
Our comprehensive experiments show that our models can achieve SOTA performance on the challenging
APPS benchmark [Hendrycks et al., 2021].
Specifically, our models reach more than 2% pass@1 , 6% pass@5 , and 19% pass@1000 .
Since our RL method is model-agnostic, we apply it to various large-scale models and achieve consistent performance gains.
We further test its zero-shot transfer ability on a simpler MBPP benchmark [Austin et al., 2021], where it sets a new SOTA result of 63.0% pass@80 over a finetuned GPT-137B’s 61.4%.
We release the improved CodeT5-large ( 770 M) model which outperforms many pretrained LMs of much larger sizes.
Related Work 2.1 Program Synthesis
Program synthesis tasks can date back as early as the early adoption of machine learning research [Waldinger and Lee, 1969, Manna and Waldinger, 1971].
Earlier tasks include problem specifications in the form of input-output (IO) examples [Summers, 1977, Gulwani et al., 2012] and synthesis methods are limited to probabilistic approaches [Liang et al., 2010] or simple programming concepts [Joulin and Mikolov, 2015, Kurach et al., 2015].
As deep learning methods became popular, later approaches adopt neural models to induce output programs, assuming an inductive bias given a large number of program samples [Parisotto et al., 2016, Balog et al., 2016, Devlin et al., 2017].
More recently, we witnessed the emergence of program synthesis tasks in which output programs are extended to general-purpose programming languages [Yin and Neubig, 2017, Xu et al., 2018,
Chen et al., 2021a] and program specifications are fully described in natural English text [Hendrycks et al., 2021, Austin et al., 2021, Poesia et al., 2022].
These extensions have encouraged a rising number of applications of pretrained language models (LMs) to program synthesis to exploit the contextual representations learned from massive data of codes and natural languages [Feng et al., 2020, Clement et al., 2020, Wang et al., 2021, Wang and Komatsuzaki, 2021, Chen et al., 2022].
Recently, Nijkamp et al.
[2022] proposed a conversational program synthesis approach with large pretrained language models.
Despite impressive results in basic programming problems and initial commercial deployment 3 , existing models still perform poorly against complex problems such as those from programming competitions on Codeforces [Hendrycks et al., 2021, Li et al., 2022]. Program Synthesis in Visual Context. Another related line of research is program synthesis in computer vision domains such as images and videos. Early papers such as [Kulkarni et al., 2015,
Yang et al., 2015] introduce inverse graphics networks to infer visual properties such as pose, shape, and lighting, of visual objects.
Wu et al. [2017], Liu et al. [2019], Ellis et al. [2018] study the problem of image rendering, which transforms an image to structured and compact representations, i.e. scene programs .
Tian et al.
[2019] extends the prior work to render 3D shapes from images through shape programs , containing features to capture geometric and structural priors.
Ganin et al.
[2018] introduces an RL-based approach to render realistic images through high-level graphics programs.
Sun et al.
[2018]introduces program synthesis from demonstration synthetic videos to summarize the behaviors of the objects in the videos.
While this line of research has remarkable impacts on applications such as image/video editing, captioning, and extrapolating, these approaches are limited to programs of domain-specific languages defined for visual objects.
For instance, in [Sun et al., 2018], programming language contains basic functions for object perception, action, and control flows.
In our work, we focus on program synthesis from natural language problem specifications and the output programs are in generalpurpose languages such as Python.
This type of programming task can range from basic programming problems to competition-level programming tasks that require a high level of problem-solving skills.
Related to the program synthesis tasks are research domains of sequence generation, in which RL approaches have demonstrated remarkable achievements.
In these domains, RL approaches are used to exploit signals from non-differentiable metrics of the task at hand.
Earlier work such as [Ranzato et al., 2016] adopts this strategy with REINFORCE algorithm
[Williams, 1992] to directly optimize models for sequence-based test metrics such as BLEU [Papineni et al., 2002] and ROUGE
[Lin, 2004] scores for translation models.
In the same domain, Bahdanau et al.
[2016] introduced an actor-critic framework [Sutton, 1984, Konda and Tsitsiklis, 1999].
In visual captioning domains, Rennie et al.
[2017], Wang et al.
[2018] proposed to use RL to optimize image captioning models using variants of CIDEr scores [Vedantam et al., 2015].
Alternatively, Ren et al.
[2017] derived a new goal-oriented return estimate using visual-semantic embedding.
Johnson et al.
[2017], Trivedi et al.
[2021] introduce program generation as an auxiliary task to learn interpretable policies in question-answering and synthetic navigation tasks.
Different from prior domains, in program synthesis, Austin et al.
[2021], Chen et al. [2021a], Li et al.
[2022] demonstrated very low correlation between token-based similarity metrics and functional correctness of programs.
Hence, it is not trivial to define an appropriate optimization goal in this domain.
We propose to exploit unit test signals, which directly exhibit the functional correctness of programs, during both - model optimization and test-time generation stages.
More related to our work are RL-based program synthesis
[Guu et al., 2017, Bunel et al., 2018, Liang et al., 2018, Zhong et al., 2018] and execution-guided synthesis approaches [Ellis et al., 2019, Chen et al., 2021b].
However, these are limited to programming languages defined within a specific application domain only. Pretrained LM Problem Solution
Program Problem Actor-Critic RL Finetuning with Solution
Program Sampled program Returns Unit
Tests
Actor
Network
Critic
Network
Finetuned
LM
We
Another neural network is trained as a critic model to evaluate these synthetic samples based on their probabilities of passing tests.
The returns are estimated based on critic scores and finally factored into the learning objective L rl to finetune the actor LM network using synthetic samples.
Following a sequence-to-sequence approach, the program synthesis task contains a problem description as an input sequence D and an output sequence of program
ˆ W =
( ˆ w 1 , ..., ˆ w T ) , ˆ w t ∈ V 4 that can solve the problem.
The output at each decoding step t is a distribution over the vocabulary V , computed by the softmax function ˆ w t ∼ softmax(Linear( s t )) where s t is the contextual hidden state at decoding step t .
Conventionally, during train time, model parameters, θ , are learned by maximizing the likelihood of the ground-truth reference programs.
Denoting W =
( w 1 , ...w T ) as the ground-truth program, the objective is to minimize the cross-entropy loss:
[Equation]
where the conditional probability p θ is parameterized following the above softmax function.
During test time, models generate sequences of programs by autoregressively sampling token ˆ w t from the distribution p θ ( .
| ˆ w 1: t − 1 , D ) .
Models are evaluated against unit tests corresponding to the problem.
Each test includes a pair of input and ground-truth output.
In real-world program synthesis tasks [Hendrycks et al., 2021], example unit tests are often given as parts of the problem specification.
3.2 Pretraining Language Models on Code
We adopt Transformer models as the backbone of our program synthesis systems.
Specifically, this paper extends the CodeT5 model [Wang et al., 2021] as a foundation model for CodeRL.
CodeT5.
CodeT5 [Wang et al., 2021] is a multi-lingual code-aware language model pretrained on large-scale source code corpora curated from Github.
With a unified encoder-decoder architecture,
CodeT5 achieves state-of-the-art performance in a wide range of code intelligence tasks in the
CodeXGLUE benchmark [Lu et al., 2021] including both code understanding and generation tasks.
Improving Pretraining Data.
We enlarge the Python pretraining dataset using the recently released large-scale Github Code dataset 5 .
We have compiled public, non-personal information from GitHub consisting of permissively licensed Python code (“mit”, “apache-2”, “bsd-3-clause”, “bsd-2clause”, “cc0-1.0”, “unlicense”, “isc”).
The resulting Python dataset (GCPY) has 10.5B tokens and is 10x larger than the CodeSearchNet (CSN) corpus [Husain et al., 2019] used in the original [-CodeT
-]CodeT
[Wang et al., 2021].
https://huggingface.co/datasets/lvwerra/github-code
Improving Pretraining Objective.
While pretraining tasks in CodeT5 like masked span prediction (MSP) benefit code understanding tasks, they have a large discrepancy with program synthesis objectives.
To mitigate this gap, we introduce a pretraining task of next-token prediction (NTP) into
CodeT5.
Specifically, we uniformly sample a pivot location for each code sample, then pass the content preceding the pivot to the encoder and the remaining to the decoder.
To control the length of input and output sequences, we restrict the pivot within 10% to 90% of the original sequence.
3.3 Program Synthesis as an RL Problem
We propose to formulate the Program Synthesis as an RL problem and apply an actor-critic RL approach to improve the performance of a pretrained LM by exploiting the unit test signals in both model optimization (see Figure 2) and generation procedure (see Figure 3).
More formally, we can view the learned parameters of an LM model, θ as a stochastic policy , which decides an action as the prediction of each token.
Following each action, an LM model updates its hidden state representations which are used by the policy to determine the next action in the next decoding step.
At the end of the generation episode (i.e. an <endoftext> token is observed), the LM model receives a return r measured by the functional correctness of the generated program.
The goal of RL finetuning is to minimize the expected return:
[Equation]
where W s = ( w s 1 , ..., w sT ) is a synthetic sample in which each token w st is sampled by the LM model at decoding time step t .
Following the REINFORCE algorithm
[Williams, 1992, Sutton and Barto, 2018] and policy gradient theorem [Sutton et al., 1999] we can define an estimate of the gradient ∇ θ L ( θ ) of the non-differentiable return r as:
[Equation]
Defining Return by Unit Test Signals.
For each sample sequence W s , the return r can be defined heuristically by checking its functional correctness.
We pass generated programs together with the corresponding unit tests to a compiler.
From the outputs of the tests, we can determine the return r :
[Equation]
if W s cannot be compiled (i.e. compile error)
[Equation]
, if W
[Equation]
cannot be executed with unit tests
[Equation]
i.e. runtime error)
[Equation] [Equation]
any
[Equation] [Equation]
, if
[Equation]
passed all unit tests [Equation] [Equation] However, in related domains such as text-to-SQL research [Zhong et al., 2018, Xu et al., 2018], we note that this approach to estimate returns can lead to an unstable training process with high variance of the gradient estimate following Eq.
(3) with mini-batches in training.
Return with a Baseline.
In order to alleviate this variance, we adopt a “baseline” [Sutton and Barto, 2018].
Specifically, we use a greedy decoding strategy as a baseline and any generated samples that outperform this baseline are given positive return estimation, and negative return estimation otherwise.
This relative normalization technique allows models to explore imperfect programs, as long as their returns are better than the baseline’s.
Given a training sample, we denote the return of the baseline r ( W b ) and the expected gradient is computed as:
[Equation]
Note that at each decoding step t , our greedy decoding baseline is independent from the action w st and hence, the expected gradient term ∇ θ L rl ( θ ) from Eq.
(3) remains the same in Eq.
Intermediate Return by Critic as Error Predictor.
We observe that the above gradient estimate is only based on a final return at the end of the decoding process.
However, programs often follow fixed syntactical rules in which a single token such as an additional white-space character can render a program erroneous.
Therefore, Eq.
(8) becomes too restrictive. A straightforward solution [Figure]
Seed [Figure] [Figure] is to use token-based similarity scores [Papineni et al., 2002, Ren et al., 2020]) between each subsequence W s 1: t and the ground truth.
However, code matching is not an ideal return measure due to its poor correlation with program correctness [Hendrycks et al., 2021, Chen et al., 2021a, Austin et al., 2021] which can only be measured against fully complete programs.
Alternatively, we introduce a critic model (see Appendix A for an overview).
The critic model is parameterized as a neural network with parameters ϕ that receives inputs as the problem description D and a sampled program W s = { w s 1 , . . .
, w sT } .
The critic is trained to infer the unit test outcome; one of { CompileError , RuntimeError , FailedTest , PassedTest } as described in the return definitions in
Eq.
(4) to (7).
The training objective of the critic ϕ can be expressed as:
[Equation]
where u denotes the ground-truth unit test outcome given by the compiler.
We use Transformer models of smaller sizes than the actor model as the base architecture for the critic model.
The contextual hidden states of the program tokens ( h 1 , . .
. , h T ) obtained from the critic model decoder are max-pooled along the sequence length dimension h pool =
Pooling( h 1 , . .
. , h T ) .
The critic’s prediction on the unit test outcome is computed as ˆ u =
softmax(Linear( h pool )) .
Given a learned critic, we use the probability distribution ˆ v t =
softmax(Linear( h t )) to estimate the token-level value ˆ q of w st in relation to the ground-truth unit test output (note that we use the token level contextual representation h t here, before the pooling operation).
Specifically, ˆ q ϕ ( w st ) =
ˆ v t [ u ] where ˆ v [ . ] denotes the probability of a specific unit test outcome from the four possible ones.
We use this estimate to train the actor LM model with intermediate returns:
[Equation]
Generating Programs with Example Unit Tests and Critic.
We leverage the unit tests provided in the input problem description to improve the generation procedure during inference too (see Figure 3 for an overview and Appendix B for step-by-step explanation).
For each problem, we generate N programs, out of which we only select programs that pass example tests (leading to a set F ) and filter out the rest.
To improve sample quality, we perform another round of generation where we use sub-sequences from these filtered samples as prompts (or “seed” sequences) to the actor LM.
We employ a separate critic model ( ϕ test ) to guide our choice of sub-sequences from these filtered samples.
This critic model is trained with a similar objective as Eq.
(9), but in a binary classification setup with { FailedTest , PassedTest } labels.
Let W filter =
{ w 1 , . .
. , w T } denote a generated sample that passes the example unit tests.
We use the critic model to assign a value to each token ˆ q ϕ test ( w t ) =
p ϕ test (ˆ u = PassedTest | w 1: t , D ) corresponding to the critic’s predicted probability of the sub-sequence till t passing the unit tests.
We split the sequence at position t max corresponding to the highest critic assigned value and use the left split as the seed for the next stage.
If this seed sequence till t max contains a token with p ϕ test (FailedTest) > p ϕ test (PassedTest) , we further chop it at this token by removing tokens on the right.
This is done to pick prompts that are likely to generate successful programs in the next round.
We use these seeds to initialize and condition the (actor) LM to resample new tokens till we encounter the <endoftext> token.
In this round, each seed sequence can be stacked N/ |F| times for upsampling.
This results in the same number of output programs N .
We call this generation procedure “Critic
Sampling” (CS).
We use mini-batch generating to improve efficiency during inference and employ nucleus sampling with a batch size of N = 200 .
While we do incur additional costs to re-sample using the seed sequences, we are only required to generate partial programs in the re-generation stage, making this stage less expensive than conventional generating procedures.
Experiments 4.1 Experimental Setups and Datasets
Pretraining Setup.
We pretrain a CodeT5-large model (770M) from scratch following T5-large’s architecture [Raffel et al., 2020].
We follow the pretraining setups in CodeT5 [Wang et al., 2021] with the modifications as proposed in §3.2.
We evaluate this new pretrained CodeT5 model on
CodeXGLUE [Lu et al., 2021] and achieve new SOTA results (see Appendix D).
APPS Benchmark.
We choose the challenging APPS program synthesis benchmark [Hendrycks et al., 2021], as it has large coding problems of varying difficulties collected from multiple coding websites.
It includes training and test splits, each of which has 5000 samples of programming tasks with diverse levels of difficulty, including “Introductory”, “Interview”, and “Competition” levels.
Each sample includes 20 unit tests on average to validate the functional correctness of programs.
Finetuning Setup.
Due to the potential large number of trajectories (i.e. V T ) to generate a sequence and the unstable feedback loop between actor and critic [Lillicrap et al., 2015, Wang et al., 2018], we applied imitation learning to first warm-start a pretrained LM model with L ce only for up toepochs.
We then sampled sequences of programs from this actor network to train the critic while keeping the parameters of the actor network frozen.
For experiments with CodeT5 actor models, we use the CodeT5-small architecture for the critic model, and GPT2-small critic architecture when the actor models are GPT variants.
After training the critic, we then apply both L ce and L rl with equal weights to finetune the actor network.
Evaluation.
We follow [Hendrycks et al., 2021, Chen et al., 2021a] and evaluate the models using the pass@k metric, which is the percentage of problems solved by using k generated programs per problem.
We also follow Li et al.
[2022] and use n@k metric which only considers a subset of n candidates from k generated programs per problem.
The subset of n candidates are typically selected by a filtering method by passing generated programs through example tests given as part of the problem description [Chen et al., 2021a, Li et al., 2022].
For more details of the experimental setup, please refer to Appendix C.
4.2 Experimental Results on APPS
Baselines.
As reported by Hendrycks et al.
[2021], we compared our models with several baselines, including GPT2 [Radford et al., 2019],
GPT-Neo
[Black et al.], and GPT3 [Brown et al., 2020].
We also compare the results with Codex [Chen et al., 2021a] and AlphaCode
[Li et al., 2022].
Note that by default, results of pretrained LMs (except for Codex and GPT3) are from models finetuned on
APPS using the standard loss L ce only.
In our ablations, since CodeRL is model-agnostic, we can also integrate it with GPT variants such as GPT-J [Wang and Komatsuzaki, 2021] and GPT-Neo.
Overall Results.
Firstly, Table 1a shows that the CodeRL with the CodeT5 model can achieve significant performance gains, outperforming many pretrained LMs of much larger sizes.
Specifically, our approach achieved new SOTA results of 2 .
57 % pass@1 , 6 . 21 % pass@5 , and 19 .
36 % pass@1000 .
Table 1b shows that when evaluating on a subset of filtered code samples, our CodeRL+CodeT5 can achieve SOTA results of 7 . 83 % 1@k and 11 .
61 % 5@k .
Note that while CodeRL incurs additional computation cost during inference with CS, our approach only requires much lower k to achieve comparable performance with other models.
Specifically, with k = 1000 only, our model performance is as good as AlphaCode with a much larger generation budget of k = 50000 . [Table] [Table] [Table]
[Table]
[Table] [Table] 4.3 Ablation Studies
In this section [Table] for a fair comparison between variants of return estimates and learning objectives, we report the results of pass@k where k = { 1 , 5 } with beam search decoding.
For larger k , we report the results with and without the CS procedure.
Impacts of Return Estimates.
Table 2 show the results of CodeT5-770M trained by different approaches to estimate returns of code samples.
Overall, we report that the CodeRL objective with relative token-level return estimates by our critic model (Model D) can achieve the best performance on pass@1 and pass@5 .
Secondly, we note that using absolute returns without a baseline (Model B) could lead to the most performance drop, as this approach heavily penalizes all incorrect samples (even though they might still be better than a naive baseline).
Thirdly, without a critic model, simply assigning identical rewards to all tokens in a code sample (Model A) is disadvantageous as these return estimates are too restrictive to be used as feedback signals for RL training.
Finally, we experimented with a distance-based critic which assumes that token values decay linearly from t = 1 to t =
T (Model C).
The lower performance suggests the benefit of training a critic network to compute the returns rather than relying on rule-based approaches.
Impacts of Learning Objectives.
Table 3 shows the results with different combinations of L ce and L rl .
We experiment with using only L rl and note the problem of vanishing gradients during finetuning [Ranzato et al., 2016, Bahdanau et al., 2016].
Secondly, we note that by using only L ce for further finetuning, despite improvement in losses during training time, the model performance indeed degrades during test time.
We expect these models are overfitting to the training data.
Interestingly, a naive approach of L ce with synthetic samples W s , all of which are treated as correct codes with r ( W s )
= 1 , still leads to some performance improvement with GPT-Neo on pass@5 (but not in [Table]
[Table] other cases). Finally, we found that using both L ce and L rl results in a more consistent performance improvement overall on pass@1 and pass@5 for the GPT-Neo and CodeT5 models.
Impact of Critic Sampling .
Table 4 shows the ablation results of critical sampling (CS) during inference.
Overall, we found positive impact of CS for improving pass@200 and pass@1000 metrics.
In addition, the positive impacts of critic sampling on pass@1 and pass@5 are indicated by comparing the results of Table 1a and Table 2 (row D, in which we only used conventional beam search decoding without critic sampling).
We can observe that, using critic sampling, model performance increases from 2 . 2 % pass@1 ( 3 . 1 % pass@5 ) to 2 .
57 % pass@1 ( 6 . 21 % pass@5 ). Interestingly, from Table 4, we observe that CS does not provide a significant gain on the n@k metric.
Note that n @ k measures the solving rate among the subset F filtered from k samples.
As CS will technically increase the size of this subset, the n@k metric will consider an exponentially larger number of options of n samples than before.
This will normalize n @ k by a larger pool of n candidate set, resulting in less impact of CodeRL on the results.
We recommend additional post-processing steps such as candidate ranking [Cobbe et al., 2021] to improve the n@k performance.
Impacts of Pretraining Approaches for CodeT5.
Table 5 reports the results of CodeT5 with different configurations of model sizes, pretraining data, and pretraining objectives.
For a fair comparison, all models are only finetuned with L ce on APPS.
As observed in prior work [Chen et al., 2021a, Austin et al., 2021], scaling up the number of model parameters or the size of the pretraining data can significantly improve the model performance of downstream synthesis tasks.
We also find that enhancing the pretraining objectives with next token prediction (NTP) is vital for generation tasks, surpassing just masked span prediction (MSP) from the original CodeT5.
4.4 Zero-shot Evaluation on MBPP Benchmark
Finally, we test the zero-shot transfer ability of CodeRL on another smaller and simpler program synthesis benchmark MBPP [Austin et al., 2021].
[Table] [Table]
results of our
CodeRL+CodeT on MBPP benchmark compared with finetuned GPT models of up to 137B size.
Our CodeRL+CodeT5 (ZS) was trained on APPS and then evaluated on MBPP in a zero-shot setting.
We observe that CodeRL with CodeT5 of a much smaller model size yields surprisingly good zero-shot performance, setting a new SOTA result of 63.0% pass@ over GPT-137B’s 61.4% pass@80 .
This validates the strong zero-shot transfer ability of CodeRL for unseen tasks.
For additional experiments, analysis, and qualitative results, please refer to Appendix D.
Limitations and Broader Impacts
One limitation of our approach is the computation cost of training critic model to estimate returns in addition to the original LM (actor network).
However, in practice, we found that training a good critic model does not require large-scale models to attain a decent performance.
For instance, a finetuned critic model initialized from a pretrained GPT-2 (small) can achieve over 75% error prediction accuracy on synthetic samples.
Program synthesis can lead to substantial positive social impacts, e.g., transforming future software developing tools, increasing the productivity of developers, and improving the accessibility and quality of programming courses.
Yet, some risks and bias issues are still worth considering before deploying such models at scale.
For example, training data from public GitHub code repos may contain vulnerabilities and the resulting synthesis models may generate programs with weak security measures [Hammond Pearce et al., 2021].
Conclusion
We present CodeRL, a novel framework for program synthesis, using deep reinforcement learning to improve pretrained LMs, by exploiting unit test signals in both training and inference stages.
Specifically, we introduce an actor-critic training approach to optimize pretrained LMs with dense feedback signals on synthetic code samples.
During inference, we propose a new generation procedure with critical sampling, which enables the model to automatically regenerate programs based on feedback from unit tests and critic scores.
We integrate CodeRL with the improved CodeT5-large model (770M) and achieve new SOTA results on both the APPS and MBPP benchmarks, surpassing the prior SOTA by massive pretrained LMs of much larger model sizes.
Our comprehensive analysis shows that CodeRL achieved consistent improvement upon the conventional pretrained LMs for code generation tasks.
CodeRL is a general framework that integrates pretrained LMs and RL holistically for program synthesis, and can be extended and improved in various ways.
For example, it can be easily integrated with other better pretrained LMs and can be improved with more fine-grained feedback from the environment, such as feedback received from a static code analyzer.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] See Section 3 and(b)
Did you describe the limitations of your work?
[Yes] See Section 5.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section 5.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See Appendix D and code at https://github.com/salesforce/CodeRL (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] See Appendix C. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[No] As it is very expensive to experiment with large-scale language models, we did not try different random seeds due to the limitation of computation resources.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See the configurations in Appendix C. 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] See Appendix C. (b) Did you mention the license of the assets?
[Yes] See Appendix C. (c) Did you include any new assets either in the supplemental material or as a URL?
[No] We do not curate any new dataset in this paper.
We will release the code and models.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] All datasets evaluated in our experiments are publicly available for use.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A]
The data we are using are code samples from public programming competitions which do not include personally identifiable information or offensive content. If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Overview of Critic Model
Figure 4 shows an overview of our critic model. In our CodeRL framework, besides the actor LM network θ , we introduce a critic model that is trained as an error predictor and parameterized as a neural network with parameters ϕ . The critic receives as inputs a problem description D and a corresponding synthetic program W S sampled from the actor. The critic is required to output a prediction of the unit test outcome of the input program.
We define 4 possible outcomes u :
CompileError , RuntimeError , FailedTest , and PassedTest .
The critic model is trained by minimizing the following loss:
[Equation]
The ground-truth outcome of a synthetic sample is obtained by passing it to the unit tests corresponding to the problem.
Note that since our critic model is applied in a supervised learning environment with available ground truth, we also use the training samples from the original dataset with ground truth output u = PassedTest to train the critic.
The learned hidden state representations of program tokens when passed through the critic are then used to measure their return estimates for our RL optimization objective.
The return estimates are incorporated as intermediate returns at decoding steps to compute the expected gradient of the actor network ∇ θ L rl ( θ ) .
B Critic Sampling Procedure
Refer to Algorithm 1 for a step-by-step explanation of our critic sampling procedure.
[Figure]
C Additional Experimental Setup Details
Pretraining Setup.
For CodeT5, we adopt the code-specific tokenizer as described by Wang et al.
Note that we employ 6 programming languages (PLs) in CodeSearchNet
[Husain et al., 2019] (CSN) instead of 8 PLs in CodeT5 as C/C
# datasets are not publicly available.
We employ only the pretraining task of masked span prediction (MSP) in CodeT5 and hence, we do not have to parse programs into abstract syntax trees (ASTs) to obtain the identifier information.
This preprocessing step was required in other original pretraining tasks like masked identifier prediction [Wang et al., 2021].
To further speed up training, we concatenate data samples to batch size 512 for pretraining with MSP and the resulting number of tokens is
1.1B.
APPS Benchmark.
We follow the same preprocessing step as in [Hendrycks et al., 2021] to formulate the input sequences from problem descriptions.
APPS consists of 10,000 coding problems with a 50-50 train-test split.
Each problem is accompanied by 23.2 correct Python programs and 21.2 unit tests on average.
The average length per problem is 293.2 words and the average length per program is 18.0 lines.
The dataset is categorized into three levels of difficulty: Introductory (3639, train/test=2639/100), Interview (5000, train/test=2000/3000), and Competition (1361, train/test=361/1000).
Similarly to [Hendrycks et al., 2021], we employ the strict accuracy metric to evaluate the functional correctness of a program, where it is counted as correct if it can pass all the unit tests corresponding to the problem.
MBPP Benchmark.
We additionally include another smaller and simpler Python program synthesis dataset called MBPP [Austin et al., 2021]
(Mostly Basic Programming Problems) for evaluation.
The dataset contains 974 instances with 374/90/500 instances for training/validation/testing respectively and 10 reserved for few-shot learning.
The problems are typically short, usually one sentence of natural language descriptions each.
Each problem is accompanied by 1 correct solution (6.8 lines of code on average) and 3 unit tests in the form of assert statements for validating the functional correctness.
Unlike APPS, unit tests in MBPP are not hidden and are explicitly incorporated into the source sequences for program synthesis models.
This might encourage models to be overfitting to these assert statements via hard-coding an if-expression very occasionally.
However, for a fair comparison with the baselines, we construct the source sequences in the same way as in prior work.
Specifically, we adopt the same prompt format as in [Austin et al., 2021] to prepare the input sequence as: problem descriptions + “Your code should satisfy these tests:” + 3 assert statements.
Finetuning Setup.
Following [Bahdanau et al., 2016], since our RL method is applied in a supervised learning task, in addition to synthetic programs, we also use the ground-truth programs of training samples to train the critic.
These samples are considered perfect programs and always have a label of PassedTest .
To optimize the LM actor network, in practice, following previous work
Input: Problem D , Language model θ , Critic model ϕ
Output: A set of N generated solution programs 1 begin /* step 1: generate program solutions by LM */ a set of output solutions S ←− ∅ for N times do /* use sampling decoding to generate programs */
Program ˆ W i ←−
Sampling p θ ( W i | D )
←− Sampling p θ ( w it | w i 1: t − 1 , D )
S ←− S ∪ { ˆ W i } /
* step 2: extract example unit tests from problem */ example unit tests I ←− ExtractExampleInputOutput
( D ) /
* step 3: filter for passed programs and pass to the critic model to select sub-sequences */ a set of filtered programs F ← ∅ for each generated program ˆ W i ∈ S do
Testoutcomes u
W i , I ) samples return all current programs for evaluation */
a do
PassedTest |
4: subsequence as seeds to regenerate programs
/ length of subsequence m
|
N times do
Sampling p
| W subi
D )
Sampling p θ ( w jt | w jm t − 1 , W subi , D )
S ′
S
W j } / return the regenerated programs for evaluation
/ return S ′
[Bahdanau et al., 2016, Rennie et al., 2017, Wang et al., 2018], in each training optimization step, we can simply approximate the expected gradient with a single sample W s ∼ p θ :
[Equation]
Configurations.
For pretraining, we perform our experiments on a Kubernetes with 16 A100-40G
GPUs on Google Cloud Platform and the total pretraining duration is around 21 days.
In the first pretraining stage with MSP, we employ a corruption rate of 15%, a peak learning rate (LR) of 2e-4, and a batch size of 2048 .
We pretrain on CSN for 150 epochs ( 10 days) and then on GCPY for 10 epochs ( 5 days).
For the second stage pretraining with NTP, we adopt a peak LR of 1e-4 and a batch size of 256 , and pretrain for 10 epochs (6 days).
We set the maximum length to 768 andfor source and target sequences respectively for this objective.
For all experiments, we employ an
AdamW optimizer [Loshchilov and Hutter, 2019] with a 0 .
05 weight decay and a linear decay LR scheduler with a warmup step of 1000 .
For finetuning on APPS, we adopt a batch size of 64 and warmup LR from 0 to 2e-5 for the firststeps and polynomially (power= 0 .
5 ) decay to 1e-5 until the end of 10 epochs, which takes around
[Table]
hours on one A100 GPU.
We set the maximum source and target sequence length to 600 andrespectively.
To train the critic model (either a GPT2-small or a CodeT5-base), we train the models with synthetic sampled programs for up to 10 epochs.
We found that in this phase, the critic model usually converges quite earlier than the 10-epoch training limit.
Once the critic model is trained, we continue to finetune the actor language model for 10 more epochs using both the conventional NTP loss and RL-based loss.
For MBPP, in finetuning experiments with CodeT5, due to its small training set, we finetune the models for 60 epochs with a constant LR of 2e-5 and a batch size of 32 , which takes less thanmins on one A100.
We set its maximum source and target length to 382 and 306 respectively.
Critic Sampling.
Note that in CS, many rounds of regeneration can be run to keep improving the output programs.
However, in practice, we found that one round of regeneration is sufficient to obtain good programs.
In the initial generation, if no samples pass example unit tests, we do not apply critic sampling and simply use all available programs for evaluation.
D Additional Experimental Results D.1 CodeXGLUE Benchmark Results
To validate the effectiveness of our simplified pretraining strategies of CodeT5-large, we extensively evaluate it on a variety of generation tasks in CodeXGLUE [Lu et al., 2021], including code-to-text generation (i.e. summarization, see Table 7), text-to-code generation (see Table 8), and code-to-code generation (i.e., code translation and code refinement, see Table 9).
Different from APPS [Hendrycks et al., 2021] and MBPP [Austin et al., 2021], we follow the default similarity-based evaluation metrics in the CodeXGLUE benchmark, including BLEU
[Papineni et al., 2002] and CodeBLEU [Ren et al., 2020], and exact match (EM) scores.
Table 7, 8, and 9 show that our simplified pretrained
CodeT5-large sets new SOTA results on a large majority of the tasks, and hence, can be served as a better foundation model for other code-related generation tasks.
Note that in these experiments, we employ the conventional finetuning objective with L ce and there might be potential to improve the performance further with our CodeRL framework.
D.2 MBPP Benchmark Results
Following Austin et al.
[2021], we adopt temperature sampling to generate multiple candidate solutions.
We empirically find that CodeT5 benefits from a higher temperature of 1 . 2 (less greedy decoding or more diverse) than their GPT’s temperature of 0 . 5 on this benchmark. Table 10 reports the pass@80 and pass@1000 results for both finetuned and zero-shot settings. For baselines, we compared with GPT models with sizes ranging from 224M to 137B [Austin et al., 2021], which are pretrained on 2.93B web documents (13.8M containing source code) using standard language modeling objective.
Results of GPT models are obtained from the original authors.
From the comparison among various CodeT5 variants, we again confirm that larger model sizes and pretraining data, and better pretraining objective of NTP all lead to a performance boost.
Particularly,
[Table] [Table]
our CodeT5-770M yields a pass@80 of 46.8%, surpassing GPT-8B’s 40.6% with a much smaller model size.
In addition, we find CodeT5 models finetuned on APPS can achieve a surprisingly good zero-shot performance on MBPP with a pass@80 of 60.2% and further improved to 63.0% with the help of CodeRL, which even outperforms the largest GPT-137B’s performance of 61.4%.
This indicates APPS is a comprehensive program synthesis benchmark and CodeT5+CodeRL models trained on it are able to generalize to other simpler coding tasks.
If we further increase the budget of attempts up to 1000, all models witness a consistent and significant boost in solving rate, especially our CodeT5+CodeRL yielding a new SOTA result of 81.8% pass@1000 .
A common concern about transfer learning is that the source (APPS) and target (MBPP) tasks might overlap in their training data, which could result in the source model tending to memorize these substantially similar data when applied to the target task.
To address this concern, we analyze how many lines of code appear in both the training set of APPS and the programs of MBPP following Austin et al.
For this analysis, we discard code comments and normalize the whitespaces for each line, and then exclude lines that appear more than twice anywhere in MBPP, as these are likely to be common Python keywords such as return and break .
Figure 5 illustrates the number of absolute duplicated lines (Left) and the relative fraction of duplicated lines (Right) in the MBPP programs.
As can be seen, the overlap between APPS and MBPP seems to be minimal.
Only 12.6% MBPP programs have more than half of their lines matched somewhere in the APPS training data. Besides, more than half (514 out of 974) of programs have a zero overlap and 90.9% have only no more than 3 lines overlapping with the APPS training set. If we further require the lines to be consecutive, there are no more than 2 consecutive duplicated lines. D.3 APPS Benchmark Results on Competition-level Tasks
Figure 6 shows the results of pass@k and n@k with k ranging from 1 to 200 and n = { 1 , 5 } , for
CodeRL+CodeT5 and CodeT5 only.
We choose to investigate a subset of the APPS test split, which contains the test samples of the highest difficulty level (i.e. competition programming tasks).
Since
CodeRL is model-agnostic, we also integrate it with GPT-J [Wang and Komatsuzaki, 2021] and report
[Table]
the results.
To focus on the impact of our RL optimization, during test time, we compare models using only nucleus sampling and without the CS procedure.
Figure 6 shows that the performance gains are quite consistent on both GPT-J and CodeT5.
In particular, as k increases, the performance gain of CodeRL is more significant on both GPT-J and CodeT5 models.
We attribute these gains to the CodeRL learning objective L rl that encourages models to explore code solutions drawn from the model’s sampling distribution.
During test time with an increasing k sampling budget, models are allowed to generate diverse code solutions and the impact of L rl becomes more significant.
D.4 CodeT5 Ablation by Training Epochs
Figure 7 shows the performance of CodeT5 model variants by finetuning epochs and by difficulty levels of programming tasks.
Note that in these experiments, we only need to compare among
CodeT5 model variants by pretraining strategies, and hence, only engage L ce (imitation learning) in the finetuning stage on APPS.
Consistent with our prior analysis (see Section 4 of the main paper), enhancing both pretraining data (with larger data of GCPY) and pretraining objectives (with NTP objective) improves model performance across training epochs in general.
Moreover, as noted by our analysis of learning objectives, only using L ce often leads to overfitting performance, typically after epoch 10 in our case.
Hence, to further finetune large-scale LMs, we recommend adopting our RL objective
L rl to utilize synthetic training samples and avoid overfitting models.
D.5 Impacts of CodeRL on Program Quality by Unit Test Signals
Figure 8 demonstrates the average percentages of generated programs by their test signals.
Specifically, we use CodeT5 or CodeRL+CodeT5 to generate programs and randomly select 200 generated programs per test sample in the APPS test split.
We pass programs to either example unit tests or hidden unit tests corresponding to the problem and group the output programs by their output signals, including CompileError, RuntimeError, FailedTest, and PassedTest.
We observe that integrating
CodeRL can increase the likelihood that a program can pass unit tests and reduces the likelihood that it fails one or more unit tests, or whether it contains compiling or runtime errors.
However, we note that there are significant gaps in performance by test signals between example unit tests and hidden unit tests.
This observation suggests that example tests are not as comprehensive as hidden tests and hence, applying our CS procedure might lead to false positive samples for regeneration.
We recommend additional methods to improve example unit tests, such as through data augmentation by mutating input/output pairs [Austin et al., 2021].
D.6 Example Generated Programs and Qualitative Analysis
Figure 9 to 11 show examples of programming problems from the APPS benchmark and corresponding programs generated by CodeT5 variants.
Specifically, based on the same foundation pretrained
CodeT5 (pretrained with GCPY data and NTP objective), we compare the CodeT5 model that is finetuned by L ce only and another model that follows our CodeRL framework.
In CodeRL+CodeT5, we further show programs before and after applying the CS procedure.
The example programs show that applying CodeRL can improve the quality of generated programs and the CS procedure further improves the functional correctness of the programs.
For instance, in Figure 11, CodeT5 model misunderstands the problem and focuses on finding the greatest common divisor between a and b only.
Instead, the CodeRL model avoids this mistake and tackles the problem based on the factorials of a and b .
In Figure 10, we note that the CS procedure improves the program by reordering the elif code blocks.
The resulting program is more correct and is able to pass all hidden unit tests.
We also found that CodeRL can improve the efficiency of the programs, an important quality in complex programming problems.
For instance, in the interview-level programs in Figure 11, we note that without applying CS, the generated program is functionally correct but fails during execution due to a timeout error.
Applying the CS procedure can condition models on parts of the prior program and (re)generates new tokens to produce a more efficient program.
Hence, the resulting final program is able to pass all hidden unit tests (including tests with extremely large values) without timeout errors.
Input: 3 4 5 Output:
QUESTION:
right
=
BC and CA , find right is an integer. -----Constraints----- - 1 \leq |AB|,|BC|,|CA| \leq- All values in input are integers.
- The area the triangle ABC is an integer.
-----Input-----Input is given from Standard Input in the following format: |AB| |BC| |CA| -----
[Figure]
-----Print the area of the triangle ABC.
[Figure] [Figure]
-----Sample
Output----
[Figure]
This triangle has an area of 6.
Use Standard Input format ANSWER: Input Problem CodeT5 Output Program
CodeRL+CodeT5 Output Program (after CS)
CodeRL+CodeT5 Output Program (before CS)
Example Unit Tests
Ground-truth Program
Input: 5 12 13 Output:
Input: 24 7 25 Output:...
Hidden Unit Tests
PASSED TESTS FAILED TESTS [Figure] [Figure]
[Table]
[Figure]
[Table]
Use
[Table]
format
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
Example
[Table] [Figure] [Table]
Example
[Table] programs and refine the code. Specifically, we observe the model can simply reorder the elif blocks between line 11 and 15 to fix the error. The resulting program is functionally correct and passes all hidden tests.
[Table]
QUESTION:Holidays have finished. Thanks to the help of the hacker Leha, Noora managed to enter the university of her dreams which is located in a town Pavlopolis. It's well known that universities provide students with dormitory for the period of university studies. Consequently Noora had to leave Vičkopolis and move to Pavlopolis.
Thus Leha was left completely alone in a quiet town Vičkopolis. He almost even fell into a depression from boredom! Leha came up with a task for himself to relax a little. He chooses two integers A and B and then calculates the greatest common divisor of integers "A factorial" and "B factorial".
Formally the hacker wants to find out GCD(A!,
B!). It's well known that the factorial of an integer x is a product of all positive integers less than or equal to x. Thus x! = 1·2·3·...·(x - 1)·x. For example 4!
{+= 1·2·3·
 = 24. Recall that GCD(x, y) is+} the largest positive integer q that divides (without a remainder) both x and y. Leha has learned how to solve this task very effective.
You are able to cope with it not worse, aren't you? -----Input-----
The first and single line contains two integers A and B (1 ≤ A, B ≤ 10^9, min(A, B) ≤ 12).
-----Output-----
Print a single integer denoting the greatest common divisor of integers A! and B!.
{+-----Example----- Input
 Output -----Note----- Consider the sample.+}
{+= 1·2·3·
 = 24. = 1·2·3 = 6. The greatest common divisor of integers 24 and 6 is exactly 6.+}
Use Standard Input format ANSWER: [Table] CodeRL+CodeT5 Output Program (after CS)
CodeRL+CodeT5 Output Program (before CS)
[Table] {+Input: 
 [Figure]+}