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

We introduce DLLM-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 by 33% 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 pp on 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 pretrained 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.