Multi-task Learning based Pre-trained Language Model for Code Completion

Statistical language models capture the statistical patterns in languages by assigning occurrence probabilities to a sequence of words in a particular sequence, which will score an utterance high, if it sounds “natural” to a native speaker, and score low the unnatural (or wrong) sentences. Programming languages are kind of languages that contain predictable statistical properties (Hindle et al., 2012), which can be modeled by statistical language models. Given a token sequence S = $s_{1},s_{2},...,s_{t}$, the probability of the sequence is computed as:

The probabilities are hard to estimate when the number of the context tokens $s_{1},s_{2},...,s_{t-1}$ is tremendous. The N-gram model based on the Markov assumption is proposed to address this challenge, where the probability of a token is dependent only on the $n-1$ most recent tokens. N-gram based models have been generally applied to code completion (Hindle et al., 2012; Tu et al., 2014; Hellendoorn and Devanbu, 2017). These models have been proved to capture the repetitive regularities in the source code effectively. In recent years, deep recurrent neural networks, including Long Short-Term Memory (LSTM) (Hochreiter and Schmidhuber, 1997) and Gate Recurrent Unit (GRU) (Cho et al., 2014), have shown great performance on modeling programming languages (Liu et al., 2016; Bhoopchand et al., 2016; Li et al., 2018). By using recurrent connections and gate mechanisms, information can cycle inside these networks for a long time, which loosens the fixed context size and can capture longer dependencies than the N-gram model.

However, the introduction of the gating mechanism in LSTMs and GRUs might not be sufficient to address the gradient vanishing and explosion issue fully. To ease this issue, attention mechanisms (Bahdanau et al., 2015; Vaswani et al., 2017), which add direct connections between long-distance word pairs, are proposed. For example, the Transformer (Vaswani et al., 2017) is an architecture based solely on attention mechanism. It uses a multi-headed self-attention mechanism to replace the recurrent layers to reduce sequential computation and capture longer-range dependency. Later, Transformer-XL (Dai et al., 2019) is proposed by introducing the notion of recurrence into the deep self-attention network. Thus it enables the Transformer networks to capture the very long-term dependency during language modeling.

Multi-task learning is an approach for knowledge transfer across related tasks. It improves generalization by leveraging the domain-specific information contained in the training signals of related tasks (Caruana, 1997). Through sharing hidden layers among tasks, the model can capture the common features among all the tasks. Furthermore, by preferring the representation that all tasks prefer, the risk of over-fitting is reduced, and the model can be more general to new tasks in the future. Multi-task learning has been successfully used in many fields including natural language processing (Liu et al., 2015; Guo et al., 2018; Devlin et al., 2019), speech recognition (Deng et al., 2013) and computer vision (Long and Wang, 2015; Lu et al., 2017).

Language model pre-training has shown to be effective for NLP, and has achieved the state-of-the-art results across many NLP tasks (Dai and Le, 2015; Devlin et al., 2019; Peters et al., 2018; Radford et al., 2018; Howard and Ruder, 2018). The advantages of the pre-trained model can be summarized as follows: (1 By pre-training on the huge corpus, the model can learn universal representations and help with the target tasks; 2) The pre-trained model can provide a better model initialization, which leads to a better generalization performance on the downstream tasks. 3) Pre-training can be regarded as a kind of regularization to avoid over-fitting on small data. To apply the pre-trained language representations to downstream tasks, the feature-based approaches use the pre-trained representations as additional features (Peters et al., 2018), and the fine-tuning approaches directly adapt the model on the downstream tasks by simply fine-tuning the pre-trained parameters (Radford et al., 2018; Devlin et al., 2019). Generative Pre-trained Transformer (GPT) (Radford et al., 2018) and Bidirectional Encoder Representations from Transformers (BERT) (Devlin et al., 2019) are the widely used fine-tuning approach, where BERT has significantly improved the performance of a wide range of natural language understanding tasks. However, the bidirectionality nature of BERT makes it difficult to be applied to natural language generation tasks. To overcome this limitation, UNIfied pre-trained Language Model (UNILM) (Dong et al., 2019) that can be applied to both natural language understanding (NLU) and natural language generation (NLG) tasks was proposed. Inspired by these models, we build a pre-trained language model for code understanding and generation, and then fine-tune it on code completion.

