Compilable Neural Code Generation with Compiler Feedback

As shown in Figure 2(a), we adopt CodeGPT as the generator, which uses GPT-2 Radford et al. (2019) as the starting point and is continually pre-trained on the large-scale code corpus. Our generator is then fine-tuned on the target task to minimize the cross-entropy loss:

$$\mathcal{L}_{\rm G}=-\frac{1}{|\mathcal{M}|}\sum^{|\mathcal{M}|}_{i}\sum^{|\mathcal{V}|}_{j}\ Y_{ij}\ {\rm log}\ P_{ij}\,,$$

(2)

where $\mathcal{M}$ denotes the set of the generated code tokens, $\mathcal{V}$ represents the vocabulary, $Y_{ij}$ denotes the label of the code token $i$ in class $j$, and $P_{ij}$ is the predicted probability of token $i$ in class $j$.

During training, the generator takes $x=\{\texttt{<BOS>},c,\texttt{<EOS>}\}$ as the input in the code completion task, and $x=\{d,\texttt{<BOS>},c,\texttt{<EOS>}\}$ as input in the text-to-code generation task, correspondingly. Special tokens <BOS> and <EOS> indicate the start and end symbols of code sequences. After several epochs of supervised fine-tuning on the target task dataset, we save the trained generator, which will be used in the next stage.

Reinforcement Learning (RL) is a method of learning the optimal policy by obtaining reward signals from the real environment Sutton and Barto (1998); Wan et al. (2018). As shown in Figure 2(b), we use the fine-tuned generator $\rho$ (after Stage 1) as the reference model. Then we initialize a policy $\pi=\rho$. Given an input sequence $s\in\mathcal{S}$, our goal is to find a policy $\pi$ that generates an output sequence $t\in\mathcal{T}$ with the objective of maximizing the compilability-based reward. We use RL (specifically PPO2 version of Proximal Policy Optimization Schulman et al. (2017)) to directly optimize the expected reward as:

$$\mathbb{E}_{\pi}\left[r\right]=\mathbb{E}_{s\sim\mathcal{S},t\sim\pi(.|s)}\left[r(s,t)\right],$$

(3)

where the policy $\pi$ is rewarded by the compiler (Eq. 1), $r$ is the reward function. We define $r(s,t)=1.0$ iff the code can be compiled by the program compiler and $r(s,t)=-1.0$ otherwise.

It is worth mentioning that code compilability constraints can be strong or weak. Strong constraint is defined that a long piece of code snippet may not be correctly compiled if a certain token is changed. And weak constraint means a blank string consisting of whitespace characters can be correctly compiled by the compiler. Concretely, in the text-to-code generation task, if the generator generates a string composed of whitespace characters, the compiler will consider it as a good case. In the code completion task, if the previous code snippet is compilable, the generator can fool the compiler easily. The RL is good at making use of this, resulting in the generated code can be compiled, but seriously deviating from the generation likelihood objective.

To avoid active model $\pi$ being too far away from reference model $\rho$, we add a Kullback-Leibler (KL) penalty with expectation, e.g., $\beta$KL$(\pi,\rho)$ Ziegler et al. (2019). Therefore, the modified reward will be reformulated as follows:

$$r(s,t)=r(s,t)-\beta\ {\rm log}\frac{\pi(t|s)}{\rho(t|s)}\,,$$

(4)

where $\beta$ is a constant, which plays the role of an entropy bonus, preventing the policy from moving too far from the range where $r$ is valid.

To alleviate the imbalance between the reward term and the KL penalty term and improve the stability of training, we use autoregressive fine-tuning (Causal Language Modeling) Radford et al. (2019) to make the KL penalty term fluctuate within a small range after RL training. This fine-tuning process incorporates a compilability-aware discriminator that will be introduced in the next stage.

Figure 3 shows an example of code completion. We mask the last five tokens of a Python function and ask the generator to complete them. The generator generates five candidates with high probabilities. Some minor mistakes prevent four of them from being successfully compiled. We hope the generator can have more perception power to explicitly distinguish compilable and non-compilable code generated by itself. Therefore, at this stage, we design a compilability-aware discriminator to deal with this issue.

Concretely, we add a discriminator (a two-layer MLP equipped with the ${\rm tanh}$ activation function between layers) after the final hidden layer of the generator. As shown in Figure 2(c), given the input sequence ($s$), we perform beam search on the generator to generate top-$k$ candidates ($t$). Each entire code $c\in\mathcal{Q}$ ($c=[s;t]$ in the code completion task) is labeled by the program compiler as positive (1) or negative (0), depending on whether it can be successfully compiled (see Eq. 1).

