Project.md

Project: Code-JEPA

This document is the scientific source of truth for the project. Code, experiments, and paper drafts should follow this document. Update it when the research understanding changes, not when implementation details move around.

See also: docs/diagrams/code-jepa-project-map.pdf.

Main idea

Train a code-native Joint-Embedding Predictive Architecture (Code-JEPA) that learns code-space geometry useful for code agents and code LLM post-training.

The representation should be:

• invariant to harmless program changes;
• sensitive to tiny behavior-changing edits;
• usable without executing code;
• useful as a frozen judge, memory, and ranker for downstream training.

The downstream name we have been using is RLJF: Reinforcement Learning from JEPA Feedback. The broad paper is bigger than only RLJF: it also includes hindsight feedback and failure-memory exploration for code agents.

Core hypothesis

A Code-JEPA trained with positive behavior-preserving transformations and hard negative behavior-impacting mutations can learn a more useful code similarity space than generic code embeddings or token-level similarity.

Target geometry:

renaming / formatting / safe refactor        -> close
< vs <= / wrong variable / swapped args      -> far in semantic/local space
same task, different algorithm, both correct -> semantically close, strategically maybe far
unrelated code                               -> far

The model is not a correctness oracle. It ranks, clusters, detects duplicate failures, and supplies preference signals when there is an anchor such as a reference solution, verifier success, accepted patch, or hindsight final solution.

Model shape

Use a shared-encoder Siamese setup as the default paper path:

code view -> shared transformer encoder -> hidden states H
          -> mask-aware mean pool        -> h
          -> projection head p(h)        -> z

anchor/positive/negative views -> same shared encoder/projection -> z_anchor/z_pos/z_neg
train z_anchor close to z_pos and far from z_neg
SIGReg(z outputs)

Use one shared encoder for all views and all languages. The default learning signal is masked-span latent prediction (data2vec / I-JEPA style): encode the full anchor as a stop-gradient target, encode a span-masked anchor as context, and use a predictor conditioned on the masked positions to predict the target's per-token representations in latent space. This is the JEPA mechanism that learns code semantics without input-space token reconstruction. It replaces the earlier MLM auxiliary, which is explicitly rejected as off-thesis: predicting tokens is input-space reconstruction and contradicts the JEPA premise. Prediction is in representation space only — never reconstruct tokens. Because the predictor is conditioned on the mask/context it is a conditioned predictor; a bare unconditioned z_anchor -> predictor -> z_pos remains only an ablation. Use one shared encoder for the context and the target branch — never a separate or EMA target network. Collapse prevention is SIGReg, applied at three sites: the per-token latent targets, the pooled embedding h, and the projected space z. This is the LeJEPA position and it is confirmed in-repo (2026-07-06/07 collapse diagnostic): the hybrid latent run, which regularized only z, showed monotonic latent-target variance decay (0.45 -> 0.22) and pooled-h cone collapse (effective rank 21 vs 66 at random init), while the pure LeJEPA run with token+pool SIGReg kept a healthy, retrievable h. Stop-gradient on the target branch is an implementation detail to ablate, not a load-bearing mechanism, and EMA teachers are off the recipe entirely. Use a small SIGReg coefficient around 0.03-0.05 per site; 0.10 should be treated as high/likely too much, not a normal setting.

Masked-span latent prediction (default semantic objective)

