基于轨迹的在策略蒸馏用于掩码扩散语言模型

# 基于轨迹的在线蒸馏方法用于掩码扩散语言模型  
来源：https://arxiv.org/html/2607.16872  

Haolin Ren, Ziyang Huang∗, Chenhao Yuan, Jun Zhao, Kang Liu  
中国科学院自动化研究所，北京，中国  
中国科学院大学，北京，中国  
renhaolin2026@ia\.ac\.cn  

###### 摘要  

扩散大语言模型（dLLMs）是自回归生成的一种有前景的替代方案。然而，面向推理的 dLLM 后训练仍然具有挑战性。dLLM 的监督微调（SFT）需要密集但往往是非策略（off-policy）的掩码状态，而强化学习（RL）则依赖稀疏奖励或价值建模。本文提出**基于轨迹的在线蒸馏方法（TOPD）**，一种教师监督框架，无需奖励估计即可将推理能力迁移至目标 dLLM。关键思想是在目标 dLLM 自身的去噪轨迹上进行监督，重点关注构成最终响应的轨迹对齐（trace-aligned）词元决策。具体来说，TOPD 从目标 dLLM 中采样在线（on-policy）扩散轨迹，从教师模型获取对应部分去噪状态下的教师词元分布，并使用词元级别的反向 KL 散度（Reverse-KL）目标更新目标 dLLM。这种设计在保持密集教师监督的同时，使训练与模型自身的去噪状态对齐。在数学推理基准上，TOPD 使 SDAR-4B-Chat 的 MATH500 准确率与其 RL 训练版本 TraDo-4B-Instruct 相当，静态评估提升 +5.7，动态评估提升 +4.5。与 RL 训练版本相比，TOPD 只需 4 倍的 rollout 轮数，对应估算的 96.0 倍的模型计算加速比（to-accuracy）。  

---

# 基于轨迹的在线蒸馏方法用于掩码扩散语言模型  

Haolin Ren††Equal contribution., Ziyang Huang∗, Chenhao Yuan, Jun Zhao, Kang Liu  
中国科学院自动化研究所，北京，中国  
中国科学院大学，北京，中国  
renhaolin2026@ia\.ac\.cn  

## 1 引言  

参考图1  
图1：TOPD 与随机掩码 RL 在 dLLM 中的概念对比。  
问题：一张票80美元，折扣25%，然后加税10%。最终价格是多少？  
完整答案：80×(1-25%)=60，60×(1+10%)=66，最终答案：66美元  
在线状态：80×(1-25%)=__，__×(1+10%)=__，最终答案：$__  
保留因果方向：小计由前面计算填充  
随机掩码状态：__×(1-25%)=__，60×(1+__%)=66，最终答案：$__  
破坏因果推理：60 由后续证据监督  
图2：一个合成的问答示例，说明轨迹不匹配问题。随机掩码可能暴露后续结果而隐藏早期推理变量，迫使模型进行反向重构而非正向计算。

后训练已在推理任务上为自回归大语言模型（LLMs）带来了巨大提升（Ouyang 等人，2022；Shao 等人，2024；Guo 等人，2025；Chu 等人，2025）。扩散大语言模型（dLLMs）基于迭代去噪提供了一种有竞争力的非自回归替代方案（Li 等人，2022；Sahoo 等人，2024；Nie 等人，2025；Cheng 等人，2025；Arriola 等人，2025；Gong 等人，2025；Ye 等人，2025）。然而，面向推理的 dLLM 后训练研究仍然较少，特别是从更强的 dLLM 向较弱的学生 dLLM 蒸馏推理能力。本文研究如何在迭代去噪生成过程中，利用更强的 dLLM 作为监督，高效地适配学生 dLLM 用于推理。

