Better Language Models of Code through Self-Improvement
Year: 2023

Abstract:
Pre-trained language models for code (PLMCs) have gained attention in recent research. These models are pre-trained on large-scale datasets using multi-modal objectives. However, finetuning them requires extensive supervision and is limited by the size of the dataset provided. We aim to improve this issue by proposing a data augmentation framework using knowledge distillation. Our framework utilizes knowledge gained during the pre-training and finetuning stage to augment training data, which is then used for the next step. We incorporate this framework into the state-of-the-art language models, such as CodeT5, CodeBERT, and UnixCoder. The results show that our framework significantly improves PLMCs’ performance in sequence-generation tasks, such as code summarization and code generation in the CodeXGLUE benchmark.

heading: 1 Introduction
text: Pre-trained models for code (PLMCs), such as CodeBERT (Feng et al., 2020), PLBART (Ahmad et al., 2021), CodeT5 (Wang et al., 2021), UniXCoder (Guo et al., 2022), and DISCO (Ding et al., 2022), have become the foundation to solve many practical software engineering tasks such as code summarization, code translation, program repair. Those PLMCs, like large language models (LLMs), are typically first pretrained on very large-scale datasets with a variety of multi-modal objectives under a self-supervised training style. They can then be fine-tuned using task-specific datasets in a supervised training style.
We hypothesise that, while fine-tuned models may not achieve peak performance, PLMCs can produce reasonable outputs that can be regarded as high quality data because they have been pretrained on large scale datasets, and that such data
∗Equal contribution. Listing order is based on the alphabetical ordering of author surnames.
can be leveraged as additional high-quality training data. Our framework utilizes the self-improvement capability of PLMCs through an simple data augmentation step. This approach is particularly useful for tasks involving code-related sequence generation, such as code summarization and code generation. Our method involves fine-tuning a PLMC on a downstream dataset, allowing the model to gain knowledge about the task. The model then generates an augmented version of the original training data, which are used to further fine-tuning. Our framework is similar to sequence-level knowledge distillation (Kim and Rush, 2016), but our approach focuses on improving model performance without compressing the model by utilizing the same technique.
Our empirical evaluation results show that our framework significantly improves the state-of-thearts PLMCs, including CodeBERT, CodeT5, UniXCoder with significant margins. In short, we summarize our contributions as follows.
• We present a simple self-improvement framework and show how it can be easily adapted to PLMCs for the task of code-related sequence generation.
• We conduct extensive evaluation on two tasks: code summarization and code generation, and compare it with the well-known, state-of-the-art PLMCs. The results show that our framework consistently improvesover all PLMCs by a significant margin in those tasks.
• We provide analysis and explanations on how utilizing a simple framework consistently improves the performance of PLMCs.
Our work is publicly available. 1
1 https://github.com/Fsoft-AIC/
Code-LM-Self-Improvement
12994

heading: 2 Related Work
text: Exposure bias and hallucination in Sequence Generation Tasks The exposure bias problem is regarded as the difference between the training and inference phases for auto-regressive sequence generation models. Previous work has attempted to reduce exposure bias in training phase (Bengio et al., 2015; Ranzato et al., 2015; Wiseman and Rush, 2016; Wang and Sennrich, 2020). In the sense that our self-improvement step involves training model on its own prediction, the exposure bias is close to our approach.
Code understanding and generation Code learning problems have recently emerged as one of the primary tasks for assessing the capability of language models. Most recent code models are pretrained on multi-modal objectives before being fine-tuned on specific downstream tasks (Feng et al., 2020; Ahmad et al., 2021; Wang et al., 2021; Guo et al., 2022; Ding et al., 2022).
Knowledge Distillation Knowledge distillation is the process of transferring knowledge from a large unwieldy model or set of models to a single smaller model that can be practically deployed under real-world constraints, and such smaller model can usually keep the same performance or even better than the original model (Hinton et al., 2015;
Kim and Rush, 2016; Wang et al., 2020; Chen et al., 2020; Mukherjee et al., 2021). We perform an additional self-improvement step to improve the original model without using external resources, our work is relevant to knowledge distillation.

