Contrastive Code Representation Learning

In plain words

Language models for code are often trained to fill in masked tokens. They learn how code is written, so a small change such as renaming a variable can throw them off, even when the program still does exactly the same thing.

ContraCode takes the opposite route. It takes unlabelled JavaScript functions and rewrites each one with compiler-style tools: renaming identifiers, removing dead code, computing constant expressions in advance and similar changes. Two rewrites of the same function count as a matching pair. The model is then trained to pick the matching rewrite out of a large pile of unrelated functions. To succeed, it has to represent what the code does rather than how it looks.

The trained encoder is then reused for tasks such as detecting code clones, guessing a function's name from its body and predicting types. The paper reports gains on each of them.

TL;DR

ContraCode is a self-supervised pre-training method for source code. It makes functionally equivalent JavaScript variants with a source-to-source compiler and trains an encoder with a MoCo-style InfoNCE loss, so programs that do the same thing get similar embeddings §3.1 §3.3. The paper reports robustness to adversarial edits in zero-shot clone detection and accuracy gains on summarization and TypeScript type inference (author-reported) §4. Relevant if you need code embeddings that survive renaming and restyling, but note the heavy pre-processing and several unstated details (see Pitfalls).

Problem

Masked-token pre-training (RoBERTa on code) yields representations that are sensitive to implementation style. In zero-shot clone detection, a few label-preserving edits (three in the introduction, four in the evaluation) push RoBERTa below random guessing §1 fig. 1 §4.1. Supervised code models have a similar weakness, and labelled code is scarce, so the paper asks for label-free pre-training that targets program functionality §1 §2. Creating equivalent program pairs from GitHub is hard because equivalence is undecidable and checking it by running code is costly §3.1. The method sidesteps this by constructing equivalent variants by design §3.1.

Method

The method has three parts: compiler-based augmentation §3.1, transform dropout to diversify variants §3.2, and contrastive pre-training of the encoder §3.3.

1. Compile variants. A JavaScript method is parsed into an abstract syntax tree and rewritten by Babel and Terser passes. The paper lists 11 transforms in three groups: code compression, identifier modification and regularisation §3.1 tab. 1 app. A. Ten preserve semantics by construction. Line subsampling (keep each line with probability 0.9) is lossy and is kept as a regulariser app. A tab. 1.

2. Transform dropout. Each transform is applied with its own probability, in a fixed order, and the results are deduplicated §3.2 alg. 1. The paper precomputes up to 20 variants for each of 1.8 million methods, and reports that 89% of methods have more than one distinct alternative §3.2. Subword regularisation is applied as a last step each batch, so two variants of one method still tokenise differently §3.2.

# Transform dropout  [alg. 1]
# Input:  source program, transforms with probabilities, count N
# Output: a set of program variants, the source included

def make_variants(source, transforms, count):
    # the set starts with the untouched program  [alg. 1]
    variants = {source}
    for _ in range(count - 1):
        # each variant starts again from the source  [alg. 1]
        current = source
        # transforms run in a pre-specified order  [§3.2]
        for transform, prob in transforms:
            # apply each transform with its own probability  [§3.2]
            if bernoulli(prob):
                # convert to AST or source as the transform needs  [alg. 1]
                current = convert_for(transform, current)
                current = transform(current)
        # store the result as source text; the set removes duplicates  [alg. 1]
        variants.add(to_source(current))
    return variants

3. Contrastive pre-training. A query encoder f_q and a key encoder f_k embed programs into 128-dimensional vectors app. C; the code normalises them to unit length before comparing them code_moco.py. Two variants of the same program form a positive pair; keys of other programs fill a queue of negatives §3.3. The loss is InfoNCE eq. 1: the query should score its positive key higher than all queued negatives, with a softmax temperature t §3.3 eq. 1. The key encoder is an exponential moving average of the query encoder, following MoCo §3.3 app. E. The encoder can be a 6-layer Transformer or a 2-layer BiLSTM §3.3 app. C.

# One contrastive pre-training step  [§3.3] [eq. 1] [app. C]
# query_batch, key_batch: two variants per program (positive pairs)
# queue: cached key embeddings of other programs (negatives)

