Few-Step Diffusion Language Models via Trajectory Self-Distillation

会议: NeurIPS 2026
arXiv: 2602.12262
代码: https://github.com/Tyrion58/T3D
领域: 模型压缩
关键词: 扩散大语言模型, 自蒸馏, 轨迹蒸馏, 标记分解误差, 直接判别优化

一句话总结

针对掩码扩散大语言模型在步数压缩时因标记独立分解假设导致的质量塌缩，本文提出轨迹自蒸馏框架 T3D，通过拟合教师完整去噪轨迹并结合基于反向 KL 的直接判别优化（DDO）与早期标记路径一致性正则，在极少步数下实现了高保真推理并大幅缩小与全步数扩散解码的性能差距。

研究背景与动机

扩散大语言模型（DLLM）凭借并行动作生成与非自回归特质，展现出极高吞吐潜能，但其实际推理受制于漫长的反向扩散采样链。当前主流模型通常需要数十甚至上百步去噪迭代，导致计算开销巨大；而一旦直接大幅减少步数（如将每步生成的标记数成倍提高），生成文本的语法连贯性与复杂逻辑推理能力便会急剧劣化。

这种性能崩溃的核心根源在于掩码扩散语言模型（MDLM）的平均场参数化假设。在反向去噪时，模型为了处理离散标记将多标记联合转移概率强制近似为独立因式分解的乘积。当推理步数充足时，每一步仅需填补极少量的掩码标记，跨标记间的依赖主要由细粒度的迭代步骤逐步捕获；而一旦进行大跨度的少步数解码（Few-Step Jump），单步内需要同时确定的大量标记之间存在极强依赖，强制因式分解会导致严重的条件全相关（Conditional Total Correlation, TC）爆炸与标记分解误差（Token Factorization Error）。此前尝试加速扩散模型的端点自蒸馏方案（如基于修正流思想的 ReDi）在离散掩码设定下彻底失效：由于初始完全掩码的先验分布是退化的确定性分布，终点监督下的条件全相关退化为数据分布本身的常数，根本无法压缩中间转移的分解误差。

为了打破这一理论与经验瓶颈，本文主张必须从终点监督迈向轨迹监督。核心 idea：通过利用教师模型完整反向轨迹中的成对中间去噪状态训练少步数学生模型，结合直接判别优化（DDO）诱导向教师高概率模式聚集的反向 KL 效应，并在路径一致性中对早期解码标记实施重点加权，从根本上压低大步长去噪的条件全相关并消除误差累积。

方法详解

整体框架

T3D 框架由离线教师轨迹收集与学生模型轨迹蒸馏两部分构成。输入是预训练教师模型在任务数据上的多步去噪推演轨迹，输出是能够在极少步数（如每步生成 2 到 8 个 token）下直接完成高质量去噪的学生模型。教师模型在静态解码配合低置信度重新掩码（Remasking）下生成完整去噪轨迹，记录每个标记解码的时间步与中间掩码状态 \(({\mathbf{x}}_0, {\mathbf{x}}_t)\)。少步数学生模型初始化自该教师模型，通过直接判别优化目标（DDO）以及路径一致性正则化直接学习教师轨迹的跨步转移映射。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["预训练教师模型<br/>全步数反向去噪推演"] --> B["轨迹收集与状态对提取<br/>记录解码顺序与中间状态对"]
    B --> C["轨迹级自蒸馏<br/>中间转移匹配降低条件全相关"]
    C --> D["直接判别优化 DDO<br/>隐式判别器诱导模式寻求行为"]
    D --> E["路径一致性正则化<br/>早期解码标记线性衰减加权"]
    E --> F["少步数扩散语言模型<br/>跨步并行跳跃推理"]

关键设计

1. 轨迹自蒸馏：在非平庸去噪区间压低条件全相关