heading: 3 Method
text: Algorithm 1 Data Augmentation Process Input:
• θfine−tuned, the fine-tuned model checkpoint on a specific task T ∈ {code summarization, code generation, etc. }. • D = {(xi, yi) ∣ i = 1, n}, the train dataset on which the θfine−tuned is fine-tuned. • Bk denotes the beamsearch algorithm with beam size of k. It returns a list of k best sequences as prediction.
Output: • Augmented dataset D̃
1: procedure DATAAUGMENTATIONPROCESS 2: D̃ ← ∅ 3: for each datapoint (xi, yi) ∈ D do: 4: LK ← Bk(Pθfine−tuned(y ∣ xi)) 5: In other words, LK =[ŷi1, ŷi2, ..., ŷik] 6: ỹi ← argmax
ŷij∈LK (sim(ŷij , yi))
7: Adding (xi, ỹi) → D̃ 8: end for 9: return D̃
10: end procedure
Our method utilizes three distinct sets of model parameters: θpre−trained, θfine−tuned, and θimproved. Each corresponds to the stage of the model parameters after pre-trained, fine-tuned, and self-improved, respectively. The model generates tokens in autoregressive manner, progressing token-by-token.
Typically, models are pretrained on large scale corpora, resulting in a pre-trained checkpoint θpre−trained. These pre-trained models are subsequently fine-tuned on a targeted downstream dataset D using a supervised learning approach, yielding in a set of fine-tuned parameters θfine−tuned. Our investigation has revealed that model’s performance can be further enhanced if we continue to fine-tuned these parameters on an augmented version of D. Figure 1 portrays our proposed self-improvement step within the broader training framework. This step encompasses a data
augmentation technique and an additional round of fine-tuning, in addition to the pre-traingn and fine-tuning paradigm.
The process of augmenting the dataset is illustrated in Figure 2. We provide a detailed algorithm for this procedure in Algorithm 1. For each training pair of sequences (xi, yi) in the train dataset D, we employ beam search to generate a list of K-best predictions LK . This list comprises k predictions, where k represents the beam size. Subsequently, we evaluate the similarity between each prediction ŷij and its corresponding ground truth sequence yi using a similarity function sim based on BLEU score. The prediction with highest similarity is then selected ỹi = argmaxŷij∈LK(sim(ŷij , yi)). Finally, we add the sequence pair (xi, ỹi) to a new empty dataset D̃, which we refer as the augmented dataset. Essentially, the augmented dataset contains an equal number of datapoints as the original training dataset due to an one-by-one mapping during augmetation. Moreover, the augnmentation process occurs offline, with each newly augmented datapoint being saved to storage before being used for training in the self-improving phase.
The subsequent step involves fine-tuning θfine−tuned on D̃ until convergence. This results in a new set of model parameters denoted as θimproved. It is important to note that the index j in ŷij denotes the jth prediction in the beam, rather than the j th token of the predicted sequence. Additionally,
only the training dataset D is augmented, while the validation and test dataset remain unchanged for evaluation purpose.

heading: 4 Experimental Setup
text: Our goal is to show that for any of the fine-tuned model for a sequence generation task (F-PLMC), after applying our self-improvement method (SPLMC), the result improves.
Dataset and Downstream Tasks To achieve our goal of enhancing the code-related sequence generation task, we selected code summarization and code generation as our experimental areas. To evaluate these tasks, we utilized the CodeXGLUE benchmark (Lu et al., 2021), which comprises various datasets for various code understanding and code generation tasks. Specifically, we utilized the code summarization and code generation datasets from CodeXGLUE and disregarded the other ones. The statistics for each dataset is reported in Appendix.
Baseline Models We chose CodeBERT (Feng et al., 2020), CodeT5 (Wang et al., 2021), and UniXCoder (Guo et al., 2022) as baseline models. Each model represents a distinct neural architecture style. CodeBERT abides to the Bidirectional Transformer architecture, which is a well-known PLMCs, despite the fact that it may not produce the best results for the tasks in CodeXGLUE. CodeT5
and UniXCoder are the two PLMCs that achieve state-of-the-arts performance on the CodeXGLUE benchmark. CodeT5 is an encoder-decoder architecture that follows the Seq2Seq learning style by following T5. UniXCoder, on the other hand, is a unified language model. It can behave as an encoder-only, decoder-only, or encoder-decoder model by modifying the masked attention matrices inside each transformer layer. Note that while CodeT5 and UniXCoder are proposed to address both code summarization and code generation, CodeBERT only considers the first problem in there paper. So we only consider CodeBERT for code summarization in our work.
Evaluation Metric We follow CodeXGLUE benchmark in employing evaluation metrics for each task. Smoothed BLEU-4 (Lin and Och, 2004) is used as the evaluation metric for code summarization. For code generation, smoothed BLEU-4, CodeBLEU (Ren et al., 2020), and exact match (EM) are employed. We aim to improve all of these metrics after apply our self-improvement method into the PLMCs.
Implementation We carefully selected the checkpoints for CodeBERT 2, CodeT5 3, and UniXCoder 4 based on the corresponding tasks. It is important to note that not all of these methods have released fine-tuned checkpoints. CodeT5 stands out as the only model to have released checkpoints for code summarization and code generation tasks. Conversely, CodeBERT and UniXCoder only offer
2 https://github.com/microsoft/CodeBERT/tree/
master/CodeBERT 3 https://github.com/salesforce/CodeT5
4 https://github.com/microsoft/CodeBERT/tree/
master/UniXcoder
training scripts. When checkpoints were not available, we employed the provided training scripts to fine-tune the pretrained models until we obtained results comparable to those reported in the original research. Additionally, CodeT5 and UniXCoder may have different checkpoint options with varying model sizes, such as small, base, and large. We selected the base version for both CodeT5 and UniXCoder. The same validation set is employed self-improving as in fine-tuning. The best checkpoint is selected according to BLEU score evaluated on this set among the training epochs.