def contrastive_step(query_batch, key_batch, queue):
    # query embedding from the trained encoder  [§3.3]
    query = normalize(encoder_q(query_batch))         # (B, 128)
    # EMA key encoder: theta_k <- m theta_k + (1 - m) theta_q  [app. E]
    update_key_encoder(momentum=0.999)
    # key embedding from the EMA encoder, no gradient  [§3.3]
    key = normalize(encoder_k(key_batch)).detach()    # (B, 128)
    # similarity to the positive key, one value per query  [eq. 1]
    positive = sum(query * key, dim=1, keepdim=True)  # (B, 1)
    # similarity to every cached negative key  [eq. 1]
    negatives = query @ queue                         # (B, K)
    # divide all logits by the temperature t = 0.07  [app. C]
    logits = cat([positive, negatives], dim=1) / 0.07 # (B, 1 + K)
    # InfoNCE: the positive is always at index 0  [eq. 1]
    loss = cross_entropy(logits, target=0)
    # push the new keys into the queue  [§3.3]
    enqueue(queue, key)
    return loss

Encoder heads. For the Transformer, token embeddings are averaged over the sequence and passed through a two-layer MLP (512 hidden units, ReLU) to a 128-dimensional vector app. C. For the BiLSTM, the four terminal hidden states (two directions, two layers; 1024 values) go through the same MLP head. Contrasting these terminal states worked better than contrasting mean token states app. C app. E tab. 8.

Transfer. For zero-shot clone detection the pre-trained encoder is used as is §3.3 §4.1. For tasks whose output differs from the encoder, a task-specific head is added and the whole network is fine-tuned end-to-end §3.3. For type inference a 2-layer MLP head reads the token embeddings produced by the DeepTyper BiLSTM §4.2. Transferring the first MLP layer together with the encoder is reported to help type inference in one BiLSTM setting app. E tab. 7.

Scale. Pre-training uses CodeSearchNet's JavaScript methods: 1,843,099 programs in total, of which 81,487 have both a name and a docstring §4. Only the augmented, unlabelled pool is used for pre-training; labels come from the supervised datasets §4 §3.2.

Hyperparameters

Results

Author-reported. Setup: JavaScript clone detection on 274 HackerRank programs for 33 problems (2,065 positive pairs, 2,065 negatives) §4.1 app. C; TypeScript type inference on DeepTyper's data, regenerated, 187 training projects §4.2 app. C; JavaScript method-name prediction on 81,487 CodeSearchNet methods §4.3 app. C.

Zero-shot clone detection (AUROC and AP, natural code; adversarial AUROC with N = 4 edits) tab. 2:

With N = 16 adversarial edits, ContraCode reaches 58.32 AUROC and RoBERTa 4.51 tab. 2. Table 2 has a fourth, unlabelled column group that the text does not explain.

Type inference, top-1 and top-5 accuracy tab. 3:

The paper reports +8.9 points top-1 for the best model over CheckJS §4.2 tab. 3.

Method-name prediction (precision, recall, F1) tab. 4:

Ablations. Compiler augmentations used only in supervised training lower accuracy: the Transformer's F1 drops from 16.86 to 15.65, and DeepTyper's Acc@1 from 51.73 to 50.33 tab. 5. In contrastive pre-training (DeepTyper BiLSTM, 20k steps), removing code compression costs about 2 points of Acc@1 (52.65 to 50.69), and removing line subsampling costs 1.6 points tab. 6 §4.4.

Paper vs code

The official repository is parasj/contracode, linked in the paper's abstract. Commit 831ae5aa, Apache-2.0 LICENSE. I checked the MoCo model, the pre-training and fine-tuning entry points and the Slurm scripts. I did not check the clone-detection pooling or the BiLSTM encoder in full.

