自蒸馏轨迹感知玻尔兹曼建模：弥合扩散语言模型中的训练-推理差异

# 自蒸馏轨迹感知玻尔兹曼建模：弥合扩散语言模型中的训练-推理差异

来源: https://arxiv.org/html/2605.11854  
Kecheng Chen${}^{1,\star}$, Ziru Liu${}^{2,\star,\spadesuit}$, Xijia Tao${}^{3}$, Hui Liu${}^{1}$, Yibing Liu${}^{1}$, Xinyu Fu${}^{2}$, Shi Wu${}^{2}$, Suiyun Zhang${}^{2}$, Dandan Tu${}^{2,\clubsuit}$, Lingpeng Kong${}^{3}$, Rui Liu${}^{2,\clubsuit}$, Haoliang Li${}^{1,\clubsuit}$  
${}^{1}$香港城市大学, ${}^{2}$华为研究, ${}^{3}$香港大学  
邮箱: [email protected]; [email protected]; [email protected]

###### 摘要

扩散语言模型（DLMs）最近作为自回归语言模型的一种有前景的替代方案出现，提供了更强的全局感知能力和高度并行的生成能力。然而，使用标准的基于负证据下界（NELBO）的监督微调（SFT）对预训练后的 DLM 进行训练仍然效率低下：训练是在单步中重建随机掩码的标记，而推理则遵循置信度引导的多步“由易到难”的去噪轨迹。最近的基于轨迹的自蒸馏方法主要利用这种推理轨迹来进行采样步压缩和加速，通常在不实质性增强模型底层能力的情况下提高解码效率，甚至可能在完全扩散解码下导致性能下降。在这项工作中，我们探讨自蒸馏轨迹是否不仅用于更快的推理，还能用于真正的知识获取。尽管这些轨迹位于预训练 DLM 自身的分布流形上，从而可能提供更低的优化障碍，但我们发现，使用标准的 NELBO 目标在它们上面进行简单的微调仅能带来微小的增益。为了解决这一限制，我们提出了**基于玻尔兹曼建模的轨迹对齐优化（TABOM）**，这是一种基于自蒸馏轨迹的后训练框架，旨在使训练与推理的“由易到难”结构保持一致。TABOM 将推理时的去掩码偏好建模为预测熵上的玻尔兹曼分布，并推导出一个可处理的对成排序目标，以将模型的确定性排序与观察到的解码轨迹对齐。实验表明，与标准 SFT 相比，TABOM 在新领域取得了显著的增益，扩展了 DLM 的有效知识边界，并显著缓解了灾难性遗忘。

${}^{11}$脚注文本: ${}^{\star}$对本工作贡献同等。${}^{\clubsuit}$通讯作者。${}^{\spadesuit}$项目负责人。

## 1 引言

参见标题 **图 1: 提出的 TABOM 框架概览**。标准 SFT 使用均匀随机掩码训练 DLM，这导致了与推理期间使用的“由易到难”去掩码轨迹的不匹配。相反，TABOM 通过将轨迹感知重建与基于能量的对成排序损失相结合，从自蒸馏轨迹中学习，使模型的熵景观与推理时的玻尔兹曼去掩码分布保持一致。

