DLLM-JEPA: Joint Embedding Predictive Architectures for Masked Diffusion Language Models

# Joint Embedding Predictive Architectures for Masked Diffusion Language Models
Source: [https://arxiv.org/html/2606.00091](https://arxiv.org/html/2606.00091)
###### Abstract

We introduceDLLM\-JEPA, a JEPA formulation for masked diffusion language models\. JEPA objectives have so far been applied to autoregressive LMs at the cost of explicit paired views and multiple gradient passes per training step\. By leveraging the diffusion noise schedule, DLLM\-JEPA constructs two views from a single input without requiring paired data, and reduces JEPA training FLOPs by33%relative to LLM\-JEPA’s two\-gradient\-view design through a single gradient pass per step\.

Across four tasks and two diffusion backbones, DLLM\-JEPA consistently improves over diffusion\-only fine\-tuning, with modest gains in stable settings \(e\.g\.,\+1\.8 ppon GSM8K\) and larger improvements under more aggressive fine\-tuning, while also tightening seed\-to\-seed variance on the high\-variance LLaDA\-8B GSM8K cells\. In addition, it does not degrade base\-model performance on a held\-out diffusion\-loss probe nor on a small MMLU sanity check\.

We further analyze the representation dynamics induced by the objective and observe a consistent empirical pattern: models trained with DLLM\-JEPA exhibit larger geometric drift from their pre\-trained initialization while maintaining comparable or lower functional forgetting\.

These results suggest that DLLM\-JEPA provides an efficient way to incorporate representation\-level objectives into diffusion language model fine\-tuning\.

Joint Embedding Predictive Architectures, Masked Diffusion Language Models, Representation Learning, Fine\-Tuning

## 1Introduction

The dominant paradigm for training large language models relies on input\-space reconstruction: autoregressive \(AR\) next\-token prediction\(Brown et al\.,[2020](https://arxiv.org/html/2606.00091#bib.bib4); Touvron et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib18)\)or masked token reconstruction\(Devlin et al\.,[2019](https://arxiv.org/html/2606.00091#bib.bib6)\)\. In contrast, vision has increasingly adopted Joint Embedding Predictive Architectures \(JEPAs\)\(Assran et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib1); Bardes et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib3)\), which learn representations by predicting embeddings of one view from another—operating entirely in latent space\. JEPAs have been shown to learn richer, more abstract representations by avoiding pixel\-level reconstruction biases\(LeCun,[2022](https://arxiv.org/html/2606.00091#bib.bib11); Littwin et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib12)\)\.

The recent LLM\-JEPA\(Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\)represents a first step toward bringing JEPA objectives to language models\. By treating \(text, code\) pairs as two views of the same underlying knowledge, LLM\-JEPA adds a JEPA loss alongside the standard next\-token prediction objective\. Despite promising results, LLM\-JEPA faces two fundamental limitations rooted in the autoregressive architecture:

1. 1\.Explicit view requirement\.LLM\-JEPA requires datasets with natural two\-view structures \(e\.g\., natural language paired with code\)\. The authors themselves note this as a key limitation:*“developing a mechanism akin to data\-augmentation in vision would enable JEPA objectives to be used on any dataset”*\(Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\)\.
2. 2\.Computational overhead from unidirectionality\.AR models process tokens left\-to\-right with causal masking\. To obtain independent embeddings for two views, LLM\-JEPA requires a custom block\-causal attention mask and two forward passes—both carrying gradients—resulting in∼\\sim2×\\timesthe training compute of standard fine\-tuning\.

We observe that*masked diffusion language models*\(Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15); Sahoo et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib16); Shi et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib17)\)resolve both limitations naturally\. Models such as LLaDA\(Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15)\)employ bidirectional attention and learn to denoise randomly masked tokens—a process that is structurally analogous to JEPA’s view prediction\. This insight leads to our proposed method,DLLM\-JEPA, whose contributions can be summarized as \(1\) an*architecturally efficient*adaptation of JEPA to masked\-diffusion LMs, \(2\) a*consistent empirical pattern*in how JEPA reshapes representations in bidirectional LMs, and \(3\) consistent task gains together with maintained base\-capability:

