T3D：通过轨迹自蒸馏与直接判别优化实现少步扩散语言模型 T3D: Few-Step Diffusion Language Models via Trajectory Self-Distillation with Direct Discriminative Optimization

T3D 的训练分为三个清晰的步骤。第一步是轨迹收集：用预训练的教师模型 $p_\phi$ 在训练集上进行全步解码（例如 block size 4、4 步/block 的静态解码），记录每个 token 的解码顺序，然后根据这个顺序构造中间状态对 $(x_0, x_t)$。具体来说，给定最终生成的干净序列 $x_0$ 和其解码顺序，在时间步 $t$ 处将尚未解码的 token 替换为掩码 token，得到对应的含噪序列 $x_t$。第二步是轨迹蒸馏训练：学生模型 $p_\theta$ 从教师模型初始化，使用 DDO 目标函数进行训练。DDO 损失 $\mathcal{L}_{\text{traj-DDO}}$ 包含两项——第一项鼓励学生对教师生成的样本赋予比参考模型 $p_{\theta_{ref}}$ 更高的似然，第二项惩罚学生对参考模型样本的过度估计。参考模型每 10 个全局步从当前学生更新一次（多轮自博弈更新）。同时计算路径一致性正则化损失 $\mathcal{L}_{\text{path}} = -\mathbb{E}\left[\sum_i w_i \log p_\theta(x_0^i | x_t^{(i)})\right]$，其中权重 $w_i = \frac{B - \pi_i + 1}{B}$，$\pi_i$ 是 token $x_0^i$ 的解码步骤，越早解码的 token 权重越大。第三步是推理：训练完成后，学生模型可以直接用少步解码（例如 2 步/block）进行推理，跳过教师需要的中间步骤。此外，训练时还混入随机 token 以提高鲁棒性（Mixture of Random Tokens 技术）。

T3D 的技术新颖性体现在三个维度。第一，它首次为 MDLM 建立了轨迹蒸馏的严格理论基础，证明了轨迹监督 vs 端点监督的本质区别不是经验性的而是理论性的——端点监督在 MDLM 中存在不可逾越的天花板（Theorem 4.1 和 Corollary 4.2）。第二，它将 DDO 引入扩散语言模型的蒸馏场景，这是一种全新的目标函数选择。与 ReDi 使用的前向 KL 不同，DDO 产生反向 KL 风格的模式寻求行为，论文用 Figure 4 的熵分析实验展示了这种行为差异：DDO 在全掩码阶段保持探索性，在后续阶段实现锐利精炼。第三，路径一致性正则化是一个轻量但有效的设计，它基于一个直觉观察——在压缩的解码过程中，早期 token 的预测错误会传播到后续步骤，因此应该给它们更大的训练权重。实验表明，这种早 token 加权策略一致优于均匀加权和晚 token 加权（Figure 5）。这三者结合使 T3D 在方法论上与 ReDi、dParallel 等前工作形成了清晰的区分。

从独立分析的角度看，T3D 存在几个值得改进的弱点。首先，DDO 的参考模型更新策略（每 10 步更新一次）是一个较为粗糙的启发式设计，在训练早期参考模型尚未稳定时可能导致训练不稳定。一个改进方向是采用指数移动平均（EMA）或渐进式更新策略，使参考模型的变化更平滑。其次，路径一致性正则化的权重设计 $w_i = \frac{B - \pi_i + 1}{B}$ 假设所有早期 token 同等重要，但在实际推理中，不同位置的 token 对最终答案的影响差异很大（例如数学推理中最终答案 token 比中间推导 token 更关键）。可以考虑基于 token 的信息量或梯度敏感度自适应调整权重。第三，T3D 的轨迹收集使用静态解码策略，但训练后的评估覆盖了动态解码场景，两者之间存在分布偏移。如果在轨迹收集阶段也引入动态解码或多样化的解码策略，可能进一步提升动态解码场景下的表现。第四，论文没有探索混合蒸馏策略——例如先用端点蒸馏快速初始化，再用轨迹蒸馏精细调优，这种两阶段方案可能在训练效率和最终性能之间取得更好的平衡。

论文的成果为多个研究方向打开了大门。首先，作者提出的一个自然延伸是将 T3D 应用于更大的模型规模（如 70B 参数级别），验证轨迹自蒸馏在模型规模扩展时是否仍能保持类似的相对收益。其次，可以探索 T3D 与系统级优化（如 KV 缓存、Tensor Parallel）的结合，因为论文的贡献是正交于这些技术的。第三，DDO 的模式寻求特性使其天然适合推理任务，但在开放域生成任务中，可能需要开发自适应的模式覆盖/模式寻求切换机制。第四，论文的理论分析（Theorem 4.1）为研究「什么样的轨迹监督最有效」提供了框架，未来可以探索更细粒度的轨迹采样策略——例如对因子化误差大的区间给予更多监督。第五，T3D 的自蒸馏范式可以扩展到多轮蒸馏（iterative distillation），即用 T3D 的输出作为新的教师进行下一轮蒸馏，逐步压缩到更极端的少步设置。最后，将轨迹蒸馏思想推广到其他离散扩散模型（如图像离散化、音频 token 化）也是一个有前景的方向。

论文承诺在正式接收后开源代码（GitHub 仓库已创建：https://github.com/Tyrion58/T3D），但目前提交时代码尚未公开。复现所需的关键资源包括：（1）预训练模型：SDAR-1.7B-Chat、SDAR-4B-Chat 和 LLaDA-8B-Instruct，这些都是已开源的模型；（2）训练数据：MATH 数据集（数学推理）和 PrimeIntellect 数据集（代码生成），均为公开数据集；（3）硬件需求：8 块 NVIDIA A100-40GB GPU，这是目前大多数研究团队可获得的资源；（4）推理加速：JetEngine 推理框架已开源。总训练成本约 9.5 小时（轨迹收集 1.5 小时 + 训练 8 小时），在 8 块 A100 上属于中等开销。复现难度中等——核心算法简洁（Algorithm 1 仅 7 行），但轨迹收集和 DDO 的多轮自博弈更新需要仔细实现。建议复现时先在 SDAR-1.7B-Chat + MATH500 上验证，因为这是计算开销最小的设置。