Contrastive Code Representation Learning

TL;DR

This work tackles the fragility of reconstruction-based code representations by introducing ContraCode, a contrastive pretraining framework that uses compiler-based, semantics-preserving transformations to learn functional code representations. By forming positive pairs from variants of the same program and negatives from others, ContraCode learns invariances to surface-level code edits and enhances robustness across clone detection, type inference, and extreme code summarization. The approach yields significant improvements over MLM baselines and demonstrates robustness to adversarial edits, including a new zero-shot JavaScript clone-detection dataset. The results underscore the value of function-focused, data-augmentation-driven self-supervision for scalable, language-agnostic code understanding.

Abstract

Recent work learns contextual representations of source code by reconstructing tokens from their context. For downstream semantic understanding tasks like summarizing code in English, these representations should ideally capture program functionality. However, we show that the popular reconstruction-based BERT model is sensitive to source code edits, even when the edits preserve semantics. We propose ContraCode: a contrastive pre-training task that learns code functionality, not form. ContraCode pre-trains a neural network to identify functionally similar variants of a program among many non-equivalent distractors. We scalably generate these variants using an automated source-to-source compiler as a form of data augmentation. Contrastive pre-training improves JavaScript summarization and TypeScript type inference accuracy by 2% to 13%. We also propose a new zero-shot JavaScript code clone detection dataset, showing that ContraCode is both more robust and semantically meaningful. On it, we outperform RoBERTa by 39% AUROC in an adversarial setting and up to 5% on natural code.

Figures (16)

• Figure 1: Robust code clone detection: On source code, RoBERTa is not robust to simple label-preserving code edits like renaming variables. Adversarially selecting between possible edits lowers performance below random guessing (dashed line). Contrastive pre-training with ContraCode learns a more robust representation of functionality, consistent across code edits.
• Figure 2: For many analyses, programs with the same functionality should have similar representations. ContraCode learns such representations by pre-training an encoder to retrieve equivalent, transformed programs among many distractors.
• Figure 3: A UMAP visualization of JavaScript method representations learned by RoBERTa and ContraCode, in $\mathbb{R}^2$. Programs with the same functionality share color and number. RoBERTa's embeddings often do not cluster by functionality, suggesting that it is sensitive to implementation details. For example, many different programs overlap, and renaming the variables of Program 19 significantly changes the embedding. In contrast, variants of Program 19 cluster in ContraCode's embedding space.
• Figure 4: A JavaScript method from our unlabeled training set with two automatically generated semantically-equivalent programs. The method is from the StackEdit Markdown editor.
• Figure 5: Histogram of the number of unique transformed variants per JavaScript method during pre-training.