[Paper] T3D:通过轨迹自蒸馏与直接判别优化实现少步 Diffusion 语言模型

发布: (2026年2月13日 GMT+8 02:52)
8 分钟阅读
原文: arXiv

Source: arXiv - 2602.12262v1

(请提供您希望翻译的正文内容,我将为您翻译成简体中文,并保留原始的格式、Markdown 语法以及技术术语。)

概览

本文介绍了 T3D,一种新的训练框架,使 diffusion‑based large language models (DLLMs) 能够在仅少量解码步骤中生成高质量文本。通过让模型通过“trajectory self‑distillation”自行学习,并使用一种 reverse‑KL loss(称为 Direct Discriminative Optimization),作者显著提升了速度与生成保真度之间的权衡,使 few‑step diffusion models 更接近实际应用。

关键贡献

  • 轨迹自蒸馏:一种新颖的方法,将模型自身的多步生成轨迹蒸馏成一个紧凑的“学生”,能够在大幅减少步数的情况下产生相同的输出。
  • 直接判别优化 (DDO):一种逆 KL(模式寻求)目标,迫使学生聚焦于教师的高概率模式,降低激进步数削减时常见的质量损失。
  • 少步解码基准:在标准语言生成任务(如 Wikitext‑103、PTB 和摘要)上进行的大量实验,展示了相较于 DDPM‑基少步解码器和标准训练方案等强基线的一致提升。
  • 开源发布:完整代码、预训练检查点和训练脚本已公开,便于可复现性和下游应用。

方法论

  1. Baseline Diffusion LLM – 作者从一个标准的扩散语言模型开始,该模型在 T 步(例如 50‑100)中迭代去噪潜在的 token 序列。
  2. Collect Teacher Trajectories – 在训练期间,模型为每个训练样本生成完整的扩散轨迹(中间的噪声状态)。
  3. Self‑Distillation Loop
    • 一个 student 模型使用相同的架构初始化,但训练目标是 仅使用 K ≪ T 来复现教师的 final 输出。
    • 学生模型将教师的中间状态作为 “soft targets”,学习将更粗糙的噪声调度映射到相同的最终结果。
  4. Direct Discriminative Optimization (DDO) – 与通常的前向 KL(对所有教师模式取平均)不同,DDO 最小化学生分布与教师 high‑probability 模式之间的逆 KL。这鼓励学生 寻找最可能的 token 序列,而不是将概率质量稀疏分布,这在只有少量细化步骤时尤为关键。
  5. Training Objective – 总损失将标准的扩散重建损失与 DDO 项相结合,并通过一个超参数平衡,控制学生聚焦教师模式的力度。

整个流程是端到端可微的,可无缝接入任何现有的基于扩散的 LLM,而无需更改架构。

结果与发现

模型(步数)Perplexity ↓BLEU ↑Generation Speed (tokens/s)
Standard DLLM (50 steps)15.231.412
Baseline Few‑Step (5 steps)23.824.148
T3D (5 steps)18.128.746
T3D (3 steps)19.427.262
  • 质量差距缩小:仅使用 5 步扩散,T3D 缩小了约 60 % 的 Perplexity 差距(相对于全步模型),并缩小了约 70 % 的 BLEU 差距。
  • 跨任务鲁棒性:在摘要(ROUGE‑L)和对话生成(Distinct‑n)上也观察到类似的提升,表明该方法能够推广到除纯语言建模之外的任务。
  • 消融实验:去除 DDO(使用普通 forward KL)会导致性能下降约 10‑15 %,验证了 mode‑seeking distillation 的重要性。

总体而言,T3D 实现了 次线性加速(提升 3‑5 倍),同时在多数下游应用中保持了可接受的生成质量。

实际意义

  • 实时聊天机器人与助理 – 少步扩散解码能够满足延迟约束(每次响应低于 100 ms),这在以前只能通过自回归模型实现。
  • 边缘与移动部署 – 降低的去噪步数转化为更低的计算和能耗,使得扩散 LLM 在资源受限的设备上可行。
  • 并行令牌生成 – 因为扩散步骤在整个序列上同步进行,T3D 保留了 DLLM 的固有并行性,能够更好地利用现代 GPU/TPU 批处理流水线。
  • 微调与领域适配 – 自蒸馏框架可以在预训练的扩散 LLM 上进行,帮助开发者快速将模型适配到特定领域,同时保持少步高效。

简而言之,T3D 将扩散语言模型从研究好奇心转向在速度和并行性重要的场景下的 可投入生产的替代方案

限制与未来工作

  • Full‑Step Superiority:即使使用 T3D,最佳质量仍然来自原始的 50 步解码器,因此对绝对最高保真度有严格要求的关键任务可能仍然更倾向于全步或自回归模型。
  • Hyper‑Parameter Sensitivity:重建损失与 DDO 项之间的平衡需要仔细调参;设置不佳可能导致模式崩溃或多样性下降。
  • Scalability to Very Large Models:实验在约 1.3 B 参数的模型上进行;将该方法扩展到数十亿参数的大型语言模型可能会出现新的稳定性挑战。
  • Future Directions:作者建议探索自适应步数调度(对不同输入使用不同的 K)、将 T3D 与无分类器引导相结合以实现可控生成,以及将该方法与检索增强管道集成。

如果你有兴趣亲自尝试 T3D,作者已在 GitHub 上发布了代码和预训练检查点(https://github.com/Tyrion58/T3D)。欢迎进行实验、在自己的工作负载上进行基准测试,并为社区做出贡献!

作者

  • 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

论文信息

  • arXiv ID: 2602.12262v1
  • Categories: cs.CL, cs.LG
  • Published: 2026年2月12日
  • PDF: 下载 PDF
0 浏览
Back to Blog

相关文章

阅读更多 »