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