Contrastive Code Representation Learning

Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 5954–5971 November 7–11, 2021. c⃝2021 Association for Computational Linguistics 5954 Contrastive Code Representation Learning Paras Jain∗ and Ajay Jain∗ and Tianjun Zhang and Pieter Abbeel and Joseph E. Gonzalez and Ion Stoica Department of EECS, UC Berkeley {parasj, ajayj, tianjunz, pabbeel, jegonzal, istoica}@berkeley.edu Abstract Recent work learns contextual representations of source code by reconstructing tokens from their context. For downstream semantic under- standing tasks like code clone detection, these representations should ideally capture pro- gram 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 Con- traCode: a contrastive pre-training task that learns code functionality, not form. Con- traCode pre-trains a neural network to iden- tify functionally similar variants of a pro- gram 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 adver- sarial code clone detection benchmark by 39% AUROC. Surprisingly, improved adversarial robustness translates to better accuracy over natural code; ContraCode improves summa- rization and TypeScript type inference accu- racy by 2 to 13 percentage points over com- petitive baselines. All source is available at https://github.com/parasj/contracode. 1 Introduction Programmers increasingly rely on machine-aided programming tools that analyze or transform code automatically to aid software development (Kim et al., 2012). Traditionally, code analysis uses hand-written rules, though the wide diversity of programs encountered in practice can limit their generality. Recent work leverages machine learn- ing for richer language understanding, such as learning to detect bugs (Pradel and Sen, 2018) and predict performance (Mendis et al., 2019). Still, neural models of source code are suscep- tible to adversarial attacks. Yefet et al. (2020) and Schuster et al. (2021) 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 se- lecting between possible edits lowers performance be- low random guessing (dashed line). Contrastive pre- training with ContraCode learns a more robust repre- sentation of functionality, consistent across code edits. function (len) { for (i = 0; i < len, i++) { ... } } function (n) { while (i < n) { ... } } function (str, len) { return str.slice(0, len); } function f(n) { return n<2 ? 1 : f(n-1) + f(n-2); } function (arr) { for (i of arr) { ... } } Maximize similarity with equivalent programs Minimize similarity with functionally different programs Given a program, Figure 2: For many analyses, programs with the same functionality should have similar representations. Con- traCode learns such representations by pre-training an encoder to retrieve equivalent, transformed programs among many distractors. significantly under adversarial perturbations for both discriminative and generative code models. In our work, we investigate adversarial attacks on code clone detection. Successful adversarial at- tacks could circumvent malware detectors. While self-supervision can improve adversar- ial robustness …

Content truncated to 500 words