针对前序工作仅监督起始掩码到最终终点（\(\mathbf{x}_T \to \mathbf{x}_0\)）导致无法在掩码先验下降低因式分解误差的理论缺陷，本文直接从教师反向采样轨迹中提取成对状态 \((\mathbf{x}_0, \mathbf{x}_t)\) 进行蒸馏。在掩码语言扩散模型中，由于全掩码先验 \(p(\mathbf{x}_T) = \delta_{\mathbf{m}}\) 是确定性的常数分布，终点蒸馏所优化的条件全相关 \(TC_J(\mathbf{x}_0 \mid \mathbf{x}_T) = \mathrm{KL}(q(\mathbf{x}_0) \,\|\, \prod_{i=1}^L q(\mathbf{x}_0^i))\) 恒等于真实数据分布的内在相关性，无法被参数优化削减。相比之下，轨迹自蒸馏作用于中间扩散时间 \(t \in (0, 1)\)，利用勾股定理（Pythagorean Inequality for KL Divergence）在理论上严格证明了：

\[\mathbb{E}_t \left[ TC_{J_{\theta^*}}(\mathbf{x}_s \mid \mathbf{x}_t) \right] \leq \mathbb{E}_t \left[ TC_{J_{\phi}}(\mathbf{x}_s \mid \mathbf{x}_t) \right]\]

这表明将学生模型与教师在轨迹各个真实中间状态对齐，能系统性降低大跨度反向转移下的标记因式分解误差，使得学生模型在少步数下能够稳健跳过中间步骤而不会引发预测分布的崩塌。

2. 基于直接判别优化（DDO）的模式寻求轨迹对齐

标准前向 KL 散度目标具有模式覆盖（Mode-covering）特性，在离散语言建模的多步压缩中极易导致平均化平滑预测，缺乏对复杂推理分支的判决敏锐度。T3D 引入无需额外判别器网络的直接判别优化机制，通过冻结参考模型 \(p_{\theta_{\text{ref}}}\) 构建隐式似然比判别器，诱导具有反向 KL 特性的模式寻求（Mode-seeking）行为。单步 DDO 优化目标形式为：

\[l(\theta) = -\log \sigma\left( \mathbb{E}_{\mathbf{x}_0 \sim p_{\phi}^{\text{Tra}}(\mathbf{x}_0 \mid \mathbf{x}_t)}\left[\log \frac{p_{\theta}(\mathbf{x}_0 \mid \mathbf{x}_t)}{p_{\theta_{\text{ref}}}(\mathbf{x}_0 \mid \mathbf{x}_t)}\right] \right) - \log\left(1 - \sigma\left( \mathbb{E}_{\mathbf{x}_0 \sim p_{\theta_{\text{ref}}}(\mathbf{x}_0 \mid \mathbf{x}_t)}\left[\log \frac{p_{\theta}(\mathbf{x}_0 \mid \mathbf{x}_t)}{p_{\theta_{\text{ref}}}(\mathbf{x}_0 \mid \mathbf{x}_t)}\right] \right)\right)\]

该目标在第一项鼓励学生赋予教师真实生成样本高于参考模型的似然，在第二项惩罚参考模型生成伪样本的高估。这一机制赋予了解码过程阶段性的探索-利用特性：在全掩码初始阶段保持较高熵值进行广泛分支探索，在去噪后半程迅速降低预测熵，精准收敛到教师模型的高置信模式，极大增强了少步数推理的确定性与代码/数学公式生成的结构完整性。

3. 早期标记加权的路径一致性正则化

在少步数扩散解码中，每一步预测的标记都将作为后续去噪的硬条件，因此越早被确定的标记一旦发生偏差，其错误将随着残余掩码的逐步填充而发生不可逆的滚雪球式级联扩散。T3D 设计了显式区分标记解码时序的路径正则项：若标记 \(\mathbf{x}_0^i\) 在整个推理预算 \(B\) 中于第 \(\pi_i \in [B]\) 步被揭开，则赋予其与时序反比的加权系数 \(w_i = \frac{B - \pi_i + 1}{B}\)：

\[\mathcal{L}_{\text{path}}(\theta) = -\mathbb{E}_{p_{\phi}(\mathbf{x}_t)} \mathbb{E}_{\mathbf{x}_0 \sim p_{\phi}(\cdot \mid \mathbf{x}_t)}\left[\sum_{i=1}^L w_i \log p_{\theta}(\mathbf{x}_0^i \mid \mathbf{x}_t^{(i)})\right]\]

