PaperlayerView exact revision on arXiv

Source-linked AI summary

Contrastive Code Representation Learning

Paras Jain, Ajay Jain, Tianjun Zhang, Pieter Abbeel, Joseph E. Gonzalez, Ion Stoica

arXiv:2007.04973v4cs.LGcs.AIcs.PLcs.SEstat.ML

TL;DR

Existing reconstruction-based code representations can be sensitive to semantics-preserving edits, leaving a gap in functionality-oriented representation learning. ContraCode uses compiler-generated equivalent programs in a contrastive pre-training task, and improves robustness and downstream JavaScript and TypeScript understanding. The paper reports gains across type inference, summarization, and clone detection, while noting scope boundaries from potentially semantics-changing transformations and the JavaScript-to-TypeScript transfer setting.

  • Problem

    Neural code representations can be sensitive to adversarial or semantics-preserving edits, despite the need for representations that capture program functionality.

  • Method

    ContraCode uses compiler-based data augmentations to generate functionally equivalent programs and trains an encoder contrastively to distinguish them from non-equivalent distractors.

  • Results

    ContraCode improves type inference top-1 accuracy by 9%, learned inference by 2%–13%, summarization F1 score by up to 8%, and clone detection AUROC by 2%–46%.

  • Takeaways & Limitations

    ContraCode learns effective functionality-oriented representations that are robust to adversarial code edits and improve three downstream JavaScript code understanding tasks.

  • Takeaways & Limitations

    Some transformations can potentially modify program semantics, and transfer to TypeScript is constrained by grammar differences, file-level data differences, and unavailable dependencies.

Abstract

from arXiv · show

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.

1 Introduction

Existing code representations can be highly sensitive to semantics-preserving edits, limiting robustness for code understanding. ContraCode addresses this by learning functionality-invariant representations through compiler-based transformations and contrastive pre-training, improving multiple downstream tasks.

  • Motivation: Neural source-code models can lose accuracy under adversarial perturbations, while RoBERTa is sensitive to stylistic implementation choices and minor syntax edits.With three minor adversarial edits, RoBERTa underperforms a random classifier on code clone detection.
  • Approach: ContraCode learns representations intended to capture program semantics by bringing functionally equivalent programs together.The approach is based on the hypothesis that programs with the same functionality should have similar underlying representations.
  • Approach: Compiler-based transformations generate syntactically diverse but functionally equivalent programs for contrastive pre-training.The pretext task requires identifying similar programs among many distractors, encouraging semantic rather than syntactic representations.
  • Findings: ContraCode improves robustness to adversarial code edits and transfers that robustness to natural code understanding.The method uses compiler transformations as data augmentations rather than reconstructing masked tokens like BERT.
  • Findings: Type inference top-1 accuracy improves by 9%, learned inference by 2%–13%, summarization F1 score by up to 8%, and clone detection AUROC by 2%–46%.These results are reported across the paper’s downstream code understanding evaluations.

2 Related work

Related work spans self-supervised and contrastive representation learning, code understanding tasks, and adversarial attacks on code models. These studies motivate learning representations that are less dependent on implementation details and more aligned with program functionality.

  • Self-supervised learning: Self-supervised learning predicts attributes of a datapoint from its remaining parts, while contrastive learning brings similar examples together and separates dissimilar negatives.BERT and RoBERTa reconstruct masked tokens; methods such as SimCLR and MoCo use many negatives for contrastive learning.
  • Code representation learning: Prior code representation learning studies address clone detection, type inference, and summarization across multiple programming languages.The related work includes graph-based and context-prediction approaches for code representations.
  • Adversarial attacks: Prior research finds code models highly sensitive to adversarial edits and reports attacks on code autocompletion tools.Code models may also be vulnerable because synthetic labels and duplication can degrade generalization.
  • Compiler-based augmentation: The paper illustrates code augmentation with automatically generated semantically equivalent JavaScript variants.The examples include transformed methods from an unlabeled training set and histograms of unique transformed variants per method.

3 Approach

ContraCode uses compiler transformations to generate functionally equivalent code variants, then applies contrastive pre-training to learn representations invariant to those transformations. The resulting encoder transfers to downstream code understanding tasks.

  • Compilation as data augmentation: ContraCode uses compiler transforms as data augmentations to generate equivalent program variants without searching for equivalences.The approach applies correct-by-construction transformations to unlabeled code, including dead-code elimination and other source-to-source passes.
  • Compilation as data augmentation: 10 of 11 automated source-to-source compiler transforms preserve operational semantics by construction.The implementation uses Babel and Terser infrastructures for JavaScript; transformations may be composed when each preserves semantics.
  • Compilation as data augmentation: The augmentation pipeline combines code compression, identifier modification, and regularizing transforms such as line subsampling.Transforms are applied stochastically in a pre-specified order, with intermediate conversion between AST and source representations as needed.
  • Compilation as data augmentation: The pipeline precomputes up to 20 augmentations for 1.8M JavaScript methods, producing syntactically distinct alternatives for 89% of methods.Variants are deduplicated before pre-training because some transformations leave programs unchanged.
  • Contrastive pre-training: Contrastive pre-training pairs variants of the same program as positives and different programs as negatives, using InfoNCE to separate them in representation space.The query encoder is trained against a queue of cached negatives, while a slowly updated key encoder stabilizes embeddings and enables over 100K negatives.
  • Transfer learning: After pre-training, the encoder transfers directly to clone detection or is combined with a task-specific decoder and fine-tuned for other tasks.ContraCode supports encoder architectures that produce global program embeddings, including Transformers and LSTMs.

