T3D: Few-Step Diffusion Language Models via Trajectory Self-Distillation with Direct Discriminative Optimization
作者: Tunyu Zhang, Xinxi Zhang, Ligong Han, Haizhou Shi, Xiaoxiao He, Zhuowei Li, Hao Wang, Kai Xu, Akash Srivastava, Hao Wang, Vladimir Pavlovic, Dimitris N. Metaxas
分类: cs.CL, cs.LG
发布日期: 2026-02-12
🔗 代码/项目: GITHUB
💡 一句话要点
提出T3D,通过轨迹自蒸馏和直接判别优化加速扩散语言模型的少步解码。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 少步解码 自蒸馏 轨迹学习 直接判别优化
📋 核心要点
- 扩散语言模型需要大量迭代步骤,导致生成速度慢,直接减少步骤数会严重影响生成质量。
- T3D通过轨迹自蒸馏,让模型学习自身高质量的生成轨迹,并结合直接判别优化,提升少步解码的性能。
- 实验表明,T3D在少步解码场景下,显著优于现有基线方法,缩小了与全步解码的差距。
📝 摘要(中文)
扩散语言模型(DLLM)有潜力通过并行解码多个token来实现快速文本生成。然而,在实践中,它们的推理效率受到大量细化步骤需求的限制,而大幅减少步骤数会导致生成质量的显著下降。为了缓解这个问题,我们提出了一个轨迹自蒸馏框架,通过蒸馏模型自身的生成轨迹来改进少步解码。我们结合了直接判别优化(DDO),这是一种反向KL目标,它促进了寻求模式的蒸馏,并鼓励学生专注于高概率的教师模式。在多个基准测试中,我们的方法始终优于强大的少步基线和严格步骤预算下的标准训练。虽然全步解码仍然更优越,但我们大大缩小了差距,为实际的少步DLLM奠定了坚实的基础。源代码可在https://github.com/Tyrion58/T3D 获取。
🔬 方法详解
问题定义:扩散语言模型(DLLM)虽然具有并行生成token的潜力,但其推理效率受限于需要大量的迭代步骤。直接减少迭代步骤会导致生成文本质量显著下降,如何在保证生成质量的前提下,减少推理步骤,加速DLLM的文本生成是本文要解决的核心问题。
核心思路:本文的核心思路是利用自蒸馏,让模型学习自身在多步迭代过程中产生的优质生成轨迹。具体来说,将多步迭代后的结果作为“教师”,指导单步或少数几步迭代的“学生”模型,从而使学生模型能够更快地生成高质量的文本。同时,采用直接判别优化(DDO)来引导蒸馏过程,鼓励学生模型专注于教师模型的高概率模式,避免陷入局部最优。
技术框架:T3D框架主要包含以下几个阶段:1) 使用标准的扩散模型训练方法训练一个teacher模型。2) 使用teacher模型生成一系列的轨迹数据,这些轨迹数据代表了从噪声到最终文本的迭代过程。3) 使用trajectory self-distillation方法,将teacher模型的轨迹数据作为监督信号,训练一个student模型,student模型的目标是尽可能地模仿teacher模型的生成轨迹。4) 在训练student模型时,使用Direct Discriminative Optimization (DDO)方法,DDO是一种reverse-KL objective,可以促进mode-seeking distillation,鼓励student模型专注于高概率的teacher模式。
关键创新:T3D的关键创新在于:1) 提出了trajectory self-distillation框架,通过蒸馏模型自身的生成轨迹来改进少步解码。2) 结合了Direct Discriminative Optimization (DDO),这是一种反向KL目标,它促进了寻求模式的蒸馏,并鼓励学生专注于高概率的教师模式。与传统的蒸馏方法不同,T3D关注的是整个生成轨迹的学习,而不仅仅是最终的生成结果。
关键设计:T3D的关键设计包括:1) 如何有效地表示和利用生成轨迹数据。2) 如何设计DDO损失函数,以促进mode-seeking distillation。3) 如何平衡teacher模型和student模型之间的差异,避免student模型过度拟合teacher模型。具体的参数设置和网络结构细节在论文中有详细描述,例如,DDO损失函数的具体形式,以及teacher和student模型的网络结构选择。
📊 实验亮点
实验结果表明,T3D在多个文本生成基准测试中,显著优于现有的少步解码方法。例如,在特定数据集上,T3D仅用少量步骤即可达到与传统扩散模型多步解码相近的生成质量,大大缩短了推理时间。与标准训练方法相比,T3D在相同步骤数下,生成质量有显著提升。
🎯 应用场景
T3D技术可应用于需要快速文本生成的场景,例如实时对话系统、机器翻译、文本摘要等。通过减少扩散模型的推理步骤,可以显著降低计算成本和延迟,提高用户体验。该研究为实际部署高效的扩散语言模型奠定了基础,并可能推动扩散模型在更多领域的应用。
📄 摘要(原文)
Diffusion large language models (DLLMs) have the potential to enable fast text generation by decoding multiple tokens in parallel. However, in practice, their inference efficiency is constrained by the need for many refinement steps, while aggressively reducing the number of steps leads to a substantial degradation in generation quality. To alleviate this, we propose a trajectory self-distillation framework that improves few-step decoding by distilling the model's own generative trajectories. We incorporate Direct Discriminative Optimization (DDO), a reverse-KL objective that promotes mode-seeking distillation and encourages the student to concentrate on high-probability teacher modes. Across benchmarks, our approach consistently outperforms strong few-step baselines and standard training under tight step budgets. Although full-step decoding remains superior, we substantially narrow the gap, establishing a strong foundation towards practical few-step DLLMs. The source code is available at https://github.com/Tyrion58/T3D.