anchor tokens
  -> span-mask ~40% (contiguous spans, mask id = vocab_size widened row)
  -> context branch  E(masked)                  -> H_context
  -> target  branch  E(full), same shared encoder -> H_target   (instance-normed; SIGReg'd; stop-grad optional)
  -> predictor(H_context, mask-position embed)   -> H_pred
  L_latent = smooth_l1( H_pred , H_target ) over masked positions only

The predictor is a small transformer (~2 layers), conditioned on which positions are masked, and is discarded downstream like any pretraining head. This is the dense, understanding-forcing signal that MLM provided, but in representation space. The full objective combines it with the contrastive and behavioral terms:

L = lambda_latent L_latent + lambda_pos L_pos + lambda_neg L_rank
  + lambda_inbatch L_inbatch + lambda_sigreg L_SIGReg

where L_latent is contextual semantics, L_pos is invariance (anchor vs behavior-preserving transform), and L_rank is behavioral sensitivity (anchor vs behavior-impacting mutation). See docs/latent-jepa-objective.md for the full spec, collapse-prevention plan, and build steps.

The default global projection head should be:

p(h) = Dense(8H) -> split 4H/4H -> SwiGLU -> RMSNorm -> Dense(D)

Keep h = pool(H) as the frozen downstream/search embedding candidate, and train JEPA/ranking losses on z = p(h). SIGReg applies to both spaces: on z to shape the training space, and (at a small coefficient) on h and the per-token states — the collapse diagnostic showed that regularizing only z leaves h free to collapse into a low-rank cone, which is what killed frozen retrieval. The current no-head pooled-output path is only a baseline and throughput path. See docs/design-notes/embedding-pooling-and-projection-head.md.

On top of the shared code encoder, use small projection/readout heads, not separate full transformers. The first multi-head implementation should be a frozen-backbone two-head post-training stage:

code -> shared transformer encoder -> hidden states H -> h = pool(H)
  -> semantic head  Dense(8H) -> split 4H/4H -> SwiGLU -> RMSNorm -> Dense(D) -> z_semantic
  -> lexical head   Dense(8H) -> split 4H/4H -> SwiGLU -> RMSNorm -> Dense(D) -> z_lexical

Use the pretrained Code-JEPA encoder as the frozen backbone first. Train the semantic head and lexical head as separate phases; because the encoder is frozen and heads do not share parameters, order does not matter. Do not use SIGReg for this first supervised head-training stage. Add top-layer encoder unfreezing only after the frozen-head baseline is measured.

The same code pair can be positive for one head and negative for another. Example: i < n vs i <= n should be lexical-close but semantic-far. Same-task accepted solutions should be semantic-close even when lexically far.

Data object hierarchy

Do not flatten everything into isolated snippets. Keep whole-file context and derive trainable units from it:

repo
  -> file
     -> imports / constants / top-level definitions
     -> class
        -> method
     -> function
     -> AST/local spans

Whole files are the canonical source and context. Functions/methods are the main training units. Local AST/token spans are needed for tiny semantic changes.

Recommended records:

• files: repo/path/license/source/imports/top-level defs/parse status;
• units: function/method/nested function/class summary/file window/span window/task solution with parent file, line range, context, identifiers, calls, and AST node-type sequence;
• spans: AST node spans and changed spans;
• views: anchor/positive/negative/context/local-span/task-solution transformed code;
• relations: file contains unit, class contains method, unit has AST auxiliary view, unit uses import, function calls function;
• triples: anchor-positive-negative training relations with per-head labels;
• semantic_pairs: same-task accepted-solution pairs, including same-language different implementations and cross-language solutions.

The canonical data-prep pipeline is tokenizer-agnostic and writes reproducible Parquet + zstd segments by dataset and transformation stage. Keep raw code, spans, context, AST side channels, transform metadata, rough length buckets, repo/task split, and sampling weights. Do not tokenize until the model/backbone tokenizer is chosen.

Code block sizes

Prepare buckets instead of one global cutoff:

Bucket	Rough size	Use
tiny	3-10 LOC / 32-128 tokens	local operator sensitivity
short	10-40 LOC / 128-512 tokens	main hard-negative training
medium	40-120 LOC / 512-1536 tokens	realistic function embeddings
long	120-250 LOC / 1536-4096 tokens	later long-context evaluation
file/class	250+ LOC	context/file JEPA, not single-vector semantic sensitivity

First serious training should focus on 10-120 LOC functions/methods while retaining file context.

Transformations

Positive transformations

Positive views should preserve behavior with high confidence:

• variable/argument renaming with references updated;
• local helper renaming when safe;
• formatting, whitespace, quote normalization;
• comment/docstring removal or rewriting when safe;
• trivial import split/merge/reorder when no side-effect risk;
• equivalent syntax rewrites only under conservative rules;
• conservative control-flow rewrites such as boolean-return simplification and if-return conditional expressions;
• guarded independent-statement reordering;
• guarded refactor-like rewrites such as simple loop-to-comprehension, range-for-to-while, and accumulator-to-sum rewrites;
• alternate structural views: AST, DFG, CFG, call graph, dependency graph, call-site context.

Riskier positives such as broad statement reordering, algebraic rewrites, top-level definition reordering, import style normalization, type annotation changes, and function extraction/inlining should be delayed or marked with confidence flags.

Hard negatives

Hard negatives are compile-valid behavior-impacting mutations relative to the original, not guaranteed test failures:

• < <-> <=, > <-> >=, == <-> !=;
• +1 <-> -1, loop-bound changes, off-by-one edits;
• and <-> or;
• wrong variable from the same scope;
• swapped call arguments;
• wrong default value or wrong API order;
• missing return / missing await;
• wrong sort direction;
• mutate copy vs original;
• missing edge-case branch or wrong exception handling.

Always record changed byte/AST spans. The local head depends on this metadata.

Dataset transform stages

These names describe transformation families only, not data-pipeline versions. The canonical prep pipeline emits each stage as a separate reproducible delta segment to avoid duplicating training triples when the whole dataset is processed. Training recipes are cumulative by selecting multiple segments, e.g. v0 + v1 + v2.

• v0: first conservative synthetic set. Positives: AST normalization, docstring removal, local variable/argument renaming. Negatives: comparison flip, boolean-op flip, call-argument swap, wrong variable, small integer flip.
• v1: delta over v0. Positives: independent assignment reorder, boolean-return simplification, if-return to conditional expression, unreachable-else removal, and same-block import sorting. Negatives: membership/identity flip, condition negation, arithmetic-op flip, subscript-index flip, default-value flip, sort reverse flip, return-value removal, and await removal.
• v2: delta over v1. Positives: simple range-for to while-loop, list-append loop to comprehension, accumulator loop to sum, De Morgan boolean rewrite, and broader independent statement reorder. Negatives: loop-bound off-by-one, missing guard/edge-case branch, wrong exception type, dropped keyword argument, copy-vs-alias mutation, and dropped context manager/resource handling.

Per-head training labels

Pair type	Strategy head	Semantic head	Local head
formatting / rename / comments	close	close	low change
behavior-preserving refactor	close	close	aligned
< vs <=, off-by-one	close	far	changed span important
wrong variable / swapped args	close	far	changed span important
same task, different accepted solution	far/maybe	close	diffuse
unrelated code	far	far	high difference

Possible loss structure:

L = L_JEPA + lambda_pos L_pos + lambda_neg L_rank + lambda_local L_local + lambda_sigreg L_SIGReg

Ranking example:

max(0, margin + sim(E(y), E(y_neg)) - sim(E(y), E(y_pos)))

The project should avoid the failure mode where a global embedding treats i < n and i <= n as merely almost identical.

Two-head post-training data

The semantic head needs enough supervised task-equivalence data despite being smaller than the backbone. POJ-104 is a useful first source, but one positive and one negative per anchor is only a smoke cache. Use many same-problem positives and different-problem negatives per anchor, then move to hard-mined pairs:

semantic easy stage:
  close = two accepted solutions with the same problem id
  far   = accepted solutions with different problem ids

semantic hard-mining stage:
  close = same-problem pairs that the current h/semantic head ranks too far apart
  far   = different-problem pairs that the current h/semantic head ranks too close

The hard-mining stage is important because random different-problem negatives are often too easy, and same-problem positives that are already close add little new signal. CodeNet/APPS/CodeContests should later provide larger same-task accepted-solution groups while POJ-104 remains a small controlled benchmark.

The lexical head can use transformed CodeSearchNet pairs:

close = anchor vs behavior-preserving transform
close = anchor vs hard semantic mutation, because it is a small edit
far   = anchor vs unrelated function

This explicitly separates edit-distance/surface similarity from semantic equivalence.

Downstream uses

1. Reference-known RLJF / reranking

Given prompt x, reference solution y_ref, and sampled candidates y_1...y_k:

score_i = sim(z_semantic(y_i), z_semantic(y_ref))

Use top/bottom candidates as preference pairs for DPO-style post-training or as reranking output. Tests stay evaluation-only in the cleanest RLJF setup.

2. Hindsight agent training when a solution is eventually reached

In pure self-training, y_ref is not known upfront. If an agent eventually reaches a verified/accepted solution:

y_1 fails, y_2 fails, ..., y_T succeeds
hindsight reference y_ref := y_T

Then earlier attempts can be ranked by similarity to the final solution and by whether they repeat known failure clusters. This is hindsight Code-JEPA feedback.

3. No-solution-yet failure memory

If all current candidates fail, Code-JEPA still helps, but only as exploration memory:

M_bad = clusters of failed attempts in strategy/semantic/local space

For a new candidate:

high strategy similarity + high semantic similarity to bad cluster -> duplicate failure / rephrase
high strategy similarity + low semantic similarity                 -> possible meaningful fix
low strategy similarity                                            -> new solution family

This rejects candidates that are only rephrasings of the same failed program. It does not certify that a novel candidate is correct. Novelty must be constrained by prompt fit, compile/static checks, type checks, or later verification.

4. Semantic duplicate rejection during self-training

When a generator samples many candidates for the same task, most may be rephrasings of the same algorithm or the same bug. Code-JEPA can cluster generated candidates before testing/judging:

20 samples -> 4 semantic/strategy clusters -> keep representatives

Use cases:

• reject candidates that are just surface rewrites of known bad attempts;
• improve candidate-set diversity before expensive verification;
• estimate unique solution-family coverage rather than raw sample count.

Metrics:

• unique failure/solution clusters per token budget;
• pass@k per unique cluster;
• same solve rate with fewer tests/rollouts;
• duplicate-rejection precision: do not reject small real fixes as rephrases.

5. Out-of-the-box embedding tasks

Code search and clone detection should be treated as cheap downstream checks that come almost for free from the Siamese embedding setup:

code search:     query/code -> embeddings -> nearest neighbors
clone detection: code/code  -> embeddings -> threshold or retrieval ranking

These tasks are not the core claim, but they are useful sanity checks and baseline comparisons because they require no agent loop, no verifier, and no decoder.

6. Cross-language retrieval and code translation

Use cross-language code-to-code retrieval as a bridge to code translation. A simple future task is Python-to-Java translation in latent space:

Python context code -> Python context encoder -> z_py
Java target code    -> Java target encoder    -> z_java
z_py -> predictor -> predicted z_java

The predictor learns the Python-to-Java semantic mapping. This is a downstream/ablation setup, not a replacement for the default shared-encoder Code-JEPA training. Compare it against the simpler Siamese shared-encoder retrieval setup where Python and Java snippets are embedded directly and matched by cosine similarity.

Evaluation should start with CodeNet-style same-problem Python/Java pairs and measure retrieval MAP/MRR before adding any decoder. If generation is added later, condition a Java generator on retrieved neighbors or the predicted Java-space embedding and compare against standard translation models.

Ablations/baselines:

• shared Siamese Code-JEPA retrieval vs Python-context/Java-target predictor mapping;
• frozen vs finetuned encoders;
• h embedding vs projected z space;
• lexical and AST baselines;
• CodeBERT / GraphCodeBERT / UniXcoder / CodeT5-style embeddings;
• supervised seq2seq code translation where available.

Evaluation and baselines

Target workshop: REALM workshop call for papers, https://realm-workshop.github.io/call_for_papers/. Fit under the theme: “Coding Agents: Specialized architectures, training strategies, and datasets for software engineering tasks such as code generation, debugging, refactoring, and unit testing.”

The first scientific bar is not DPO. It is proving that Code-JEPA is a better candidate judge/ranker than obvious baselines.

Required comparisons:

• base generator;
• supervised fine-tuning baseline;
• lexical/reference heuristics;
• CodeBERTScore / CodeBERT / UniXcoder-style embedding rankers;
• generic embedding rankers;
• Code-JEPA reranking;
• DPO with generic/code-embedding preferences;
• RLJF with frozen Code-JEPA preferences.

Useful evaluation axes:

• candidate reranking quality;
• pass@1 improvement after reranking;
• robustness to refactors/renames/formatting;
• discrimination of hard negatives;
• failure-cluster duplicate detection;
• whether small meaningful fixes are kept instead of rejected as rephrases.

Expected gains may be modest in pass@1. The stronger story is representation quality: semantic reranking, hard-negative sensitivity, and agent failure-memory usefulness.

Current empirical status: the main pretraining path is now a 25M RoBERTa-style no-predictor Siamese JAX model over six-language CodeSearchNet with custom bpe16k tokenization, hard negatives, in-batch loss, and small SIGReg. The full tokenized transform cache contains about 59.4M triplets. A one-epoch 2-H200 pretrain checkpoint learned local hard-negative structure but transferred weakly to POJ-104: zero-shot h MAP@R was 8.97% valid / 5.30% test, and a 2-epoch POJ triplet fine-tune moved it to 11.72% valid / 7.19% test. This supports the diagnosis that transform triples teach local semantic sensitivity, while task-level accepted-solution equivalence needs separate semantic-head data. Fair baselines should be matched 25M CodeBERT/UniXcoder/GraphCodeBERT-style controls trained on the same bpe16k CodeSearchNet data, not only published 125M off-the-shelf numbers.

Data sources

Start with full six-language CodeSearchNet for function-level multilingual pipeline validation: Python, Java, JavaScript, Go, PHP, and Ruby.

Practical order:

1. Full CodeSearchNet for function/method-level synthetic transform training and cross-language sanity checks.
2. CodeParrot clean Python as the public non-gated whole-file fallback when The Stack / StarCoderData auth is unavailable.
3. Larger permissive whole-file corpora from The Stack / StarCoderData / similar once license filtering and storage are decided.
4. Task/reference corpora for real semantic positives and cross-language retrieval: HumanEval, MBPP, APPS, CodeContests, CodeNet-like multi-solution datasets if accessible.

Split by repository/source/task, never by transformed view. A unit and all derived views must remain in the same split. Task datasets should emit accepted-solution semantic pairs when multiple correct implementations or multiple languages exist for the same problem.

Risks and settled constraints

• No reference or verifier means no correctness signal. Code-JEPA can only avoid known-bad regions and encourage diversity.
• Hard negatives are behavior-impacting mutations, not always proven wrong for an unknown spec.
• False positives are dangerous: unsafe “equivalent” transforms can poison invariance.
• Pure novelty can drift into irrelevant code; keep sanity/prompt-fit gates.
• A single global vector is too blunt for one-character semantic changes; keep local span training.
• Do not spend CPU time on exact tokenization until the backbone/tokenizer is chosen.

Related papers injected into this repo

Paper summaries live in docs/paper_summaries/:

• LeJEPA: docs/paper_summaries/summary_2511.08544_lejepa/summary.md
• LeWorldModel: docs/paper_summaries/summary_2603.19312_leworldmodel/summary.md
• LLM-JEPA: docs/paper_summaries/summary_2509.14252_llm_jepa/summary.md
• UniXcoder: docs/paper_summaries/summary_2203.03850_unixcoder/summary.md

The duplication/literature-risk note is in docs/literature/exa-duplication-review.md.