与均匀加权或着重监督后序标记的策略相比，这种线性衰减加权将梯度的重心精准倾斜向推理路径的前沿决策点，为整个多步生成链奠定稳固的基础语义锚点。

损失函数 / 训练策略

T3D 的总训练目标为联合优化轨迹 DDO 损失与路径一致性正则化项：

\[\mathcal{L}_{\text{T3D}}(\theta) = \mathcal{L}_{\text{traj-DDO}}(\theta) + \lambda \mathcal{L}_{\text{path}}(\theta)\]

其中正则化系数 \(\lambda\) 经消融验证设定为 \(0.2\)。训练时学生模型从预训练教师初始化，参考模型 \(p_{\theta_{\text{ref}}}\) 随训练推进每 10 个全局步使用学生的最新参数进行硬更新（Stop-Gradient），构成多轮自我博弈式蒸馏。此外，为了增强对任意去噪中间状态的泛化能力并防止学生过度拟合教师的固定轨迹伪影，训练时以概率 \(p_{\text{rand}}\) 将部分掩码替换为词表中随机采样的 token，显著提升了蒸馏鲁棒性。

实验关键数据

主实验

在代表性的块级扩散模型（SDAR-1.7B-Chat、SDAR-4B-Chat）与全局全扩散模型（LLaDA-8B-Instruct）上评估推理能力。下表汇报 SDAR-4B-Chat 在块大小 8 下不同每步生成标记数（TokPS）的主实验结果：

模型配置 / TokPS	方法	MATH500	GSM8K	MBPP	HumanEval	平均准确率	相对基线提升
TokPS = 2 (BS=8)	原始模型	49.60	72.33	33.40	46.95	50.57	-
TokPS = 2 (BS=8)	SFT	54.44	77.41	25.60	29.88	46.83	-7.40%
TokPS = 2 (BS=8)	ReDi	23.60	71.87	19.20	23.17	34.46	-31.86%
TokPS = 2 (BS=8)	dParallel	51.20	75.97	18.20	28.66	43.51	-13.96%
TokPS = 2 (BS=8)	T3D (本文)	61.60	81.96	37.00	56.10	59.17	+17.01%
TokPS = 4 (BS=8)	原始模型	16.80	41.02	10.00	16.46	21.07	-
TokPS = 4 (BS=8)	SFT	40.20	55.42	8.80	11.59	29.00	+37.64%
TokPS = 4 (BS=8)	ReDi	20.20	47.84	6.80	6.71	20.39	-3.23%
TokPS = 4 (BS=8)	dParallel	40.80	53.83	9.60	20.12	31.09	+47.56%
TokPS = 4 (BS=8)	T3D (本文)	44.80	63.99	21.20	23.17	38.29	+81.73%

在 LLaDA-8B-Instruct 极端激进的 TokPS = 8（8 个 token 每步）设置下，T3D 在 MATH500 达到 25.20%，GSM8K 达到 70.13%，平均准确率达 47.13%，相比原始模型的 17.11% 提升了 +175.47%，显著击败 CDLM（27.43%）与 dParallel（36.58%）。

消融实验

在 SDAR-4B-Chat 上使用 TokPS = 4、块大小 8 评估各模块贡献（四项基准平均准确率）：

配置	核心机制	平均准确率 (%)	相比原始模型提升
原始模型	无微调直接步数压缩	16.80	-
(a) + 轨迹蒸馏 (TD)	仅采用前向 KL 拟合轨迹对 \((\mathbf{x}_0, \mathbf{x}_t)\)	38.80	+130.95%
(b) + 直接判别优化 (DDO)	替换为隐式似然比 DDO 目标	43.20	+157.14%
(c) + 路径一致性正则 (T3D 全量)	引入早期解码标记线性衰减加权	45.00	+167.86%

关键发现