heading: 5 Evaluation Results
text: The results of our code summarization task are presented in Table 1. The "Beam sizes" column indicates the beam size used in the beam search algorithm, while the "Methods" column indicates whether or not our self-improved algorithm was utilized. We also included other models as references to compare the relative improvement of our model. On average, we observed an average of 0.76 BLUE score increase in performance across all languages. This improvement was consistent across various beam sizes (1, 5, 10), which confirms the effectiveness of our self-improved approach across a wide range of PLMCs. When comparing our model to other strong baselines, we found that our method improved the performance of CodeBERT for JavaScript from 15.78 to 16.39, surpassing the performance of PolyglotCodeBERT (15.80). This highlights the benefit of our self-improved method in improving weak models. The results of our code generation study are presented in Table 2, the performance increase by 0.81 BLUE scores on average. When using EM and CodeBLEU, the improvement also increases consistently.
While conducting our experiments, it is important to note that we did not selectively choose the most favorable random seed to optimize the performance of each entry. Instead, we utilized the default seed provided in each model repository to ensure fairness and consistency. Our code summarization experiments encompassed six different programming languages, and both code summarization and generation experiments were evaluated using three distinct beam sizes. In total, we conducted 60 runs to gather comprehensive results. The numbers reported consistently demonstrate the improvement achieved in each individual run, thereby affirming the robustness of the proposed method.

heading: 6 Ablation Study
text: Improvement Study In this section, we examine the factors that influence the improvement achieved by θimproved as compared to θfine−tuned through code summarization. We define r1 as the difference in performance measured by BLEU between inferencing with a beam size of 10 and inferencing with a beam size of 1. Additionally, we define r2 as the improvement in BLEU when inferencing with the same beam size of 1 between θfine−tuned and θimproved. By evaluating these values across a variety of beam sizes and programming languages in the code summarization dataset, we are able to visualize the results in Figure 3. Additionally, we have calculated the Pearson Correlation score, which is 0.77, indicating a strong correlation between r1 and r2. Our analysis demonstrates that a larger r1 is correlated with a better r2, suggesting that our method is more likely to yield better overall performance when r1 is large. We believe this insight is a crucial finding as it provides a simple indicator of the model’s fully trained capability.
CodeBLEU as Similarity Function Our primary findings in code generation are presented using BLEU as the similarity function. However, for a more comprehensive assessment of the correctness of the generated code, we consider CodeBLEU, which incorporates the specific characteristics of source code. CodeBLEU, therefore, aligns better with the objective of measuring similarity in data augmentation compared to BLEU, which relies on n-gram matching. This section examines the impact of using CodeBLEU as a similarity
function (CodeBLEU-augmentation) compared to BLEU (BLEU-augmentation) in the context of data augmentation for code generation.
We present the results of our UniXCoder selfimproving model with both BLEU-augmentation and CodeBLEU-augmentation in Table 3. The results indicate that CodeBLEU-augmentation enhances both BLEU and CodeBLEU scores compared to BLEU-augmentation. This suggests that using CodeBLEU as a similarity function improves the generated code at a local level, encompassing aspects such as fluency, semantics, and syntax. However, it does have a negative impact on exact match (EM). As code problems may not have a unique solution, when EM is used as an evaluation metric, it should allow for a more lenient assessment. Consequently, we argue that a slight decrease in EM would have minimal impact on the actual correctness of the generated solution. Thus, we propose placing greater emphasis on CodeBLEU as an evaluation metric for code generation.

