用扩散模型改进自监督语言学习,训练更快更准。
DLLM-JEPA: Joint Embedding Predictive Architectures for Masked Diffusion Language Models

- 将JEPA与掩码扩散语言模型结合,无需成对数据
- 仅需一次前向传播,训练计算量减少33%
- 在多个任务上显著提效,且保持知识不遗忘
联合嵌入预测架构(JEPAs)重塑了视觉自监督表征学习。近期的LLM-JEPA将JEPA引入自回归语言模型,但继承了因果注意力结构带来的双重代价:需要显式多视图数据(如文本-代码对),且每步需两次梯度传递前向传播。本文提出DLLM-JEPA,将JEPA与掩码扩散语言模型结合,一次性消除两项成本。扩散模型的双向注意力通过不同掩码率生成同一输入的语义各异的两视图,无需显式配对数据;同时支持单次梯度传递前向传播,相较LLM-JEPA降低33%训练浮点运算量。DLLM-JEPA在所有评估的任务与架构组合中均优于纯扩散微调:在LLaDA-8B GSM8K上提升最高达+18.7个百分点,在Dream-7B GSM8K上提升+11.4个百分点,对Spider、NL-RX-SYNTH和Django任务亦有持续正向增益。除精度提升外,还展现出双重优势:在LLaDA-8B Wide-t配置下,其GSM8K准确率从65.2提升至67.1(+1.8 pp),未见训数据上的Wikitext损失低于预训练基线,且在三个微调种子下保持MMLU准确率不变——而参数锚定方法仅匹配基线无提升。层间探查揭示机制:微调后的主干网络偏离预训练权重更远,但对未见数据遗忘更少,放大效应集中于中间变换器层。该现象在Dream-7B上同样出现,表明非单一主干特有。
原文摘要 · Abstract (English)
Joint Embedding Predictive Architectures (JEPAs) have reshaped self-supervised representation learning in vision. The recent LLM-JEPA ported JEPA to autoregressive language models but inherited two steep costs from the causal-attention substrate: it demands explicit multi-view data (e.g., text-code pairs), and it requires two gradient-carrying forward passes per step. We introduce DLLM-JEPA, which pairs JEPA with masked-diffusion language models to eliminate both costs at once. The bidirectional attention of diffusion models yields two semantically distinct views of the same input via different masking rates -- no explicit pairs needed -- and supports a single gradient-carrying forward pass, cutting training FLOPs by 33% relative to LLM-JEPA. DLLM-JEPA improves over diffusion-only fine-tuning in every (task, architecture) combination we evaluate: up to +18.7 pp on LLaDA-8B GSM8K and +11.4 pp on Dream-7B GSM8K, with consistent positive gains on Spider, NL-RX-SYNTH, and Django. Beyond accuracy, DLLM-JEPA exhibits a dual-win property: on LLaDA-8B with the Wide-t configuration, it simultaneously raises GSM8K accuracy (67.1 vs. 65.2, +1.8 pp), drives held-out Wikitext loss below the pre-trained base, and preserves MMLU accuracy at base level across three fine-tuning seeds -- whereas an L2-to-base parameter anchor matches baseline accuracy with no task gain. Layer-wise probing reveals the mechanism: a geometric-functional drift dissociation in which the fine-tuned backbone moves further from the pre-trained weights than the baseline yet forgets less on held-out Wikitext, with the amplification concentrated in middle transformer layers. The pattern appears on Dream-7B as well, indicating the phenomenon is not specific to a single backbone.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。