• 轨迹级监督是性能飞跃的核心主导因素：从原始模型的 16.80% 跃升至 TD 的 38.80%（绝对提升 +22.00%），证明只有学习教师反向过程中的非平庸中间流形，才能真正消除少步数跳转导致的条件全相关爆炸。
• 全步数扩散性能完全保留（Zero Forgetting）：将少步数微调后的模型切回每步 1 token 的全步数传统去噪时，dParallel 几乎完全崩塌（MATH500 仅剩 0.40%），而 T3D 在 SDAR-4B-Chat 保持 70.00% 准确率，甚至超越了原始未经蒸馏的模型（68.00%）。
• 与动态自适应解码无缝协同：在 GSM8K 上施加置信度阈值 0.9 的动态解码时，T3D 比原始模型吞吐量提升 1.45 倍，同时准确率提高 10.8%，每步平均生成的 token 数明显增多，体现出更加自信的概率聚集特性。

亮点与洞察

• 离散扩散蒸馏的理论死穴被揭示：严谨指出并证明了在离散掩码模型中因全掩码先验确定性导致的“终点监督无效性定理”，从理论上终结了直接将连续扩散领域的 Rectified Flow 蒸馏硬搬到掩码语言建模的尝试。
• 反向 KL 模式寻求与探索-利用权衡：利用 DDO 赋予离散扩散反向 KL 行为，实现在扩散早期保持宽广搜索（高熵）而在后期强力收敛（低熵），有效解决了标准前向 KL 容易导致的平滑模式崩溃与逻辑推理不坚定问题。
• 早期标记敏感性洞察：敏锐捕捉到离散扩散非自回归解码中的时间不对称性——早解码标记的误差传播破坏力远超末尾标记，提出的时间衰减加权设计兼具实现简炼与性能增益。

局限与展望

• 教师天花板与离线滚雪球开销：作为自蒸馏方法，学生性能不可避免地受制于教师模型的上限；此外教师全步数轨迹的生成与采样存储构成了额外的离线预计算负担。
• 通用开放长文本评估相对有限：实验重心偏向逻辑密集的数学与代码基准，在开放式发散创意写作与极长上下文对话中的少步数语言平滑性仍需更广泛的定性定量检验。
• 未来改进路径：探索结合多步数动态混合蒸馏（Multi-Step Joint Trajectory Matching），或进一步与轻量级在线强化学习（RL）结合，突破教师先验的固有瓶颈。

相关工作与启发

• vs ReDi：ReDi 将修正流概念简单套用至离散扩散，仅依靠 \((\mathbf{x}_T, \mathbf{x}_0)\) 终点对进行流拉直；T3D 证明了其在掩码先验下无法压缩条件全相关，转而在整个反向轨迹状态对上实施蒸馏，彻底解决了 ReDi 在大幅压缩步数时的严重掉点。
• vs dParallel：dParallel 通过监督全掩码到目标序列的最大似然加速并行解码，但破坏了扩散模型的精细去噪能力，导致在切回全步数时性能彻底崩溃（Math 仅 0.4%）；T3D 兼顾极少步数与全步数无损切换。
• vs Continuous Consistency Models：不同于连续高斯流形的确定性 ODE 一致性映射，T3D 针对离散标记转移的因式分解瓶颈，建立了贴合掩码语言特性的轨迹自蒸馏新范式。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 首次从条件全相关理论上解释并解决了离散掩码扩散少步数蒸馏的失效瓶颈，设计了结合 DDO 的优雅方案。
• 实验充分度: ⭐⭐⭐⭐⭐ 涵盖块级与全扩散两大主流架构，多基准评测、全步数复原、动态解码测试与细致消融实验一应俱全。
• 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰有力，从瓶颈定位、动机剖析到方法落地的逻辑链条紧凑完整。
• 价值: ⭐⭐⭐⭐⭐ 显著推动了扩散大语言模型从理论优势走向实用化毫秒级并行推理，工程与学术价值显著。

相关论文

• [NeurIPS 2026] Continuous-Time Distribution Matching for Few-Step Diffusion Distillation
• [NeurIPS 2026] Procedural Memory Distillation: Online Reflection for Self-Improving Language Models
• [NeurIPS 2026] Rethinking On-Policy Self-Distillation for Thinking Models
• [NeurIPS 2026] SHRED: Retain-Set-Free Unlearning via Self-Distillation with Logit Demotion
• [ICLR 2026] π-Flow: Policy-Based Few-Step Generation via Imitation Distillation