Given an input program token sequences $x=x_{1},x_{2},...,x_{n}$, CugLM obtains a contextualized vector representation for each token. The model architecture is shown in Figure 2. We adopt an $L$-layer Transformer as the language model to encode the input vectors $x=x_{1},x_{2},...,x_{n}$ into contextual representations at different levels $H^{l}=[h_{1}^{l},h_{2}^{l},...,h_{n}^{l}]$, where $H^{l}=\text{Transformer}_{l}(H^{l-1}),l\in[1,L]$. In Figure 2 and later sections, we omit the superscript $L$ for the hidden vectors of the final layer $h_{i}^{L}$ to make the illustration less cluttered. For each transformer layer (block), multi-attention heads are used to aggregate the output of the previous layer, and the output of a self-attention head $A_{l}$ is computed as:

where $H^{i}\in\mathbb{R}^{|x|\times d_{h}}$ denotes the $i$-th layer’s output. The queries $Q$, keys $K$, and values $V$ are computed by linearly projecting the previous layer’s output $H^{l-1}$ using parameter matrices $W_{l}^{Q},W_{l}^{K},W_{l}^{V}$. $M\in\mathbb{R}^{|x|\times|x|}$ is the mask matrix that determines whether a pair of tokens can be attended to each other. For different pre-training objectives, we use different mask matrices $M$ to control how many contextual tokens can a token attend to when computing its contextualized representations, as illustrated in Figure 2. For bidirectional LM, the elements of the mask matrix are all 0s, which means that all the tokens have access to each other. For unidirectional LM, the upper triangular part of the mask is set to $-\infty$, indicating that each token can only access the leftward context tokens and itself. The output of CugLM includes (1) contextual vector representation of each input token, and (2) the representation of [CLS], which is short for “CLaSsification” and works as the aggregated sequence representation and can be used for classification tasks.

During the pre-training period, the model’s parameters are shared and optimized with several objectives, namely, Masked bidirectional LM, Next Code segment Prediction, and Unidirectional LM. After the model is pre-trained, we can then fine-tune it for downstream tasks. In this paper, we fine-tune CugLM on code completion.

The input $x$ is a token sequence, which is a pair of segments packed together. As shown in Figure 3, for a given token, its vector representation is computed by summing the corresponding token, segment and position embeddings.

• •

  For token embeddings, the embedding matrix is randomly initialized, and then is adjusted as part of the training process. Two special tokens [CLS], [SEP] are defined, where [CLS], which is short for “CLaSsification”, always appears at the beginning of the input. The final hidden state corresponding to [CLS] can be used as the aggregate sequence representation for classification tasks, for example, in next sentence prediction task. [SEP], which is short for “SEPeration”, is used to separate the sentence pairs.

• •

  The segment embeddings, i.e., $E_{A}$ and $E_{B}$ are also used to differentiate the code segment pairs. For each token of the first code segment, a learned embedding $E_{A}$ is added, and a learned embedding $E_{B}$ is added to each token of the second code segment. The embedding matrix for the segment embeddings is also randomly initialized

• •

  To make use of the order of the sequence, we use learned positional embeddings with sequence lengths up to 128 tokens.

To pre-train CugLM, we adopt multi-task learning to learn three tasks jointly, as shown in Figure 2, including Masked bidirectional Language Modeling (MLM), Next Code segment Predicting (NCP), and Unidirectional Language Modeling (ULM). For the first two objectives, the Transformer network is under the bidirectional settings, and for the last objective, the Transformer network is unidirectional.

The final hidden vector corresponding to [CLS], which works as the aggregated sequence representation, is fed into the output softmax layer to produce the probability distribution of classification results.

The pre-training procedure follows the existing language model pre-training approaches. The parameters of CugLM are learned to minimize the sum of the cross-entropy losses of the three pre-training tasks, and are shared among all the tasks. The final loss function is given below:

When the model is pre-trained, we fine-tune it on code completion task. In code completion, the context of the predicted token should only consist of all the token on its left. Thus, the representation of each token can encode only the leftward context tokens and itself. During the fine-tuning procedure, the following two objectives are optimized:

where $E_{type}$ is the embedding of the predicted type, $W^{o}\in\mathbb{R}^{H_{token}\times H}$, $W^{y}\in\mathbb{R}^{V_{token}\times H_{token}}$, $b^{y}\in\mathbb{R}^{V_{token}}$ are trainable parameters. $V_{token}$ is the vocabulary size of the token, and “;” denotes the concatenation operation.

During the fine-tuning procedure, the parameters of CugLM are learned to minimize the sum of the cross-entropy losses of the two fine-tuning tasks and are shared among all the tasks. The final loss function is given below:

Through learning these two objectives jointly, we hope the model can make better predictions on both the identifiers and the other tokens.

