Year: 2021

Abstract:
Large Transformer models achieved the state-of-the-art status for Natural Language Understanding and are increasingly the baseline architecture for source code generation models. Transformers are usually pre-trained on a large unsupervised corpus, learning token representations and transformations relevant to modeling generally available text, and then fine-tuned on a particular task of interest. While fine-tuning is a tried-and-true method for adapting a model to a new domain, for example question-answering on a given topic or a source code generation model, generalization remains an on-going challenge. Here we explore the ability of various levels of model fine-tuning to improve generalization by personalized fine-tuning. In the context of generating unit tests for Java methods, here we evaluate learning to personalize to a specific project using several methods to personalize transformer models for unit test generation for a specific Java project. We consider three fine-tuning approaches: (i) custom fine-tuning, which allows all the model parameters to be tuned; (ii) lightweight fine-tuning, which freezes most of the model’s parameters, allowing a tuning of the token embeddings and softmax layer or the final layer alone; (iii) prefix tuning, which keeps language model parameters frozen, but optimizes a small project-specific prefix vector. Each of these techniques offers a different trade-off in total compute cost and prediction performance, which we evaluate by code and task-specific metrics, training time, and total computational operations. We compare these fine-tuning strategies for code generation and discuss the potential generalization and cost benefits of each in deployment scenarios.

heading: 1 INTRODUCTION
text: It is well-known that even the best models can fail to generalize properly to new domains, and even to new users of said models. For example, a model trained to answer questions in general may not answer StackOverflow questions as well as the questions in the training domain, or a software developer in an Enterprise environment with private code may have libraries and attribute name which differ from public source code used to train a code synthesis model.
The current dominant paradigm in Natural Language Processing (NLP) modeling is to pre-train a large transformer model (Vaswani et al., 2017a) on a large corpus and then fine-tune it on a particular task of interest. For example, a question-answering (Q&A) model is generally first pre-trained on a large corpus of textual data for the specific language (e.g., Wikipedia, and news articles in English), then fine-tuned on a task-specific dataset of paired questions and corresponding answers. The pre-training process aims at learning semantic vector representation of the language and words, while the fine-tuning process specializes the model for a specific domain.
Transformer models are also increasingly the baseline architecture used for code generation tasks, such as writing methods from natural language description (Clement et al., 2020; Austin et al., 2021; Chen et al., 2021), or generating test cases from the focal method under test (Tufano et al., 2021). Similarly for NLP tasks these models are pre-trained on a large corpus of natural text and publicly available source code and then fine-tuned on a specific code-related task. Further, these models also may not generalize to new domains of interest, and can benefit from task or even user-specific fine-tuning, here called customization or personalization. Customization is particularly relevant for code generation models since it provides several benefits:
• allows fine-tuning on source code data that may not be available when training a base model (e.g., private repositories or internal codebases), enabling improved overall performances on codebases with proprietary dependencies and code styles;
• the opportunity to improve data privacy by considering private or sensitive data only during the customization process on the client side;
• the opportunity to reduce deployment cost as customized models can offer better user performance without increasing model size.
Custom models can provide clear benefits to users and model providers. We envision serving tens or hundreds of thousands of custom models, but doing so presents several logistical hurdles, including the costs of training, storing, and loading these models into GPU memory for inference. Worse, memory costs will only be exacerbated when working with ever larger and more powerful models.
For these reasons, we investigate several customization approaches, some of which can dramatically reduce the memory footprint and amortized computational cost introduced by custom models. Specifically, we consider three fine-tuning approaches: (i) custom fine-tuning, which allows all the model parameters to be tuned; (ii) lightweight fine-tuning, which only optimizes the token embedding representations or the final softmax layer; (iii) prefix tuning, which keeps language model parameters frozen, but optimizes a small project-specific vector prefix.
In our extensive empirical evaluation we found that all the customization strategies lead to significant model improvements on a target project in terms of both intrinsic and task-specific metrics. While there is no unambiguous winner among the customization strategies, each approach can provide specific benefits in particular deployment scenarios. This paper provides insights on these customization strategies, their benefits and drawbacks, as well as providing guidelines and suggestions on which one to use based on the training cost, memory and storage, number of users, and deployment scenarios.

heading: 2 APPROACH
text: This section describes the proposed customization approach for code generation models. We begin by formally defining the customization process, then we provide details for each of the fine-tuning strategies.

