[Paper Review] Pre-trained Contextual Embedding of Source Code

Aditya Kanade, Petros Maniatis+2 more|arXiv (Cornell University)|Sep 25, 2019

Software Engineering Research28 references41 citations

TL;DR

This paper introduces Code Understanding BERT (CuBERT), a pre-trained contextual embedding model for source code, trained on a large-scale corpus of Python programs from GitHub. By adapting BERT's self-attention mechanism to code, CuBERT achieves state-of-the-art performance on five downstream classification tasks, outperforming baseline models by 2.9–22%, especially with limited fine-tuning data and epochs.

ABSTRACT

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 human-readable form. To facilitate this, developers use meaningful identifier 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 fine-tuned effectively for a variety of downstream supervised tasks. Inspired by these developments, we present the first 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 five classification tasks and compare fine-tuned 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 fine-tuned with smaller datasets, and over fewer epochs. We further evaluate CuBERT's effectiveness on a joint classification, localization and repair task involving prediction of two pointers.

Motivation & Objective

• To explore whether pre-trained contextual embeddings, successful in natural language processing, can be effectively applied to source code.
• To address the challenge of capturing both syntactic and semantic information in code for downstream NLP-style tasks.
• To develop a model that leverages natural-language-like patterns in code, such as meaningful identifiers and documentation, for improved representation learning.
• To evaluate the model's performance on diverse code understanding tasks, including classification, localization, and repair.
• To demonstrate the effectiveness of pre-training on large-scale code corpora for transfer learning in code intelligence applications.

Proposed method

• Curated a massive corpus of 1.6 million Python files from GitHub to pre-train the model on real-world code.
• Fine-tuned a BERT-based architecture with masked language modeling and next sentence prediction objectives adapted for code.
• Pre-trained Word2Vec embeddings on the same dataset as a baseline for comparison.
• Designed a benchmark with five distinct code classification tasks to evaluate model performance.
• Applied transfer learning by fine-tuning CuBERT on downstream tasks with limited labeled data and fewer training epochs.
• Evaluated CuBERT on a joint task involving classification, pointer prediction for code localization, and repair prediction.

Experimental results

Research questions

• RQ1 Can pre-trained contextual embeddings like BERT be effectively adapted to source code representation learning?
• RQ2 How does CuBERT compare to baseline models using Word2Vec embeddings and other sequence models on code classification tasks?
• RQ3 To what extent does pre-training on large-scale code improve performance on downstream tasks with limited labeled data?
• RQ4 Can CuBERT generalize effectively to multi-task learning scenarios involving classification, localization, and repair?
• RQ5 Does CuBERT maintain performance advantages when fine-tuned with fewer epochs or smaller datasets?

Key findings

• CuBERT outperformed all baseline models on the five classification tasks, with improvements ranging from 2.9% to 22% in F1-score.
• The model demonstrated superior performance even when fine-tuned on smaller datasets, indicating strong generalization from pre-training.
• CuBERT achieved better results with fewer training epochs, highlighting its efficiency in fine-tuning.
• The model showed strong effectiveness on a joint classification, localization, and repair task, successfully predicting two pointers.
• Pre-training on a large-scale code corpus significantly enhanced representation quality, outperforming models without such pre-training.
• The integration of natural-language-like patterns in code (e.g., meaningful identifiers and documentation) was crucial for the success of the approach.