当前 dLLM 后训练流程主要遵循两条路径：监督微调（SFT）和强化学习（RL）（Zhao 等人，2025b；Wang 等人，2025；Zhu 等人，2025；He 等人，2025；Ou 等人，2026；Wang 等人，2026；Zhao 等人，2025a；Zhong 等人，2026）。SFT 在固定目标（如人工标注或预生成样本）上训练学生，使得监督本质上是非策略（off-policy）的。这样，学生是在静态目标或理想化教师状态下优化，而不是在其自身生成策略诱导的状态上优化，导致暴露偏差和泛化能力较弱（Agarwal 等人，2024；Chu 等人，2025）。另一方面，RL 通常依赖稀疏、延迟且高方差的奖励，这使得长程信用分配困难且优化成本高昂。这些局限性催生了在线蒸馏（OPD）（Agarwal 等人，2024；Lu and Lab，2025），即学生从其自身策略采样，同时教师对这些学生生成的输出提供密集反馈。然而，将 OPD 适配到 dLLM 并非自回归模型的简单移植。在自回归 LLM 中，每个动作是一个确定的下一词元决策，因此教师反馈自然地附着在生成的词缀上。在 dLLM 中，模型迭代更新一个部分掩码的序列，这使得轨迹对齐成为核心问题。一个关键障碍是许多近期 dLLM RL 流程中的随机掩码不匹配问题：训练状态通常是通过随机破坏模型生成响应的子集来构建的（Zhao 等人，2025b；Zhu 等人，2025；Ou 等人，2026；Wang 等人，2026；Zhong 等人，2026）。如图2所示，这种状态可能暴露后续答案而隐藏早期推理变量，创建了与学生正向去噪轨迹不同的反向重构上下文。揭示顺序和推理调度会实质性地影响掩码扩散行为（He 等人，2025；Wang 等人，2025）。因此，在随机或轨迹无关状态上的教师监督会削弱与构成最终响应的决策之间的对齐。

为此，本文提出**基于轨迹的在线蒸馏方法（TOPD）**，一种面向 dLLM 的教师监督后训练框架。TOPD 采样学生自身的推理时去噪轨迹，并通过识别直接决定最终响应的词元预测，将其转换为密集的蒸馏目标。这种基于轨迹的公式避免了监督任意随机掩码配置，因为随机掩码可能暴露后续答案词元而隐藏早期推理词元，从而产生与推理时去噪不一致的条件依赖关系。TOPD 随后使用冻结的教师提供学生在策略部分去噪状态上的词元分布，并通过逐步词元分布监督来优化学生。在此过程中，利用词元级别的反向 KL 散度。结果，TOPD 保留了轨迹对齐的状态覆盖，同时用密集的教师反馈替代了稀疏奖励信用分配和价值模型开销。图1提供了 TOPD 与随机掩码 RL 的概念对比。

在数学推理基准上的实验表明，TOPD 能够以显著更低的训练成本将推理能力从 dLLM 教师迁移到更小的学生。具体来说，使用 TraDo-8B-Instruct（Wang 等人，2025）作为教师，TOPD 使基础 SDAR-4B-Chat（Cheng 等人，2025）学生的 MATH500（Hendrycks 等人，2021）准确率达到与 TraDo-4B-Instruct（即 SDAR-4B-Chat 的 RL 训练版本）相当的水平，静态评估提升 +5.7，动态评估提升 +4.5。此外，与 TraceRL（Wang 等人，2025）相比，TOPD 只需 4 倍的 rollout 轮数，在我们的参数缩放核算下实现了估算的 96.0 倍模型计算加速比（to-accuracy）。

我们的贡献总结如下：  
- **** 我们将教师监督的 dLLM 后训练形式化为一个基于轨迹的在线蒸馏问题，表明随机掩码监督可能引入与学生推理时去噪轨迹不一致的条件依赖关系。  
- **本文提出 TOPD**，一种高效的 dLLM 后训练范式，它从学生采样在线扩散轨迹，并识别出直接决定最终响应的轨迹对齐去噪决策，用逐步的教师词元分布监督替代基于奖励的更新。  

## 2 背景  

### 2.1 掩码扩散语言模型  