We use the hidden representation of the last token (<EOS>) as the final representation of the entire code $c$. Finally, the hidden representation of the last token (<EOS>) is fed into the discriminator for prediction:

	$\displaystyle h_{\rm\texttt{<EOS>}}$	$\displaystyle=$	$\displaystyle{\rm CodeGPT}(s,t)\,,$		(5)
	$\displaystyle h^{{}^{\prime}}_{\rm\texttt{<EOS>}}$	$\displaystyle=$	$\displaystyle{\rm Discriminator}(h_{\rm\texttt{<EOS>}})\,,$		(6)
	$\displaystyle P(.|t,s)$	$\displaystyle=$	$\displaystyle{\rm softmax}(h^{{}^{\prime}}_{\rm\texttt{<EOS>}})\,,$		(7)

where $h_{\rm\texttt{<EOS>}}$ denotes the representation of the last token <EOS>. The training loss of the discrimination process can be defined as:

	$\displaystyle\mathcal{L}_{\rm D}$	$\displaystyle=-\frac{1}{|\mathcal{Q}^{+}\cup\mathcal{Q}^{-}|}\left[\sum_{c\in\mathcal{Q}^{+}}\log\ P{(1|t,s)}\right.$		(8)
		$\displaystyle\left.{+\sum_{c\in\mathcal{Q}^{-}}\log\ P{(0|t,s)}}\right]\,,$

where $\mathcal{Q}^{+}$ and $\mathcal{Q}^{-}$ represent positive and negative sets respectively. The parameters of the generator and discriminator will be jointly updated.

At this stage, we jointly train the generator and discriminator, including a generating objective (to learn the generator only) and a discriminating objective (to learn the generator and discriminator together), as shown in Figure 2(c). The joint training loss is defined as follows:

$$\mathcal{L}=\mathcal{L}_{\rm G}+\mathcal{L}_{\rm D}\,.$$

(9)

We conduct experiments on two tasks: code completion and text-to-code generation. To investigate the compilability of the generated code, we need to preserve the indentation and newline operations in code. We also need to make sure that the code and its version belong to the scope of the compiler. Existing datasets on both of the two tasks usually do not serve these considerations. For convenience, we choose Python for experiments, as it is very popular and used in many projects. We conduct all experiments based on Python 3 environment and adopt the codeop¹¹1https://docs.python.org/3.6/library/codeop.html module to simulate the program compiler. We remove code that could not be compiled correctly by the compiler.

To evaluate the quality of the generated code, we adopt two widely-used evaluation metrics: Levenshtein Edit Similarity (ES) Svyatkovskiy et al. (2020); Lu et al. (2021) and Compilation Rate (CR) Kulal et al. (2019). Levenshtein Edit Similarity measures the number of single-character edits required to transform one string into another. It is a critical evaluation metric for the code generation scenario, as it measures the effort required for the developer to correct the code. Compilation Rate measures how many code can be correctly compiled by the program compiler. For both of these metrics, bigger values indicate better performance.

We compare our approach with various state-of-the-art models in the code completion task and the text-to-code generation task:

• •

  BiLSTM is a Seq2Seq model based on Bidirectional LSTM with an attention mechanism Luong et al. (2015).

• •

  Transformer Vaswani et al. (2017) is the base architecture of CodeGPT. We use 6-layer Transformer decoder to conduct experiments.

• •

  GPT-2 Radford et al. (2019) is an auto-regressive pre-trained model trained on large-scale text corpus.

• •

  CodeGPT Lu et al. (2021) is pre-trained with source code corpus on the basis of GPT-2 vis causal language modeling objective Radford et al. (2019).

• •

  PLBART Ahmad et al. (2021) is based on the BART Lewis et al. (2020) architecture, which is pre-trained on large-scale Java and Python corpora via denoising autoencoding.

• •

  CodeT5 Wang et al. (2021b) is based on the T5 Raffel et al. (2020) architecture, which employs denoising sequence-to-sequence pre-training on multiple programming languages.

In the code completion task, we set the learning rate as $1.5e\text{-}5$, the batch size as 32, the maximum fine-tuning epoch as 20, the maximum code sequence length as 96. We mask different numbers of code tokens (25, 30, 35, 40, and 45) and ask the model to complete them. We set the minimum generation length as 25, 30, 35, 40, and 45 accordingly. In the text-to-code generation task, we set the learning rate as $1.5e\text{-}5$, the batch size as 16, the maximum fine-tuning epoch as 20, the maximum text and code sequence length as 32 and 170. We set the minimum generation length as 96 (the generated code is slightly shorter than the ground-truth is allowed). In these two tasks, the generated sequence consisting of whitespace characters will be considered as a bad case.

