Exploring and Evaluating Personalized Models for Code Generation

1 INTRODUCTION . 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 . 2 APPROACH . 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 . 2.1 CUSTOMIZATION PROCESS . 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 . 2.2 CUSTOM FINE-TUNING . 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 . 2.3 LIGHTWEIGHT FINE-TUNING - EMBEDDINGS AND OUTPUT LAYER ( L-EO ) . 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 ) . 2.4 LIGHTWEIGHT FINE-TUNING - LAST DECODER BLOCK ( L-LDB ) . 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 . 2.5 PREFIX TUNING . 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 . 2.6 TRAINABLE PARAMETERS DURING FINE-TUNING . 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 . 3 EXPERIMENTAL DESIGN . 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 . 3.1 DATASET . 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 .

Summary: 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.