DLLM-JEPA：面向掩码扩散语言模型的联合嵌入预测架构

# 用于掩码扩散语言模型的联合嵌入预测架构
来源：https://arxiv.org/html/2606.00091

###### 摘要

我们提出 **DLLM-JEPA**，一种面向掩码扩散语言模型的 JEPA 公式。此前，JEPA 目标函数已应用于自回归语言模型，但代价是需要显式配对视图以及每个训练步骤多次梯度传递。通过利用扩散噪声调度，DLLM-JEPA 从单一输入中构造两个视图，无需配对数据，并将 JEPA 训练的计算量（FLOPs）相对于 LLM-JEPA 的双梯度视图设计降低了 33%——每个步骤只需一次梯度传递。在四个任务和两个扩散骨干网络上，DLLM-JEPA 始终优于仅使用扩散的微调：在稳定设置下取得适度增益（例如 GSM8K 上 +1.8 pp），在更激进的微调下取得较大改进（例如高方差的 LLaDA-8B GSM8K 单元上，种子间方差从 ±8.9 pp 收紧至 ±3.9 pp）。此外，该方法不会降低基模型在保留的扩散损失探测任务上的性能，也不会影响小规模 MMLU 基准测试的表现。我们进一步分析了该目标函数诱导的表示动态，并观察到一致的实证模式：使用 DLLM-JEPA 训练的模型相对于预训练初始化表现出更大的几何漂移，同时保持相当或更低的函数遗忘。这些结果表明，DLLM-JEPA 提供了一种将表示级目标函数高效融入扩散语言模型微调的方法。

联合嵌入预测架构，掩码扩散语言模型，表示学习，微调

## 1 引言

