DLLM-JEPA:掩码扩散语言模型的联合嵌入预测架构
DLLM-JEPA: Joint Embedding Predictive Architectures for Masked Diffusion Language Models
摘要
联合嵌入预测架构(JEPA)重塑了视觉领域的自监督表示学习。最近的 LLM-JEPA 将 JEPA 移植到自回归语言模型,但继承了因果注意基质的两个巨大成本:它需要显式的多视图数据(例如文本代码对),并且每步需要两次梯度进位前向传递。我们引入了 DLLM-JEPA,它将 JEPA 与掩码扩散语言模型配对,以一次性消除这两种成本。扩散模型的双向注意力通过不同的掩蔽率产生同一输入的两个语义上不同的视图(不需要显式对),并支持单个梯度携带前向传递,相对于 LLM-JEPA 将训练 FLOP 减少 33%。在我们评估的每个(任务、架构)组合中,DLLM-JEPA 都比仅扩散 微调 有所改进:在 LLaDA-8B GSM8K 上高达 +18.7 pp,在 Dream-7B GSM8K 上高达 +11.4 pp,在 Spider、NL-RX-SYNTH 和 Django 上具有一致的正增益。除了准确性之外,DLLM-JEPA 还表现出双赢的特性:在具有 Wide-t 配置的 LLaDA-8B 上,它同时提高了 GSM8K 准确性(67.1 vs. 65.2,+1.8 pp),将保留的 Wikitext 损失驱动到预训练基础以下,并在三个 微调 种子的基础水平上保留 MMLU 准确性 - 而 L2 到基础参数锚在没有任务的情况下匹配基线准确性增益。逐层探测揭示了这一机制:几何功能漂移解离,其中微调后的主干比基线更远离预先训练的权重,但对保留的维基文本的遗忘较少,放大集中在中间的 Transformer 层。该模式也出现在 Dream-7B 上,表明该现象并非特定于单个骨干网。