We pre-train and fine-tune our model across two programming languages: Java and TypeScript. The programs in the corpus are collected from publicly available open-source GitHub repositories by removing duplicate files and project forks. Each program is tokenized into token sequence. The detailed information is shown in Table 1. We use 60% of the projects for pre-training, and 40% of the projects for fine-tuning on code completion task. During the fine-tuning, we split the projects into train/validation/test sets in the proportion 8:1:1. For the other baselines, all the programs used in pre-training and the training programs used in fine-tuning are used as the training set, and the validation and test sets are the same as in our fine-tuning procedure. We also randomly sample 200 program files from both Java and TypeScript test sets as the small test sets for Byte Pair Encoding based Neural Language Model (BPE NLM) (Karampatsis et al., 2020) evaluation since when performing completion (testing) in their model, they use a variation of the beam search algorithm to combine the sub-units to complete tokens, which is very time-consuming. It takes several minutes to complete a single program file and will take tens of days to perform completion on the large test sets (e.g., the Java test set contains 14,600 files). Thus, we create small test sets.

For Java programs, we extract the identifiers’ type information through static analysis. For TypeScript programs, we apply the approach in Hellendoorn et al. (2018) to extract type annotations of the identifiers. We filter the programs to make sure at least 10% of type annotations are user-defined types in each TypeScript file. Figure 5 shows the examples for Java and TypeScript code, where the identifiers that have type are marked with underlines, and the green tokens next to the identifiers are the corresponding types. To generate each training input sequence for pre-training, we sample two spans of tokens from the corpus, which we refer to as segments $S_{1}$ and $S_{2}$. Each segment contains several lines of source code tokens. For the first segment $S_{1}$, we sample the first $N$ lines from one program file, where $N$ is randomly sampled from 1 to the length of the code lines of the program file. 50% of the time, the second segment $S_{2}$ is the rest of the lines from the same program file that follows $S_{1}$, and 50% of the time it is a random code segment sampled from other program files of the corpus, which is done for the “next code segment prediction (NCP)” task. They are sampled such that the combined length is $\leq$ 128 tokens. For the “Masked bidirectional Language Modeling (MLM)” task, we only mask those identifiers that have type information. For example, the underlined tokens in Figure 5.

To evaluate our proposed approach, in this section, we conduct experiments to investigate the following research questions:

1) Comparison with LSTM based closed vocabulary models (the first two baselines): To compare with Pointer Mixture Network, we downloaded their publicly available source code¹¹1https://github.com/jack57lee/neuralCodeCompletion. In their model, the programs in the datasets are parsed into ASTs, and they build the model to perform code completion on AST node sequences. Although the ASTs can provide more information, representing the programs as AST node sequences is not the natural order of typing, and the precision does not directly reflect the productivity gain of the code completion tool. More importantly, in practice, the code is incomplete, so the software project might not be compilable (code is not parsable into ASTs, or parsed ASTs miss a lot of information). Thus, representing programs as token sequences and performing code completion on the token-level might be more practical. In this paper, we focus on token-level code completion. In our corpus, the programs are tokenized into token sequences. To compare with them, we train their model within our tokenized programs using the command line arguments given in the artifact’s README file ²²2Since the Pointer Mixture Network also makes use of the additional information derived from ASTs, the results of using token sequence as input might understate the accuracy of the plain Pointer Mixture Network.. Their base model is a single layer LSTM network with an unrolling length of 50 and hidden unit size of 1500. The initial learning rate is 0.001 and is decayed by multiplying 0.6 after every epoch. The gradients’ norm is clipped to 5. The size of the attention window is 50. Since the Pointer Mixture Network is based on LSTM language model, we also list the results of the vanilla LSTM, where the parameter configuration of the vanilla LSTM network is set the same as the Pointer Mixture Network.

As shown from the results, our model outperforms the two LSTM-based models on both Java and TypeScript datasets by a large margin, especially in identifier completion. On the Java large test set, our model achieves the accuracy of 84.06% and 55.19% on token’s completion and identifier’s completion, respectively, which outperforms Pointer Mixture Network by 23.07% and 43.69%, in terms of relative improvement. On the TypeScript large test set, our model achieves the accuracy of 82.14% and 41.85% on token’s completion and identifier’s completion, respectively, which outperforms Pointer Mixture Network by 19.47% and 23.96%. The results on small test sets are similar to the large test set. We can find that the improvements on the TypeScript dataset are smaller than Java, especially in identifier completion. The reason lies in that, the (masked) identifier proportion in TypeScript (9.74%) is smaller than Java (21.04%) because the type information in TypeScript is annotated by developers, and only a part of the identifiers are annotated. In the MLM pre-training task, these identifiers are masked and are predicted based on their contextual tokens aiming at generating better contextual and informative representations for these identifiers as well as other tokens. During fine-tuning, the type information of these identifiers is used to assist the identifiers’ prediction. Due to the lower masked proportion, the pre-training and fine-tuning procedure can offer less information than Java, thus resulting in smaller improvements.

