Contrastive Code Representation Learning (2021)

Abstract: Recent work learns contextual representations of source code by reconstructing tokens from their context. For downstream semantic understanding tasks like code clone detection, these representations should ideally capture program functionality. However, we show that the popular reconstruction-based RoBERTa 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 outperforms RoBERTa on an adversarial code clone detection benchmark by 39% AUROC. Surprisingly, improved adversarial robustness translates to better accuracy over natural code; ContraCode improves summarization and TypeScript type inference accuracy by 2 to 13 percentage points over competitive baselines. All source is available at https://github.com/parasj/contracode .

## introduction programmers increasingly rely on machine-aided programming tools that analyze or transform code automatically to aid software development @xcite . traditionally, code analysis uses hand-written rules, though the wide diversity of programs encountered in practice can limit their generality. recent work leverages machine learning for richer language understanding, such as learning to detect bugs @xcite and predict performance @xcite . still, neural models of source code are susceptible to adversarial attacks. @xcite and @xcite find accuracy degrades * equal contribution 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 pretraining with contracode learns a more robust representation of functionality, consistent across code edits. significantly under adversarial perturbations for both discriminative and generative code models. in our work, we investigate adversarial attacks on code clone detection. successful adversarial attacks could circumvent malware detectors. while self-supervision can improve adversarial robustness @xcite , we find that roberta is sensitive to stylistic implementation choices of code inputs. fig. 1 plots the performance of roberta and contracode, our proposed method, on a code clone detection task as small label-preserving perturbations are applied to the input code syntax. with just three minor adversarial edits to code syntax, roberta underperforms the random classifier (in gray). in fig. 3 , we show that roberta's representations of code are sensitive to code edits in agreement with prior work @xcite @xcite . to address this issue, we develop contracode: a self-supervised representation learning algorithm that captures program semantics. we hypothesize that programs with the same functionality should have similar underlying representations for downstream code understanding tasks. contracode generates syntactically diverse but functionally equivalent programs using source-tosource compiler transformation techniques (e.g., dead code elimination, obfuscation and constant folding). it uses these programs in a challenging discriminative pretext task that requires the model to identify similar programs out of a large dataset of distractors (fig. 2 ). to solve this task, the model must embed code semantics rather than syntax. contracode improves adversarial robustness in fig. 1 . surprisingly, adversarial robustness transfers to better natural code understanding. our novel contributions include: ## related work self-supervised learning (ssl) is a learning strategy where some attributes of a datapoint are predicted from remaining parts. bert @xcite is a ssl method for nlp that reconstructs masked tokens as a pretext task. roberta @xcite further tunes bert. contrastive approaches minimize distance between learned representations of similar examples (positives) and maximize distance between dissimilar negatives @xcite . cpc @xcite encodes segments of sequential data to predict future segments. simclr @xcite and moco @xcite ) use many negatives for dense loss signal. roberta embeddings umap 1 umap 2 program 19 variant b program 19 variant a … … contracode embeddings umap 1 figure 3 : a umap visualization of javascript method representations learned by roberta and contracode, in 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 contra-code's embedding space. ## approach our core insight is to use compiler transforms as data augmentations, generating a dataset of equivalent functions ( §3.1, 3.2). we then use a contrastive objective to learn a representation invariant to these transforms ( §3.3). ## evaluation in order to evaluate whether contracode defend against adversarial code inputs, we benchmark adversarial code clone detection accuracy @xcite . we evaluate results over natural and adversarial edits. we then evaluate how improvements to adversarial robustness translate to improvements on established in-the-wild code benchmarks. while improvements on adversarial benchmarks would not be expected to translate to real code, we find significant improvements in extreme code summarization @xcite and type inference @xcite tasks. clone detection experiments show that contrastive and hybrid representations with our compiler-based augmentations are predictive of program functionality in-the-wild, and that contrastive representations are the most robust to adversarial edits ( §4.1). contrastive pre-training outperforms baseline supervised and self-supervised methods on all three tasks ( §4.1-4.3). finally, ablations suggest it is better to augment unlabeled programs during pre-training rather than augmenting smaller supervised datasets ( §4.4). experimental setup models are pre-trained on codesearchnet, a large corpus of methods extracted from popular github repositories @xcite . codesearchnet contains 1,843,099 javascript programs. only 81,487 methods have both a documentation string and a method name. the asymmetry between labeled and unlabeled programs stems from javascript coding practices where anonymous functions are widespread. the pre-training dataset described in section 3.1 is the result of augmenting all 1.8m programs. as our approach supports any encoder, we evaluate two architectures: a 2-layer bidirectional lstm with 18m parameters, similar to the supervised model used by @xcite , and a 6-layer transformer with 23m parameters. for a baseline self-supervised approach, we pretrain both architectures with the roberta mlm objective, then transfer it to downstream tasks. ## conclusion large-scale code repositories like github are a powerful resource for learning machine-aided programming tools. however, most current code representation learning approaches need labels, and popular label-free self-supervised methods like roberta are not robust to adversarial inputs. instead of reconstructing tokens like bert, learning what code says, we learn what code does. we propose contracode, a contrastive self-supervised algorithm that learns representations invariant to transformations via compiler-based data augmentations. in experiments, contracode learns effective representations of functionality, and is robust to adversarial code edits. we find that contraminimum edit distance between positive pairs in the precomputed pre-training dataset, i.e. transformed variants of the same source method. for comparison, we also compute the edit distance between negative pairs: transformed variants of different programs. the edit distance d(x q , x k ) computes the minimum number of token insertions, deletions or substitutions needed to transform the tokenized query progrm x q into the key program x k . to normalize by sequence length | • |, let