现代大型语言模型（LLMs）在各种任务中展现了卓越的能力，包括数学推理、代码生成和多轮对话\[26 (https://arxiv.org/html/2605.11854#bib.bib98), 18 (https://arxiv.org/html/2605.11854#bib.bib99), 6 (https://arxiv.org/html/2605.11854#bib.bib100)\]。尽管当前的范式依赖于自回归 Transformer 架构\[27 (https://arxiv.org/html/2605.11854#bib.bib101)\]，但扩散语言模型（DLMs）已成为一种引人注目的替代方案\[31 (https://arxiv.org/html/2605.11854#bib.bib84), 17 (https://arxiv.org/html/2605.11854#bib.bib83)\]。通过利用双向上下文并同时纳入所有位置的信息，DLM 提供了更优越的全局感知能力和高度并行标记生成的潜力。为了进一步增强预训练的 DLM，通常应用监督微调（SFT）\[31 (https://arxiv.org/html/2605.11854#bib.bib84)\]，即训练模型在单次前向传递中预测随机掩码的标记。然而，与自回归训练相比，该目标的样本效率显著较低，并引入了明显的训练-推理不匹配，因为推理需要沿解码轨迹进行迭代去噪。因此，一个自然的补救措施是直接学习高质量的解码轨迹，这需要可扩展的轨迹生成管道和高效的轨迹级学习目标。虽然强化学习（RL）似乎与此目标相关\[19 (https://arxiv.org/html/2605.11854#bib.bib3), 30 (https://arxiv.org/html/2605.11854#bib.bib4), 34 (https://arxiv.org/html/2605.11854#bib.bib6), 29 (https://arxiv.org/html/2605.11854#bib.bib7)\]，但它并非直接解决方案。结果奖励 RL 仅提供粗略的序列级监督，使得在中间去噪决策上进行信用分配变得困难。它还需要重复的回放和奖励评估，这对于 DLM 来说尤其昂贵，因为每次回放已经涉及多步去噪。

最近的工作探索了用于 DLM 后训练的基于轨迹的自蒸馏\[16 (https://arxiv.org/html/2605.11854#bib.bib97), 24 (https://arxiv.org/html/2605.11854#bib.bib87), 32 (https://arxiv.org/html/2605.11854#bib.bib9)\]，其中模型从自身生成的解码轨迹中学习。然而，现有方法主要设计用于采样步压缩。通过学习从后期扩散状态到早期状态的捷径转换，它们主要针对**推理效率**，使 DLM 能够以更少的步骤进行解码，同时保持可接受的权衡，如 Seed Diffusion\[24 (https://arxiv.org/html/2605.11854#bib.bib87)\]和 dInfer\[16 (https://arxiv.org/html/2605.11854#bib.bib97)\]所示。尽管取得了这些效率提升，最近的证据表明（参见 Zhang 等人\[32 (https://arxiv.org/html/2605.11854#bib.bib9)\]论文中的表 2），在这些轨迹上微调的 DLM 可能在完全扩散解码下遭受性能下降，即每步解码一个标记\[32 (https://arxiv.org/html/2605.11854#bib.bib9)\]。这一观察结果激发了一个更根本的问题：**除了提高推理效率外，自蒸馏轨迹能否实现真正的知识获取和性能提升？**

直观地说，由于自蒸馏轨迹是从预训练 DLM 自身的分布流形生成的，与外部构建的目标相比，它们可能提供更低的优化障碍，从而有助于在微调过程中更顺利地吸收新知识。然而，我们的初步调查显示，使用标准的负证据下界（NELBO）目标\[30 (https://arxiv.org/html/2605.11854#bib.bib4)\]在自蒸馏轨迹上简单微调 DLM 仅能带来微小的改进。为了充分利用自蒸馏轨迹来提升性能，我们提出了**基于玻尔兹曼建模的轨迹对齐优化（TABOM）**，这是一种针对 DLM 的新型后训练框架。TABOM 利用嵌入在自蒸馏轨迹中的结构化解码模式，使模型的预测分布与其实际推理时的行为保持一致。我们在理论上将这种目标行为公式化为每个标记的理想预测熵上的玻尔兹曼分布，这作为“由易到难”归纳偏置的代理，有效地缓解了训练-推理差异。

本工作的主要贡献总结如下：

- **（第3.1节 (https://arxiv.org/html/2605.11854#S3.SS1)）** 理论上，我们证明了在“由易到难”解码调度下的推理去掩码分布可以显式地建模为每个标记的理想预测熵上的玻尔兹曼分布。这使我们能够将训练-推理对齐公式化为直接的 Kullback-Leibler（KL）散度最小化问题。
- **（第3.2节 (https://arxiv.org/html/2605.11854#S3.SS2)）** 方法上，我们引入了一个基于对成排序（Pairwise Ranking）的可处理替代目标，以优化不可处理的全局 KL 散度。该机制强制局部熵梯度严格遵循“由易到难”的解码调度，平滑地将新知识转移到模型中，而不会破坏其固有的预测流形。
- **（第4.1节 (https://arxiv.org/html/2605.11854#S4.SS1)）** 实验上，我们展示了 TABOM 解决了本文确定的中央 SFT 困境：它将自蒸馏轨迹转化为显著的性能增益，同时保留了预训练模型的分布外能力，并避免了标准监督微调通常引起的灾难性遗忘。
- **（第4.2节 (https://arxiv.org/html/2605.11854#S4.SS2)）** 我们引入了轨迹判别分数（TDS）来量化模型是否在解码轨迹上保留了标记级的不确定性差异。定性（参见图2 (https://arxiv.org/html/2605.11854#S1.F2)）和定量（参见表6 (https://arxiv.org/html/2605.11854#S4.T6)）TDS 结果表明，TABOM 将熵景观重塑为朝向“由易到难”的推理偏置，而不仅仅是重用自蒸馏样本。

参见标题 **(a) MBPP**  
参见标题 **(b) GSM8K**  
**图 2: Dream 解码过程中的轨迹判别分数。** 我们计算解码时被掩码标记的预测熵的方差，在 64 个采样轨迹上取平均值，并排除每个轨迹中第一个 EOS 标记之后的步骤。较高的曲线表示沿轨迹更强的标记级不确定性判别。“Base”表示没有经过 SFT 的原始模型。

## 2 对自蒸馏轨迹的经验观察

**表 1: Dream 上的性能比较。** SFT-SD 避免了灾难性遗忘，但与 SFT-GT 相比，域内增益有限。TABOM 实现了两全其美。

为了理解自蒸馏对模型优化的影响及其在标准 SFT 范式下的局限性，我们对自蒸馏（SD）数据和离线真实（GT）数据进行了比较分析。具体而言，我们在两个领域对 Dream 模型\[31 (https://arxiv.org/html/2605.11854#bib.bib84)\]进行了实验：代码生成和数学推理。对于代码生成，我们从 Ling-Coder-SFT\[11 (https://arxiv.org/html/2605.11854#bib.bib2)\]中随机采样了 1.7 万个查询。对于数学推理，我们利用了 MixChain-Z-PRM12K\[9 (https://arxiv.org/html/2605.11854#bib.bib8)\]中的查询。对于每个查询，离线真实（GT）数据由数据集最初提供的标准问题-答案对组成。相比之下，自蒸馏（SD）数据是由基础模型使用基于熵的解码生成答案产生的，分别为代码生成和数学推理产生了约 3.8K 和约 5.1K 个有效的 SD 轨迹。

参见标题 **(a) 代码生成**  
参见标题 **(b) 数学推理**  
**图 3: 不同掩码比率下 GT 和 SD 数据的交叉熵损失比较。**

**较低的优化障碍。** 首先，我们评估模型在不同掩码比率下使用 GT 数据或 SD 数据在相同查询上的表现。如图3 (https://arxiv.org/html/2605.11854#S2.F3)所示，SD 数据的交叉熵（CE）损失始终低于 GT 数据。这表明自蒸馏使目标分布与模型的内在预测流形保持一致，提供了更平滑的优化景观并降低了优化障碍。

**困境：遗忘 vs. 边际增益。** 为了研究这种较低的障碍如何转化为下游性能，我们评估了在这些数据集上微调的模型。如表1 (https://arxiv.org/html/2605.11854#S2.T1)所示，直接在 GT 数据上微调（SFT-GT）提高了域内性能（例如，HumanEval 增加到 61.55），但在域外遭受严重的灾难性遗忘（例如，GSM8K 降至 52.33）。相反，在自蒸馏轨迹上微调（SFT-SD）有效地防止了灾难性遗忘（GSM8K 保持在 81.81）。这是合理的，因为自蒸馏轨迹是由模型自身生成的，从而保留了其固有的预测流形。然而，域内性能提升有限（HumanEval 仅为 53.66）。这一经验困境暗示了以下观察结果：尽管提供了较低的优化障碍，但在自蒸馏轨迹上进行简单微调不足以释放显著的性能增益。

## 3 自蒸馏轨迹感知玻尔兹曼建模

**动机。** 在本文中，我们认为，当前基于 NELBO 的 SFT 范式中固有的根本训练-推理差异导致了性能增益有限，尽管使用了具有较低优化障碍的自蒸馏轨迹。我们提供如下具体分析。

令 $\overline{\mathbb{X}} = \mathbb{X} \cup \{\operatorname{M}\}$ 为扩展词汇表，其中 $\operatorname{M}$ 表示吸收掩码标记。令 $\mathbf{s} \in \mathbb{X}^L$ 为提示序列，并令 $\mathbf{x}_t = \{x_t^1, \dots, x_t^N\} \in \overline{\mathbb{X}}^N$ 表示时间步 $t \in \{0, \dots, T\}$ 处的响应序列，其中 $x_t^r$ 是位置 $r$ 处的标记。该过程从完全掩码状态 $\mathbf{x}_T = \{\operatorname{M}\}^N$ 演变到从训练数据集 $\mathcal{D}$ 中采样的完全去掩码状态 $\mathbf{x}_0 \in \mathbb{X}^N$。我们使用 $\mathcal{I} = \{1, \dots, N\}$ 表示标记索引集，保留 $t, t'$ 表示解码时间步，并保留 $r, s, k$ 表示标记索引。令 $U_t \subseteq \mathcal{I}$ 和 $M_t = \mathcal{I} \setminus U_t$ 分别表示时间步 $t$ 时的去掩码和掩码索引集。为方便起见，我们使用 $\mathbf{x}_0^{U_t}$ 来表示由 $U_t$ 中的去掩码标记形成的可见上下文。

DLM 的标准训练通常使用带有均匀掩码的 NELBO 目标，其中去掩码集 $U$ 被均匀采样以作为可见上下文，并且 $M = \mathcal{I} \setminus U$ 中的掩码标记在单步中重建：

$$
\mathcal{L}_{\mathrm{NELBO}} = \mathbb{E}_{\mathbf{x}_0 \sim \mathcal{D}} \, \mathbb{E}_{U \sim q_{\mathrm{unif}}(U)} \left[ \frac{1}{|M|} \sum_{r \in M} -\log p_{\theta}(x_0^r \mid \mathbf{x}_0^U, \mathbf{s}) \right] \quad (1)
$$

###### 命题 1（NELBO 的均匀归纳偏置）

在 NELBO 目标下，模型被优化以在给定均匀采样上下文的情况下同时重建所有掩码标记。这种优化本质上灌输了一种均匀归纳偏置，即平等地对待所有标记，而不考虑其固有的预测难度，从而无法捕捉连贯生成所需的序列依赖性和确定性梯度。

命题1 (https://arxiv.org/html/2605.11854#Thmproposition1) 主要说明了现有 NELBO 目标的均匀归纳偏置。相比之下，现有 DLM 模型\[14 (https://arxiv.org/html/2605.11854#bib.bib73), 31 (https://arxiv.org/html/2605.11854#bib.bib84), 12 (https://arxiv.org/html/2605.11854#bib.bib85)\]采用的实际推理解码是一个由置信度或熵引导的多步过程。具体而言，从空去掩码集 $U_T = \emptyset$ 开始，每个反向解码步 $t$ 从掩码集 $M_t$ 中选择大小为 $b_t$ 的标记子集 $\mathcal{J}_t$ 进行去掩码，通过最小化预测熵：

$$
\mathcal{J}_t = \arg\min_{S \subset M_t, \|S\|=b_t} \sum_{r \in S} H