Exploring Continual Learning for Code Generation Models (2023)

Abstract: Large-scale code generation models such as Copilot and CodeT5 have achieved impressive performance. However, libraries are upgraded or deprecated very frequently and re-training large-scale language models is computationally expensive. Therefore, Continual Learning (CL) is an important aspect that remains under-explored in the code domain. In this paper, we introduce a benchmark called CodeTask-CL that covers a wide range of tasks, including code generation, translation, summarization, and refinement, with different input and output programming languages. Next, on our CodeTask-CL benchmark, we compare popular CL techniques from NLP and Vision domains. We find that effective methods like Prompt Pooling (PP) suffer from catastrophic forgetting due to the unstable training of the prompt selection mechanism caused by stark distribution shifts in coding tasks. We address this issue with our proposed method, Prompt Pooling with Teacher Forcing (PP-TF), that stabilizes training by enforcing constraints on the prompt selection mechanism and leads to a 21.54% improvement over Prompt Pooling. Along with the benchmark, we establish a training pipeline that can be used for CL on code models, which we believe can motivate further development of CL methods for code models.

## introduction code generation models @xcite @xcite can increase the productivity of programmers by reducing their cognitive load. these models require significant computation to train as they have billions of parameters trained on terabytes of data. hence, they are trained once and are figure 1 : we show the process of prompt selection for prompt pooling with teacher forcing when learning multiple tasks sequentially. first, we initialize the prompt pool with (key, prompt) pairs (denoted by rectangles). next, each (key, prompt) pair is assigned to either a single task or is shared by two tasks (denoted by colors). when learning task 1 (green color), we obtain the query (green circle) for a given example and select the top-k (@xmath0, prompt) pairs, highlighted in the figure. these selected pairs are then trained for the example. a similar process is followed for subsequent tasks. during inference, we remove task assignments and select the top-k pairs across all the pairs. then used repeatedly for several downstream applications. however, as software development constantly evolves with new packages, languages, and techniques @xcite , it is expensive to retrain these models. therefore, it is essential to continually improve these models to avoid errors, generate optimized code, and adapt to new domains and applications. we explore continual learning (cl) @xcite abilities of code-generation models and aim to improve them. specifically, we present a codetask-cl benchmark for code-based cl and aim to train a model on sequentially presented tasks with different data distributions without suffering from catastrophic forgetting (cf) @xcite . this occurs when the model overfits the current task, resulting in a decline in performance on previously learned tasks. given the lack of cl benchmarks for the code domain, we create a benchmark called codetask-cl using existing datasets. it consists of tasks like code completion @xcite @xcite , code translation @xcite , code summarization @xcite , and code refinement @xcite . this benchmark presents a new and challenging scenario as it necessitates the adaptation of the model to varying input and output programming languages. along with this benchmark, we also present a training framework to easily apply cl methods to code generation models. next, we evaluate the effectiveness of popular cl methods from nlp and vision domains in the context of code generation models. we consider prompting methods @xcite and experience-replay (de lange et al., 2019) due to their good performance for pre-trained models @xcite . we also experiment with prompt pooling (pp) @xcite , an effective prompting-based method for cl in the vision domain. our results show that prompt pooling suffers from catastrophic forgetting on our proposed codetask-cl benchmark because of the complex distribution shift from varying input and output programming languages across tasks. with further investigation, we find that the unconstrained prompt selection mechanism leads to an unstable training problem. to address this, we propose our method prompt pooling with teacher forcing (pp-tf), which imposes constraints on prompt selection during training by assigning certain prompts to fixed tasks during training (see figure 1 ). this results in stable training and better performance. interestingly, we find when a replay buffer is available, the simple experience-replay @xcite ) method outperforms other cl methods and achieves performance similar to a multitask baseline @xcite where all tasks are provided at once. in summary, our contributions include: (1) being the first study on cl for code generation tasks, (2) establishing a benchmark and a novel pipeline that supports cl for code generation to motivate future work, (3) identifying and addressing the unstable training issue of prompt pooling through our proposed method pp-tf, and (4) discussion on the best cl methods to use in different use cases. ## related work code generation models. code generation and language modeling for source code is an emerging research field experiencing active growth. several model architectures have been examined recently, including encoder-only models @xcite , encoder-decoder models @xcite , and decoder-only models @xcite @xcite . however, none of these models have been studied in the context of continual learning. continual learning. there are various methods for continual learning (cl) and they fall into three categories: regularization, replay, and parameter isolation methods. regularization methods @xcite @xcite ) assign importance to model components and add regularization terms to the loss function. replay methods @xcite @xcite retain a small memory buffer of data samples and retrain them later to avoid catastrophic forgetting (cf). parameter isolation methods, such as prompting-based methods @xcite @xcite , introduce or isolate network parameters for different tasks. for a more comprehensive overview of all cl methods, we refer the reader to @xcite ; @xcite . to the best of our knowledge, there are currently no studies or benchmarks for cl on code generation models. therefore, we evaluate the effectiveness of prompting @xcite and experience replay @xcite based methods, which have demonstrated strong performance in cl on large pretrained models @xcite . we do not consider regularization methods as they are not effective in continually learning large-scale pretrained models @xcite . next, we discuss our proposed benchmark and methods. ## codetask-cl benchmark we present the codetask-cl benchmark to assess the cl abilities of code generation models. we also provide a novel training pipeline that can be used to continually train and evaluate code generation models. all of the datasets used to create the codetask-cl benchmark are available under the mit license and more details on the dataset splits and input-output domains are in ## prompt pooling with teacher forcing prompt pooling @xcite is a highly effective technique that possesses two key benefits. firstly, the number of prompts required does not increase linearly with the number of tasks. secondly, the prompts within the pool can be utilized across multiple tasks, thereby enabling the reuse of previously acquired knowledge. these abilities are advantageous in real-world scenarios, particularly when a model needs to be continually adjusted to accommodate a large number of users/tasks. in prompt pooling (pp), a set of learnable prompts @xmath1. we follow @xcite and utilize a query and key-matching process to select the prompts for each task. this process has four steps: (1) a learnable key, represented as k @xmath2, is defined for each prompt, resulting in a prompt pool of the form {(k i , p i )} m @xmath3; ## experiments we focus on the scenario of known task identities for continual learning. this is commonly the case in code-related domains and task identities can also be determined through input and output analysis in certain situations. in the field of nlp and vision, methods utilizing experience replay and prompting have been highly effective for cl on large pretrained models @xcite @xcite . moreover, regularization methods are shown to not work well in conjunction with pre-trained models @xcite , and hence, we skip them from our study. next, we present these methods along with some baseline methods. ## conclusion in conclusion, we have introduced a novel benchmark, codetask-cl, tailored to cover a broad spectrum of tasks in the code domain, aiming to fuel advancements in continual learning (cl) for large-scale code generation models. our study underscores the shortfalls of popular cl methods like prompt pooling when applied to coding tasks, predominantly due to catastrophic forgetting. however, we demonstrate that our proposed method, prompt pooling with teacher forcing (pp-tf), can effectively mitigate this issue, leading to a significant improvement of 21.54% over the baseline. furthermore, we establish a comprehensive training pipeline catering to cl on code models. we believe that our contributions, both in the form of the codetask-cl benchmark and the pp-tf method, will ignite further exploration and innovation in cl techniques specifically designed for the dynamic and evolving realm of code generation.