4 Evaluation

Evaluation across clone detection, type inference, and summarization shows that contrastive pre-training improves functionality-oriented code representations, including robustness to adversarial edits. Ablations indicate that contrastive learning and pre-training augmentation matter more than applying augmentations directly to smaller supervised datasets.

  • 4.4 Understanding augmentation importance: Ablations show that augmenting unlabeled programs during ContraCode pre-training is more effective than augmenting smaller supervised datasets.The ablation indicates that augmentations alone do not suffice; contrastive learning is important.
  • 4.1 Robust Zero-shot Code Clone Detection: Contrastive representations are most robust to adversarial code transformations while improving zero-shot clone detection on natural HackerRank programs.The cosine-similarity probe reports +8% AUROC for contrastive representations and +10% for the hybrid over a heuristic textual similarity probe.
  • 4.1 Robust Zero-shot Code Clone Detection: With only four semantics-preserving edits, RoBERTa falls below random guessing, whereas ContraCode retains much of its performance with +39% AUROC over RoBERTa.The edits preserve ground-truth labels, making the degradation reflect sensitivity to implementation details.
  • 4.2 Fine-tuning for Type Inference: Contrastive pre-training outperforms learned baselines for TypeScript type inference, with top-1 gains of +1.2%, +6.3%, and +2.3% over Transformer, RoBERTa, and DeepTyper.The corresponding top-5 gains are +1.8%, +5.7%, and +2.8%.
  • 4.3 Code Summarization: Contrastive pre-training improves code summarization over code2seq by +8.2% precision, +7.3% recall, and +7.9% F1, and exceeds RoBERTa by +4.8% F1.It also performs slightly better than a Transformer trained from scratch.
  • 4.4 Understanding augmentation importance: Semantics-preserving compression passes are the most important augmentation group, improving type-inference top-1 accuracy by 1.95% when included.Line subsampling acts as a regularizer despite changing semantics, while identifier modifications preserve semantics but alter naming information.

5 Conclusion

ContraCode learns code representations intended to capture functionality rather than surface form, using compiler-based augmentations and contrastive self-supervision. Across three JavaScript understanding tasks, it improves performance and robustness to adversarial code edits.

  • 5 Conclusion: ContraCode learns representations invariant to compiler-generated transformations instead of reconstructing tokens like BERT.The method uses compiler-based data augmentations to learn what code does rather than what it says.
  • 5 Conclusion: ContraCode significantly improves performance on three downstream JavaScript code understanding tasks.The conclusion also reports robustness to adversarial code edits.

A Program transformation details

The transformation pipeline combines compiler and tokenization operations that diversify JavaScript code while generally preserving program behavior. These transformations include formatting, optimization, identifier changes, no-op insertion, and subword variation.

  • Compiler infrastructure: Babel and Terser provide compiler-based transformations, including variable renaming, dead-code insertion, and other source-to-source operations.Terser supplies mangling and compression modes with controls for formatting, comments, logging, and dead-code elimination.
  • Syntax and formatting transformations: Formatting, beautification, and compression normalize coding style or remove unnecessary syntax without changing code semantics.Compression can eliminate redundant variable declarations and brackets.
  • Compiler transformations: Dead-code elimination removes unused side-effect-free code, including statements that are stale, unneeded, or safely inlineable.It is an analysis-based compiler pass.
  • Semantic-preserving transformations: Type upconversion replaces some JavaScript values with equivalent representations, such as encoding booleans as true or 1.The transformation exploits JavaScript’s polymorphic types.
  • Compiler transformations: Constant folding evaluates compile-time expressions, replacing expressions such as (2 + 3) * 4 with 20.The pass inlines values that can be pre-computed during compilation.
  • Identifier transformations: Variable renaming and identifier mangling alter naming information while preserving program behavior, improving robustness to naming choices.Arguments receive random word sequences and identifiers can become short tokens.
  • Syntax transformations: Dead-code insertion adds common no-ops such as comments and logging to create syntactic variation.These insertions are not intended to change the program’s functional behavior.
  • Tokenization transformations: Subword regularization tokenizes text in multiple ways, allowing different token sequences for the same word.It is applied as a tokenization-level transformation.

B How similar are transformed programs?