训练大语言模型的主流范式依赖于输入空间重建：自回归（AR）下一词预测（Brown et al., 2020 (https://arxiv.org/html/2606.00091#bib.bib4); Touvron et al., 2023 (https://arxiv.org/html/2606.00091#bib.bib18)）或掩码词重建（Devlin et al., 2019 (https://arxiv.org/html/2606.00091#bib.bib6)）。相比之下，视觉领域越来越多地采用联合嵌入预测架构（JEPA）（Assran et al., 2023 (https://arxiv.org/html/2606.00091#bib.bib1); Bardes et al., 2024 (https://arxiv.org/html/2606.00091#bib.bib3)），这种架构通过从一个视图预测另一个视图的嵌入来学习表示——完全在潜在空间中操作。JEPA 已被证明可以通过避免像素级重建偏差来学习更丰富、更抽象的表示（LeCun, 2022 (https://arxiv.org/html/2606.00091#bib.bib11); Littwin et al., 2024 (https://arxiv.org/html/2606.00091#bib.bib12)）。最近的 LLM-JEPA (Huang et al., 2025 (https://arxiv.org/html/2606.00091#bib.bib10)) 是将 JEPA 目标引入语言模型的第一步。通过将 (文本, 代码) 对视为同一底层知识的两个视图，LLM-JEPA 在标准下一词预测目标之外添加了 JEPA 损失。尽管取得了有希望的结果，LLM-JEPA 面临着源于自回归架构的两个根本限制：

1. 1. **显式视图需求**。LLM-JEPA 需要具有自然双视图结构的数据集（例如，自然语言与代码配对）。作者本人也指出这是关键限制：“开发一种类似于视觉中数据增强的机制，将使 JEPA 目标能够用于任何数据集”（Huang et al., 2025 (https://arxiv.org/html/2606.00091#bib.bib10)）。
2. 2. **单向性带来的计算开销**。AR 模型使用因果掩码从左到右处理词元。为了获得两个视图的独立嵌入，LLM-JEPA 需要自定义的块因果注意力掩码和两次前向传播——两者都携带梯度——导致训练计算量约为标准微调的 ∼2× 倍。

我们观察到 **掩码扩散语言模型**（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)）自然地解决了这两个限制。诸如 LLaDA (Nie et al., 2025 (https://arxiv.org/html/2606.00091#bib.bib15)) 之类的模型采用双向注意力，并学习去噪随机掩码的词元——这一过程在结构上类似于 JEPA 的视图预测。这一见解引出了我们提出的方法 **DLLM-JEPA**，其贡献可概括为：（1）将 JEPA 高效适配到掩码扩散语言模型的**架构性**方法；（2）JEPA 如何在双向语言模型中重塑表示的**一致实证模式**；（3）一致的任务增益与基础能力保持：

- • **节省 33% 训练计算量的 JEPA 公式**。掩码扩散模型的双向注意力使我们能够通过不同的掩码率（无需显式视图对）从同一输入生成两个 JEPA 视图，并在单次前向传播中处理上下文视图，梯度同时产生扩散 logits 和 JEPA 嵌入；只有目标视图需要额外的无梯度前向传播。结果是，每个 JEPA 训练步骤的计算量比 LLM-JEPA 少 33%（表1 (https://arxiv.org/html/2606.00091#S3.T1)），使用标准的双向注意力（无需自定义块因果掩码），并且适用于任何文本数据集。
- • **一致的实证模式：几何-函数漂移分离**。逐层探测表明，DLLM-JEPA 目标似乎并没有最小化表示变化——而是重新引导了变化方向。在 GSM8K 上，DLLM-JEPA 隐藏状态相对于预训练权重的漂移 **大于** 基线（各种配置下漂移比 1.3–3.6×），但其 Wikitext 函数遗忘 **更小**。这种放大集中在中层 Transformer 层——与先前的解释（中层编码组合结构）一致——并且在 Dream-7B 上复现（比率 1.28×）。我们还通过实证排除了表示坍塌（§3 (https://arxiv.org/html/2606.00091#S3)，附录 A.7 (https://arxiv.org/html/2606.00091#A1.SS7)）。我们在全文使用“分离”一词描述这种观察到的现象。
- • **任务增益与基础保持的联合**。在 LLaDA-8B 上使用 Wide-tt 配置（t_L=0.1, t_H=0.9, lr 1.4×10⁻⁶，比主任务表中使用的激进 (0.2, 0.7) 调度更温和），DLLM-JEPA 同时提高了 GSM8K 准确率 **并** 将 Wikitext 扩散损失降至预训练基础之下，而 L2 参数锚定于基线则只能实现弱基础保持且无任务增益。
- • **4个任务 × 2个骨干网络上的一致增益**。在单一 4-shot 评估协议下，DLLM-JEPA 改善了每个 (任务, 架构) 组合（表2 (https://arxiv.org/html/2606.00091#S4.T2)）。3 种子的平均提升：在 LLaDA-8B GSM8K 的 Wide-tt 配置下为 +1.8 pp（±0.4 vs 基线 ±0.9），在激进调度任务上为 +2.6–3.0 pp，其中 DLLM-JEPA 将 ±8.9 pp 的基线 spread 收紧至 ±3.9 pp。最佳种子结果（最高 +18.7 pp）及可用的多种子统计信息见附录 A.6 (https://arxiv.org/html/2606.00091#A1.SS6)。

#### 比较范围。
我们将 LLM-JEPA 定位为结构动机而非直接比较对象：两种方法处于不同的注意力基底上，因此我们在同一骨干网络上将 DLLM-JEPA 与仅扩散微调进行基准比较。我们关注的是 JEPA 目标如何自然地实例化到掩码扩散 LM 中，以及它们如何影响微调动态。

## 2 相关工作

#### 掩码扩散语言模型。
用于文本生成的离散扩散模型作为自回归解码的替代方案已获得关注（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)) 引入了一个简单的掩码扩散框架，在 8B 参数规模上取得了与 AR 模型相当的性能。MDLM (Sahoo et al., 2024 (https://arxiv.org/html/2606.00091#bib.bib16)) 和 SEDD (Shi et al., 2024 (https://arxiv.org/html/2606.00091#bib.bib17)) 提供了进一步的理论和实证基础。这些模型采用双向注意力，并学习逆转前向掩码过程，从而在推理时实现并行词元生成。

#### 联合嵌入预测架构。
I-JEPA (Assran et al., 2023 (https://arxiv.org/html/2606.00091#bib.bib1)) 证明了预测掩码图像块的潜在表示能产生优于像素级重建（MAE）的视觉特征。V-JEPA (Bardes et al., 2024 (https://arxiv.org/html/2606.00091#bib.bib3)) 将其扩展到视频。LLM-JEPA (Huang et al., 2025 (https://arxiv.org/html/2606.00091#bib.bib10)) 将该框架适配到自回归语言模型，在代码生成和数学推理任务上均取得改进。然而，LLM-JEPA 需要显式视图对，并且由于因果注意力的限制而产生了显著的计算开销。

#### 微调扩散语言模型。
掩码扩散语言模型的微调动态仍未得到充分探索。Nie 等人 (2025 (https://arxiv.org/html/2606.00091#bib.bib15)) 证明了对 LLaDA 在指令数据上进行监督微调能产生强大的指令跟随行为，但未分析表示动态或任务适应与基础保持之间的权衡。我们的工作使用 DLLM-JEPA 作为探究该权衡的透镜，并表明表示级正则化在性质上重塑了这一权衡。

## 3 方法

### 3.1 预备知识：掩码扩散语言模型

LLaDA (Nie et al., 2025 (https://arxiv.org/html/2606.00091#bib.bib15)) 定义了一个逐步掩码词元的前向过程。给定一个干净序列 x₀ = (x₀¹, ..., x₀ᴸ)，前向过程在时间步 t ∈ [0, 1] 以概率 t 独立地将每个词元替换为特殊的 [MASK] 词元：

q(x_tⁱ ∣ x₀ⁱ) = \begin{cases} x₀ⁱ & \text{概率为 } 1-t \\ \texttt{[MASK]} & \text{概率为 } t \end{cases}      (1)

模型 f_θ 学习从掩码序列预测原始词元。令 M_t = {i: x_tⁱ = \texttt{[MASK]}} 表示时间步 t 时被掩码位置集合，扩散训练目标是对掩码位置的逐词元交叉熵：

L_diff = E_{t∼U(0,1), x_t∼q(⋅∣x₀,t)} \left[ -\frac{1}{|M_t|} \sum_{i∈M_t} \log p_θ(x₀ⁱ ∣ x_t) \right]      (2)

关键在于，f_θ 使用 **双向注意力**：每个词元可以关注到其他所有词元，从而单次前向传播即可获得丰富的上下文表示。

### 3.2 DLLM-JEPA：面向扩散语言模型的 JEPA

请参阅图注
图 1: DLLM-JEPA 概览。*第一行（训练流程）*。单一干净输入 x₀ 在两个掩码率（t_L=0.2, t_H=0.7）下加噪，形成上下文视图和目标视图——无需配对数据集。在线骨干 f_θ 在单次带梯度的前向传播中处理上下文视图，产生扩散 logits（得到 L_diff）和池化嵌入 z_{t_L}；目标视图由 EMA 副本 f_{θ'}（衰减 τ=0.996）在无梯度（no_grad）下处理，产生 z_{t_H}。预测器 g_φ 将 z_{t_L} 映射为 \hat{z}_{t_H}，L_JEPA 是到 z_{t_H} 的余弦差距；L_total = L_diff + λ L_JEPA（成本 4F/步，表1 (https://arxiv.org/html/2606.00091#S3.T1)）。中左部分显示了一个具体的掩码文本示例。*底部面板（观察到的机制）*。使用 DLLM-JEPA 微调 **增加** 了隐藏态漂移（GSM8K 上为基线的 1.36–3.60×，集中在中层 Transformer 层），却 **减少了** Wikitext 函数遗忘（43–58%）。3D 密度图在逐种子点上可视化了相同效应：基线质量集中在高遗忘区域，而 DLLM-JEPA 质量沿漂移轴延伸，但更接近零遗忘底面（§5 (https://arxiv.org/html/2606.00091#S5)）。

#### 通过噪声调度生成视图。
我们通过从扩散调度中采样两个掩码率来生成每个输入的两个视图。给定 t_L > 1) 在三个 GSM8K 单元上（LLaDA 的两个配置和 Dream-7B）是明显的；NL-RX 和 Django 接近或略低于 1，表明是任务特定的放大，而非均匀漂移。（C）DLLM-JEPA 在每个微调任务上的 Wikitext 函数遗忘 **减少**（43–58%），与漂移趋势相反。（D）每区域漂移比（早期/中期/晚期第三部分）。在 GSM8K 上，DLLM-JEPA 的放大集中在中层，与先前关于中层组合结构的解释一致。所有单元均为多种子，但 LLaDA-8B GSM8K Wide-tt 除外（单种子，受检查点存储限制）。

### 5.1 GSM8K 上的几何-函数漂移分离

图 2 (https://arxiv.org/html/2606.00091#S5.F2) (A–B) 揭示了一个反直觉的发现：在 GSM8K 上，DLLM-JEPA 相对于基线的隐藏态漂移 **大于** 基线——而非更小。在 LLaDA-8B 上，DLLM-JEPA / 基线漂移比在使用温和 Wide-tt 配置时为 3.60×，在使用激进配置时为 1.36×。在 NL-RX 和 Django（激进配置）上，该比率分别为 0.94× 和 0.99×——接近 1，意味着 DLLM-JEPA 的漂移放大是 **任务特定的**，而非均匀的。然而，图 2 (https://arxiv.org/html/2606.00091#S5.F2) (C) 显示，在测量 Wikitext 遗忘的三个任务上，DLLM-JEPA **减少** 了 43–58% 的函数遗忘。这种几何-函数漂移分离可以用描述性解释：JEPA 目标似乎并没有最小化表示变化，而是将其重新引导到对任务有用的轴上，从而保留了基模型在保留集输出函数上的大部分能力。

#### 漂移模式的跨架构复现。
在 Dream-7B（不同的参数量、分词器和预训练语料）上出现了相同的效应方向：DLLM-JEPA 的漂移比为 1.28×（图 2 (https://arxiv.org/html/2606.00091#S5.F2) B，最后一列）。在一个额外骨干网络上的复现具有提示性而非结论性。

#### 区域特定放大（GSM8K）。
额外的漂移集中在 **中间** Transformer 层（图 2 (https://arxiv.org/html/2606.00091#S5.F2) D；激进调度下中层比 1.64× 对比早期 1.38× / 晚期 1.18×），这与先前关于组合推理的工作相关联。该模式在 GSM8K 上明显，而在其他任务上较弱。

### 5.2 组件消融：非对称视图与预测器

图 3 (https://arxiv.org/html/2606.00091#S5.F3) 展示了在 LLaDA-8B GSM8K（激进调度，种子 42）上移除两个 DLLM-JEPA 设计选择的效果。(i) **对称视图**（t_L = t_H = 0.2，两个视图使用相同掩码率）移除了我们方法的非对称噪声结构。(ii) **无预测器**（g_φ 为恒等映射）移除了 JEPA 的预测头，因此