heading: 2.1 CUSTOMIZATION PROCESS
text: We use the term customization to refer to the process of fine-tuning a model m, previously trained on a generic dataset for a task t, with the goal of improving its performance on a specific dataset p. The performance of a machine learning model m on a dataset p is measured by one or more evaluation functions f(m, p), where f can be either a maximization (e.g., BLEU, top-k accuracy) or minimization (e.g., perplexity) function. The customization process is designed to modify the trainable parameters of the model m, obtaining the model m′, such that the performance of m′ on p is better than what was attained by m. Specifically, f(m′, p) > f(m, p) for maximization functions, or f(m′, p) < f(m, p) for minimization functions.
In this work, m is an encoder-decoder transformer model, t is a code generation task, and p is a target software project to which we intend to customize m.

heading: 2.2 CUSTOM FINE-TUNING
text: Custom fine-tuning is the most straightforward customization approach. The model to be customized is taken as is and trained on a selected project. All parameters are trainable during this process. Figure 1a shows the model during fine-tuning, where all the parameters from the encoder and decoder blocks, as well as embeddings and output layers can be modified.

heading: 2.3 LIGHTWEIGHT FINE-TUNING - EMBEDDINGS AND OUTPUT LAYER (L-EO)
text: Fully fine-tuning a model for every project or user may be prohibitive in terms of storage and memory costs. As a result, we explore ways to mitigate these costs by reducing the number of parameters that vary from one custom model to another. In our lightweight fine-tuning experiments, we achieve this
by freezing most parameters in the baseline model, and only keeping a small subset trainable. Figure 1b shows the Lightweight fine-tuning - Embeddings and Output Layer (L-EO) design, where most of the model parameters are frozen (displayed in gray), and we allow only the embedding and output layers parameters to be fine-tuned, following the approach in Lu et al. (2021).

heading: 2.4 LIGHTWEIGHT FINE-TUNING - LAST DECODER BLOCK (L-LDB)
text: In this lightweight fine-tuning strategy, shown in Figure 1c (L-LDB), most of the model’s parameters are kept frozen, while only the parameters in the last decoder block are trainable, this includes: self-attention, encoder-decoder attention, layernorm and feedforward layers. This design decision of training only the last decoder block is motivated by experimental results analyzing the model’s parameter changes during custom fine-tuning. Figure 2 reports the average absolute changes, during fine-tuning, in the parameters belonging to different Encoder and Decoder blocks for a BART model. We observe that, as we go through the transformer model, the average change in parameter values tends to increase, with the last decoder block showing the highest changes in parameter values. As a result, we hypothesize that it could be sufficient to tune the last decoder block and obtain performance improvements similar to the fully custom fine-tuned model.

heading: 2.5 PREFIX TUNING
text: Prefix tuning was first introduced by Li and Liang (2021), with the goal of fine-tuning a general model for different tasks. The technique concatenates a sequence (prefix) of virtual tokens (trainable parameters) to the front of the input of every encoder and decoder block. In our context, the intuition behind this approach is that the prefix embeds the properties of a specific project, which allows the model to generate customized responses for that repository. Practically, we set the prefix length to 200 tokens, and thus with an embedding size of 1024, this gives a total of 1024× 200× 24× 2 ≈ 10M trainable parameters. The prefix is initialized to the most frequent words in the repository for which the model is customized.

heading: 2.6 TRAINABLE PARAMETERS DURING FINE-TUNING
text: Table 1 provides an overview of the number of total and trainable parameters involved in each customization process, in the case of a BART Transformer model with 406M parameters. Custom fine-tuning allows to train 100% of the 406M available parameters in the model. During L-EO finetuining, instead, only 13% (53M) parameters are trained. The L-LDB fientuning reduces the
number of trainable parameters to 4.2% (17M). Finally, Prefix tuning has the lowest number of trainable parameters, only 2.4% (10M) of the total, but these are additional parameters added to the model, which reaches a total of 416M.