The paper measures token-level dissimilarity between transformed variants of the same method and variants of different methods using normalized Levenshtein distance. Transformations make positive pairs substantially different in token sequence while retaining their functional relationship.

  • Results: The histogram shows that transformed variants of the same method form contrastive positives despite having substantially different token sequences.The transformations therefore create syntactic diversity without making positive pairs identical in token space.
  • Distance measure: The analysis compares Levenshtein edit distance for positive pairs from the same source method with negative pairs from different programs.The distance counts token insertions, deletions, and substitutions needed to transform one tokenized program into another.
  • Distance measure: Dissimilarity normalizes edit distance by the maximum sequence length, ranging from 0% for identical token sequences to 100% for sequences without shared tokens.Whitespace changes do not affect the metric because the tokenizer collapses repeated whitespace.
  • Results: Positive pairs have 65% mean token dissimilarity, compared with 86% for negative pairs.Positive-pair dissimilarity percentiles are 59%, 66%, and 73% at the 25th, median, and 75th percentiles; negative-pair values are 82%, 87%, and 90%.

C Experimental setup

The experiments evaluate ContraCode across code summarization, type prediction, and clone detection using Transformer and BiLSTM encoders. The setup standardizes datasets, pooling choices, training schedules, and clone-detection metrics.

  • Architectures: The Transformer encoder uses six layers and 23M parameters, with four causal decoder layers added for code summarization.The architecture uses sinusoidal positional embeddings and is paired with sequence-to-sequence decoding for summaries.
  • Contrastive pre-training: ContraCode represents a program by averaging token embeddings and projecting the result through a two-layer MLP into a 128-dimensional contrastive embedding.The pre-training objective uses InfoNCE with temperature t = 0.07 and a momentum-updated key encoder.
  • Representation choices: The BiLSTM experiments compare global terminal-state representations with local mean-pooled token representations before the contrastive projection.The BiLSTM uses a 1024-dimensional concatenated terminal representation or averaged bidirectional hidden states, with batch size 512 and weight decay.
  • Clone detection: Clone detection samples 2,065 same-problem pairs and 2,065 different-problem pairs, evaluating AUROC and Average Precision.Whole-program embeddings are obtained by mean or max pooling token representations before cosine-similarity classification.
  • Training schedules: DeepTyper-style models use 20K pre-training iterations, whereas Transformer models generally use 240K steps before downstream fine-tuning.The ContraCode summarization Transformer converges in 20K iterations, substantially faster than RoBERTa’s 240K iterations.
  • Tasks and datasets: The study evaluates summarization, type prediction, and code clone detection across JavaScript, TypeScript, and HackerRank programs.The summarization task predicts method names from method bodies and comments, while type prediction uses separate training, validation, and test projects.

D Baselines

The baselines are reimplemented on the JavaScript dataset to normalize dataset effects, spanning AST-based models, DeepTyper variants, and RoBERTa masked-language-model pre-training.

  • Baseline normalization: The study reimplements diverse state-of-the-art baselines on JavaScript to normalize differences in programming languages and datasets.This addresses inconsistencies in the original training settings of competing models.
  • AST-based models: The code2vec and code2seq baselines use AST paths and are extended from Java to JavaScript using a generated AST path dataset.Their published F1 scores are not directly used because path-mining embeddings are sensitive to datasets.
  • DeepTyper: The DeepTyper baseline replaces its two-layer GRU with a parameter-matched bidirectional LSTM using 512-dimensional hidden states and subword tokenization.The modified architecture is used to improve performance while keeping parameter counts comparable.
  • RoBERTa: The RoBERTa baseline pre-trains an encoder with masked language modeling on the same augmented CodeSearchNet data used for ContraCode.The model is subsequently fine-tuned on downstream datasets.

Code clone detection ROC, PR curves

The clone-detection analyses compare natural and adversarial performance across representation-learning strategies, showing that contrastive pre-training produces more robust and semantically clustered representations.

  • Adversarial robustness: ContraCode preserves much of its AUROC and Average Precision as adversarial transformations intensify, unlike untrained and RoBERTa-MLM models.The adversarial evaluation applies transformations to one input program in a clone-detection pair.
  • Natural-code detection: The hybrid model combining ContraCode and MLM achieves better ROC and precision-recall tradeoffs than the other zero-shot classifiers.The curves compare true-positive rate versus false-positive rate and precision versus recall using cosine similarity or token-level edit distance.
  • Transfer ablations: Transferring part of the contrastive MLP head improves type-inference top-5 accuracy by 2% over transferring only the encoder.The ablation uses local mean-hidden-state representations and a two-layer BiLSTM encoder.
  • Representation ablations: +2.38% acc@1 is achieved by global-feature pre-training over local-feature pre-training in the BiLSTM ablation.This comparison occurs after 10K pre-training iterations and before convergence; the global strategy achieves the best reported ablation result.
  • Decoder ablations: Longer pre-training significantly improves method-name prediction when the summarization decoder has one layer.The decoder-depth ablation examines how encoder pre-training benefits shallow sequence-to-sequence decoders.
  • Queue refresh: Refreshing the negative queue more frequently leads to faster pre-training convergence.The queue stores negative program representations, and the comparison varies its fill rate.
  • Representation visualization: t-SNE shows ContraCode clustering transformed variants of the same program closely together, indicating invariance to compiler-based transformations.Variants share colors in the visualization, whereas masked-language-model representations are less tightly grouped.
Loading 2007.04973v4…