P RE - TRAINED C ONTEXTUAL E MBEDDING OF S OURCE C ODE
Anonymous authors
Paper under double-blind review
A BSTRACT
The source code of a program not only serves as a formal description of an executable task, but it also serves to communicate developer intent in a humanreadable form.
To facilitate this, developers use meaningful identiﬁer names and natural-language documentation.
This makes it possible to successfully apply sequence-modeling approaches, shown to be effective in natural-language processing, to source code.
A major advancement in natural-language understanding has been the use of pre-trained token embeddings; BERT and other works have further shown that pre-trained contextual embeddings can be extremely powerful and can be ﬁnetuned effectively for a variety of downstream supervised tasks.
Inspired by these developments, we present the ﬁrst attempt to replicate this success on source code.
We curate a massive corpus of Python programs from GitHub to pre-train a BERT model, which we call Code Understanding BERT (CuBERT).
We also pre-train Word2Vec embeddings on the same dataset.
We create a benchmark of ﬁve classiﬁcation tasks and compare ﬁnetuned CuBERT against sequence models trained with and without the Word2Vec embeddings.
Our results show that CuBERT outperforms the baseline methods by a margin of 2.9–22%.
We also show its superiority when ﬁnetuned with smaller datasets, and over fewer epochs.
We further evaluate CuBERT’s effectiveness on a joint classiﬁcation, localization and repair task involving prediction of two pointers.
I NTRODUCTION
Modern software development places a high value on writing clean and readable code.
This helps other developers understand the author’s intent so that they can maintain and extend the code.
Developers use meaningful identiﬁer names and natural-language documentation to make this happen (Martin, 2008).
As a result, source code contains substantial information that can be exploited by machine-learning algorithms.
Sequence modeling on source code has been shown to be successful in a variety of software-engineering tasks, such as code completion (Hindle et al., 2012; Raychev et al., 2014), source code to pseudocode mapping (Oda et al., 2015), API-sequence prediction (Gu et al., 2016), program repair (Pu et al., 2016; Gupta et al., 2017), and natural language to code mapping (Iyer et al., 2018), among others.
The distributed vector representations of tokens, called token (or word) embeddings, are a crucial component of neural methods for sequence modeling.
Learning useful embeddings in a supervised setting with limited data is often difﬁcult.
Therefore, many unsupervised learning approaches have been proposed to take advantage of large amounts of unlabeled data that are more readily available.
This has resulted in ever more useful pre-trained token embeddings (Mikolov et al., 2013a; Pennington et al., 2014).
However, the subtle differences in the meaning of a token in varying contexts are lost when each word is associated with a single representation.
Recent techniques for learning contextual embeddings (McCann et al., 2017; Peters et al., 2018;
Radford et al., 2018; 2019; Devlin et al., 2019;
Yang et al., 2019) provide ways to compute representations of tokens based on their surrounding context, and have shown signiﬁcant accuracy improvements in downstream tasks, even with only a small number of task-speciﬁc parameters.
Inspired by the success of pre-trained contextual embeddings for natural languages, we present the ﬁrst attempt to apply the underlying techniques to source code.
In particular, BERT (Devlin et al., 2019) produces a bidirectional Transformer encoder (Vaswani et al., 2017) by training it to predict values of masked tokens and whether two sentences follow each other in a natural discourse.
The pre-trained model can be ﬁnetuned for downstream supervised tasks and has been shown to produce state-of-the-art results on a number of NLP benchmarks.
In this work, we derive contextual embedding of source code by training a BERT model on source code.
We call our model CuBERT, short for Code Understanding BERT .
In order to achieve this, we curate a massive corpus of Python programs collected from GitHub.
GitHub projects are known to contain a large amount of duplicate code.
To avoid biasing the model to such duplicated code, we perform deduplication using the method of Allamanis (2018).
The resulting corpus has 6.6M unique ﬁles with a total of 2 billion words.
We also train Word2Vec embeddings (Mikolov et al., 2013a;b), namely, continuous bag-of-words (CBOW) and Skipgram embeddings, on the same corpus.
For evaluating CuBERT, we create a benchmark of ﬁve classiﬁcation tasks, ranging from classiﬁcation of source code according to presence or absense of certain classes of bugs, to mismatch between a function’s natural language description and its body, to predicting the right kind of exception to catch for a given code fragment.
These tasks are motivated by prior work in this space, but unfortunately, the associated datasets come from different languages and varied sources.
We want to ensure that there is no overlap between pre-training and ﬁnetuning datasets, and that all of the tasks are deﬁned on Python code.
We therefore create new datasets for the ﬁve tasks after carefully separating the pre-training and ﬁnetuning corpora.
To evaluate CuBERT’s effectiveness on a more complex task, we create a task for joint classiﬁcation, localization and repair of variable misuse bugs (Vasic et al., 2019), which involves predicting two pointers.
We ﬁnetune CuBERT on each of the classiﬁcation tasks and compare the results with multi-layered bidirectional LSTM (Hochreiter & Schmidhuber, 1997) models.
We train the LSTM models from scratch and also using pre-trainined Word2Vec embeddings.
Our results show that CuBERT consistently outperforms these baseline models by 2.9–22% across the tasks.
We perform a number of additional studies by varying the sampling strategies used for training Word2Vec models, by varying program lengths, and by comparing against Transformer models trained from scratch.
In addition, we also show that CuBERT can be ﬁnetuned effectively using only 33% of the task-speciﬁc labeled data and with only 2 epochs, and that it attains results competitive to the baseline models trained with the full datasets and much larger number of epochs. CuBERT when ﬁnetuned on the variable misuse localization and repair task, produces high classiﬁcation, localization and localization+repair accuracies.
The contributions of this paper are as follows:
• We present the ﬁrst attempt at pre-training a BERT contextual embedding of source code.
• We show the efﬁcacy of the pre-trained contextual embedding on ﬁve classiﬁcation tasks.
Our results show that the ﬁnetuned models outperform the baseline LSTM models supported by Word2Vec embeddings, and Transformers trained from scratch.
Further, the ﬁnetuning works well even for smaller datasets and fewer training epochs. We also evaluate CuBERT on a multi-headed pointer prediction task.
• We plan to make the models and datasets publicly available for use by others.
R ELATED W ORK
Given the abundance of natural-language text, and the relative difﬁculty of obtaining labeled data, much effort has been devoted to using large corpora to learn about language in an unsupervised fashion, before trying to focus on tasks with small labeled training datasets.
Word2Vec (Mikolov et al., 2013a;b) computed word embeddings based on word co-occurrence and proximity, but the same embedding is used regardless of the context.
The continued advances in word embeddings (Pennington et al., 2014) led to publicly released pre-trained embeddings, used in a variety of tasks.
To deal with varying word context, contextual word embeddings were developed (McCann et al., 2017; Peters et al., 2018;
Radford et al., 2018; 2019), in which an embedding is learned for the context of a word in a particular sentence, namely the sequence of words preceding it and possibly following it.
BERT (Devlin et al., 2019) improved natural-language pre-training by using a denoising autoencoder.
Instead of learning a language model, which is inherently sequential, BERT optimizes for predicting a noised word within a sentence.
Such prediction instances are generated by choosing a word position and either keeping it unchanged, removing the word, or replacing the word with a random wrong word.
It also pre-trains with the objective of predicting whether two sentences can be next to each other.
These pre-training objectives, along with the use of a Transformer-based architecture, gave BERT an accuracy boost in a number of NLP tasks over the state-of-the-art.
BERT has been improved upon in various ways, including modifying training objectives, utilizing ensembles, combining attention with autoregression (Yang et al., 2019), and expanding pre-training corpora and time (Liu et al., 2019).
However, the main architecture of BERT seems to hold up as the state-of-the-art, as of this writing.
In the space of programming languages, attempts have been made to learn embeddings in the context of speciﬁc software-engineering tasks.
These include embeddings of variable and method identiﬁers using local and global context (Allamanis et al., 2015), abstract syntax trees or ASTs (Mou et al., 2016), paths in ASTs (Alon et al., 2019), memory heap graphs (Li et al., 2016), and ASTs enriched with data ﬂow information (Allamanis et al., 2018).
These approaches require analyzing source code beyond simple tokenization.
In this work, we derive a pre-trained contextual embedding of tokenized source code without explicitly modeling source-code-speciﬁc information, and show that the resulting embedding can be effectively ﬁnetuned for downstream tasks.
E XPERIMENTAL S ETUP 3. C ODE C ORPUS FOR F INETUNING T ASKS
We use the ETH Py150 corpus (Raychev et al., 2016) to generate datasets for the ﬁnetuning tasks.
The ETH Py150 corpus consists of 150K Python ﬁles from GitHub, and is partitioned into a training split (100K ﬁles) and a test split (50K ﬁles).
We held out 10K ﬁles from the training split as a validation split.
We deduplicated the dataset in the fashion of Allamanis (2018), resulting in a slightly smaller dataset of 85K, 9.5K, and 47K ﬁles in train, validation, and test, respectively.
3. T HE G IT H UB P YTHON P RE -T RAINING C ODE C ORPUS
We used the public GitHub repository hosted on Google’s BigQuery platform (the github repos dataset under BigQuery’s public-data project, bigquery-public-data ).
We extracted all ﬁles ending in .py , under open-source, redistributable licenses, removed symbolic links, and retained only ﬁles reported to be in the refs/heads/master branch.
This resulted in about 16.1M ﬁles.
To avoid duplication between pre-training and ﬁnetuning data, we removed ﬁles that had high similarity to the ﬁles in the ETH Py150 dataset, using the method of Allamanis (2018).
In particular, two ﬁles are considered similar to each other if the Jaccard similarity between the sets of tokens (identiﬁers and string literals) is above 0.8 and in addition, it is above 0.7 for multi-sets of tokens.
This brought the dataset to 13.5M ﬁles.
We then further deduplicated the remaining ﬁles, by clustering them into equivalence classes holding similar ﬁles according to the same similarity metric, and keeping only one exemplar per equivalence class.
This helps avoid biasing the pre-trained embedding.
Finally, we removed ﬁles that could not be tokenized.
In the end, we were left with 6.6M Python ﬁles containing over 2 billion words.
This is our Python pre-training code corpus.
3. S OURCE C ODE M ODELING
We ﬁrst tokenize a Python program using the standard Python tokenizer (the tokenize package).
We leave language keywords intact and produce special tokens for syntactic elements that have either no string representation (e.g., DEDENT tokens, which occur when a nested program scope concludes), or ambiguous interpretation (e.g., new line characters inside string literals, at the logical end of a Python statement, or in the middle of a Python statement result in distinct special tokens).
We split identiﬁers according to common heuristic rules (e.g., snake or Camel case).
Finally, we split string literals using heuristic rules, on whitespace characters, and on special characters.
We limit all thus produced tokens to a maximum length of 15 characters.
We call this the program vocabulary .
Our Python pre-training code corpus contained 10.2M unique tokens, including 12 reserved tokens.
We greedily compress the program vocabulary into a subword vocabulary (Schuster & Nakajima, 2012) using the SubwordTextEncoder from the Tensor2Tensor project (Vaswani et al., 2018), resulting in slightly over 50K tokens.
All words in the program vocabulary can be losslessly encoded using one or more of the subword tokens.
We encode programs ﬁrst into program tokens, as described above, and then encode those tokens one by one in the subword vocabulary.
The objective of this encoding scheme is to preserve syntactically meaningful boundaries of tokens.
For example, the identiﬁer “ snake case ” could be encoded as “ sna ke ca se ”, preserving the snake case split of its characters, even if the subtoken “ e c ” were very popular in the corpus; the latter encoding might result in a smaller representation but would lose the intent of the programmer in using a snake-case identiﬁer.
Similarly, “ i=0 ” may be very frequent in the corpus, but we still force it to be encoded as separate tokens i , = , and 0 , ensuring that we preserve the distinction between operators and operands.
Both the BERT model and the Word2Vec embeddings are built on the subword vocabulary.
3. F INETUNING T ASKS
To evaluate CuBERT, we design ﬁve classiﬁcation tasks and a multi-headed pointer task.
These are motivated by prior work, but unfortunately, the associated datasets come from different languages and varied sources.
We want the tasks to be on Python code, and for accurate results, we ensure that there is no overlap between pre-training and ﬁnetuning datasets.
We therefore create all the tasks on the ETH Py150 corpus (see Section 3.1).
As discussed in Section 3.2, we ensure that there is no duplication between this and the pre-training corpus.
We hope that our datasets for these tasks will be useful to others as well.
The ﬁnetuning tasks are described below.
A more detailed discussion is presented in Appendix A.
Variable Misuse Classiﬁcation Allamanis et al.
(2018) observed that developers may mistakenly use an incorrect variable in the place of a correct one.
These mistakes may occur when developers copy-paste similar code but forget to rename all occurrences of variables from the original fragment, or when there are similar variable names in contexts that can be confused with each other.
These can be subtle errors that remain undetected during compilation.
The task by Allamanis et al.
(2018) is to predict a correct variable name at a location within a function and was devised on C# programs.
We take the classiﬁcation version restated by Vasic et al.
(2019), wherein, given a function, the task is to predict whether there is a variable misuse at some location in the function, without specifying a particular location to consider.
In this setting, the classiﬁer has to consider all variables and their usages to make the decision.
In order to create negative (buggy) examples, we replace a variable use at some location with another variable that is deﬁned within the function.
Wrong Binary Operator Pradel & Sen (2018) proposed the task of detecting whether a binary operator in a given expression is correct.
They use features extracted from limited surrounding context.
We use the entire function with the goal of detecting whether any binary operator in the function is incorrect.
The negative examples are created by randomly replacing some binary operator with another type-compatible operator.
Swapped Operand Pradel & Sen (2018) propose the wrong binary operand task where a variable or constant is used incorrectly in an expression, but that task is quite similar to the variable misuse task we already use.
We therefore deﬁne another class of operand errors where the operands of noncommutative binary operators are swapped.
The operands can be arbitrary subexpressions, and are not restricted to be just variables or constants.
To simplify example generation, we restrict examples for this task to those in which the binary operator and its operands all ﬁt within a single line.
Function-Docstring Mismatch Developers are encouraged to write descriptive docstrings to explain the functionality and usage of functions.
This provides parallel corpora between code and natural language sentences that have been used for machine translation between the two (Barone & Sennrich, 2017), detecting uninformative docstrings (Louis et al., 2018) and to evaluate their utility to provide supervision in neural code search (Cambronero et al., 2019).
We prepare a sentence-pair classiﬁcation problem where the function and its docstring form two distinct sentences. Similar to the other ﬁnetuning tasks, we use the ETH Py150 corpus to create this dataset.
The positive examples come from the correct function-docstring pairs.
We create negative examples by replacing correct docstrings with docstrings of other functions, randomly chosen from the dataset.
For this task, the existing docstring is removed from the function body. [Table]
Exception Type While it is possible to write generic exception handlers (e.g., “ except
Exception ” in Python), it is considered a good coding practice to catch and handle the precise exceptions that can be raised by a code fragment.
We identiﬁed the 20 most common exception types from the GitHub dataset, excluding the catch-all Exception (full list in Table 6).
Given a function with an except clause for one of these exception types, we replace the exception with a special “hole” token.
The task is the multi-class classiﬁcation problem of predicting the original exception type. Variable Misuse Localization and Repair
As an instance of a non-classiﬁcation task, we consider the joint classiﬁcation, localization and repair version of the variable misuse task from Vasic et al.
Given a function, the task is to predict one pointer (called the localization pointer) to identify a variable misuse location and another pointer (called the repair pointer) to identify a variable from the same function that is the right one to use at the faulty location.
The model is also trained to classify functions that do not contain any variable misuse as bug-free by making the localization pointer point to a special location in the function.
We create negative examples using the same method as used in the Variable Misuse Classiﬁcation task.
Table 1 lists the sizes of the resulting benchmark datasets extracted from the (deduplicated) ETH
Py150 corpus.
The Exception Type task contains fewer examples than the other tasks, since examples for this task only come from functions that catch one of the chosen 20 exception types.
3. BERT FOR S OURCE C ODE
The BERT model (Devlin et al., 2019) consists of a multi-layered Transformer encoder.
It is trained with two tasks: (1) to predict the correct tokens in a fraction of all positions, some of which have been replaced with incorrect tokens or the special [MASK] token (the Masked Language Model task) and (2) to predict whether the two sentences separated by the special [SEP] token follow each other in some natural discourse (the Next Sentence Prediction task).
Thus, each example consists of one (for MLM) or two (for NSP) sentences , where a sentence is the concatenation of contiguous lines from the source corpus, sized to ﬁt the target example length.
To ensure that every sentence is treated in multiple instances of both MLM and NSP, BERT by default duplicates the corpustimes, and generates independently derived examples from each duplicate.
With 50% probability, the second example sentence comes from a random document (for NSP).
With 15% probability, a token is chosen for an MLM prediction (up to 20 per example), and from those chosen, 80% are masked, 10% are left undisturbed, and 10% are replaced with a random token.
CuBERT is similarly formulated, but a CuBERT sentence is a logical code line, as deﬁned by the
Python standard.
Intuitively, a logical code line is the shortest sequence of consecutive lines that may constitute a legal statement, e.g., it has correctly matching parentheses.
We count example lengths by counting the subword tokens of both sentences (see Section 3.3).
We train the BERT Large model, consisting of 24 layers with 16 attention heads and hidden size of 1024 units.
Sentences are created by parsing our pre-training dataset.
Task-speciﬁc classiﬁers pass the embedding of a special start-of-example [CLS] token through feedforward and softmax layers.
For the pointer prediction task, the pointer is computed over the sequence of outputs generated by the last layer of the BERT model.
3. B ASELINES 3.6. W ORD 2V EC
We train Word2Vec models using the same pre-training corpus as the BERT model.
To maintain parity, we generate the dataset for Word2Vec using the same pipeline as BERT but by disabling masking and generation of negative examples for NSP.
The dataset is generated without any duplication.
We train both CBOW and Skipgram models using GenSim ( ˇReh˚uˇrek & Sojka, 2010).
To deal with the large vocabulary, we use negative sampling and hierarchical softmax (Mikolov et al., 2013a;b) to train the two versions.
In all, we obtain four Word2Vec embeddings.
. B IDIRECTIONAL LSTM AND T RANSFORMER
In order to obtain context-sensitive encodings of input sequences for the ﬁnetuning tasks, we use multi-layered bidirectional LSTMs (Hochreiter & Schmidhuber, 1997) (BiLSTMs).
These are initialized with the pre-trained Word2Vec embeddings.
Additionally, to further evaluate whether
LSTMs alone are sufﬁcient without pre-training, we try initializing the BiLSTM with an embedding matrix that is trained from scratch.
We also trained Transformer models (Vaswani et al., 2017) for our ﬁnetuning tasks.
We used BERT’s own Transformer implementation, to ensure comparability of results.
For comparison with prior work, we use the unidirectional LSTM and pointer model from (Vasic et al., 2019) for the Variable Misuse Localization and Repair task.
E XPERIMENTAL R ESULTS 4. T RAINING D ETAILS
As stated above, CuBERT’s dataset generation duplicates the corpus 10 times, whereas Word2Vec is trained without duplication.
To compensate for this difference, we trained Word2Vec for 10 epochs and CuBERT for 1 epoch.
We pre-train CuBERT with the default conﬁguration of the BERT Large model.
For sequences of length 128, 256 and 512, we use batch sizes of 8192, 4096 andrespectively.
For Word2Vec, when training with negative samples, we choose 10 negative samples.
The embedding sizes for all the pre-trained models are set at 1024.
For the baseline BiLSTM models, we did extensive experimentation on the Variable Misuse task by varying the number of layers (1–3) and the number of hidden units (128, 256, 512).
We also tried
LSTM output dropout probability (0.1, 0.5), optimizers (Adam (Kingma & Ba, 2014) and AdaGrad (Duchi et al., 2011)), and learning rates (1e-3, 1e-4, 1e-5).
The most promising combination was a 3-layered BiLSTM with 512 hidden units per layer, LSTM output dropout probability of 0.
and Adam optimizer with learning rate of 1e-3.
We use this set of parameters for all the tasks except the Exception Type task.
Due to the much smaller dataset size of the latter (Table 1), we did a separate search and chose a single-layer BiLSTM with 256 hidden units.
We used the batch size of 8192 for the larger tasks and 64 for the Exception Type task.
For the baseline Transformer models, we originally attempted to train a Transformer model of the same conﬁguration as CuBERT.
However, the size of our training dataset seemed too small to train that large a Transformer.
Instead, we performed a hyperparameter search over transformer layers (1–6), hidden units (128, 256, 512), learning rates (5e-5, 1e-4, 5e-4, 1e-3) and batch sizes (64, 256, 1024, 2048, 4096, 8192) on the Variable Misuse task.
The best architecture (4 layers, 512 hidden units, 16 attention heads, learning rate of 5e-4, batch size of 4096) is used for all the tasks except the Exception Type task.
A separate experimentation for the smaller Exception Type dataset resulted in the best conﬁguration of 3 layers, 512 hidden units, 16 attention heads, learning rate of 5e-5, and batch size of 2048. Finally, for our baseline pointer model (referred to as LSTM+pointer below) we searched over the following hyperparameter choices: hidden sizes of 512 and 1024, token embedding sizes of 512 and 1024, learning rates of 0.1, 0.01, and 0.001, and the AdaGrad and Gradient Descent optimizers. In contrast to the original work, we generated one pair of buggy/bug-free examples per function (rather than one per variable use, per function, which would bias towards longer functions), use CuBERT’s subword-tokenized vocabulary of 50K subtokens (rather than a limited full-token vocabulary, which leaves many tokens out of vocabulary).
[Table] 4. R ESEARCH Q UESTIONS
We set out to answer the following research questions.
We will address each with our results.
Do contextual embeddings help with source-code analysis tasks, when pre-trained on an unlabeled code corpus?
We compare CuBERT to BiLSTM models with and without pretrained Word2Vec embeddings on the classiﬁcation tasks (Section 4.3).
Does ﬁnetuning actually help, or is the Transformer model behind CuBERT the main power behind the approach?
We compare ﬁnetuned CuBERT models to Transformer-based models trained from scratch on the classiﬁcation tasks (Section 4.4).
How does the performance of CuBERT on the classiﬁcation tasks scale with the amount of labeled training data?
We compare the performance of ﬁnetuned CuBERT models when ﬁnetuning with one third, two thirds, or the full training dataset for each task (Section 4.5).
How does example length affect the beneﬁts of CuBERT?
We compare ﬁnetuning performance for different example lengths on the classiﬁcation tasks (Section 4.6).
How does CuBERT perform on complex tasks?
We implemented and ﬁnetuned a model for a multi-headed pointer prediction task, namely, the Variable-Misuse Localization and Repair task (Section 4.7).
We compare it to the model from Vasic et al. Except for Section 4.6, all the results are presented for sequences of length 512.
We give examples of classiﬁcation instances in Appendix B and include visualizations of attention weights for them.
4. C ONTEXTUAL VS . W ORD E MBEDDINGS
The purpose of this analysis is to understand how much pre-trained contextual embeddings help, compared to word embeddings.
For each classiﬁcation task, we trained BiLSTM models starting with each of our baseline Word2Vec embeddings, namely, continuous bag of words (CBOW) and Skipgram trained with negative sampling or hierarchical softmax.
In all the models, the Word2Vec embeddings can be reﬁned during training.
Within the ﬁrst 100 epochs, the performance of the
BiLSTM models stopped improving.
The best model weights per task were selected by ﬁnding the minimum validation loss on the corresponding dataset (Table 1) over the ﬁrst 100 epochs.
On the
CuBERT side, we ﬁnetuned the pre-trained model for 20 epochs, with similar model selection.
The resulting test-split accuracies are shown in Table 2.
CuBERT consistently outperforms BiLSTM (with the best task-wise Word2Vec conﬁguration) on all tasks, by a margin of 2.9–22%.
Thus, the pre-trained contextual embedding provides superior results even with a smaller budget ofepochs, compared to the 100 epochs used for BiLSTMs.
The Exception Type classiﬁcation task is an interesting case since it has an order of magnitude less training data than the other tasks (see
[Table] Table 1).
The difference between the performance of BiLSTM and CuBERT is the highest for this task.
Thus, ﬁnetuning is of much value for tasks with training
We analyzed the performance of CuBERT with the reduced ﬁnetuning budget of only 2 andepochs (see Table 2).
Except for the Operand task, CuBERT outperforms BiLSTM within 2 ﬁnetuning epochs.
On the Operand task, the performance difference between CuBERT with 2 orﬁnetuning epochs and BiLSTM is about 1%.
For the rest of the tasks, CuBERT with only 2 ﬁnetuning epochs outperforms BiLSTM (with the best task-wise Word2Vec conﬁguration) by a margin of 0.7–12%.
This shows that CuBERT can reach accuracies that are comparable to or better than those of BiLSTMs trained with Word2Vec embeddings within only a few epochs.
We also trained the BiLSTM models from scratch, that is, without using the Word2Vec embeddings.
The results are shown in the ﬁrst row of Table 2.
Compared to those, the use of Word2Vec embeddings performs better by a margin of 1.5–10.5%.
Though no single Word2Vec conﬁguration is the best, CBOW trained with negative sampling gives the most consistent results overall.
4. I S T RANSFORMER A LL Y OU N EED ?
One may wonder if CuBERT’s promising results derive more from using a Transformer-based model for its classiﬁcation tasks, and less from the actual, unsupervised pre-training.
Here we compare our results on the classiﬁcation tasks to a Transformer-based model trained from scratch, i.e., without the beneﬁt of a pre-trained embedding.
All the models were trained for 100 epochs during which their performance stopped improving.
We selected the best model per task using least validation loss.
As seen from the last row of Table 2, the performance of CuBERT is substantially higher than the Transformer models trained from scratch.
We therefore conclude that pre-training is crucial to CuBERT’s success.
4. T HE E FFECTS OF L ITTLE S UPERVISION
The big draw of unsupervised pre-training followed by ﬁnetuning is that some tasks have small labeled datasets.
We study here how CuBERT fares when the size of its training split is reduced.
We sampled uniformly the training split of ETH Py150 down to 2 / 3 rds and 1 / 3 rd of its original size, and produced training datasets for each of the classiﬁcation tasks from each sub-split.
We then ﬁnetuned the pre-trained CuBERT model with each of the 3 different training splits.
Validation and testing were done with the same original datasets.
Table 3 shows the results.
The Function Docstring task seems robust to the reduction of the training dataset, both early and late in the ﬁnetuning process (that is, within 2 vs. 20 epochs), whereas the Exception Classiﬁcation task is heavily impacted by the dataset reduction, given that it has relatively few training examples to begin with.
Interestingly enough, for some tasks, even ﬁnetuning for only 2 epochs and only using a third of the training data outperforms the baselines.
For example, for both Variable Misuse and
Function Docstring, CuBERT at 2 epochs and 1 / 3 rd training data outperforms the BiLSTM with
Word2Vec and the Transformer baselines. [Table] [Table]
4. T HE E FFECTS OF R EDUCING C ONTEXT
Context size is especially useful in code tasks, given that some relevant information may lie many “sentences” away from its locus of interest.
Here we study how reducing the context length (i [Table] e., the length of the examples used to pre-train and ﬁnetune) affects performance.
We produce data with shorter example lengths by following the standard BERT mechanism.
Table 4 shows the results [Table]
Although context seems to be important to most tasks, the Function Docstring task seems to improve with reduced context.
This may be because the task primarily depends on comparison between the docstring and the function signature [Table] and including more context dilutes the model’s focus.
For comparison, we also evaluated the BiLSTM model on sequences of length 128 and 256 for the Variable Misuse task.
We obtained accuracies of 71.34% and 73.63% respectively, which are lower than the best BiLSTM accuracy on sequence length 512 and also lower than the accuracies of
CuBERT for the corresponding lengths (see Table 4).
4. E VALUATION ON A M ULTI - HEADED P OINTER P REDICTION T ASK
We now discuss the results of ﬁnetuning CuBERT to predict the localization and repair pointers for the variable misuse task. For this task, we implement the multi-headed pointer model from Vasic et al. (2019) on top of CuBERT.
The baseline consists of the same pointer model on a unidirectional LSTM as used in Vasic et al.
We refer to these respectively as CuBERT+pointer and
LSTM+pointer models, respectively.
Due to limitations of space, we omit the details of the pointer model and refer the reader to the above paper.
As reported in Section 4 of Vasic et al.
(2019), to enable comparison with an enumerative approach, the evaluation was performed only on 12K test examples.
In comparison, we report the numbers on all 430K test examples (Table 1) for both the models. Similar to other tasks, we trained the baseline model for 100 epochs and ﬁnetuned CuBERT for up to 20 epochs. Table 5 gives the results along the same metrics as Vasic et al. The metrics are deﬁned as follows: 1) True Positive is the percentage of bug-free functions classiﬁed as bug-free.
2) Classiﬁcation Accuracy is the percentage of correctly classiﬁed examples (between bug-free and buggy).
3) Localization Accuracy is the percentage of buggy examples for which the localization pointer correctly identiﬁes the bug location.
4) Localization+Repair Accuracy is the precentage of buggy examples for which both the localization and repair pointers make correct predictions. As seen from Table 5, the CuBERT+pointer model outperforms the LSTM+pointer model consistently across all the metrics, and even within 2 and 10 epochs.  C ONCLUSIONS AND F UTURE W ORK We present the ﬁrst attempt at pre-trained contextual embedding of source code by training a BERT model, called CuBERT, which we ﬁnetuned on ﬁve classiﬁcation tasks and compared against BiLSTM with Word2Vec embeddings and Transformer models.
As a more challenging task, we also evaluated CuBERT on a multi-headed pointer prediction task.
CuBERT outperformed the baseline models consistently.
We evaluated CuBERT with less data and fewer epochs, highlighting the beneﬁts of pre-training on a massive, unsupervised code corpus.
We see this as a promising step towards source-code understanding, and plan to explore its utility on other programming languages and tasks.
We use a tokenized representation of source code and leave it to the underlying Transformer model to infer any structural interactions between the input tokens through its self-attention mechanism.
However, the literature on deep learning for source code also demonstrates the utility of encoding explicit structural information such as data-ﬂow information (Allamanis et al., 2018).
The Transformer model has been extended to incorporate explicitly provided structural information (Shaw et al., 2018).
Using such relation-aware Transformers for pre-training source code representations will be an important direction to explore in future. A D ATA P REPARATION FOR F INETUNING T ASKS A. L ABEL F REQUENCIES All four of our binary-classiﬁcation ﬁnetuning tasks had an equal number of buggy and bug-free examples. The Exception task, which is a multi-class classiﬁcation task, had a different number of examples per class (i.e., exception types).
We show the breakdown of example counts per label for our ﬁnetuning dataset splits in Table 6.
A. F INETUNING T ASK D ATASETS
In this section, we describe in detail how we produced our ﬁnetuning datasets (Section 3.4).
A common primitive in all our data generation is splitting a Python module into functions.
We do this by parsing the Python ﬁle and identifying function deﬁnitions in the Abstract Syntax Tree that have no other function deﬁnition between themselves and the root of the tree.
The resulting functions include functions deﬁned at module scope, but also methods of classes and subclasses.
Not included are functions deﬁned within other function and method bodies, or methods of classes that are, themselves, deﬁned within other function or methods bodies.
We do not ﬁlter functions by length, although task-speciﬁc data generation may ﬁlter out some functions (see below).
When generating examples for a ﬁxed-length pre-training or ﬁnetuning model, we prune all examples to the maximum target sequence length (in this paper we consider 128, 256, and 512 subtokenized sequence lengths).
Note that if a synthetically generated buggy/bug-free example pair differs only at a location beyond the target length (say on the 600-th subtoken), we still retain both examples.
For instance, for the Variable Misuse Localization and Repair task, we retain both buggy and bug-free examples, even if the error and/or repair locations lie beyond the end of the maximum target length.
Commutative
Non-Commutative
Arithmetic +, * -, /, %
Comparison =
=, !
=, is, is not < , < =, > , > =
Membership in, not in
Boolean and, or
A.2. R EPRODUCIBLE D ATA G ENERATION
We make pseudorandom choices at various stages in ﬁnetuning data generation.
It was important to design a pseudorandomness mechanism that gave (a) reproducible data generation, (b) nondeterministic choices drawn from the uniform distribution, and (c) order independence.
Order independence is important because our data generation is done in a distributed fashion (using Apache
Beam), so different pseudorandom number generator state machines are used by each distributed worker.
More speciﬁcally, pseudorandomness is computed based on an experiment-wide seed, but is independent of the order in which examples are generated.
Speciﬁcally, to make a pseudorandom choice about a function, we hash (using MD5) the seed and the function data (its source code and metadata about its provenance), and use the resulting hash as a uniform pseudorandom value from the function, for whatever needs the data generator has (e.g., in choosing one of multiple choices).
In that way, the same function will always result in the same choices given a seed, regardless of the order in which each function is processed, resulting in reproducible dataset generation.
To choose among multiple choices, we hash the function’s pseudorandom value along with all choices (sorted in a canonical order) and use the digest to compute an index within the list of choice.
Note that given two choices over different candidates but for the same function, independent decisions will be drawn.
We also use such order-independent pseudorandomness when subsampling datasets (e.g., to generate the validation datasets).
In those cases, we hash a sample with the seed, as above, and turn the resulting digest into a pseudorandom number in [0 , 1] , which can be used to decide given a target sampling rate.
A.2. V ARIABLE M ISUSE C LASSIFICATION
A variable use is any mention of a variable in a load scope.
This includes a variable that appears in the right-hand side of an assignment, or a ﬁeld dereference.
We regard as deﬁned all variables mentioned either in the formal arguments of a function deﬁnition, or on the right-hand side of an assignment.
We do not include in our deﬁned variables those declared in module scope (i.e., globals).
To decide whether to generate examples from a function, we parse it, and collect all variable-use locations, and all deﬁne variables, as described above.
We discard the function if it has no variable uses, or if it deﬁnes fewer than two variables; if there is only one variable deﬁned, the problem of detecting variable misuse is moot.
For any function that we do not discard, we generate a buggy and a bug-free example, as described next.
To generate a buggy example from a function, we choose one variable use pseudorandomly (see above how multiple-choice decisions are done), and replace its current occupant with a different pseudorandomly-chosen variable deﬁned in the function (with a separate multiple-choice decision).
A.2. W RONG B INARY O PERATOR
This task considers both commutative and non-commutative binary operators (unlike the SwappedArgument Classiﬁcation task). See Table 7 for the full list, and note that we have excluded relatively infrequent operators, e.g., the Python integer division operator // . If a function has no binary operators, it is discarded. Otherwise, it is used to generate a bug-free example, and a single buggy example as follows: one of the operators is chosen pseudorandomly (as described above), and a different operator chosen to replace it in the same row of the Table 7.
So, for instance, a buggy example would only swap =
= with is , but not with not in , which would not type-check if we performed static type inference on Python.
We take appropriate care to ensure the code parses after a bug is introduced.
For instance, if we swap the operator in the expression 1==2 with is , we ensure that there is space between the tokens (i.e., 1 is 2 rather than the incorrect 1is2 ), even though it was not needed before.
A.2. S WAPPED O PERAND
Since this task targets swapping the arguments of binary operators, we only consider noncommutative operators from Table 7.
Functions without eligible operators are discarded, and the choice of the operator to mutate in a function, as well as the choice of buggy operator to use, are done as above, but limiting choices only on non-commutative operators.
To avoid complications due to format changes, we only consider expressions that ﬁt in a single line (in contrast to the Wrong Binary Operator Classiﬁcation task).
We also do not consider expressions that look the same after swapping (e.g., a - a ).
A.2. F UNCTION -D OCSTRING M ISMATCH
In Python, a function docstring is a string literal that directly follows the function signature and before the main function body.
Whereas in other common programming languages, the function documentation is a comment, in Python it is an actual, semantically meaningful string literal.
We discard functions that have no docstring from this dataset, or functions that have an empty docstring.
We split the rest into the function deﬁnition without the docstring, and the docstring summary (i.e., the ﬁrst line of text from its docstring), discarding the rest of the docstring.
We create bug-free examples by pairing a function with its own docstring summary.
To create buggy examples, we pair every function with another function’s docstring summary, according to a global pseudorandom permutation of all functions: for all i , we combine the i -th function (without its docstring) with the P i -th function’s docstring summary, where P is a pseudorandom permutation, under a given seed.
We discard pairings in which i =
=
P [ i ] , but for the seeds we chose, no such pathological permuted pairings occurred.
A.2. E XCEPTION T YPE
Note that, unlike all other tasks, this task has no notion of buggy or bug-free examples.
We discard functions that do not have any except clauses in them.
For the rest, we collect all locations holding exception types within except clauses, and choose one of those locations to query the model for classiﬁcation.
Note that a single except clause may hold a comma-separated list of exception types, and the same type may appear in multiple locations within a function.
Once a location is chosen, we replace it with a special HOLE token, and create a classiﬁcation example that pairs the function (with the masked exception location) with the true label (the removed exception type).
The count of examples per exception type can be found in Table 6.
A.2. V ARIABLE M ISUSE L OCALIZATION AND R EPAIR
The dataset for this task is identical to that for the Variable Misuse Classiﬁcation task (Section A.2.2).
However, unlike the classiﬁcation task, examples contain more features relevant to localization and repair.
Speciﬁcally, in addition to the token sequence describing the program, we also extract a number of boolean input masks:
• A candidates mask, which marks as True all tokens holding a variable, which can therefore be either the location of a bug, or the location of a repair.
The ﬁrst position is always a candidate, since it may be used to indicate a bug-free program.
• A targets mask, which marks as True all tokens holding the correct variable, for buggy examples.
Note that the correct variable may appear in multiple locations in a function,
[Figure]
mask may have multiple True positions
-free
have an
• An error-location mask, which marks as True the location where the bug occurs (for buggy examples) or the ﬁrst location (for bug-free examples).
All the masks mark as True some of the locations that hold variables.
Because many variables are subtokenized into multiple tokens, if a variable is to be marked as True in the corresponding mask, we only mark as True its ﬁrst subtoken, keeping trailing subtokens as False.
B A TTENTION V ISUALIZATIONS
In this section, we provide sample code snippets used to test the different classiﬁcation tasks. Further, Figures 1–5 show visualizations of the attention matrix of the last layer of the ﬁnetuned CuBERT model (Coenen et al., 2019) for the code snippets.
In the visualization, the Y-axis shows the query tokens and X-axis shows the tokens being attended to.
The attention weight between a pair of tokens is the maximum of the weights assigned by the multi-head attention mechanism. The color changes from dark to light as weight changes from 0 to 1.
[Figure] [Figure] [Figure]