heading: 3 EXPERIMENTAL DESIGN
text: The goal of our experimental design is to investigate whether custom models outperform the baseline model, leading to performance improvements in terms of intrinsic metrics (RQ1), as well as extrinsic task-specific metrics (RQ2). Next, we analyze and compare the different customization approaches in terms of training and compute costs (RQ3) as well as model size and required storage for deployment.
In our case study, we chose Unit Test Case generation as our code generation task t, and AthenaTest by Tufano et al. (2021) as our baseline model m, which is a BART transformer model pre-trained on source code and English, and fine-tuned on Java unit test generation. The task is modeled as a translation task, where the input is a focal method (i.e., method under test), and the output is a test case which tests the focal method’s correctness. We randomly sample 20 projects from the test set, each of those representing the dataset p on which a custom model is fine-tuned. Specifically, for each project p, we start from the baseline model m and fine-tune four different custom models according to the four proposed fine-tuning strategies. For each project and fine-tuning strategy (e.g., L-EO), we fine-tune and evaluate the models using 4-fold cross-validation. The models are trained until the best validation loss is reached, independently for every fold, every repository, and every customization approach. In total, we fine-tune and evaluate 20(projects) × 4(approaches) × 4(folds) = 320 models.

heading: 3.1 DATASET
text: Table 2 reports information about the 20 GitHub repositories sampled from the test set, which will be used to customize our models. The table shows (i) the Project ID, which will be used in the paper to
reference a specific project; (ii) the project name; (iii) the project size in terms of disk usage; (iv) the popularity of the project in terms of number of stars obtained on GitHub; (v) and the dataset size, which corresponds to the number of data points for the unit test generation task (i.e., pair of focal method and test case). The list of projects represent a diverse set of repositories with different size, domain, and popularity. They span from small personal projects (e.g., Tutorials with 6 stars), to open source projects developed by large organizations such as Apache and Google.
3.2 RQ1: INTRINSIC EVALUATION METRICS
RQ1: Do custom models obtain better performances on intrinsic metrics, such as BLEU and perplexity, w.r.t. the baseline? To begin, we investigate how the different model customization approaches described in Sec. 2 score on intrinsic metrics such as BLEU and perplexity. All approaches entail fine-tuning the baseline model to the dataset of a specific project, with the choice of parameters being tuned depending on the approach taken. The four variants are trained independently until the best validation loss is achieved. We report the BLEU4 score and the mean perplexity per token on the test fold, for all the 20 projects. Next, we perform statistical tests to investigate whether the observed differences between the baseline and custom models are significant, as well as differences among the customization approaches. Specifically, we rely on the Kruskal-Wallis test, a non-parametric statistical test.
3.3 RQ2: TASK-SPECIFIC PERFORMANCES
RQ2: Do custom models improve on performance metrics specific to unit test generation? We want to investigate how the different customization approaches compare with respect to the downstream task of generating unit tests. Beyond BLEU score and perplexity, we would like to see if custom models can produce the correct target code, how closely their unit tests mimic the repository style, or even if they can perfectly match the desired output.
• Perfect Matches: We compare the model’s output string with the target developer-written unit test. If the two strings are identical, this is considered a perfect match. We do not take spacing and indentation into account as we are using a Java dataset (where indentation is not required). We report the proportion of perfect matches among the top 5 model predictions.
• Abstracted Code Matches: We pass the model output and target output through the src2abs tool (Tufano, 2018), to obtain an abstracted version, masking variable names, method names, etc. We also do not distinguish between different objects of the same type.
• Coding Style: For each project’s custom model, we would like to determine how closely the model learns the developer’s personal programming style and preferences. To this end, we extract the collection of all identifiers (i.e., variables and functions’ names) from the unit tests written by the developer as well as those generated by the models. We then pass these text outputs through a tf-idf vectorizer and compute the cosine similarity between them. This allows us to compare the developer’s and the models’ word usage. We examine the similarity between the developer’s unit tests and the baseline and custom models generated tests. This scores the vocabulary similarity of the unit tests with the model generated code.
3.4 RQ3: TRAINING COST COMPARISON
RQ3: Given the same amount of compute, which custom models achieve the biggest performance improvement? Since our four training regimes tune a different number of parameters, simply comparing the training time or number of optimization steps to reach the best validation loss may not be appropriate. For a model with N parameters, we approximate the computation cost of a forward pass to be C ≈ 2N floating point operations per training token, with an additional correction for embedding layers. The backward pass takes roughly twice the amount of compute, but it is unnecessary for layers that are frozen. For additional details, we refer to Table 1 in Kaplan et al. (2020). We report the resulting compute in petaFLOPS-seconds.

