Exploring Distributional Shifts in Large Language Models for Code Analysis (2023)

Abstract: We systematically study how three large language models with code capabilities - CodeT5, Codex, and ChatGPT - generalize to out-of-domain data. We consider two fundamental applications - code summarization, and code generation. We split data into domains following its natural boundaries - by an organization, by a project, and by a module within the software project. We establish that samples from each new domain present all the models with a significant challenge of distribution shift. We study how established methods adapt models to better generalize to new domains. Our experiments show that while multitask learning alone is a reasonable baseline, combining it with few-shot finetuning on examples retrieved from training data can achieve very strong performance. Moreover, this solution can outperform direct finetuning for very low-data scenarios. Finally, we consider variations of this approach to create a more broadly applicable method to adapt to multiple domains at once. We find that for code generation, a model adapted to multiple domains simultaneously performs on par with those adapted to a single domain.

## introduction since the late 2000s, researchers have been reporting poor generalization of statistical learning models to new software systems @xcite , a phenomenon that has become important with the rise of large language models (llms) for code, such as github copilot, amazon codewhisperer, replit, etc. thus, it is crucial to understand when pretrained large language model performance on a private software system will differ from the performance obtained on a benchmark. prior work has studied some aspects of this problem, among others studying generalization from older to newer code, large software projects, and small competition problems, authors, and code representations @xcite @xcite . however, the challenges of distribution shifts stemming from the hierarchical nature of software data, as depicted in figure 1 , have not been systematically studied with regard to large language models for code. motivated by that, in this work, we probe the generalization capacity of large language models with code capabilities, specifically codex @xcite , codet5 @xcite and chatgpt, in code generation and summarization tasks, examining three scenarios: generalization across companies, projects, and project components. these scenarios are routinely considered for analyzing software systems @xcite @xcite due to the careful consideration that goes into combining or separating such entities. figure 1 : organization of a software system by the granularity of its components first, we want to understand how models perform on new domains -if models struggle with out-of-domain generalization, they should be used with caution. at the same time, we empirically establish the legitimacy of our definitions for outof-domain scenarios by demonstrating that these examples present a distributional shift. to answer this question, we compare the performance of the models without any additional adaptation with that of the models that have been adapted on limited data from a random domain or from the test domain. adaptation with labeled examples from the test domain is the proxy for model performance if there were no distributional shift. we find that all three models suffer from a drop in performance when applied out-of-domain. in this experiment, the difference is more pronounced for code summarization, where adapting models with few in-domain examples, on average, leads to an improvement of over 10 bleu @xcite score points. next, we explore ways to improve the out-ofdomain generalization of large language models with code capabilities, recognizing that relying on labeled in-domain data for every new domain is impractical. instead, we investigate the use of labeled out-of-domain data and small amounts of unlabelled in-domain data to enhance generalization. we test methods known to be successful in other transfer learning scenarios, such as metalearning @xcite and multitask learning @xcite . we also leverage unlabeled in-domain data to retrieve similar labeled examples from an out-of-domain corpus for adapting to the new domain. we find that while meta-learning and multitask learning do not solve the out-of-domain generalization problem, domain adaptation with retrieved examples is a good technique for low-data domains. in our evaluation on codesearchnet dataset we find that models supervised with retrieved examples perform on par, or better, than models that have been adapted using a few samples (e.g., 8 or 16) of in-domain labeled data. we are particularly interested in scenarios with an extreme scarcity of labeled data -ranging from a few labeled instances to no labeled data at all. this is due to how new data emerges in software engineering domains -it is not difficult to imagine a new repository, or a new module, with fewer than 32 functions, let alone -32 labeled functions. lastly, we study if we can make the code models more broadly applicable and retain their generalization capacities, rather than having to adapt them to every new domain? depending on the approach to model adaptation (e.g. weight update vs in-context demonstrations) we vary the set of retrieved examples for each new domain, or for each test input individually. we compare performance obtained this way with that of the models that are adapted simultaneously to multiple domains (or instances, correspondingly). we find that codex is very sensitive to these changes, so it is best to retrieve similar instances for each test data point. on the other hand, codet5 has a minor drop in code summarization and a negligible drop in code generation. this makes it feasible to adapt and apply codet5 to multiple domains simultaneously with minimal tradeoff, eliminating the need to store separate copies of the model for each domain. ## background the shifts in underlying semantics between the training and evaluation data can be one of the most impacting factors for deteriorating performance at test time. prior work in code analysis has mainly focused on cross-project shifts, i.e. training and evaluating the model on disjunct sets of code projects. additionally, the studies were mainly conducted in the context of traditional machine learning methods, such as linear classifiers, support vector machines, and later, lstms @xcite @xcite . more recent works consider shifts caused by different authors of the code, the timeline of the project, distributions of code tokens, etc @xcite @xcite . however, the abilities of large language models under distribution shift are still under-explored. we conduct a comprehensive empirical analysis to probe the large language models' capabilities in handling three different granularity of distribution shifts (company, domain, module) when different training and adaptation methods are used. in addition to directly fine-tuning vanilla llms, we experiment with enhancing pretrained models using the methods described below. meta-learning and multi-task learning. in our work, we experiment with both meta-learning and multi-task learning to get better initialization for few-shot performance on the downstream task. for meta-learning, we employ model-agnostic meta-learning (maml) @xcite which is a gradient-based method. it is a conceptually simple and model-agnostic algorithm that has been shown to outperform existing approaches in several tasks. multi-task learning (mtl) aims to learn a shared and generalized representation by jointly training on several tasks. we adopt the simplest approach to multi-task learning by jointly finetuning a shared language model on multiple tasks. parameter efficient methods. parameterefficient methods have been shown to obtain performance comparable to finetuning all model parameters with finetuning only a tiny fraction of model parameters. in our work, we experiment with low-rank adaptation (lora) @xcite , which is a low-rank update method. in-context learning. @xcite demonstrated the ability of large language models to perform few-shot predictions, where the model is given a description of the task in natural language with few examples. in our work, we conduct experiments on in-context learning on codex. retrieval based example selection. it has been shown in @xcite that in-context examples selected following a strategy may serve as more informative input to unleash gpt3's extensive knowledge. inspired by this, we leverage a similarity-based retrieval for domain adaptation. we study scenario where users seek to integrate a large language model, such as codex or codet5, into their software project. the primary focus of this study is to gain a deeper understanding of the performance characteristics exhibited by these models, particularly when confronted with source code originating from an unseen organization, an unseen project, or specific project components that have not been previously encountered. ## problem setting for every code data point in the dataset, we have information about the organization, project, and the module within the project that the data point comes from. based on this information, we can group data points into sets, and end up with three sets of sets, as illustrated in figure 2 . for example, the middle set in the figure contains multiple sets of data points. each of those sets corresponds to a unique organization to which all data points within it belong. in other words, all data points within a set belong to the same domain. appendix, section 7.3 contains additional analysis on splitting the data points in this manner. for simplicity, we refer to a set of examples from the same domain as τ i . we refer to splits of such a set into train, development, or test sections as τ i train , τ i dev , and τ i test . ## analysis in this section, we formulate the research questions that we aim to answer and give a more detailed description of the setups that we have used for analyzing and answering each question. ## conclusion we evaluate large language models for code -codet5, codex (code-cushman-001), and chat-gpt (gpt-3.5-turbo) -on two fundamental code applications -code generation and code summarization. we study how the models perform under distribution shifts that can commonly occur due to the nature of the software. we experiment with three granularities for defining domains in applications for code -organization, project, and module or folder. our experiments show that all models evaluated are susceptible to reduced performance due to domain shifts. we experiment with a number of training and domain adaptation techniques for achieving better out-of-domain generalization. we discover that retrieving similar out-of-domain examples from training data is the most effective approach for adapting to new domains in the absence of in-domain data. in addition, we experiment with adapting models to multiple new domains simultaneously and find that such models can perform very well for code generation. however, we find the generality of the model to be a tradeoff for its performance for code summarization.