掩码扩散语言模型（MDLMs）通过迭代去噪掩码词元序列来生成文本，而不是从左到右依次提交词元（Sahoo 等人，2024；Nie 等人，2025）。给定一个干净序列 x∈V^n，训练过程将词元以噪声水平 t 破坏为 [MASK]，并学习一个针对掩码位置的去噪器 π_θ(·|z)，通常使用一个调度加权掩码重建目标（Sahoo 等人，2024；Shi 等人，2024）。全注意力 dLLM 允许响应词元双向关注整个被破坏的响应，而块注意力 dLLM 将响应划分为块，并使用块因果掩码，例如，块 k 中的词元仅关注块 ≤k 的块（Arriola 等人，2025；Cheng 等人，2025）。然而，在推理时，模型遵循其自身的去噪轨迹 τ=(s_0,a_0,...,s_T)，从完全掩码的响应开始，而揭示顺序会影响生成行为（He 等人，2025；Wang 等人，2025）。由于中间决策可能后被修改或删除，我们用 **trace** 表示 τ 中那些最终保留在 s_T 中的决策子集。

### 2.2 dLLM 的后训练信号  

dLLM 的后训练通常依赖于 SFT 或 RL（Zhao 等人，2025b；Wang 等人，2025；Zhu 等人，2025；He 等人，2025；Ou 等人，2026；Wang 等人，2026；Zhao 等人，2025a；Zhong 等人，2026）。在 SFT 中，目标响应被部分破坏得到 ỹ，模型被训练来重建掩码位置：  
L_SFT(θ) = E_{q,y,ỹ} [ Σ_{i∈M(ỹ)} -log π_θ(y_i | q, ỹ) ],   (1)  
其中 M(ỹ) 是掩码位置集合，对 ỹ 的期望表示应用于 y 的随机掩码过程。  
RL 则从当前学生采样，并最大化：  
J_RL(θ) = E_{q, τ~P_θ(·|q)} [ R(q, s_T) ].   (2)  
这里 P_θ(τ|q) 表示由学生的迭代去噪策略诱导的轨迹分布，R(q, s_T) 是分配给最终响应的标量奖励。  
SFT 提供密集监督，但训练是在固定目标响应上，而非当前学生诱导的状态上。RL 使用学生样本，但其稀疏且延迟的奖励必须跨许多非自回归去噪决策进行分配，使得信用分配成本高昂且方差大。教师监督提供了一种折中方案：在学生访问的状态上提供密集的词元分布，无需延迟奖励估计。

### 2.3 将 OPD 适配到 dLLM 的挑战  

OPD 在学生当前策略采样的状态上训练学生，同时教师在这些相同状态上提供密集的目标分布（Agarwal 等人，2024；Lu and Lab，2025）。将这一思想适配到 dLLM 并非易事，因为扩散步骤是在部分掩码序列上操作，并且可能并行提出多个词元决策。与自回归 OPD 不同，自回归 OPD 中每个生成的词缀定义了一个确定的下一词元上下文，而 dLLM 蒸馏必须决定在多步揭示过程中哪些部分去噪状态和词元位置应接收教师反馈。随机掩码监督提供了密集的掩码词元目标，但这些人工状态可能与当前学生在生成过程中遇到的状态或揭示顺序不匹配（He 等人，2025；Wang 等人，2025）。因此，核心挑战是使 OPD 既是在策略（on-policy）又是轨迹感知的。学生应在其自身去噪过程采样的状态上被监督，并且教师信号应根据该过程诱导的轨迹构建来组织，而不是根据独立采样的掩码。TOPD 通过采样在线扩散轨迹并在轨迹对齐的去噪决策上蒸馏教师分布来解决这一问题，我们接下来对其进行形式化。

## 3 方法  

参考图3  
图3：TOPD 概述。学生采样一个在线去噪轨迹，教师提供逐步词元分布监督，仅对最终输出中保留的轨迹对齐决策 ã_t 进行更新。

我们引入 **TOPD**，一种用于掩码扩散语言模型后训练的基于轨迹的在线蒸馏框架。给定一个提示，TOPD 采样学生的去噪轨迹，保留最终响应中保留的词元决策，并在这些学生访问的状态上蒸馏冻结的教师分布。图3展示了该方法。