Invariant Pretraining for Robust Code Representations

Yifeng He; Yundi Xu; Christopher Castro Gaw Gonzalo; Zili Wang; Hao Chen

Source-linked AI summary

Invariant Pretraining for Robust Code Representations

Yifeng He, Yundi Xu, Christopher Castro Gaw Gonzalo, Zili Wang, Hao Chen

arXiv:2608.15412v1Aug 15, 2026cs.LGcs.AIcs.SE

TL;DR

Encoder code representations often fail on semantically equivalent programs written in different syntactic forms. InvPT uses code-only invariant pretraining with multi-positive contrastive learning, improving robustness across clone detection and classification while matching or improving standard performance.

Problem
Encoder-based code models often mishandle semantically equivalent but syntactically varied programs, despite their practical importance for low-cost discriminative tasks.
Method
InvPT continues pretraining with semantics-preserving code transformations, masked language modeling, and multi-positive supervised contrastive learning over self-contrast and invariant-contrast pairs.
Results
InvPT improves robustness in every model–dataset comparison, by median 8.07 percentage points on clone detection and 3.56 on classification, while matching or improving standard performance.
Takeaways & Limitations
Code-only invariant pretraining recovers part of robustness degradation, with invariant contrastive learning as the primary driver and self-contrast providing consistent additional improvement.
Takeaways & Limitations
The evaluation establishes robustness only to rewrite compositions built from the six pretraining operators, leaving independently generated transformations and natural refactorings untested.

Abstract

from arXiv · show

Encoder-based code representation models remain widely deployed for discriminative tasks such as clone detection and code classification, where their small size and low inference cost are decisive. Their robustness, however, is fragile: under invariant programs, semantically equivalent code written in different syntactic forms, learned representations degrade substantially even though program behavior is unchanged. We present an empirical study of this robustness gap across four encoder baselines, two downstream tasks, and four datasets, together with a minimal code-only continued pretraining recipe that closes much of it. Our method, invariant pretraining (InvPT), applies semantic-preserving transformations to the corpus and combines masked language modeling with multi-positive supervised contrastive learning that treats all augmentations of the same source function as positives, mixing self-contrast pairs (same code, different masks) with invariant-contrast pairs (transformed code) for positives of varying difficulty. Unlike prior contrastive code encoders, InvPT does not require paired natural-language data. Across our evaluation, InvPT improves robustness on transformed test sets by up to 11 percentage points on clone detection and 19 on code classification while matching or improving standard accuracy, and our ablations isolate multi-positive invariant contrast as the main source of the gains. Our aim is not a new objective but a careful measurement of where encoder robustness breaks and how far a simple, code-only recipe can recover it.

1 Introduction

Encoder-based code models remain practical for discriminative tasks but degrade on semantically equivalent rewrites. The paper measures this robustness gap and evaluates InvPT, a minimal code-only continued-pretraining recipe combining invariant transformations with multi-positive contrastive learning.

Motivation: Encoder-based models remain practical for clone detection and code classification because they are smaller than generation models while performing well.They also support pipeline tasks such as corpus deduplication, code retrieval and ranking, and near-duplicate detection.
Problem: Semantically equivalent but syntactically varied programs expose a robustness gap because model behavior degrades despite unchanged observable program behavior.The paper calls these semantically equivalent rewritings invariant programs, including transformations such as rewriting a counting for loop as an equivalent while loop.
Method: InvPT is a minimal, code-only continued-pretraining recipe that applies semantics-preserving transformations and supervised contrastive learning to programming-language data.It mixes self-contrast pairs from different masks with invariant-contrast pairs from transformed code, providing positives of varying difficulty.
Ablations: Multi-positive invariant contrastive learning is the dominant source of improvement, self-contrast adds consistent gains, and paired natural-language descriptions are unnecessary.These findings support the paper’s emphasis on measuring robustness and testing a simple code-only recipe rather than introducing a new objective.
Evaluation: Four encoder models are evaluated under semantics-preserving transformations across clone detection and code classification on POJ104, Java250, Python800, and C++1400.The study quantifies consistent degradation under semantically equivalent rewrites.
Results: 8.07 points median clone-detection robustness improvement and 3.56 classification improvement occur across transformed model–dataset comparisons, while standard clone-detection accuracy gains up to 6.57 points.InvPT improves robustness in every transformed model–dataset comparison but recovers only part of the degradation.

2 Invariant Pretraining

InvPT continues pretraining on original code and semantically equivalent syntactic variants, combining masked language modeling with multi-positive supervised contrastive learning. Its design mixes self-contrast and invariant-contrast positives of varying difficulty, with both signals contributing to performance.

Method: InvPT generates semantically equivalent, syntactically varied code variants and trains on originals and variants with masked language modeling plus invariant contrastive learning.The method’s three components are invariant transformations, code-only contrastive learning, and mixed-difficulty contrastive pairs.
Invariant transformations: The six transformation operators are implemented at the AST level with language-specific tooling, while unsupported or unparsable cases are excluded.Operator applicability varies by language; for example, Python iterator-based for-loops preclude W2F and F2W.
Invariant transformations: Transformations preserve semantics only under explicit preconditions, including restricting loop rewrites and primitive numeric operators where needed.F2W excludes loop bodies containing continue, while Python AA2EA can change aliased object behavior because augmented assignment may mutate in place.
Contrastive objective: Multi-positive supervised contrastive learning treats all augmentations of the same source function as positives, avoiding InfoNCE’s false negatives when variants co-occur.With k occurrences of a source function, each anchor can have up to 2k−1 positives.
Contrastive objective: Removing self-contrast while retaining only hard invariant pairs lowers performance on every original benchmark, showing that easy and hard positive signals are complementary.The shared encoder is trained with MLM on original and invariant code together with invariant contrastive loss.