• Key embedding timing. The paper says keys are embedded once per epoch because of the queue §3.3. The code recomputes the key embedding for every batch with the EMA encoder code_moco.py and enqueues it code_moco.py.
• Pre-training batch size. The paper uses 96 for the Transformer app. C. The pre-training entry point defaults to 256 pretrain.py, and the hybrid Slurm script defaults to 64 pretrain_hybrid.sh.
• Queue size. The paper says "over 100K" §3.3. The code's CodeMoCo default is 107,520 keys code_moco.py. The base class default is 61,440 code_moco.py. The paper's 65k and 125k queue comparison app. E matches neither.
• Momentum and temperature. m = 0.999 and t = 0.07 match the paper app. E app. C code_moco.py.
• Pre-training learning rate and schedule. The paper gives none app. C. The code is inconsistent. pretrain.py uses Adam at 8e-4 with no scheduler pretrain.py pretrain.py. pretrain_distributed.py uses 8e-4 with 5,000 warmup steps and 600,000 total steps pretrain_distributed.py. The Slurm scripts pass 4e-4 pretrain_hybrid.sh or 1e-4 pretrain_dist.sh.
• Fine-tuning learning rate. The paper gives a Transformer maximum of 1e-4 with 5,000 warmup steps app. C. The type-inference script defaults to 8e-4, 5,000 warmup and 200,000 steps type_prediction.py. Its Slurm wrapper defaults to 1e-4 with batch 16 supervised_type_finetune.sh. The paper's BiLSTM maximum of 1e-3 app. C was not compared with code.
• Adam epsilon. Not stated in the paper. The fine-tuning default is 1e-6 type_prediction.py and the pre-training default is 1e-9 pretrain.py.
• Subword regularisation. The paper applies it each batch §3.2 but gives no alpha. The pre-training default is 0 pretrain.py.
• Queue initialisation. Not stated in the paper. The code fills the queue with random vectors, normalised code_moco.py.
• Transform probabilities. Not stated in the paper. I found no per-transform probability in node_src/transform.js, node_src/javascript_augmentations.js or scripts/transform*.

Pitfalls

• Key timing: the paper's once-per-epoch key embedding does not match the code. Per-batch key embedding is the standard MoCo behaviour and is what the code does code_moco.py.
• Queue must divide evenly: the code asserts that the queue size is a multiple of the batch size code_moco.py. The defaults 107,520, 96 and 256 satisfy this; 65k or 125k with batch 96 may not.
• Unstated transform probabilities and pool size. The 1.8M-method pool is built with unstated probabilities. The paper says a single core handles 300 functions per second §3.2. My derived estimate is about 33 core-hours for 20 variants of 1.8M methods, which is an estimate and not a reported number.
• Ablation numbers differ from the headline. The row labelled "all augmentations" in the ablation table tab. 6 is 52.65 Acc@1, but the headline DeepTyper + ContraCode number is 54.01 tab. 3. The ablation uses 20k steps app. E.
• Caption and text mismatches. The clone-detection caption claims +8% AUROC over the heuristic for contrastive pre-training tab. 2. The table gives 75.73 − 69.55 = +6.2. The type-inference caption gives 1.5% to 2.28% gains tab. 3, but the Transformer row gives +1.20 Acc@1. The text's "+6.3% over RoBERTa" matches the hybrid row (47.16 − 40.85 = 6.31), not ContraCode (+6.01) §4.2 tab. 3.
• Adversarial selection is underspecified. The attack picks the "loss-maximising" transform among N samples §4.1. The loss used for zero-shot cosine similarity is not defined.
• Ambiguous checkpoint rule. The appendix says to keep the checkpoint with the "minimum" accuracy@1 app. C. This is probably meant as maximum, but the paper does not say so.
• Parameter counts. The BiLSTM is 18M in the evaluation section §4 and 17.5M in the appendix app. C.
• Batch 96 vs 256 vs 64. Use the paper's batch size for paper replication, but note the code defaults differ (see Paper vs code).
• Memory for the queue (derived). At 107,520 keys of 128 float32 values the queue is about 55 MB. This is my arithmetic, not a stated number.
• Hidden costs. Each method needs many compiler passes and an extra EMA encoder. Pre-training uses a large negative pool, so memory and per-step cost grow with the queue.

Lineage

• Builds on: MoCo, which the method extends to code §3.3; InfoNCE §3.3 eq. 1; BERT and RoBERTa masked-token pre-training as the baseline §2 app. C app. D; DeepTyper and its type-inference data §4.2 app. C.
• Compared against: edit-distance heuristic §4.1; RoBERTa MLM §4.1 §4.3; code2vec and code2seq for summarization §4.3 app. D; TypeScript CheckJS and GPT-3 Codex for type inference §4.2.