heading: 4 RESULTS
text: 4.1 RQ1: INTRINSIC EVALUATION METRICS
Table 3 presents the results of custom models in terms of the intrinsic metrics: BLEU and perplexity. Specifically, for each project, we report the average BLEU and perplexity over the four folds, achieved on the test set by the different customization strategy, as well as the baseline model. We observe notable improvements in both metrics for every project w.r.t. the baseline, with BLEU going from 16.1 achieved by the baseline model to 36-38 by custom models.
The statistical tests reported in Table 4 confirm that the improvement observed by the four customization techniques are statistically significant w.r.t. the baseline (p < 1e-7). However, we do not observe statistical significance in the differences among the customization strategies, meaning that, in terms of intrinsic metrics performances, the differences are within margin of error.
4.2 RQ2: TASK-SPECIFIC PERFORMANCES
The results in terms of task-specific performances are presented in Figure 3. The plot 3a shows the top-k accuracy for perfect matches (solid line) and abstracted matches (dotted line), aggregated over the 20 projects. The baseline model outputs the same code structure (abstracted) in roughly 3% of all cases, and virtually never produces the exact target output (<1%). Moreover, its performance does not improve as we consider more predictions. All customization processes show significant improvement compared to the baseline. Specifically, these improvements are observed for every single project (full results will be available on our online appendix). Customized models produce the correct code structure as their top prediction in ∼13-14% of instances, and a perfect match in ∼4-6% of cases. They also tend to improve as we consider their top 5 predictions. Between the different customization processes, Custom consistently performs the best, closely followed by Prefix and L-LDB. When considering abstracted code matches, these three approaches are nearly identical. L-EO, however, performs slightly worse than the others.
Plot 3b shows the distribution of tf-idf cosine similarity computed between identifiers used in the developers’ written tests and the models’ generated outputs. We observe that the distribution for custom models is skewed towards the higher values of cosine similarity. This result demonstrates that custom models tend to use variable and function names that are more similar to what developers used in their own tests.
4.3 RQ3: TRAINING COST COMPARISON
For each customization process, we plot validation loss as a function of compute, as defined in section 3.4. The results are presented in Figure 4, where the light lines represent the validation loss curve for each individual project and fold, while the bold line represents the average for each custom strategy. First note that Custom achieves very large gains during the first epoch, as evidenced by the fact that its validation loss starts much lower than L-EO and L-LDB. Custom also outperforms other customization processes when given a limited amount of compute. However, we observe that beyond a certain amount of compute, Custom and L-LDB tend to achieve similar performances. In contrast, L-EO starts at the same validation loss as L-LDB but converges much slower to the best loss, requiring 2-3 times as much compute.
Since the prefix parameters suffer from poor initialization, Prefix is the most expensive customization process. To overcome this problem, it is possible to first train the prefix on a large generic dataset. Then, given proper hyperparameter tuning, it is possible to substantially cut down compute cost for customizing the prefix.

heading: 10 1 100 101 102 103

heading: 5 DISCUSSION & LESSONS LEARNED
text: The four customization strategies considered in this work are effective in improving a code generation model’s performances on a given software project. Specifically, all custom models significantly outperform the baseline in terms of intrinsic metrics (i.e., BLEU and perplexity) as well as taskspecific metrics (i.e., abstract and raw matches). While the differences among the customization approaches are not significant (no clear winner), each strategy offers specific advantages in different circumstances and deployment scenarios.
Custom fine-tuning achieves the overall best performances and the customization process is relatively fast and efficient. This is somewhat expected, since this customization strategy allows all the model’s parameters to be tuned on the specific project. This characteristic also leads to the major disadvantage of this approach: each custom model is an entire copy of the original model. Storage and inference costs could become prohibitive when serving many users with personalized custom models.
Lightweight fine-tuning achieves good results while training fewer parameters. This allows to serve potentially many users with custom models which can be stored and loaded efficiently. Specifically, L-LDB trains fewer parameters than L-EO, however the latter could allow to deploy the embedding and output layers on the user side, with a privacy-preserving focus.
Prefix fine-tuning trains the lowest number of parameters (only 2.4% for a BART model), while improving over the baseline. However, it increases the total number of parameters of the model (prefixes are additional virtual tokens) and requires more compute time to achieve good performances, mostly due to the prefix initialization problem. On the bright side, this strategy allows to batch together requests from different users (with different prefixes), which can be processed by a single model, generating personalized outputs.