3 Experiments

Experiments evaluate InvPT across clone detection and code classification using four encoder families and transformed test sets with unseen transformation compositions. InvPT preserves standard accuracy while improving robustness, with multi-positive contrastive learning identified as the main driver.

Experimental setup: InvPT is evaluated on clone detection and code classification across POJ104 and three CodeNet subsets, comparing four baseline encoder families with corresponding InvPT variants.Clone detection uses MAP@R, while code classification uses accuracy; standard accuracy refers to original, untransformed test sets.
Standard performance: InvPT matches or slightly improves baseline accuracy on all four code-classification datasets despite near-saturated original-test performance.This tests whether robustness gains come at the cost of standard downstream performance.
Robustness: All 16 transformed model–dataset comparisons improve, including cross-language gains on held-out C/C++ data absent from invariant pretraining.On transformed C++1400, InvCodeBERT improves over CodeBERT by +5.84 pp (+22.8%).
Ablations: Removing contrastive learning leaves MLM-only training below full InvCodeBERT, whereas removing self-contrast consistently reduces original-test performance and usually robustness.On original POJ104, MLM-only training reaches 83.87 versus 85.47 for CodeBERT; removing self-contrast lowers original performance by roughly 0.5 to 1.1 points.
Ablations: Adding natural-language descriptions is unnecessary: +NL averages 0.76 pp above InvCodeBERT on transformed data but 0.18 pp below it on original data.The transformed-data comparison is 56.58 versus 55.82, while the original-data comparison is 78.49 versus 78.67.

4 Related Work

Prior code-representation and contrastive-learning methods use increasingly rich signals and semantics-preserving transformations, but do not explicitly train invariant embeddings in the same code-only continued-pretraining setting. InvPT differs through multi-positive contrastive learning, explicit transformed-set evaluation, and pretraining-level invariance rather than fine-tuning augmentation or defenses.

Pretrained code models: CodeBERT and GraphCodeBERT exemplify encoder pretraining for code representation, while later methods add AST sequences, data-flow graphs, execution signals, and transformation-based augmentation.These approaches do not explicitly train for invariant embeddings under semantics-preserving transformations.
Contrastive learning for code: Corder and ContraCode contrast semantics-preserving source variants, but their pairwise InfoNCE objectives treat co-batched variants of the same function as negatives.InvPT uses a multi-positive supervised contrastive mask keyed by a hash of the canonical source.
Contrastive learning for code: InvPT applies code-only continued pretraining to released encoders and evaluates robustness on transformed test sets and a held-out target language, rather than standard accuracy alone.The method does not rerun prior objectives on its backbone, so it reports no control comparison for those objectives.
Contrastive learning for code: ContraBERT combines MoCo with bi-modal NL-PL pretraining and primarily evaluates variable-renaming robustness, whereas InvPT uses code-only supervised contrastive pretraining across multiple transformation types and downstream tasks.InvPT transformations include loop equivalence, branch reversal, scalar-increment forms, and variable renaming.
Adversarial robustness: SPAT uses semantics-preserving transformations for fine-tuning augmentation, while adversarial defenses such as DAMP, MHM, and ALERT attack code models through variable renaming and dead-code insertion.InvPT instead builds invariance into pretraining, amortizing the benefit across downstream tasks.

5 Limitations

The evaluation establishes robustness only for a specific rewrite family and benchmark setting, while leaving broader transformations, industrial code, other languages, and generation untested. Results are also subject to run-to-run uncertainty and possible behavior changes from unverified automated rewrites.

Evaluation scope: The test-time compositions use the same six pretraining operators, so the results establish invariance to that rewrite family rather than semantics-preserving changes generally.Independently generated transformations and naturally occurring refactorings remain untested, as does comparison with downstream-only augmentation.
Experimental uncertainty: Single runs per configuration make differences of a point or less, including several Table 8 ablation margins, uncertain.The authors recommend reading these small differences with caution.
Transformation validity and language coverage: Automated, non-formally-verified rewrites may rarely alter runtime behavior, especially Python’s AA2EA, while extending the operators to Haskell or Rust requires non-trivial engineering.The affected corpus fraction is expected to be small but was not measured.
Task and model scope: The study uses competitive-programming benchmarks and 125M-parameter CodeSearchNet-pretrained encoder-only models, under-representing industrial codebases and limiting generality.Invariance applies when outputs should remain unchanged under rewriting; generation instead requires equivariance, which remains future work.

6 Conclusion

InvPT is a code-only continued pretraining method that improves robustness across clone detection and code classification while matching or improving standard performance. The gains recover only part of robustness degradation, motivating equivariance under known transformation-specific output maps as a promising direction.

Method: InvPT combines semantics-preserving transformations with supervised contrastive learning and mixed-difficulty positives, without paired natural-language data.The method mixes transformed-code positives with self-contrast pairs.
Results: InvPT improves robustness in every model–dataset comparison, by median 8.07 percentage points on clone detection and 3.56 on code classification.It matches or improves standard task performance.
Results: Up to 11.04 percentage points on clone detection and 19.16 on code classification are recovered, though these gains address only part of the degradation.Evaluation uses compositions of operator families seen during pretraining.
Ablations: Invariant contrastive learning is the primary driver of the gains, according to ablations.The passage also indicates self-contrast contributes, but the supplied text truncates its description.
Limitations and future work: Robustness to structurally different rewrites remains open, while learning equivariance under a known transformation-specific output map is identified as most promising.Equivariance is expressed as g(T(x)) = M_T(g(x)); invariance is the special case M_T = id.