- •A JEPA formulation that saves 33% of the training FLOPs\.The bidirectional attention of masked\-diffusion models lets us generate both JEPA views from the same input via different masking rates \(no explicit view pairs\) and process the context view in a single forward pass with gradients that simultaneously yields diffusion logits and the JEPA embedding; only the target view requires an additional no\-gradient forward pass\. The result is a JEPA training step with 33% fewer FLOPs than LLM\-JEPA \(Table[1](https://arxiv.org/html/2606.00091#S3.T1)\), standard bidirectional attention \(no custom block\-causal mask\), and applicability to any text dataset\.
- •A consistent empirical pattern: geometric–functional drift dissociation\.Layer\-wise probing reveals that the DLLM\-JEPA objective does not appear to minimize representation change—it redirects it\. On GSM8K, DLLM\-JEPA’s hidden\-state drift from the pre\-trained weights is*larger*than the baseline’s \(drift ratio 1\.3–3\.6×\\timesacross configurations\), yet its Wikitext functional forgetting is*smaller*\. The amplification concentrates in middle transformer layers—consistent with prior interpretations of middle layers as encoding compositional structure—and replicates on Dream\-7B \(ratio 1\.28×\\times\)\. We also empirically rule out representation collapse \(§[3](https://arxiv.org/html/2606.00091#S3), Appendix[A\.7](https://arxiv.org/html/2606.00091#A1.SS7)\)\. We use “dissociation” in a descriptive sense throughout\.
- •Joint task gain and base preservation\.On LLaDA\-8B with the Wide\-ttconfiguration \(tL=0\.1,tH=0\.9t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9, lr1\.4×10−61\.4\\times 10^\{\-6\}, gentler than the aggressive\(0\.2,0\.7\)\(0\.2,0\.7\)schedule used for the main task table\), DLLM\-JEPA simultaneously improves GSM8K accuracy*and*drives Wikitext diffusion loss below the pre\-trained base, while an L2\-to\-base parameter anchor achieves weak base preservation with zero task gain\.
- •Consistent gains across 4 tasks×\\times2 backbones\.Under a single 4\-shot evaluation protocol, DLLM\-JEPA improves every \(task, architecture\) combination \(Table[2](https://arxiv.org/html/2606.00091#S4.T2)\)\. The 3\-seed mean improvement is\+1\.8\+1\.8pp on LLaDA\-8B GSM8K under Wide\-tt\(±0\.4\\pm 0\.4vs baseline±0\.9\\pm 0\.9\) and\+2\.6\+2\.6–3\.03\.0pp on aggressive\-schedule tasks, where DLLM\-JEPA tightens a±8\.9\\pm 8\.9pp baseline spread to±3\.9\\pm 3\.9pp\. Best\-seed numbers \(up to\+18\.7\+18\.7pp\) and available multi\-seed statistics are in Appendix[A\.6](https://arxiv.org/html/2606.00091#A1.SS6)\.

#### Scope of comparison\.

We position LLM\-JEPA as a structural motivator rather than a direct comparator: the two methods sit on different attention substrates, so we benchmark DLLM\-JEPA against diffusion\-only fine\-tuning on the same backbones\. Our focus is on how JEPA objectives can be naturally instantiated in masked diffusion LMs and how they affect fine\-tuning dynamics\.

## 2Related Work

#### Masked Diffusion Language Models\.

Discrete diffusion models for text generation have gained attention as alternatives to autoregressive decoding\(Austin et al\.,[2021](https://arxiv.org/html/2606.00091#bib.bib2); Lou et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib14)\)\. LLaDA\(Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15)\)introduced a simple masked diffusion framework that achieves competitive performance with AR models at the 8B parameter scale\. MDLM\(Sahoo et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib16)\)and SEDD\(Shi et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib17)\)provide further theoretical and empirical foundations\. These models employ bidirectional attention and learn to reverse a forward masking process, enabling parallel token generation at inference time\.

#### Joint Embedding Predictive Architectures\.

I\-JEPA\(Assran et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib1)\)demonstrated that predicting latent representations of masked image patches produces superior visual features compared to pixel\-level reconstruction \(MAE\)\. V\-JEPA\(Bardes et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib3)\)extended this to video\. LLM\-JEPA\(Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\)adapted the framework to autoregressive language models, showing improvements across code generation and math reasoning tasks\. However, LLM\-JEPA requires explicit view pairs and incurs significant computational overhead due to the causal attention constraint\.

#### Fine\-tuning Diffusion Language Models\.

The fine\-tuning dynamics of masked diffusion language models remain relatively underexplored\.Nie et al\. \([2025](https://arxiv.org/html/2606.00091#bib.bib15)\)demonstrate that supervised fine\-tuning of LLaDA on instruction data yields strong instruction\-following behavior, but do not analyze representation dynamics or the trade\-off between task adaptation and base preservation\. Our work uses DLLM\-JEPA as a lens to probe this trade\-off and shows that representation\-level regularization reshapes it qualitatively\.

## 3Method

### 3\.1Preliminaries: Masked Diffusion Language Models

LLaDA\(Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15)\)defines a forward process that progressively masks tokens\. Given a clean sequencex0=\(x01,…,x0L\)x\_\{0\}=\(x\_\{0\}^\{1\},\\ldots,x\_\{0\}^\{L\}\), the forward process at timet∈\[0,1\]t\\in\[0,1\]independently replaces each token with a special\[MASK\]token with probabilitytt:

q​\(xti∣x0i\)=\{x0iwith probability​1−t\[MASK\]with probability​tq\(x\_\{t\}^\{i\}\\mid x\_\{0\}^\{i\}\)=\\begin\{cases\}x\_\{0\}^\{i\}&\\text\{with probability \}1\-t\\\\ \\texttt\{\[MASK\]\}&\\text\{with probability \}t\\end\{cases\}\(1\)
The modelfθf\_\{\\theta\}learns to predict the original tokens from the masked sequence\. Lettingℳt=\{i:xti=\[MASK\]\}\\mathcal\{M\}\_\{t\}=\\\{i:x\_\{t\}^\{i\}=\\texttt\{\[MASK\]\}\\\}denote the set of masked positions at timett, the diffusion training objective is the per\-token cross\-entropy over masked positions,

ℒdiff=𝔼t∼𝒰\(0,1\),xt∼q\(⋅∣x0,t\)​\[−1\|ℳt\|​∑i∈ℳtlog⁡pθ​\(x0i∣xt\)\]\\mathcal\{L\}\_\{\\text\{diff\}\}=\\mathbb\{E\}\_\{t\\sim\\mathcal\{U\}\(0,1\),\\,x\_\{t\}\\sim q\(\\cdot\\mid x\_\{0\},t\)\}\\left\[\-\\frac\{1\}\{\|\\mathcal\{M\}\_\{t\}\|\}\\sum\_\{i\\in\\mathcal\{M\}\_\{t\}\}\\log p\_\{\\theta\}\(x\_\{0\}^\{i\}\\mid x\_\{t\}\)\\right\]\(2\)
Crucially,fθf\_\{\\theta\}employs*bidirectional*attention: every token can attend to every other token, enabling rich contextual representations from a single forward pass\.

### 3\.2DLLM\-JEPA: JEPA for Diffusion Language Models

![Refer to caption](https://arxiv.org/html/2606.00091v1/figures/dllm_jepa_flow.png)Figure 1:DLLM\-JEPA at a glance\.*Top row \(training flow\)\.*A single clean inputx0x\_\{0\}is noised at two masking rates \(tL=0\.2t\_\{L\}\{=\}0\.2,tH=0\.7t\_\{H\}\{=\}0\.7\) to form a context view and a target view—no paired dataset required\. The online backbonefθf\_\{\\theta\}processes the context view in a single forward pass with gradients that yields diffusion logits \(givingℒdiff\\mathcal\{L\}\_\{\\text\{diff\}\}\) and a pooled embeddingztLz\_\{t\_\{L\}\}; the target view is processed by an EMA copyfθ′f\_\{\\theta^\{\\prime\}\}\(decayτ=0\.996\\tau\{=\}0\.996\) underno\_gradto produceztHz\_\{t\_\{H\}\}\. A predictorgϕg\_\{\\phi\}mapsztL→z^tHz\_\{t\_\{L\}\}\{\\to\}\\hat\{z\}\_\{t\_\{H\}\}andℒJEPA\\mathcal\{L\}\_\{\\text\{JEPA\}\}is the cosine gap toztHz\_\{t\_\{H\}\};ℒtotal=ℒdiff\+λ​ℒJEPA\\mathcal\{L\}\_\{\\text\{total\}\}=\\mathcal\{L\}\_\{\\text\{diff\}\}\+\\lambda\\mathcal\{L\}\_\{\\text\{JEPA\}\}\(cost 4F/step, Table[1](https://arxiv.org/html/2606.00091#S3.T1)\)\. A concrete masked\-text example is shown in the middle\-left\.*Bottom panels \(observed mechanism\)\.*Fine\-tuning with DLLM\-JEPA*increases*hidden\-state drift \(1\.36–3\.60×\\timesbaseline on GSM8K, concentrated in the middle transformer layers\) yet*reduces*Wikitext functional forgetting \(43–58%\)\. The 3D density landscape visualizes the same effect over per\-seed points: baseline mass peaks at high forgetting, while DLLM\-JEPA mass extends along the drift axis yet stays closer to the zero\-forgetting floor \(§[5](https://arxiv.org/html/2606.00091#S5)\)\.#### View generation via noise schedule\.

We generate two views of each input by sampling two masking rates from the diffusion schedule\. GiventL<tHt\_\{L\}<t\_\{H\}\(we use fixedtL=0\.2t\_\{L\}=0\.2,tH=0\.7t\_\{H\}=0\.7\), we construct:

xtL\\displaystyle x\_\{t\_\{L\}\}∼q​\(xt∣x0,t=tL\)\(context view: 20% masked\)\\displaystyle\\sim q\(x\_\{t\}\\mid x\_\{0\},t=t\_\{L\}\)\\quad\\text\{\(context view: 20\\% masked\)\}\(3\)xtH\\displaystyle x\_\{t\_\{H\}\}∼q​\(xt∣x0,t=tH\)\(target view: 70% masked\)\\displaystyle\\sim q\(x\_\{t\}\\mid x\_\{0\},t=t\_\{H\}\)\\quad\\text\{\(target view: 70\\% masked\)\}\(4\)
The context view retains most tokens and provides a nearly complete picture, while the target view is heavily masked and forces prediction from a more incomplete, abstract context\. Both are generated from the*same*input via different noise levels—no explicit view pairs needed\.

#### Encoding and prediction\.

The backbonefθf\_\{\\theta\}processes the context viewxtLx\_\{t\_\{L\}\}to produce hidden states, from which we obtain:

- •Diffusion logits:x^0=Classifier​\(fθ​\(xtL\)\)\\hat\{x\}\_\{0\}=\\text\{Classifier\}\(f\_\{\\theta\}\(x\_\{t\_\{L\}\}\)\), used for the denoising objective\.
- •JEPA embedding:ztL=Pool​\(fθ​\(xtL\)\)z\_\{t\_\{L\}\}=\\text\{Pool\}\(f\_\{\\theta\}\(x\_\{t\_\{L\}\}\)\), a pooled representation of the context view\. We use mean pooling over non\-masked, non\-padding tokens followed by LayerNorm; mean pooling is the standard sequence\-to\-vector reduction in vision JEPAs\(Assran et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib1); Bardes et al\.,[2024](https://arxiv.org/html/2606.00091#bib.bib3)\), and LayerNorm stabilizes the embedding scale—collapse prevention itself rests on the EMA target, stop\-gradient, predictor, and the simultaneous diffusion denoising objective rather than on LayerNorm alone\.

The target embeddingztH=Pool​\(fθ′​\(xtH\)\)z\_\{t\_\{H\}\}=\\text\{Pool\}\(f\_\{\\theta^\{\\prime\}\}\(x\_\{t\_\{H\}\}\)\)is computed with an*exponential moving average \(EMA\)*copy of the backbone,fθ′f\_\{\\theta^\{\\prime\}\}, following I\-JEPA\(Assran et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib1)\)/ BYOL\(Grill et al\.,[2020](https://arxiv.org/html/2606.00091#bib.bib7)\)conventions\. The EMA target is updated each step asθ′←τ​θ′\+\(1−τ\)​θ\\theta^\{\\prime\}\\leftarrow\\tau\\theta^\{\\prime\}\+\(1\-\\tau\)\\thetawith decayτ=0\.996\\tau\{=\}0\.996, and its forward pass runs underno\_gradso the target branch contributes no backward compute or optimizer state\. The stop\-gradient \+ EMA\-target \+ predictor design provides a stable, slowly\-moving prediction target and serves the standard role of preventing trivial\-constant collapse, complementing SimSiam\(Chen & He,[2021](https://arxiv.org/html/2606.00091#bib.bib8)\)\-style stop\-gradient\-only schemes\.

A lightweight predictorgϕg\_\{\\phi\}consisting ofkktransformer decoder layers \(chosen for architectural consistency with LLM\-JEPA’s predictor design\) maps the context embedding to predict the target:

z^tH=gϕ​\(ztL\)\\hat\{z\}\_\{t\_\{H\}\}=g\_\{\\phi\}\(z\_\{t\_\{L\}\}\)\(5\)

#### JEPA objective\.

Following best practices from vision JEPAs\(Assran et al\.,[2023](https://arxiv.org/html/2606.00091#bib.bib1)\), we use cosine similarity:

ℒJEPA=1−cos⁡\(sg​\(ztH\),z^tH\)\\mathcal\{L\}\_\{\\text\{JEPA\}\}=1\-\\cos\\\!\\big\(\\text\{sg\}\(z\_\{t\_\{H\}\}\),\\;\\hat\{z\}\_\{t\_\{H\}\}\\big\)\(6\)wheresg​\(⋅\)\\text\{sg\}\(\\cdot\)denotes stop\-gradient\. Note thatsg​\(ztH\)\\text\{sg\}\(z\_\{t\_\{H\}\}\)is redundant by construction \(the target encoder is an EMA copy withrequires\_grad=Falseand the forward is wrapped inno\_grad\); we keep the notation for consistency with the BYOL/I\-JEPA family of objectives\.

#### Combined training objective\.

The total loss balances diffusion denoising and JEPA prediction:

ℒtotal=ℒdiff\+λ⋅ℒJEPA\\mathcal\{L\}\_\{\\text\{total\}\}=\\mathcal\{L\}\_\{\\text\{diff\}\}\+\\lambda\\cdot\\mathcal\{L\}\_\{\\text\{JEPA\}\}\(7\)whereλ≥0\\lambda\\geq 0controls the JEPA contribution andkkdetermines the predictor depth\.

#### Empirical check: no representation collapse\.

Cosine\-similarity\-only JEPA objectives can in principle collapse to a trivial\-constant solution\. In our setting this is unlikely for two reasons: \(i\) the EMA target evolves slowly and provides a non\-trivial prediction target; and \(ii\) the backbone is simultaneously optimized by the diffusion denoising lossℒdiff\\mathcal\{L\}\_\{\\text\{diff\}\}, which constrains the token\-level output distribution\. We verify directly on the fine\-tuned checkpoints \(Appendix[A\.7](https://arxiv.org/html/2606.00091#A1.SS7)\): on LLaDA\-8B \(D=4096D\{=\}4096\), the pooled JEPA embeddingsztL,ztHz\_\{t\_\{L\}\},z\_\{t\_\{H\}\}retain an effective rank of 42–44 \(close to the pre\-trained base’s 42–43\) with per\-dimension standard deviation 0\.73–0\.95 and cosine diversity 0\.25–0\.28—virtually identical to the baseline fine\-tuning on the same metrics\. The pattern is the same on Dream\-7B \(D=3584D\{=\}3584, effective rank 43–44\)\. The JEPA objective therefore does not reduce the rank or shrink the variance of the representation space; it*redirects*the representation without collapsing it\.

### 3\.3Computational Advantage over LLM\-JEPA

Table 1:Per\-step training cost\.FF: one forward pass through the backbone;BB: one backward pass \(≈2​F\\approx 2Fin FLOPs\)\. LLM\-JEPA requires gradient through both view encodings; DLLM\-JEPA processes the target view without gradients\.MethodFwd \(grad\)Fwd \(no grad\)BackwardTotal \(FF\)OverheadAR Baseline1FF–1BB≈\\approx2FF3FF–LLM\-JEPA2FF–2BB≈\\approx4FF6FF\+100%Diffusion Baseline1FF–1BB≈\\approx2FF3FF–DLLM\-JEPA \(ours\)1FF1FF1BB≈\\approx2FF4FF\+33%Causal attention forces LLM\-JEPA to encode the two views through*two*forward passes with gradients under a custom block\-causal mask\. In contrast, the bidirectional attention in masked\-diffusion backbones lets a single forward pass throughxtLx\_\{t\_\{L\}\}simultaneously yield the diffusion logits and the context embeddingztLz\_\{t\_\{L\}\}; the target embeddingztHz\_\{t\_\{H\}\}needs one additional forward pass through the EMA copy of the backbone, but this forward is gradient\-free \(no backward pass, no gradient memory, no optimizer state\)\. The bookkeeping \(Table[1](https://arxiv.org/html/2606.00091#S3.T1)\) gives\(6​F−4​F\)/6​F≈𝟑𝟑%\(6F\{\-\}4F\)/6F\\approx\\mathbf\{33\\%\}fewer training FLOPs per step\. The design also uses standard bidirectional attention, applies to any text dataset rather than paired \(text, code\) data, and adds only a pooling layer, a lightweight predictor, and one no\-grad EMA forward to the vanilla diffusion training loop\.

## 4Experiments

We evaluate DLLM\-JEPA on four tasks spanning code generation, SQL, regex synthesis, and math reasoning, using*two*masked\-diffusion backbones \(LLaDA\-8B and Dream\-7B\)\. Throughout, the comparison is DLLM\-JEPA vs\. diffusion\-only fine\-tuning on the same backbone; LLM\-JEPA is used as the structural motivator rather than a direct head\-to\-head baseline\.

### 4\.1Experimental Setup

#### Models\.

LLaDA\-8B\(Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15)\)and Dream\-7B\(Ye et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib20)\), two recently released masked\-diffusion language models with different pre\-training recipes and tokenizers\.

#### Tasks and metrics\.

- •Django\(Oda et al\.,[2015](https://arxiv.org/html/2606.00091#bib.bib21)\): 18,805 Python code generation pairs\. Metric: whitespace\-normalized prefix match \(Appendix[A\.9](https://arxiv.org/html/2606.00091#A1.SS9)\); 0\-shot strict exact match also reported in Table[5](https://arxiv.org/html/2606.00091#A1.T5)\.
- •Spider\(Yu et al\.,[2018](https://arxiv.org/html/2606.00091#bib.bib19)\): 7,000 text\-to\-SQL pairs\. Metric: exact match and execution match on the Spider validation split \(the official test set is hidden behind a leaderboard\)\.
- •NL\-RX\-SYNTH\(Locascio et al\.,[2016](https://arxiv.org/html/2606.00091#bib.bib13)\): 10,000 text\-to\-regex pairs\. Metric: exact match and functional equivalence \(tested against 200 random strings\)\.
- •GSM8K\(Cobbe et al\.,[2021](https://arxiv.org/html/2606.00091#bib.bib5)\): 7,473 grade\-school math problems with chain\-of\-thought answers\. Metric: answer accuracy\.

#### Training\.

All runs use full fine\-tuning for 2 epochs on 8×\\timesNVIDIA A100 80GB GPUs with gradient checkpointing\. We use AdamW with learning rate1×10−51\{\\times\}10^\{\-5\}\(primary\) or1\.4×10−61\.4\{\\times\}10^\{\-6\}\(Wide\-ttpreservation\-focused config\), a fixed masking scheduletL=0\.2,tH=0\.7t\_\{L\}\{=\}0\.2,t\_\{H\}\{=\}0\.7for most experiments \(Wide\-ttuses\[0\.1,0\.9\]\[0\.1,0\.9\]\), and DLLM\-JEPA hyperparametersλ∈\{0\.5,1\.0,2\.0\},k∈\{1,…,5\}\\lambda\\in\\\{0\.5,1\.0,2\.0\\\},k\\in\\\{1,\\ldots,5\\\}\.

#### Evaluation\.

We report 0\-shot and 4\-shot accuracy using each model’s standard iterative unmasking generation with 128 diffusion steps\. For Dream on code\-structured tasks, whose natural output format may emit tokens after the answer, we additionally report the stopping\-robust variant \(whitespace\-normalized prefix match for Django; truncation at the first;plus standard Spider cleaning for Spider\) to enable fair cross\-architecture comparison; raw metrics are reported in the appendix\.

### 4\.2Main Results: Task Performance

Table[2](https://arxiv.org/html/2606.00091#S4.T2)presents DLLM\-JEPA results across 4 tasks×\\times2 architectures\. Under the unified 4\-shot protocol of Table[2](https://arxiv.org/html/2606.00091#S4.T2), DLLM\-JEPA improves every \(task, architecture\) cell, with gains ranging from\+1\.00\+1\.00to\+18\.73\+18\.73percentage points\. A finer\-grained 0\-shot decomposition on LLaDA\-8B Django \(Table[5](https://arxiv.org/html/2606.00091#A1.T5)\) shows a larger\+21\.89\+21\.89pp gain on the strict exact\-match metric\.

Reporting protocol \(read once\)\.Table[2](https://arxiv.org/html/2606.00091#S4.T2)reports*task\-optimal*\(λ,k\)\(\\lambda,k\)per \(task, architecture\) cell under a unified 4\-shot evaluation protocol using the*aggressive*training schedule \(tL=0\.2,tH=0\.7t\_\{L\}\{=\}0\.2,t\_\{H\}\{=\}0\.7\)\. A*single\-config*sibling \(Table[8](https://arxiv.org/html/2606.00091#A1.T8)\) repeats the comparison with\(λ=1,k=3\)\(\\lambda\{=\}1,k\{=\}3\)fixed across tasks; gains remain positive throughout but smaller\. Multi\-seed \(n=3n\{=\}3\) statistics are in Table[10](https://arxiv.org/html/2606.00091#A1.T10)\(Appendix[A\.6](https://arxiv.org/html/2606.00091#A1.SS6)\); LLaDA\-8B Spider and Dream\-7B Spider/NL\-RX are single\-seed grids over\(λ,k\)\(\\lambda,k\)\. Preservation results \(§[4\.3](https://arxiv.org/html/2606.00091#S4.SS3)\) use the*Wide\-tt*schedule \(tL=0\.1,tH=0\.9t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9\) chosen*a priori*; base\-capability probes are multi\-seed on MMLU; Wikitext rows for Baseline and DLLM\-JEPA are single\-seed point estimates \(the L2 anchor is reported over multiple seeds\)\.

Table 2:Main results\.Task performance across four tasks and two masked\-diffusion backbones, at task\-optimal\(λ,k\)\(\\lambda,k\)\. Baselines and DLLM\-JEPA share the same learning rate, training schedule, and seed within each row;\(λ,k\)\(\\lambda,k\)is not applicable to the baseline \(λ=0\\lambda\{=\}0\)\. All rows use a single 4\-shot protocol and the same metric for both models\. Spider: SQL execution match with semicolon\-truncation cleaning; Django: whitespace\-normalized prefix match \(Appendix[A\.9](https://arxiv.org/html/2606.00091#A1.SS9)\)\. Fixed single\(λ,k\)\(\\lambda,k\)comparison in Table[8](https://arxiv.org/html/2606.00091#A1.T8)\.TaskMetric \(4\-shot\)LLaDA\-8BDream\-7BBL→\\rightarrowJEPAΔ\\DeltaBL→\\rightarrowJEPAΔ\\DeltaGSM8Kaccuracy42\.61→\\rightarrow61\.33\+18\.7334\.87→\\rightarrow46\.25\+11\.38NL\-RXfunc match47\.50→\\rightarrow58\.20\+10\.7042\.00→\\rightarrow46\.80\+4\.80Spiderexec match \(cleaned\)35\.40→\\rightarrow39\.36\+3\.9720\.89→\\rightarrow25\.15\+4\.26Djangows\-prefix match74\.40→\\rightarrow75\.40\+1\.0069\.58→\\rightarrow72\.35\+2\.77Every \(task, architecture\) cell shows a positive improvement under a single 4\-shot evaluation protocol and the same metric applied to both models\. The LLaDA\-8B GSM8K baseline of42\.642\.6sits at the lower tail of a±8\.9\\pm 8\.9pp seed\-dependent spread induced by the aggressive training schedule on diffusion\-only fine\-tuning; DLLM\-JEPA collapses that spread to±3\.9\\pm 3\.9pp \(Table[10](https://arxiv.org/html/2606.00091#A1.T10)\), so the cell\-level gap reflects both peak improvement and variance reduction\. Seeds are matched across methods within each row\.

#### Cross\-architecture replication \(task performance\)\.

Dream\-7B differs from LLaDA\-8B in tokenizer, parameter count, and pre\-training recipe\. The same fine\-tuning protocol applied to both yields the same direction of improvement on every task, supporting that the DLLM\-JEPA effect is not an LLaDA\-specific artifact; we emphasize that this is replication on one additional backbone, not a demonstration across all diffusion LMs\.

### 4\.3Base Capability Preservation

A central property of DLLM\-JEPA is that task gain coexists with preservation of the base model on a held\-out diffusion\-loss probe\. Preservation of pre\-trained capability during task adaptation is the classic catastrophic\-forgetting question\(Kirkpatrick et al\.,[2017](https://arxiv.org/html/2606.00091#bib.bib9)\); parameter\-space regularizers such as EWC are one response, while our analysis below offers a representation\-space alternative\. We evaluate this using held\-out Wikitext\-103 validation text, measuring the diffusion denoising loss \(ℒdiff\\mathcal\{L\}\_\{\\text\{diff\}\}with 50% masking, deterministic seed\) of each fine\-tuned checkpoint against the pre\-trained base\. \(DLLM\-JEPA’s Wikitext loss in fact ends up slightly*below*the pre\-trained base; we read this as evidence of preservation rather than a strong improvement claim, given that absolute Wikitext deltas are at the∼10−3\\sim\\\!10^\{\-3\}scale\.\)

#### Comparison with a parameter\-anchor baseline\.

To show that DLLM\-JEPA’s benefit is not reducible to simple parameter proximity, we compare against anL2\-to\-base anchor: an SFT objective with an explicit penaltyλL​2​∥θ−θbase∥22\\lambda\_\{L2\}\\lVert\\theta\-\\theta\_\{\\text\{base\}\}\\rVert\_\{2\}^\{2\}\(gradient\-side injection for memory efficiency,λL​2=10−4\\lambda\_\{L2\}\{=\}10^\{\-4\}\)\. Both methods use the same Wide\-ttconfiguration on LLaDA\-8B GSM8K\.

Table 3:Task gain and base preservation together\.LLaDA\-8B GSM8K, Wide\-ttconfiguration\. DLLM\-JEPA is the only method that raises task accuracy while pushing Wikitext loss further below the base\. The GSM8K accuracy column is over three seeds \(per\-seed breakdown in Table[12](https://arxiv.org/html/2606.00091#A1.T12)\); Wikitext rows for Baseline and DLLM\-JEPA are single\-seed \(checkpoint\-storage constraint\), so we report a point estimate\.MethodGSM8K 0\-shotWikitextΔ\\Deltaloss \(vs\. base\)Base \(no fine\-tuning\)–0\.0000Diffusion Baseline \(λ=0\\lambda\{=\}0\)65\.23±\\pm0\.93−0\.0004\-0\.0004L2\-to\-base anchor \(λL​2=10−4\\lambda\_\{L2\}\{=\}10^\{\-4\}\)65\.18±\\pm0\.87−0\.0007±0\.0002\-0\.0007\\pm 0\.0002DLLM\-JEPA\(ours\)67\.07±\\pm0\.41−0\.0017\\mathbf\{\-0\.0017\}Three observations emerge\. \(1\) The diffusion baseline achieves modest task gain and essentially matches the base in Wikitext loss\. \(2\) The L2 anchor further reduces Wikitext drift slightly, but it does not improve task performance—its GSM8K score is statistically indistinguishable from the baseline\. \(3\) DLLM\-JEPA is the only method to achieve both: higher GSM8K accuracy \(\+1\.84 pp over baseline\) together with a WikitextΔ\\Deltabelow either competing method\. The absolute Wikitext differences are small \(∼10−3\{\\sim\}10^\{\-3\}\), so readers should weigh them in combination with the task\-accuracy column rather than in isolation\.

#### MMLU sanity check \(no catastrophic forgetting\)\.

As a sanity check that fine\-tuning does not catastrophically degrade general capability, we also evaluate each Wide\-ttcheckpoint onMMLU\(500 stratified questions; same items for every model\), across three seeds \(Table[4](https://arxiv.org/html/2606.00091#S4.T4)\)\. Both methods stay within about half a point of the base \(baseline57\.93±0\.4257\.93\\pm 0\.42, DLLM\-JEPA57\.53±0\.2357\.53\\pm 0\.23vs\. base57\.4057\.40\); the small mean gap is within seed noise, so we read MMLU as evidence*against catastrophic forgetting on either method*rather than as a comparative claim\. DLLM\-JEPA’s seed std is roughly half the baseline’s; this matches the variance\-tightening we observe on the high\-variance LLaDA GSM8K cells but not on lower\-variance cells\. Our preservation argument primarily rests on the Wikitext evidence \(Table[3](https://arxiv.org/html/2606.00091#S4.T3)\)\.

Table 4:MMLU sanity check\.500\-question stratified subset, three fine\-tuning seeds\. Both Wide\-ttfine\-tuning methods stay within about half a point of the base; the mean gap between Baseline and DLLM\-JEPA is within seed noise\.MethodMMLU accuracy \(%\)Base \(no fine\-tuning\)57\.40Diffusion Baseline \(λ=0\\lambda\{=\}0\)57\.93±\\pm0\.42DLLM\-JEPA\(ours\)57\.53±\\pm0\.23

## 5Mechanism Analysis

Having established that DLLM\-JEPA delivers both task gains and base preservation, we next ask:*what does the auxiliary objective do to the backbone internally?*We probe the fine\-tuned checkpoints by comparing their layer\-wise hidden\-state representations against the pre\-trained base on 100 Wikitext\-103 samples\. For each transformer layer we compute the mean\-pooled representation and report the geometric drift1−cos⁡⟨hℓbase,hℓFT⟩1\-\\cos\\\!\\langle h^\{\\text\{base\}\}\_\{\\ell\},h^\{\\text\{FT\}\}\_\{\\ell\}\\rangleaveraged over samples\.

![Refer to caption](https://arxiv.org/html/2606.00091v1/x1.png)Figure 2:Geometric–functional drift dissociation for DLLM\-JEPA\.\(A\) Layer\-wise geometric drift on LLaDA\-8B fine\-tuned on GSM8K \(aggressive configuration\)\. DLLM\-JEPA drifts further from the base than the diffusion baseline throughout the network\. \(B\) Mean drift ratio \(DLLM\-JEPA / baseline\) across five \(task, configuration\) cells\. Amplification \(ratio\>\>1\) is clear on the three GSM8K cells \(both LLaDA configurations and Dream\-7B\); NL\-RX and Django are near or slightly below 1, indicating task\-specific amplification rather than uniform drift\. \(C\) Wikitext functional forgetting is*reduced*for DLLM\-JEPA on every fine\-tuning task \(43–58%\), opposite to the drift trend\. \(D\) Per\-region drift ratio \(early / middle / late thirds\)\. On GSM8K, DLLM\-JEPA’s amplification concentrates in the middle layers, consistent with prior interpretations of middle\-layer compositional structure\. All cells are multi\-seed except LLaDA\-8B GSM8K Wide\-tt\(single seed, checkpoint\-storage constraint\)\.### 5\.1Geometric–Functional Drift Dissociation on GSM8K

Figure[2](https://arxiv.org/html/2606.00091#S5.F2)\(A–B\) reveals a counterintuitive finding: on GSM8K, DLLM\-JEPA’s hidden\-state drift from the base is*larger*than the baseline’s—not smaller\. On LLaDA\-8B, the DLLM\-JEPA / baseline drift ratio is3\.60×3\.60\\timeswith the gentle Wide\-ttconfiguration and1\.36×1\.36\\timeswith the aggressive configuration\. On NL\-RX and Django \(aggressive configuration\) the ratio is0\.94×0\.94\\timesand0\.99×0\.99\\timesrespectively—near 1, meaning DLLM\-JEPA’s drift amplification is*task\-specific*, not uniform\. Yet Figure[2](https://arxiv.org/html/2606.00091#S5.F2)\(C\) shows that on the three tasks where we measure Wikitext forgetting, DLLM\-JEPA*reduces*functional forgetting by 43–58%\.

This geometric–functional drift dissociation admits a descriptive interpretation: the JEPA objective does not appear to minimize representation change but instead redirects it along task\-useful axes, leaving the base model’s held\-out output function largely intact\.

#### Cross\-architecture replication of the drift pattern\.

The same direction of effect appears on Dream\-7B \(different parameter count, tokenizer, and pre\-training corpus\): DLLM\-JEPA’s drift ratio is1\.28×1\.28\\times\(Figure[2](https://arxiv.org/html/2606.00091#S5.F2)B, last column\)\. Replication on one additional backbone is suggestive rather than conclusive\.

#### Region\-specific amplification \(GSM8K\)\.

The extra drift concentrates in the*middle*transformer layers \(Figure[2](https://arxiv.org/html/2606.00091#S5.F2)D; ratios1\.64×1\.64\\timesmiddle vs1\.38×1\.38\\timesearly /1\.18×1\.18\\timeslate under the aggressive schedule\), which prior work associates with compositional reasoning\. The pattern is sharp on GSM8K and muted on other tasks\.

### 5\.2Component ablation: asymmetric views and predictor

Figure[3](https://arxiv.org/html/2606.00091#S5.F3)shows the effect of removing two DLLM\-JEPA design choices on LLaDA\-8B GSM8K \(aggressive schedule, seed 42\)\. \(i\)Symmetric views\(tL=tH=0\.2t\_\{L\}\{=\}t\_\{H\}\{=\}0\.2, same masking rate for both views\) removes the asymmetric\-noise structure of our method\. \(ii\)No\-predictor\(identity map forgϕg\_\{\\phi\}\) removes JEPA’s predictive head so the loss reduces to a plain stop\-gradient two\-view cosine consistency\. The full DLLM\-JEPA uses\(0\.2,0\.7\)\(0\.2,0\.7\)and akk\-layer predictor\.

![Refer to caption](https://arxiv.org/html/2606.00091v1/x2.png)Figure 3:Ablating two DLLM\-JEPA components\.\(A\) GSM8K accuracy on LLaDA\-8B with aggressive fine\-tuning\. Removing asymmetric views collapses 0\-shot to38\.938\.9pp —*below*the diffusion\-only baseline — and removing the predictor collapses 4\-shot by−16\.9\-16\.9pp relative to the full method\. \(B\) Fractional gain over baseline: each partial configuration captures only a small \(or negative\) fraction of the full DLLM\-JEPA gain\.Two findings deserve emphasis\. First, a symmetric\-noise two\-view setup, which would be a natural simplification, performs*worse*than the diffusion\-only baseline: the two views must probe different noise levels for the JEPA objective to produce a useful signal\. Second, retaining the asymmetric views but discarding the predictor leaves much of the 0\-shot gain intact but collapses the 4\-shot gain, indicating that DLLM\-JEPA is not reducible to a simple two\-view consistency regularizer — JEPA’s*predictive*structure contributes most of the few\-shot benefit\. Full noise\-schedule sensitivity \(tL,tHt\_\{L\},t\_\{H\}grid on GSM8K\) and additional ablations are in Appendix[A\.4](https://arxiv.org/html/2606.00091#A1.SS4)\.

Per\-seed drift–forgetting scatter, joint density landscape, Django content/termination decomposition, and a structural comparison table to LLM\-JEPA are in the appendix \(§[A\.1](https://arxiv.org/html/2606.00091#A1.SS1)–§[A\.3](https://arxiv.org/html/2606.00091#A1.SS3)\)\.

## 6Limitations

We identify three limitations that temper the strength of our conclusions:

\(L1\) Scope of comparison\.We do not claim superiority over LLM\-JEPA\. The two methods sit on different attention substrates \(causal vs\. bidirectional\), so a matched head\-to\-head would require aligning pre\-training corpora, tokenizers, scales, and inference protocols—and even then the result would confound substrate with the JEPA objective itself\. We therefore benchmark DLLM\-JEPA against diffusion\-only fine\-tuning on the same backbone, and view the two methods as complementary instantiations of JEPA on different substrates\.

\(L2\) Preservation evidence\.Our preservation argument rests primarily on held\-out Wikitext\-103 diffusion loss \(Table[3](https://arxiv.org/html/2606.00091#S4.T3)\); the Wide\-ttrow is a single surviving seed \(storage cleanup\), so absolute deltas should be read alongside the direction\-of\-effect rather than as precise quantities\. MMLU \(Table[4](https://arxiv.org/html/2606.00091#S4.T4)\) is reported as a no\-catastrophic\-forgetting check, not a comparative claim—the small Baseline / DLLM\-JEPA gap on MMLU is within seed noise\.

\(L3\) Statistical and operating\-point scope\.The drift\-dissociation and region\-specific claims rest on three seeds per cell; the effect is largest on GSM8K, muted on NL\-RX/Django, and replicated on one additional backbone \(Dream\)\. Our two configurations—aggressive\(tL=0\.2,tH=0\.7,\(t\_\{L\}\{=\}0\.2,t\_\{H\}\{=\}0\.7,lr10−5\)10^\{\-5\}\)and Wide\-tt\(tL=0\.1,tH=0\.9,\(t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9,lr1\.4×10−6\)1\.4\\times 10^\{\-6\}\)—are illustrative operating points, not the result of a fully factored\(lr,tL,tH\)\(\\text\{lr\},t\_\{L\},t\_\{H\}\)sweep, and all experiments are 2\-epoch full\-rank fine\-tunings at the 7–8B scale; pre\-training and other scales remain untested\.

## 7Conclusion

We introduced DLLM\-JEPA, a JEPA formulation for masked\-diffusion language models that cuts training FLOPs by 33% relative to LLM\-JEPA’s architectural requirements while generating its two views automatically via the diffusion noise schedule\. Under a single 4\-shot protocol, DLLM\-JEPA improves every \(task, architecture\) combination we evaluate, with gains up to \+18\.7 pp on LLaDA\-8B GSM8K and \+11\.4 pp on Dream\-7B GSM8K, and tightens run\-to\-run variance on the high\-variance LLaDA\-8B GSM8K aggressive\-schedule baseline\. Under the Wide\-ttconfiguration \(tL=0\.1,tH=0\.9t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9, lr1\.4×10−61\.4\\times 10^\{\-6\}\) the method achieves task gain together with base preservation—raising GSM8K accuracy while driving Wikitext loss below the pre\-trained base, unlike a parameter\-level L2 anchor; MMLU stays close to base level on either method, with no evidence of catastrophic forgetting—and layer\-wise probing uncovers a geometric–functional drift dissociation concentrated in middle transformer layers that replicates on Dream\-7B\. Future directions include applying DLLM\-JEPA at pre\-training, theoretical analysis of the dissociation mechanism, and adaptive scheduling ofλ\\lambdaduring training; a matched\-protocol comparison with LLM\-JEPA is the natural next step\.

## References

- Assran et al\. \[2023\]M\. Assran, Q\. Duval, I\. Misra, P\. Bojanowski, P\. Vincent, M\. Rabbat, Y\. LeCun, and N\. Ballas\.Self\-supervised learning from images with a joint\-embedding predictive architecture\.In*CVPR*, 2023\.
- Austin et al\. \[2021\]J\. Austin, D\. D\. Johnson, J\. Ho, D\. Tarlow, and R\. van den Berg\.Structured denoising diffusion models in discrete state\-spaces\.In*NeurIPS*, 2021\.
- Bardes et al\. \[2024\]A\. Bardes, Q\. Garrido, J\. Ponce, X\. Chen, M\. Rabbat, Y\. LeCun, M\. Assran, and N\. Ballas\.Revisiting feature prediction for learning visual representations from video\.*arXiv preprint arXiv:2404\.08471*, 2024\.
- Brown et al\. \[2020\]T\. Brown, B\. Mann, N\. Ryder, et al\.Language models are few\-shot learners\.In*NeurIPS*, 2020\.
- Cobbe et al\. \[2021\]K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. Schulman\.Training verifiers to solve math word problems\.*arXiv preprint arXiv:2110\.14168*, 2021\.
- Devlin et al\. \[2019\]J\. Devlin, M\.\-W\. Chang, K\. Lee, and K\. Toutanova\.BERT: Pre\-training of deep bidirectional transformers for language understanding\.In*NAACL*, 2019\.
- Grill et al\. \[2020\]J\.\-B\. Grill, F\. Strub, F\. Altché, C\. Tallec, P\. H\. Richemond, E\. Buchatskaya, C\. Doersch, B\. Ávila Pires, Z\. D\. Guo, M\. G\. Azar, B\. Piot, K\. Kavukcuoglu, R\. Munos, and M\. Valko\.Bootstrap your own latent: A new approach to self\-supervised learning\.In*NeurIPS*, 2020\.
- Chen & He \[2021\]X\. Chen and K\. He\.Exploring simple siamese representation learning\.In*CVPR*, 2021\.
- Kirkpatrick et al\. \[2017\]J\. Kirkpatrick, R\. Pascanu, N\. Rabinowitz, J\. Veness, G\. Desjardins, A\. A\. Rusu, K\. Milan, J\. Quan, T\. Ramalho, A\. Grabska\-Barwinska, D\. Hassabis, C\. Clopath, D\. Kumaran, and R\. Hadsell\.Overcoming catastrophic forgetting in neural networks\.*PNAS*, 114\(13\):3521–3526, 2017\.
- Huang et al\. \[2025\]H\. Huang, Y\. LeCun, and R\. Balestriero\.LLM\-JEPA: Large language models meet joint embedding predictive architectures\.*arXiv preprint arXiv:2509\.14252*, 2025\.
- LeCun \[2022\]Y\. LeCun\.A path towards autonomous machine intelligence\.*openreview\.net preprint*, 2022\.
- Littwin et al\. \[2024\]E\. Littwin, O\. Saremi, M\. Advani, V\. Thilak, P\. Nakkiran, C\. Huang, and J\. Susskind\.How JEPA avoids noisy features: The implicit bias of deep linear self distillation networks\.In*NeurIPS*, 2024\. arXiv:2407\.03475\.
- Locascio et al\. \[2016\]N\. Locascio, K\. Narasimhan, E\. DeLeon, N\. Kushman, and R\. Barzilay\.Neural generation of regular expressions from natural language with minimal domain knowledge\.In*EMNLP*, 2016\.
- Lou et al\. \[2024\]A\. Lou, C\. Meng, and S\. Ermon\.Discrete diffusion modeling by estimating the ratios of the data distribution\.In*ICML*, 2024\.
- Nie et al\. \[2025\]S\. Nie, F\. Zhu, Z\. You, X\. Zhang, J\. Ou, J\. Hu, J\. Zhou, Y\. Lin, J\.\-R\. Wen, and C\. Li\.Large language diffusion models\.*arXiv preprint arXiv:2502\.09992*, 2025\.
- Sahoo et al\. \[2024\]S\. S\. Sahoo, M\. Arriola, Y\. Schiff, A\. Gokaslan, E\. Marroquin, J\. T\. Chiu, A\. Rush, and V\. Kuleshov\.Simple and effective masked diffusion language models\.In*NeurIPS*, 2024\. arXiv:2406\.07524\.
- Shi et al\. \[2024\]J\. Shi, K\. Han, Z\. Wang, A\. Doucet, and M\. Titsias\.Simplified and generalized masked diffusion for discrete data\.In*NeurIPS*, 2024\.
- Touvron et al\. \[2023\]H\. Touvron, L\. Martin, K\. Stone, et al\.Llama 2: Open foundation and fine\-tuned chat models\.*arXiv preprint arXiv:2307\.09288*, 2023\.
- Yu et al\. \[2018\]T\. Yu, R\. Zhang, K\. Yang, M\. Yasunaga, D\. Wang, Z\. Li, J\. Ma, I\. Li, Q\. Yao, S\. Roman, Z\. Zhang, and D\. Radev\.Spider: A large\-scale human\-labeled dataset for complex and cross\-domain semantic parsing and text\-to\-SQL task\.In*EMNLP*, 2018\.
- Ye et al\. \[2025\]J\. Ye, Z\. Xie, L\. Zheng, J\. Gao, Z\. Wu, X\. Jiang, Z\. Li, and L\. Kong\.Dream 7B: Diffusion large language models\.*arXiv preprint arXiv:2508\.15487*, 2025\.
- Oda et al\. \[2015\]Y\. Oda, H\. Fudaba, G\. Neubig, H\. Hata, S\. Sakti, T\. Toda, and S\. Nakamura\.Learning to generate pseudo\-code from source code using statistical machine translation\.In*ASE*, 2015\.
- Merity et al\. \[2017\]S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher\.Pointer sentinel mixture models\.In*ICLR*, 2017\.
- Hendrycks et al\. \[2021\]D\. Hendrycks, C\. Burns, S\. Basart, A\. Zou, M\. Mazeika, D\. Song, and J\. Steinhardt\.Measuring massive multitask language understanding\.In*ICLR*, 2021\.

## Appendix AAppendix

### A\.1Per\-Seed Drift–Forgetting Scatter and Density Landscape

![Refer to caption](https://arxiv.org/html/2606.00091v1/x3.png)Figure 4:Per\-run drift vs\. functional forgetting\.One marker per fine\-tuning run\. Baselines show a strong positive relationship \(r=\+0\.94r\{=\}\+0\.94\); DLLM\-JEPA visibly shifts below the baseline cluster at comparable drift levels, weakening the relationship \(r=\+0\.75r\{=\}\+0\.75\) and occupying region with lower forgetting at the same geometric drift\.Figure[4](https://arxiv.org/html/2606.00091#A1.F4)plots per\-seed geometric drift against functional forgetting\. Baseline runs show a strong positive drift–forgetting correlation \(r=\+0\.94r=\+0\.94, 95% CI\[0\.76,0\.99\]\[0\.76,0\.99\],n=10n=10\); DLLM\-JEPA runs sit below the baseline envelope at comparable drift levels, weakening the correlation tor=\+0\.75r=\+0\.75\(95% CI\[0\.23,0\.94\]\[0\.23,0\.94\],n=10n=10\)\. A Fisherrr\-to\-zztest givesZ=1\.42Z=1\.42,p=0\.16p=0\.16at this sample size\. We use “dissociation” for this weakened coupling rather than its elimination—DLLM\-JEPA’srris still positive with a CI excluding zero\.

![Refer to caption](https://arxiv.org/html/2606.00091v1/x4.png)Figure 5:Joint \(drift, forgetting\) density landscape \(n=10n\{=\}10runs per method, exploratory\)\.Gaussian\-KDE surfaces over the same runs as Figure[4](https://arxiv.org/html/2606.00091#A1.F4), baseline in red and DLLM\-JEPA in blue\. At matched drift, DLLM\-JEPA’s mass sits below the baseline’s, and the baseline mass never reaches the low\-forgetting corner that DLLM\-JEPA occupies\. Intended as an exploratory visualization, not a formal density estimate\.Figure[5](https://arxiv.org/html/2606.00091#A1.F5)re\-plots the same data as a 3D density landscape\. The baseline distribution is peaked at high forgetting \(mass at WikitextΔ​loss≳0\.04\\Delta\\text\{loss\}\\gtrsim 0\.04\); DLLM\-JEPA’s mass stays closer to theΔ​loss=0\\Delta\\text\{loss\}=0plane and extends further along the drift axis without a corresponding rise in forgetting\.

### A\.2Content vs\. Termination Decomposition on Django

A finer\-grained view of where the task gains come from is obtained by decomposing Django outputs into three categories: \(i\) exact match, \(ii\) “stopping artifact” \(the correct code followed by trailing tokens\), and \(iii\) genuine content error\. For this decomposition we use the 0\-shot exact\-match setting on LLaDA\-8B Django \(single representative run\)\.

Table 5:Django 0\-shot decomposition\.Single\-seed illustrative run on LLaDA\-8B\. DLLM\-JEPA simultaneously reduces stopping artifacts and reduces real content errors\.CategoryBaselineDLLM\-JEPAΔ\\DeltaExact match34\.40%56\.29%\+21\.89Stopping artifact \(correct \+ trailing\)29\.81%11\.30%−18\.52\-18\.52Content error35\.79%32\.41%−3\.38\-3\.38Reasoning\-correct total64\.21%67\.59%\+3\.38The exact\-match improvement \(\+21\.89pp\) decomposes into an 18\.52pp reduction in stopping artifacts and a 3\.38pp improvement in underlying content quality\. DLLM\-JEPA’s predictor objective provides an additional gradient signal about what comes next in representation space, translating into both sharper answer endings and cleaner content\.

### A\.3Structural Comparison with LLM\-JEPA

Table[6](https://arxiv.org/html/2606.00091#A1.T6)summarizes the architectural differences between LLM\-JEPA\[Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\]and DLLM\-JEPA\. The comparison is purely structural: these properties follow directly from the choice of substrate \(causal vs\. bidirectional attention\), not from an empirical head\-to\-head\. All task\-performance claims in the main text are measured against diffusion\-only fine\-tuning\.

Table 6:Structural comparison: LLM\-JEPA vs\. DLLM\-JEPA\.Architectural differences determined by the underlying substrate\.PropertyLLM\-JEPADLLM\-JEPA \(ours\)Base architectureAutoregressiveMasked DiffusionAttentionCausalBidirectionalView constructionExplicit \(Text↔\\leftrightarrowCode\)Automatic \(noise schedule\)Requires view pairsYesNoForward passes \(w/ grad\)21Custom attention maskRequiredNot neededTraining FLOP overhead\+100%\+33%Applicable datasetsMulti\-view onlyAny text
### A\.4Ablations

#### λ×k\\lambda\\times kgrid\.

Table[7](https://arxiv.org/html/2606.00091#A1.T7)shows GSM8K accuracy across the grid on LLaDA\-8B; every non\-zero\(λ,k\)\(\\lambda,k\)configuration outperforms the diffusion baseline\. The best configuration is\(λ=2\.0,k=4\)\(\\lambda\{=\}2\.0,k\{=\}4\)\.

Table 7:GSM8Kλ×k\\lambda\\times kgrid\.0\-shot accuracy on LLaDA\-8B\. Baseline \(λ=0\\lambda\{=\}0\): 0\.402\.λ\\lambda\\kk123450\.50\.4450\.4500\.4570\.4590\.4381\.00\.4450\.4560\.4670\.4670\.4442\.00\.4370\.4600\.4730\.4880\.453
#### Single\-configuration comparison\.

With a single fixed\(λ=1\.0,k=3\)\(\\lambda\{=\}1\.0,k\{=\}3\)applied across all tasks \(Table[8](https://arxiv.org/html/2606.00091#A1.T8)\), DLLM\-JEPA still improves over the diffusion baseline on every row\.

Table 8:Single fixed configuration across tasks\.TaskShotMetricBaselineDLLM\-JEPAΔ\\DeltaDjango0\-shotexact34\.4045\.82\+11\.42NL\-RX4\-shotexact11\.2518\.74\+7\.49Spider4\-shotexact25\.7329\.21\+3\.48GSM8K \(Wide\-tt\)0\-shotacc65\.2866\.87\+1\.59
#### Noise schedule ablation\.

We use two illustrative training configurations: \(i\) the*aggressive*schedule with view\-generation\(tL=0\.2,tH=0\.7\)\(t\_\{L\}\{=\}0\.2,t\_\{H\}\{=\}0\.7\)and lr10−510^\{\-5\}, used for the main task table; and \(ii\) the*Wide\-tt*schedule with view\-generation\(tL=0\.1,tH=0\.9\)\(t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9\)paired with a lower lr \(1\.4×10−61\.4\\times 10^\{\-6\}\), used for the preservation/joint\-improvement experiment in §[4\.3](https://arxiv.org/html/2606.00091#S4.SS3)\. The Wide\-ttschedule increases DLLM\-JEPA’s drift ratio on GSM8K from1\.36×1\.36\\timesto3\.60×3\.60\\times\(Figure[2](https://arxiv.org/html/2606.00091#S5.F2)B\) while simultaneously achieving deeper base preservation \(Table[3](https://arxiv.org/html/2606.00091#S4.T3)\)\. These two schedules are illustrative operating points rather than the result of a fully factored\(lr,tL,tH\)\(\\text\{lr\},t\_\{L\},t\_\{H\}\)sweep; see Appendix[A\.4](https://arxiv.org/html/2606.00091#A1.SS4.SSS0.Px4)for our partial sensitivity grid \(within the aggressive lr\)\.

#### \(tL,tH\)\(t\_\{L\},t\_\{H\}\)sensitivity\.

Table[9](https://arxiv.org/html/2606.00091#A1.T9)sweeps the view\-generation pair on LLaDA\-8B GSM8K \(aggressive schedule, seed 42\)\. The\(0\.2,0\.7\)\(0\.2,0\.7\)default dominates other pairs sharply on 4\-shot \(by\+14\+14to\+19\+19pp\) and more modestly on 0\-shot\. Wider \(0\.3,0\.90\.3,0\.9\) and narrower \(0\.1,0\.50\.1,0\.5\) pairs both degrade; the symmetric case\(0\.2,0\.2\)\(0\.2,0\.2\)is reported in §[5\.2](https://arxiv.org/html/2606.00091#S5.SS2)and drops below the diffusion\-only baseline entirely\. These results corroborate the component\-decomposition finding that*asymmetric*noise levels are load\-bearing for DLLM\-JEPA\.

Table 9:\(tL,tH\)\(t\_\{L\},t\_\{H\}\)sensitivity on LLaDA\-8B GSM8K \(aggressive, seed 42\)\.\(tL,tH\)\(t\_\{L\},t\_\{H\}\)0\-shot4\-shot\(0\.1,0\.5\)\(0\.1,0\.5\)43\.5944\.05\(0\.2,0\.5\)\(0\.2,0\.5\)42\.0046\.40\(0\.2,0\.7\)\(0\.2,0\.7\)\(default\)44\.8861\.33\(0\.3,0\.9\)\(0\.3,0\.9\)40\.1142\.76
#### Multi\-seed robustness\.

Table[10](https://arxiv.org/html/2606.00091#A1.T10)reports 3\-seed mean±\\pmstd\. The direction matches Table[2](https://arxiv.org/html/2606.00091#S4.T2), with smaller means due to seed variance in the diffusion\-only baseline\. On LLaDA\-8B GSM8K \(aggressive\), the baseline’s seed variance is±8\.9\\pm 8\.9pp, which DLLM\-JEPA narrows to±3\.9\\pm 3\.9pp while remaining within the same upper envelope\.

Table 10:Multi\-seed robustness\.Mean±\\pmstd over three fine\-tuning seeds where available; std values are population standard deviations \(NN\-denominator\)\.ModelTaskBaseline→\\rightarrowDLLM\-JEPAΔ\\DeltaLLaDA\-8BGSM8K 4\-shot \(aggressive\)55\.14±\\pm8\.90→\\rightarrow58\.12±\\pm3\.91\+2\.98LLaDA\-8BGSM8K 0\-shot \(Wide\-tt\)65\.23±\\pm0\.93→\\rightarrow67\.07±\\pm0\.41\+1\.84LLaDA\-8BNL\-RX 4\-shot \(func\)45\.93±\\pm2\.70→\\rightarrow48\.52±\\pm4\.55\+2\.58LLaDA\-8BDjango 4\-shot \(ws\-prefix\)74\.63±\\pm0\.16→\\rightarrow74\.83±\\pm0\.56\+0\.20Dream\-7BGSM8K 4\-shot37\.65±\\pm3\.77→\\rightarrow40\.21±\\pm4\.34\+2\.56Dream\-7BDjango 4\-shot \(ws\-prefix\)68\.72±\\pm1\.72→\\rightarrow69\.81±\\pm1\.82\+1\.09

### A\.5Full Grid Search Results

Table 11:NL\-RX\-SYNTHλ×k\\lambda\\times kgrid\.4\-shot functional match on LLaDA\-8B\. Baseline: 0\.475\.λ\\lambda\\kk12340\.50\.5120\.4630\.4970\.5011\.00\.4650\.4320\.5820\.5342\.00\.4710\.4140\.4910\.518Complete LLaDA\-8B grids for Spider \(both exact and execution match\), NL\-RX \(exact, functional, and DFA\-based functional match\), and Django are released with the code\.

### A\.6Multi\-Seed Results

Table 12:Multi\-seed Wide\-ttGSM8K\.LLaDA\-8B 0\-shot accuracy across three seeds\.Methodseed 42seed 123seed 777mean±\\pmstdBaseline FT65\.2864\.0666\.3465\.23±\\pm0\.93L2\-to\-base anchor65\.2864\.0666\.1965\.18±\\pm0\.87DLLM\-JEPA67\.6366\.9466\.6467\.07±\\pm0\.41DLLM\-JEPA’s mean is higher than both baselines, with lower variance on this Wide\-ttGSM8K cell specifically; aggregate variance behavior across all multi\-seed cells is mixed \(Table[10](https://arxiv.org/html/2606.00091#A1.T10)\)\. Multi\-seed runs on the aggressive configuration across GSM8K, NL\-RX, and Django are included in the released artifacts; the direction of improvement is consistent with the headline numbers reported in Table[2](https://arxiv.org/html/2606.00091#S4.T2), though absolute values differ by task and configuration as expected\.

### A\.7Representation\-Collapse Diagnostic

To empirically rule out representation collapse in our stop\-gradient JEPA setting, we compute pooled context/target embeddingsztL=Pool​\(fθ​\(xtL\)\)z\_\{t\_\{L\}\}=\\mathrm\{Pool\}\(f\_\{\\theta\}\(x\_\{t\_\{L\}\}\)\)andztH=Pool​\(fθ​\(xtH\)\)z\_\{t\_\{H\}\}=\\mathrm\{Pool\}\(f\_\{\\theta\}\(x\_\{t\_\{H\}\}\)\)on 100 Wikitext\-103 samples for each checkpoint, then report: \(i\) mean per\-dimension standard deviation, \(ii\) effective rankreff=\(∑iσi\)2/∑iσi2r\_\{\\mathrm\{eff\}\}=\(\\sum\_\{i\}\\sigma\_\{i\}\)^\{2\}/\\sum\_\{i\}\\sigma\_\{i\}^\{2\}, and \(iii\) cosine diversity1−\|cos⁡⟨zi,zj⟩\|¯1\-\\overline\{\|\\cos\\langle z\_\{i\},z\_\{j\}\\rangle\|\}on 5,000 random pairs\. Collapse would manifest asreff→1r\_\{\\mathrm\{eff\}\}\\to 1and cos\-div→0\\to 0\.

Table 13:Collapse diagnostic\.Measured on 100 Wikitext samples\. For LLaDA\-8B \(D=4096D\{=\}4096\) and Dream\-7B \(D=3584D\{=\}3584\) across base / baseline\-FT / DLLM\-JEPA\-FT checkpoints\. Effective rank remains in the 41–44 range across all checkpoints, essentially matching the pre\-trained base; cosine diversity and std are also preserved\. No collapse is observed\.Checkpointviewstdreffr\_\{\\mathrm\{eff\}\}cos\-divLLaDA baseztLz\_\{t\_\{L\}\}0\.7642\.30\.28LLaDA baseztHz\_\{t\_\{H\}\}0\.9541\.80\.27LLaDA baseline\-FT \(Wide\-tt\)ztLz\_\{t\_\{L\}\}0\.7742\.80\.28LLaDA baseline\-FT \(Wide\-tt\)ztHz\_\{t\_\{H\}\}0\.9641\.20\.28LLaDA DLLM\-JEPA \(Wide\-tt\)ztLz\_\{t\_\{L\}\}0\.7342\.10\.25LLaDA DLLM\-JEPA \(Wide\-tt\)ztHz\_\{t\_\{H\}\}0\.9541\.80\.27LLaDA baseline\-FT \(aggressive\)ztLz\_\{t\_\{L\}\}0\.7744\.00\.31LLaDA DLLM\-JEPA \(aggressive\)ztLz\_\{t\_\{L\}\}0\.7542\.40\.28Dream baseztLz\_\{t\_\{L\}\}1\.4344\.10\.47Dream baseline\-FT \(GSM8K\)ztLz\_\{t\_\{L\}\}1\.3344\.00\.39Dream DLLM\-JEPA \(GSM8K\)ztLz\_\{t\_\{L\}\}1\.3143\.00\.39
### A\.8Reference Values from LLM\-JEPA on Comparable Benchmarks

Table[14](https://arxiv.org/html/2606.00091#A1.T14)lists the improvements reported in LLM\-JEPA\[Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\]alongside the improvements we obtain with DLLM\-JEPA on the same task names\. This isnota matched comparison—the two methods use different backbones \(AR vs\. masked\-diffusion\), different*model scales*\(1B vs\. 7–8B\), different pre\-training corpora, different baseline recipes, and in one case a different metric variant\. We present the table solely as a reference point for readers familiar with the LLM\-JEPA paper; reading it as a direct DLLM\-JEPA\-vs\-LLM\-JEPA comparison would be incorrect\. Dream\-7B is not shown here because LLM\-JEPA did not evaluate it\.

Table 14:LLM\-JEPA reference numbers \(not a matched comparison\)\.Values from LLM\-JEPA\[Huang et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib10)\]\(Tables 2–4 therein\) and DLLM\-JEPA \(this work\), each measured against its own baseline\. Note the model\-scale gap \(1B vs\. 7–8B\), different substrates, and the more lenient “startswith” metric used by LLM\-JEPA on NL\-RX\-SYNTH\.TaskMethod \(backbone, scale\)MetricBaseline→\\rightarrowmethodΔ\\Delta\(pp\)GSM8KLLM\-JEPA \(Llama\-3\.2\-1B\-Instruct, 1B\)accuracy \(3\-seed\)32\.4→\\rightarrow36\.4\+4\.0DLLM\-JEPA \(LLaDA\-8B, Wide\-tt\)0\-shot acc \(3\-seed\)65\.2→\\rightarrow67\.1\+1\.8DLLM\-JEPA \(LLaDA\-8B, aggressive\)4\-shot acc \(3\-seed\)55\.1→\\rightarrow58\.1\+3\.0SpiderLLM\-JEPA \(Llama\-3\.2\-1B\-Instruct, 1B\)accuracy \(3\-seed\)47\.5→\\rightarrow50\.6\+3\.0DLLM\-JEPA \(LLaDA\-8B\)4\-shot exec \(1 seed\)35\.4→\\rightarrow39\.4\+4\.0NL\-RX\-SYNTHLLM\-JEPA \(Llama\-3\.1\-8B\-Instruct, 8B\)startswith†35\.8→\\rightarrow63\.6\+27\.8DLLM\-JEPA \(LLaDA\-8B\)4\-shot func \(3\-seed\)45\.9→\\rightarrow48\.5\+2\.6†*startswith*counts a prediction as correct if it begins with the gold regex string, irrespective of anything that follows; this is substantially more lenient than the functional\-equivalence metric \(our default\), which tests the predicted regex against 200 randomly sampled strings and requires exact agreement with the gold regex’s output\. Neither baseline accuracy \(35\.8\) nor final accuracy \(63\.6\) under startswith is directly comparable to our functional\-equivalence numbers\. For orientation, on our LLaDA\-8B checkpoints the startswith metric gives consistently higher absolute numbers than the functional metric, but this is not a substitute for evaluating LLM\-JEPA under a matched protocol\.

Setting the metric caveat aside, effect sizes on GSM8K and Spider fall in the same general range \(single\-digit percentage points over each paper’s own baseline\)\. Model\-scale effects \(LLM\-JEPA at 1–8B, DLLM\-JEPA at 7–8B\) also factor into the absolute numbers; at a given scale, a larger model typically shows smaller absolute pp gains from the same auxiliary loss because its baseline is already closer to its ceiling\.

### A\.9Stopping\-Robust Evaluation Protocol

Dream\-7B’s natural generation for structured outputs \(Python code, SQL\) occasionally emits additional tokens after completing the answer, reflecting a generative style inherited from its pre\-training\. Under strict exact\-match, this appears as a task failure despite the model having produced a correct answer\. Because this is a model\-style artifact rather than a quality difference, we adopt stopping\-robust metrics for the tasks and models where this matters, and apply them*uniformly*to both baseline and DLLM\-JEPA, and to*both*LLaDA\-8B and Dream\-7B in the main table \(so column\-to\-column comparison is apples\-to\-apples\):

- •Django \(ws\-prefix\):whitespace\-normalized prefix match between prediction and gold\. On LLaDA\-8B the ws\-prefix and raw exact metrics are similar; on Dream\-7B, ws\-prefix avoids false negatives from trailing whitespace/docstrings\.
- •Spider \(exec\-cleaned\):truncate prediction at the first semicolon, then run the standard Spider execution match against the database\.

#### Raw vs\. cleaned numbers for Dream\.

On Dream\-7B Django 4\-shot, raw exact match drops after DLLM\-JEPA fine\-tuning \(e\.g\., baseline 12\.1%→\\rightarrowJEPA 7\.9% on a representative run\), which is entirely attributable to the stopping artifact: the model now more often emits a longer continuation including correct code plus additional tokens\. Under ws\-prefix, the same checkpoints show consistent positive gains \(Table[2](https://arxiv.org/html/2606.00091#S4.T2)\)\. On Dream\-7B Spider 4\-shot, raw exact/exec match are near zero for both baseline and JEPA \(baseline 0\.3% / 0\.6%, JEPA 2\.3% / 5\.2%\) due to the same mechanism; after SQL\-cleaning, both rise into the 20–25% range with DLLM\-JEPA providing a clean \+4\.26 pp improvement\.

#### 0\-shot results for completeness\.

The main\-table 4\-shot protocol is chosen for direct LLaDA–Dream comparability\. 0\-shot results include LLaDA\-8B Django exact\-match 34\.4%→\\rightarrow56\.3% \(\+21\.9 pp\), LLaDA\-8B GSM8K Wide\-tt0\-shot 65\.3%→\\rightarrow67\.6% \(also reported in §[4\.3](https://arxiv.org/html/2606.00091#S4.SS3)\), and Dream\-7B NL\-RX 0\-shot functional match \(DFA\) 34\.7%→\\rightarrow42\.4% \(\+7\.7 pp\)\. These are not included in the main table only to preserve the unified 4\-shot framing\.

### A\.10Implementation and Reproducibility Details

#### Backbones\.

LLaDA\-8B\[Nie et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib15)\]: 32 transformer layers, hidden dim 4096\. Dream\-7B\[Ye et al\.,[2025](https://arxiv.org/html/2606.00091#bib.bib20)\]: 28 layers, hidden dim 3584\. Both use bidirectional attention\. We use the public HuggingFace checkpoints and trust\_remote\_code\.

#### JEPA module\.

The predictorgϕg\_\{\\phi\}is a stack ofk∈\{1,2,3,4,5\}k\\in\\\{1,2,3,4,5\\\}transformer decoder layers with the same hidden dimension as the backbone, randomly initialized\. Pooling is mean pooling over non\-masked, non\-padding tokens followed by a single LayerNorm\. The EMA target encoder is a no\-grad deepcopy of the backbone with decayτ=0\.996\\tau\{=\}0\.996, updated every optimizer step\.

#### Optimization\.

AdamW withβ1=0\.9,β2=0\.999\\beta\_\{1\}\{=\}0\.9,\\beta\_\{2\}\{=\}0\.999, weight decay0\.010\.01, cosine LR decay with 5% warmup\. Per\-device batch size 2, gradient accumulation 2 \(effective batch 32 across 8 GPUs\)\. Gradient checkpointing enabled for memory\. DeepSpeed ZeRO\-2 \(optimizer\-state \+ gradient sharding across 8 GPUs\)\.

#### Training schedules\.

- •*Aggressive*\(main task table\): lr1×10−51\{\\times\}10^\{\-5\}, 2 epochs, fixed view rates\(tL=0\.2,tH=0\.7\)\(t\_\{L\}\{=\}0\.2,t\_\{H\}\{=\}0\.7\)\.
- •*Wide\-tt*\(preservation\): lr1\.4×10−61\.4\{\\times\}10^\{\-6\}, 2 epochs, fixed view rates\(tL=0\.1,tH=0\.9\)\(t\_\{L\}\{=\}0\.1,t\_\{H\}\{=\}0\.9\)\.

DLLM\-JEPA hyperparameters:λ∈\{0\.5,1\.0,2\.0\}\\lambda\\in\\\{0\.5,1\.0,2\.0\\\}\(loss weight onℒJEPA\\mathcal\{L\}\_\{\\text\{JEPA\}\}\),k∈\{1,2,3,4,5\}k\\in\\\{1,2,3,4,5\\\}\(predictor depth\)\. Best\(λ,k\)\(\\lambda,k\)per cell reported in Table[2](https://arxiv.org/html/2606.00091#S4.T2); fullk∈\{1,…,5\}k\\in\\\{1,\.\.\.,5\\\}grid for GSM8K in Table[7](https://arxiv.org/html/2606.00091#A1.T7), availablek∈\{1,…,4\}k\\in\\\{1,\.\.\.,4\\\}grid for NL\-RX in Table[11](https://arxiv.org/html/2606.00091#A1.T11)\.

#### Seeds\.

Multi\-seed results use seeds\{42,123,777\}\\\{42,123,777\\\}\. Seeds are matched between baseline and DLLM\-JEPA within each task / configuration cell so that any difference is attributable to the objective\.

#### Hardware and wall\-clock\.

8×8\\timesNVIDIA A100 80GB on a single node, 80 GB HBM each\. Per\-run training wall\-clock: GSM8K∼\\sim45 min \(LLaDA\-8B aggressive\),∼\\sim1\.0 h \(Dream\-7B aggressive\),∼\\sim50 min \(LLaDA\-8B Wide\-tt\)\. Per\-run evaluation wall\-clock: GSM8K 0\-shot∼\\sim1\.5 h, 4\-shot∼\\sim1\.5 h \(8\-GPU DataParallel\)\. Total compute for the experiments in this paper is approximately 800 A100\-hours\.

#### Generation / evaluation\.

LLaDA’s iterative unmasking with 128 diffusion steps, block length 128, generation length 256 \(GSM8K\) or 512 \(Django/NL\-RX\)\. Greedy decoding \(temperature 0\)\. Same generation config for baseline and DLLM\-JEPA within each row\.

#### Datasets\.

GSM8K\[Cobbe et al\.,[2021](https://arxiv.org/html/2606.00091#bib.bib5)\], Spider\[Yu et al\.,[2018](https://arxiv.org/html/2606.00091#bib.bib19)\], NL\-RX\-SYNTH\[Locascio et al\.,[2016](https://arxiv.org/html/2606.00091#bib.bib13)\], Django\[Oda et al\.,[2015](https://arxiv.org/html/2606.00091#bib.bib21)\]from public sources\. Wikitext\-103 validation\[Merity et al\.,[2017](https://arxiv.org/html/2606.00091#bib.bib22)\]for held\-out diffusion loss\. MMLU\[Hendrycks et al\.,[2021](https://arxiv.org/html/2606.00091#bib.bib23)\]500\-question stratified subset \(approximately 8–9 questions per subject across all 57 subjects, fixed seed\) for general\-capability check\.

#### Code and artifacts\.

All training and evaluation scripts, configuration files, and the EMA \+ predictor \+ Wikitext probing code will be released at the paper’s accompanying GitHub repository upon publication\. The released bundle includes the exact\(λ,k\)\(\\lambda,k\)values per cell, masking\-schedule scripts, and seed\-by\-seed checkpoints used for the multi\-seed tables\.