heading: 6 RELATED WORK
text: This work is related to two areas of the existing literature: neural source code generation and model personalization. Neural code generation has generated an intense recent interest in NLP, using Transformer models Vaswani et al. (2017b) in particular for code completion Svyatkovskiy et al. (2020; 2019); Clement et al. (2020); Raychev et al. (2014); Bruch et al. (2009); Brockschmidt et al. (2018), code synthesis from examples Chen et al. (2018), natural language to code Clement et al. (2020); Chen et al. (2018); Austin et al. (2021), code feature summarization Liu et al. (2021); Moreno et al. (2013); Scalabrino et al. (2017); Wan et al. (2018); Alon et al. (2018); Moreno et al. (2014), code search Husain et al. (2019); Feng et al. (2020), unit test generation Tufano et al. (2021) and even bug fixing Drain et al. (2021) and detection Zhai et al. (2020). This paper naturally is an extension and evaluation of personalized unit test generations as studied by Tufano et al. (2021), and an important contribution to the understanding optimization in a deployment scenario.
Much of the previous literature on personalized models focuses on client-side training to keep data on device Shor et al. (2019); Popov et al. (2018), and most work is in the domain of search query completion Jaech and Ostendorf (2018), natural language completion Popov et al. (2018), or even automated speech recognition Shor et al. (2019). Naturally this work extends the domain of evaluation beyond natural language tasks and into the software engineering domain. This paper does not evaluate methods for client side training with restricted resources, however, as the most powerful large language models which enable state of the art code synthesis have 10-100 million parameters. At the time of writing such large models cannot be executed in a reasonable amount of time on most consumer laptops. We leave to future work extending these studies to models which have been pruned, quantized, distilled, and optimized to be ran in limited resource environments.

heading: 7 CONCLUSION
text: In this paper we explored different ways to customize a code generation model for a given codebase, with the goal of improving its performances on a target project. We described and analyzed four customization strategies and applied them on 20 different software projects for the task of generating unit test cases. Specifically, we considered the following strategies: (i) custom fine-tuning, which allows all the model parameters to be tuned on the target project; (ii) L-EO fine-tuning, a lightweight training which freezes most of the model’s parameters, tuning only embedding and output layers; (iii) L-LDB fine-tuning, a lightweight training which only tunes the last decoder block; (iv) prefix tuning, which keeps language model parameters frozen, but optimizes a small project-specific vector (prefix).
In our extensive empirical evaluation we found that all the customization strategies lead to significant model’s improvements on a target project, in terms of both intrinsic and task-specific metrics, with the custom models adapting to the coding style of the target project. While there is no clear winner among the customization strategies, each approach can provide specific benefits in particular deployment scenarios.

Review summaries:
- This work studies customizing models for code towards specific projects/coding standards/preferences for unit test case generation task. It specifically studies this in the context of server-side customization, where one entity would need to maintain multiple customized models, as opposed to client-side customization where each client can store its own customized model. - It is an empirical study of how 4 customization strategies -- a) modify entire model weighs, b,c) modify a chosen subset of model weights, and d) prefix tuning -- affects the performance of the fine-tuned model on the test case generation task. - This work finds no substantial difference in the final fine-tuned performance of any of the customization strategies, but argues that each strategy has its advantages and disadvantages, therefore allowing the user/implementor of this approach to make a more informed choice.

The paper proposes personalization of a baseline model to a project by fine-tuning it on the project-specific labeled examples. It demonstrates that personalization yields performance benefits over the baseline model. The baseline model (Tufano et al. 2021) is obtained from a pre-trained Java BART model after fine-tuning for the task of generating Java unit tests. The paper evaluates four fine-tuning strategies that select different subsets of weights to be fine-tuned.

Software projects tend to differ strongly from one another in terms of coding style, which significantly reduces the efficacy of even large pretrained models in new project contexts. This work explores the practical performance of several fine-tuning strategies of a model trained for a software engineering task (test-case generation) to new project contexts. The results suggest that fine-tuning can substantially improve model performance, but the specific choice of fine-tuning strategy makes very little difference.