Better Language Models of Code through Self-Improvement (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, fine-tuning 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 fine-tuning 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.

## introduction pre-trained models for code (plmcs), such as codebert @xcite , plbart @xcite , codet5 @xcite , @xcite , and disco @xcite , 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 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 @xcite , 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, unix-coder with significant margins. in short, we summarize our contributions as follows. our work is publicly available. ## method algorithm 1 data augmentation process 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 θ f ine-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 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 (x i , y i ) in the train dataset d, we employ beam search to generate a list of k-best predictions l k . this list comprises k predictions, where k represents the beam size. subsequently, we evaluate the similarity between each prediction ŷij and its corresponding ground truth sequence y i using a similarity function sim based on bleu score. the prediction with highest similarity is then selected ỹi = argmax ŷij ∈l k (sim(ŷ ij , y i )). finally, we add the sequence pair (x i , ỹ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 θ f ine-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 ŷij denotes the j th 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. ## experimental setup 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 (s-plmc), 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 @xcite , 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. ## evaluation results 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. improvement study in this section, we examine the factors that influence the improvement achieved by θ improved as compared to θ f ine-tuned through code summarization. we define r 1 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 r 2 as the improvement in bleu when inferencing with the same beam size of 1 between θ f ine-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 r 1 and r 2 . our analysis demonstrates that a larger r 1 is correlated with a better r 2 , suggesting that our method is more likely to yield better overall performance when r 1 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 code-bleu, 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 ## conclusion 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 (code-bert, 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.