2) Comparison with open vocabulary model (BPE NLM): To compare with BPE NLM, we downloaded their publicly available source code³³3https://github.com/mast-group/OpenVocabCodeNLM and train their model on our datasets. They use a single layer GRU NLM with an unrolling length of 200 built upon sub-word units learned from BPE. The embedding size and the hidden unit size are both set to 512 in their model. To keep the number of parameters comparable with our model and other baselines, we increase the hidden unit size and the embedding size of their model to 1500. There are three scenarios: static, dynamic, and maintenance, where the dynamic and maintenance settings update model’s parameters during testing. Since our model and other baselines do not update parameters during the test process, we present the results of the static scenario to make the comparison fair, and realize that evaluating dynamically may improve accuracy. As shown from the results, BPE NLM performs best on completing identifiers among all the baseline models on both datasets, which proves the power of the open vocabulary LM for predicting the identifiers. Even though, our model still outperforms the BPE NLM on completing identifiers. When evaluating on completing all kinds of tokens, the performance of BPE NLM is not as well as the identifier completion. Our model outperforms BPE NLM on completing all kinds of tokens by a large margin.

3) Comparison with transformer network based model (Transformer-XL): To find out if CugLM’s promising results derive more from using a Transformer-based model for code completion, or from the multi-task learning based pre-training and fine-tuning, we compare our results to a Transformer-based model trained from scratch, i.e., without the benefit of a pre-trained embedding. Transformer-XL is a Transformer network based language model, which introduces the notion of recurrence to model the long-term dependency of the input sequence. We use a 6-layer Transformer-XL network with 5 parallel heads. The dimension of each head is set to 64. We set the segment length to be 128, and the length of the cached segments to 256. The dimensionality of the model (hidden unit) and the embedding size is set to 800. The dimension of the feed-forward layer is set to 1024. As seen from Table 2, transformer-XL model outperforms the other baseline models that are based on the recurrent neural networks on both datasets, which demonstrates that the Transformer-based network is more powerful than recurrent neural networks on this task. The performance of our model is substantially higher than the Transformer-XL model trained from scratch. We therefore conclude that pre-training and fine-tuning are crucial to CugLM’s success.

The above results demonstrate that all of the pre-training tasks are necessary to improve the performance, and MLM contributes most to the improvements.

Except for identifiers, we also give a detailed breakdown of the accuracies for completing different types of tokens on both our model and BPE NLM (Karampatsis et al., 2020), and also present these tokens’ proportion. The results are shown in Table 4. Punctuations make up the majority of the completions, and the accuracies of both our model and BPE NLM on predicting punctuations are high, where BPE NLM performs better than CugLM. The punctuation tokens are much easier to complete than identifiers, but these completions are not that useful for developers (Karampatsis et al., 2020). For keyword completions, our model outperforms BPE LM by a large margin. The keywords are predefined, reserved words used in programming that have special meanings to the compiler, which contain the syntactic information or the attribute information of the objects. The great performance of CugLM on completing keywords further demonstrates that through multi-task learning based pre-training and fine-tuning, the representations generated by our model can capture syntactic and semantic information better. For numeral and operator completions, which are more related to the semantic of the programs, our model also outperforms BPE NLM substantially.

To analyze the complexity of our model and the baseline models, we present the number of trainable parameters for all the models, shown in Table 5. The number of trainable parameters of our model is less than all the baselines. Although we adopt multi-task learning for both pre-training and fine-tuning, the number of trainable parameters does not increase much as all of the tasks share one multi-layer transformer network. To improve training efficiency and avoid over-fitting, we do not use large parameter settings. Even though, our model still outperforms the other baselines by a large margin thanks to the pre-training and fine-tuning.

To further improve the performance of our model, we also conduct experiments on applying Byte Pair Encoding (BPE) algorithm to build up the vocabulary of sub-words as in (Karampatsis et al., 2020), where the rare tokens will be segmented into more common sub-word units, and no word is OoV. However, the performance on Java corpus is comparable with the origin model, and the accuracy decreases slightly on TypeScript corpus. We analyze the possible reasons are as follows. During pre-training, we mask the identifiers with type information. When we apply BPE algorithm, most of these masked identifiers will be split into sub-word units. Thus, all of these units will be masked, which leads to the high mask proportion and increased the difficulty of learning the semantics of embeddings. Besides, during fine-tuning, our model utilizes the predicted type information to assist the token’s prediction. After splitting the tokens into sub-word units, all of the units from one token correspond to the same type, resulting in the semantic inconsistencies between the type information and the sub-word units. For example, the same unit from different tokens might correspond to different types. Thus, applying BPE does not improve the performance of our model.