We use the Adam optimizer to update model parameters. We train our model on the basis of CodeGPT checkpoint²²2https://huggingface.co/microsoft/CodeGPT-small-py-adaptedGPT2. Our model is trained on 2 NVIDIA Tesla V100 with 32GB memory. We employ the same tokenizer as CodeGPT. To train the policy $\pi$, we use the PPO2 version of Proximal Policy Optimization Schulman et al. (2017). In each epoch, we only randomly select 5% training data for the stability of RL training (Stage 2). In other stages (Stages 1 and 3), we use the full training data. To generate candidates (at Stage 3), we set the beam size as 5 in beam search. For efficiency, we update the candidates every 5 epochs.

Table 1 shows the results of the code completion task. We mask 25 tokens at the tail of code functions and ask the generation model to complete. We can observe that: (1) The code generated by existing autoregressive models has a low Compilation Rate. CodeGPT and GPT-2 only achieve 46.84 and 43.26 scores respectively on the Compilation Rate, which means that more than half of the code generated by them cannot be correctly compiled by the program compiler. (2) CompCoder significantly improves the Compilation Rate. It obtains 94.48 scores on the Compilation Rate, which is 47.64 points higher than the closest competitor (CodeGPT). (3) When our approach significantly improves the Compilation Rate, it does not sacrifice the fluency of the generated code. CompCoder obtains a comparable and even slightly better Edit Similarity score than other baselines, indicating that it effectively preserves the code fluency.

Figure 4 presents more results of the code completion task in the setting of completing 30, 35, 40, and 45 tokens. CompCoder still effectively improves the Compilation Rate when generating longer code. As the completion length increases, our approach outperforms CodeGPT by 49.66, 47.68, 46.64, and 33.36 points in the setting of completing 30, 35, 40, and 45 tokens, respectively. On average, our approach outperforms CodeGPT by 45 points across a different number of tokens for the task of code completion.

Table 2 presents the results of the text-to-code generation task. We could see that: (1) CompCoder significantly outperforms all other models w.r.t. the Compilation Rate. E.g., CompCoder achieves 23.1 points and 24.3 points improvements when compared with PLBART and CodeT5 respectively. (2) Compared to code completion task (Table 1), all models in the text-to-code generation task have relatively higher Compilation Rate. One of the main reasons we think may be: code completion requires the generated code to be constrained by the existing (previous) code, which is a much stronger restriction than the text-to-code generation.

We compare several simplified versions of our model to understand contributions of different components, including the Reinforcement Learning (RL) component and the discriminator’s effect for both model training (D${}_{\text{train}}$) and model inference (D${}_{\text{test}}$). As a case study, we take the code completion task as an example in the setting of completing 25 tokens and present the results in Table 3.

Several meaningful observations can be drawn: First, both RL (Row 2) and D${}_{\text{train}}$ (Row 3) effectively increase the code Compilation Rate of the generation model (CodeGPT in Row 1), which confirms that the two components we designed can indeed improve the ability of the generator for compilable code generation. Second, applying RL and D${}_{\text{train}}$ together (Row 4) further improves the Compilation Rate over their individual contributions. Third, using the discriminator to select the output during model inference stage (D${}_{\text{test}}$) is beneficial. It further boosts the Compilation Rate of vanilla “D${}_{\text{train}}$” by 17.08% (Row 5 v.s. Row 2) and boosts “RL+${\text{D}}_{\text{train}}$” by 11.34% (Row 6 v.s. Row 4). Forth, these three components (RL, D${}_{\text{train}}$, D${}_{\text{test}}$) that effectively improve the Compilation Rate do not compromise the generation capability measured by the Edit Similarity.

To better understand the effectiveness of our proposed approach, we present two cases for code completion and text-to-code generation tasks respectively. For both CodeGPT and CompCoder, we present top-1 result in Figure 5. For code completion, we observe that CodeGPT can not complete code with high quality (non-compilable), while CompCoder can complete the code well, and it is exactly the same for the reference solution. For text-to-code generation, we observe that although both models can not generate exactly the same code as the reference solution, CompCoder generates a compilable code at the function level. These results reveal the effectiveness of our proposed approach for compilable code generation.