diff --git a/doc/self-contrast.md b/doc/self-contrast.md new file mode 100644 index 0000000..80a2acc --- /dev/null +++ b/doc/self-contrast.md @@ -0,0 +1,174 @@ +# Self-Contrast: Analysis + +## What Is Self-Contrast? + +Self-contrast is a data-level strategy for handling rows in the pre-training dataset +that have **no code transformation** (i.e., `transformed` is `None`). +Instead of discarding these rows, self-contrast copies the original code into the +`transformed` field, creating a pair where both the anchor and the augmentation are +the **same source code**. + +Because MLM masking is applied **independently** to each copy +(`DataCollatorForLanguageModeling` with `mlm_probability=0.15`), the model sees +two views of the same code with **different random masks**. This produces an "easy" +contrastive pair — same semantics, same syntax, different masked tokens. + +### Implementation + +The entire mechanism is a single `dataset.map` call in `pretrain.py:268-276`: + +```python +if self_contrast: + dataset = dataset.map( + lambda transformed, code: { + "transformed": transformed if transformed is not None else code + }, + input_columns=["transformed", "code"], + num_proc=num_proc, + ) +``` + +No changes are needed in the loss functions, model architecture, or training loop. +The contrastive loss (InfoNCE, SupCon, or Grouped) treats self-contrast pairs +identically to invariant-contrast pairs. + +## Why It Is Feasible + +### 1. MLM Masking Creates Distinct Views + +Even though both copies share the same token sequence, the `DataCollatorForLanguageModeling` +samples a fresh 15% mask for each copy independently. This means: + +- The `code_input_ids` has some tokens replaced with `[MASK]`. +- The `aug_input_ids` has a **different** set of tokens replaced with `[MASK]`. + +The CLS embeddings produced by the encoder therefore differ, because the model +conditions on different visible tokens. This difference is what makes the contrastive +signal non-trivial — the model must learn that both masked views correspond to the +same underlying program. + +### 2. Stochastic Masking as a Natural Data Augmentation + +MLM masking is itself a form of stochastic data augmentation. +Given a sequence of length $L$, the number of possible 15%-masks is +$\binom{L}{\lfloor 0.15L \rfloor}$, which is astronomically large for typical code +sequences ($L \approx 100$–$256$). Each training step therefore produces a genuinely +novel pair, even if the underlying code is the same. The model never sees the exact +same (mask_a, mask_b) combination twice in practice. + +### 3. Zero Implementation Overhead + +Self-contrast requires no additional model components, no separate training stage, +and no architectural changes. It is purely a data-filling operation that replaces +`None` with the original code string. The downstream data collator, tokenizer, +and loss functions are completely agnostic to whether a pair is self-contrast or +invariant-contrast. + +### 4. Consistency with the Shared-Encoder Design + +InvPT uses a **single shared encoder** for both the original and transformed code +(unlike ContraBERT's momentum contrast with a separate momentum encoder). +In this architecture, the same parameters process both copies. +Self-contrast is a natural fit: since the encoder is shared, feeding the same code +twice with different masks is equivalent to asking "do two partial observations of +the same program yield similar representations?" — a well-defined learning objective. + +## Why It Is Beneficial + +### 1. Implicit Curriculum Learning + +Self-contrast creates an implicit curriculum with two difficulty levels: + +| Pair Type | Difficulty | Source of Difference | +| ------------------ | ---------- | ---------------------------------------------- | +| Self-contrast | Easy | Different MLM masks only | +| Invariant-contrast | Hard | Syntactic transformation + different MLM masks | + +**Easy pairs** (self-contrast) produce CLS embeddings that are already relatively +close, generating small but consistent gradients that stabilize training. +**Hard pairs** (invariant-contrast) produce more distant embeddings, generating +larger gradients that push the model to learn semantic invariance. + +Training on both simultaneously provides the benefits of curriculum learning +without explicit scheduling: the easy pairs act as a regularizer, preventing the +model from being destabilized by the harder invariant pairs early in training. +As the model improves, self-contrast pairs become trivially easy and the learning +signal shifts toward the harder invariant pairs, naturally concentrating the +gradient on the most informative examples. + +### 2. Data Efficiency — No Wasted Rows + +Not every code sample in the dataset has a successful transformation. +Transformation operators may fail (e.g., no for-loops to convert, no local +variables to rename). Without self-contrast, these rows are **discarded**: + +```python +# Without self-contrast: rows without augmentation are dropped +dataset = dataset.filter( + lambda transformed: transformed is not None and transformed != "", + input_columns=["transformed"], +) +``` + +This wastes a potentially large fraction of the dataset. With self-contrast, +**every row contributes** to both the MLM and contrastive objectives. The model +sees more code, learns better MLM representations, and the contrastive batch +is denser (more positives in SupCon mode, since rows sharing the same +`function_id` all participate). + +### 3. MLM Regularization on the Full Corpus + +Even if the contrastive gradient from a self-contrast pair is small, the **MLM +loss** on that pair is still fully informative — the model must predict different +masked tokens in each copy. Retaining these rows ensures the MLM objective +covers the entire pre-training corpus, not just the transformable subset. +This is important because the transformation operators have language-specific +coverage (e.g., Python lacks `F2W` and `W2F`), so dropping untransformable rows +would bias the MLM objective toward code containing certain syntactic patterns. + +### 4. Representation Stability via Alignment + +Self-contrast directly optimizes what Chen & He (2021) call the **alignment** +property of contrastive representations: embeddings of the same instance +(under different views) should be close. By ensuring the encoder maps different +masked views of the same code to similar CLS vectors, self-contrast encourages +a smooth, stable representation manifold where small input perturbations +(missing tokens) do not cause large embedding shifts. This stability transfers +to downstream tasks where inputs may be noisy or partial. + +### 5. Compatibility with SupCon Multi-Positive Masking + +In SupCon mode, all embeddings sharing the same `function_id` are treated as +positives. Self-contrast rows share `function_id` with their corresponding +invariant-contrast rows (since `function_id` is a hash of the original code). +This means self-contrast rows add **additional positive pairs** to the SupCon +loss, increasing the effective number of positives per anchor. Khosla et al. +(2020) showed that SupCon benefits from more positives per class, making +self-contrast especially synergistic with the SupCon contrastive mode. + +## Connection to Prior Work + +Self-contrast is conceptually related to several established techniques: + +- **SimCLR** (Chen et al., 2020): Uses two random augmentations of the same + image as a positive pair. Self-contrast uses two random MLM masks of the + same code as a positive pair — the MLM mask plays the role of the + stochastic augmentation. +- **Dropout-as-augmentation** (Gao et al., 2021, SimCSE): Passes the same + sentence through the encoder twice with different dropout masks to create + positive pairs. Self-contrast is analogous but uses MLM masking instead + of dropout as the source of stochasticity. +- **BERT's own MLM objective**: Standard MLM already trains the model to be + robust to random token masking. Self-contrast extends this to the + contrastive objective, requiring not just correct token prediction but + also consistent sequence-level (CLS) representations across masks. + +## Summary + +Self-contrast is a simple, zero-cost data strategy that: + +1. Turns untransformable rows into useful training signal (data efficiency). +2. Creates an implicit easy/hard curriculum (training stability). +3. Ensures MLM coverage of the full corpus (unbiased language modeling). +4. Encourages alignment and smoothness in the representation space (downstream robustness). +5. Requires no changes to the model, loss, or training loop (implementation simplicity). diff --git a/experiments/grouped/codebert.yaml b/experiments/grouped/codebert.yaml index b414305..a1ea518 100644 --- a/experiments/grouped/codebert.yaml +++ b/experiments/grouped/codebert.yaml @@ -2,7 +2,7 @@ # Effective batch size: 16 * 16 = 256 # Usage: python modeling/cli.py run experiments/grouped/codebert.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "microsoft/codebert-base" batch_size: 32 @@ -19,3 +19,4 @@ max_seq_length: 512 contra_mode: "grouped" max_num_augs: 6 +self_contrast: true diff --git a/experiments/grouped/contrabert_c.yaml b/experiments/grouped/contrabert_c.yaml index a7456e4..7d4b740 100644 --- a/experiments/grouped/contrabert_c.yaml +++ b/experiments/grouped/contrabert_c.yaml @@ -2,7 +2,7 @@ # Effective batch size: 16 * 16 = 256 # Usage: python modeling/cli.py run experiments/grouped/contrabert_c.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "./saved_models/ContraBERT_C" tokenizer_name: "microsoft/codebert-base" @@ -20,3 +20,4 @@ max_seq_length: 512 contra_mode: "grouped" max_num_augs: 6 +self_contrast: true diff --git a/experiments/grouped/contrabert_g.yaml b/experiments/grouped/contrabert_g.yaml index d9ae460..1c7cfd1 100644 --- a/experiments/grouped/contrabert_g.yaml +++ b/experiments/grouped/contrabert_g.yaml @@ -2,7 +2,7 @@ # Effective batch size: 16 * 16 = 256 # Usage: python modeling/cli.py run experiments/grouped/contrabert_g.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "./saved_models/ContraBERT_G" tokenizer_name: "microsoft/graphcodebert-base" @@ -20,3 +20,4 @@ max_seq_length: 512 contra_mode: "grouped" max_num_augs: 6 +self_contrast: true diff --git a/experiments/grouped/graphcodebert.yaml b/experiments/grouped/graphcodebert.yaml index da9ff73..d7eb62a 100644 --- a/experiments/grouped/graphcodebert.yaml +++ b/experiments/grouped/graphcodebert.yaml @@ -2,7 +2,7 @@ # Effective batch size: 16 * 16 = 256 # Usage: python modeling/cli.py run experiments/grouped/graphcodebert.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "microsoft/graphcodebert-base" batch_size: 32 @@ -19,3 +19,4 @@ max_seq_length: 512 contra_mode: "grouped" max_num_augs: 6 +self_contrast: true diff --git a/experiments/supcon/codebert.yaml b/experiments/supcon/codebert.yaml index fb70d1a..8dca2c0 100644 --- a/experiments/supcon/codebert.yaml +++ b/experiments/supcon/codebert.yaml @@ -2,7 +2,7 @@ # Effective batch size: 64 * 4 = 256 # Usage: python modeling/cli.py run experiments/supcon/codebert.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "microsoft/codebert-base" batch_size: 256 @@ -18,3 +18,4 @@ temperature: 0.1 max_seq_length: 512 contra_mode: "supcon" +self_contrast: true diff --git a/experiments/supcon/contrabert_c.yaml b/experiments/supcon/contrabert_c.yaml index 039e83f..14d4c56 100644 --- a/experiments/supcon/contrabert_c.yaml +++ b/experiments/supcon/contrabert_c.yaml @@ -2,7 +2,7 @@ # Effective batch size: 64 * 4 = 256 # Usage: python modeling/cli.py run experiments/supcon/contrabert_c.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "./saved_models/ContraBERT_C" tokenizer_name: "microsoft/codebert-base" @@ -19,3 +19,4 @@ temperature: 0.1 max_seq_length: 512 contra_mode: "supcon" +self_contrast: true diff --git a/experiments/supcon/contrabert_g.yaml b/experiments/supcon/contrabert_g.yaml index eb262b6..f3fe39a 100644 --- a/experiments/supcon/contrabert_g.yaml +++ b/experiments/supcon/contrabert_g.yaml @@ -2,7 +2,7 @@ # Effective batch size: 64 * 4 = 256 # Usage: python modeling/cli.py run experiments/supcon/contrabert_g.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "./saved_models/ContraBERT_G" tokenizer_name: "microsoft/graphcodebert-base" @@ -19,3 +19,4 @@ temperature: 0.1 max_seq_length: 512 contra_mode: "supcon" +self_contrast: true diff --git a/experiments/supcon/graphcodebert.yaml b/experiments/supcon/graphcodebert.yaml index 8d8c835..e5fc701 100644 --- a/experiments/supcon/graphcodebert.yaml +++ b/experiments/supcon/graphcodebert.yaml @@ -2,7 +2,7 @@ # Effective batch size: 64 * 4 = 256 # Usage: python modeling/cli.py run experiments/supcon/graphcodebert.yaml -dataset_path: "data/aug_csn.jsonl" +dataset_path: "data/csn.jsonl" model_name: "microsoft/graphcodebert-base" batch_size: 256 # 64 @@ -18,3 +18,4 @@ temperature: 0.1 max_seq_length: 512 contra_mode: "supcon" +self_contrast: true diff --git a/modeling/cli.py b/modeling/cli.py index 28cef55..97c548a 100644 --- a/modeling/cli.py +++ b/modeling/cli.py @@ -24,18 +24,26 @@ def run( Optional[float], typer.Option("--sample-rate", "-sr", help="Override sample rate from config."), ] = None, + self_contrast: Annotated[ + Optional[bool], + typer.Option( + help="Override self-contrast setting from config. Use --no-self-contrast to disable.", + ), + ] = None, ) -> None: """Run pre-training from a YAML config file. All parameters are read from the config file to ensure full reproducibility. - Use --sample-rate to override the dataset sampling fraction. + Use --sample-rate / --no-self-contrast to override specific settings. - Example: python -m modeling run experiments/base.yaml --sample-rate 0.1 + Example: python -m modeling run experiments/base.yaml --no-self-contrast """ cfg = load_config(config) kwargs = asdict(cfg) if sample_rate is not None: kwargs["sample_rate"] = sample_rate + if self_contrast is not None: + kwargs["self_contrast"] = self_contrast main(**kwargs) @@ -88,6 +96,12 @@ def pretrain( max_num_augs: Annotated[ int, typer.Option(help="Max augmentations per anchor (grouped mode).") ] = 6, + self_contrast: Annotated[ + bool, + typer.Option( + help="Use self-contrast (same code, different MLM masks) for rows without augmentation. If disabled, rows without augmentation are dropped." + ), + ] = True, ) -> None: """Run pre-training with all parameters specified as CLI options. @@ -113,6 +127,7 @@ def pretrain( checkpoint=checkpoint, contra_mode=contra_mode, max_num_augs=max_num_augs, + self_contrast=self_contrast, ) diff --git a/modeling/config.py b/modeling/config.py index 4771ea2..f24b365 100644 --- a/modeling/config.py +++ b/modeling/config.py @@ -32,6 +32,7 @@ class PretrainConfig: tokenizer_name: str | None = None contra_mode: ContraMode = ContraMode.INFO_NCE max_num_augs: int = 6 + self_contrast: bool = True def load_config(path: str | Path) -> PretrainConfig: diff --git a/modeling/pretrain.py b/modeling/pretrain.py index fa92802..38ee846 100644 --- a/modeling/pretrain.py +++ b/modeling/pretrain.py @@ -223,6 +223,7 @@ def main( tokenizer_name: str | None = None, contra_mode: ContraMode = "info_nce", max_num_augs: int = 6, + self_contrast: bool = True, ): set_seed(seed) @@ -264,11 +265,23 @@ def main( } ) dataset = load_dataset("json", data_files=dataset_path, features=features)["train"] - dataset = dataset.filter( - lambda transformed: transformed is not None, - input_columns=["transformed"], - num_proc=num_proc, - ) + if self_contrast: + # For rows without a transformation, use the original code as the + # augmentation (self-contrast: same code, different MLM masks). + dataset = dataset.map( + lambda transformed, code: { + "transformed": transformed if transformed is not None else code + }, + input_columns=["transformed", "code"], + num_proc=num_proc, + ) + else: + # Drop rows that have no real augmentation. + dataset = dataset.filter( + lambda transformed: transformed is not None and transformed != "", + input_columns=["transformed"], + num_proc=num_proc, + ) if sample_rate < 1.0: dataset = dataset.shuffle(seed=seed).select(