heading: 7 Conclusion
text: We introduced a self-improvement technique as a final fine-tuning step to enhance model performance. Our experiments showed that this method, when applied to popular pre-trained code models (CodeBERT, CodeT5, and UniXCoder), significantly improves performance on code summarization and code generation tasks. We also provided insights on when this method is most effective in improving PLMCs. We intend to implement our technique in larger-scale models and other tasks, and believe it is an efficient way to optimize the capabilities of any code language model without the need for extensive architecture modifications or large-scale dataset assembly. We leave all of these investigations for the future.
Limitations
We discuss a few limitations of our work. One limitation of Self-Improved is its complexity in usage. The data augmentation process involves generating predictions for the entire training dataset with a large beam size, resulting in a time complexity of O(nk), where n is the train dataset size and k is the beam size. Additionally, the fine-tuning step to derive θimproved also adds a significant amount of computational complexity. This limitation is discussed in Section 6 to weigh the performance benefits of our method against the computational sacrifices. Another limitation is that Self-Improved has only been applied to encoder-decoder models in this work. However, it is also applicable to other types of auto-regressive models, including encoderonly models, which are commonly used for tasks such as code completion (Radford et al., 2019; Lu et al., 2021; Guo et al., 2022). A few models can be named are GPT models (Radford et al., 2019; Brown et al., 2020), CodeX (Chen et al., 2021), CodeGen (Nijkamp et al., 2022), etc. Further research into these applications is left for future work.

heading: A Data statistics
text: We include the statstics for the data used in our experiments. We access directly to CodeXGLUE page5 for downloading data. CodeXGLUE gathers datasets from multiple sources for each downstream task. For code summarization, the CodeSearchNet (Husain et al., 2019) are used with the number of examples for each train/dev/test are reported in Table 4. The dataset comprise the codetext pairs in 6 programming languages. CONCODE (Iyer et al., 2018) dataset is employed as the benchmark for code generatiom with statistics reported in Table 5. It contains pairs of Java member function and natural language description Java language.

heading: B Hyperparameter selection
text: In the fine-tuning phase, we maintain the model hyperparameter values as specified in the training script provided by the official repositories. However, we make a modification by increasing the batch size to fully utilize the memory capacity of a NVIDIA A100 80GB. In the self-improvement phase, we further decrease the learning rate by a factor of 10.
5 https://github.com/microsoft/CodeXGLUE
ACL 2023 Responsible NLP Checklist

heading: A For every submission:
text: 3 A1. Did you describe the limitations of your work?
8
3 A2. Did you discuss any potential risks of your work? 8
3 A3. Do the abstract and introduction summarize the paper’s main claims? 1
7 A4. Have you used AI writing assistants when working on this paper? Left blank.
B 3 Did you use or create scientific artifacts? 4
3 B1. Did you cite the creators of artifacts you used? 4
B2. Did you discuss the license or terms for use and / or distribution of any artifacts? Not applicable. Left blank.
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Not applicable. 4
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Not applicable. Left blank.
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? Not applicable. Left blank.
3 B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Left blank.
C 3 Did you run computational experiments? 4

heading: 3 C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? 4
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: 3 C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? 4, Appendix D

heading: 3 C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? 4

heading: 3 C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did
text: you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? 4
D 7 Did you use human annotators (e.g., crowdworkers) or research with human participants? Left blank.
D1. Did you report the full text of instructions given to participants, including e.g., screenshots, disclaimers of any risks to participants or annotators, etc.? No response.
D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students) and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? No response.
D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to crowdworkers explain how the data would be used? No response.
D4. Was the data collection protocol approved (or determined exempt) by an ethics review board? No response.
D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data? No response.

heading: Future Work
text: We leave all of these investigations for the future. Limitations
We discuss a few limitations of our work. One limitation of Self-Improved is its complexity in usage. The data augmentation process involves generating predictions for the entire training dataset with a large beam size, resulting in a time complexity of O(nk), where n is the train dataset size and k is the beam size. Additionally, the fine-tuning step to derive θimproved also adds a significant amount of computational complexity. This limitation is discussed in Section 6 to weigh the performance benefits of our method against the computational sacrifices. Another limitation is that Self-Improved has only been applied to encoder-decoder models in this work. However, it is also applicable to other types of auto-regressive models, including encoderonly models, which are commonly used for tasks such as code completion (Radford et al. , 2019; Lu et al. , 2021; Guo et al. , 2022). A few models can be named are GPT models (Radford et al. , 2019; Brown et al. , 2020), CodeX (Chen et al. , 2021), CodeGen (Nijkamp et al. , 2022), etc. Further research into these applications is left for future work