[Paper Review] Learning and Evaluating Contextual Embedding of Source Code

Aditya Kanade, Petros Maniatis+2 more|arXiv (Cornell University)|Dec 21, 2019

Software Engineering ResearchComputer Science156 citations

TL;DR

CuBERT pre-trains a contextual embedding for Python code on a large deduplicated GitHub corpus and shows strong fine-tuned performance across multiple code-understanding tasks, outperforming Word2Vec baselines, BiLSTMs, and from-scratch Transformers, with data-efficient training.

ABSTRACT

Recent research has achieved impressive results on understanding and improving source code by building up on machine-learning techniques developed for natural languages. A significant advancement in natural-language understanding has come with the development of pre-trained contextual embeddings, such as BERT, which can be fine-tuned for downstream tasks with less labeled data and training budget, while achieving better accuracies. However, there is no attempt yet to obtain a high-quality contextual embedding of source code, and to evaluate it on multiple program-understanding tasks simultaneously; that is the gap that this paper aims to mitigate. Specifically, first, we curate a massive, deduplicated corpus of 7.4M Python files from GitHub, which we use to pre-train CuBERT, an open-sourced code-understanding BERT model; and, second, we create an open-sourced benchmark that comprises five classification tasks and one program-repair task, akin to code-understanding tasks proposed in the literature before. We fine-tune CuBERT on our benchmark tasks, and compare the resulting models to different variants of Word2Vec token embeddings, BiLSTM and Transformer models, as well as published state-of-the-art models, showing that CuBERT outperforms them all, even with shorter training, and with fewer labeled examples. Future work on source-code embedding can benefit from reusing our benchmark, and from comparing against CuBERT models as a strong baseline.

Motivation & Objective

• Motivate improved representation learning for source code using contextual embeddings similar to BERT in NLP.
• Create a large, deduplicated Python corpus for pre-training CuBERT.
• Design a cohesive, multi-task Python code benchmark covering classification and program repair tasks.
• Evaluate CuBERT against strong baselines (Word2Vec, BiLSTM, Transformer) and state-of-the-art methods.
• Release models and datasets to enable future research and benchmarking.

Proposed method

• Pre-train CuBERT on a deduplicated 7.4M Python file corpus (9.3B tokens) with a Python-specific tokenization and subword vocabulary.
• Use BERT-style MLM and Next-Sentence Prediction objectives with CuBERT treated as line-based code inputs, where a line is a logical code line.
• Fine-tune CuBERT on five classification tasks and one pointer-based variable-misuse localization/repair task.
• Compare CuBERT to Word2Vec-based BiLSTMs and Transformers trained from scratch, as well as published state-of-the-art models.
• Evaluate with varying example lengths and fine-tuning data budgets to assess data efficiency and context Effects.
• Provide open-source code and datasets for benchmark reuse.

Experimental results

Research questions

• RQ1 Do contextual embeddings improve source-code analysis when pre-trained on unlabeled code?
• RQ2 Does fine-tuning a Transformer-based model yield benefits beyond training from scratch?
• RQ3 How does CuBERT performance scale with limited task-specific labeled data?
• RQ4 How does context size (example length) affect CuBERT performance on code tasks?
• RQ5 How does CuBERT perform on complex tasks like variable-misuse localization/repair compared to prior state-of-the-art?

Key findings

• CuBERT consistently outperforms BiLSTMs with best Word2Vec embeddings across all classification tasks by 3.2% to 14.7% (test accuracy).
• CuBERT achieves strong results with only 2–20 fine-tuning epochs, approaching or surpassing full-data baselines.
• Fine-tuning CuBERT on 33%–100% of task data yields competitive or superior performance vs. baselines trained with full datasets.
• CuBERT significantly outperforms state-of-the-art models on the variable-misuse localization and repair task.
• Compared to transformers trained from scratch, CuBERT (pre-trained + fine-tuned) yields substantially higher accuracy, indicating the value of pre-training for code representations.
• The authors provide open-source models and benchmarks to facilitate future research.