Temporal Self-Distillation: Faster Inference in Discrete Diffusion Language Models
作者: Shijian Xu, Andrea Miele, Metod Jazbec, Volker Roth, Eric Nalisnick, Ilija Bogunovic
分类: cs.LG
发布日期: 2026-09-14
💡 一句话要点
提出时间自蒸馏方法以加速离散扩散语言模型推理
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 离散扩散语言模型 时间自蒸馏 快速推理 并行解码 模型优化
📋 核心要点
- 现有的离散扩散语言模型在并行解码时性能下降严重，限制了其推理速度。
- 本文提出的时间自蒸馏（TSD）方法通过跨时间蒸馏来优化模型的预测，提升推理效率。
- 在七个基准测试中，TSD显著提升了模型的推理速度和质量，推动了低计算需求下的性能提升。
📝 摘要（中文）
离散扩散语言模型（dLLMs）通过并行生成多个标记来实现快速推理，但在过度并行解码时性能严重下降。本文提出了一种简单的在线策略方法——时间自蒸馏（TSD），通过跨时间蒸馏预测来训练dLLMs以实现快速推理。具体而言，TSD将模型在早期时间步的去噪分布蒸馏到最终时间步的分布，从而鼓励早期预测更好地预见模型的最终输出，支持更激进的并行解码。由于教师信号来自模型自身，TSD无需离线教师生成，能够无缝应用于基础和后训练策略。在数学、规划和代码等七个基准测试中，TSD显著推动了速度与质量的平衡，提供了一种简单的单阶段加速dLLMs的方法，取得了与离线蒸馏相当的加速效果，同时避免了复杂的两阶段流程。
🔬 方法详解
问题定义：本文旨在解决离散扩散语言模型在推理过程中因过度并行解码导致的性能下降问题。现有方法在并行生成多个标记时，往往无法保持较高的输出质量，影响了实际应用效果。
核心思路：提出的时间自蒸馏（TSD）方法通过将早期时间步的去噪分布蒸馏到最终时间步的分布，促使模型在早期阶段就能更好地预测最终输出，从而实现更高效的并行解码。
技术框架：TSD方法的整体架构包括两个主要模块：首先是模型的去噪过程，其次是蒸馏过程。模型在训练时会在多个时间步生成预测，并通过蒸馏机制将早期预测与最终预测进行对比和优化。
关键创新：TSD的主要创新在于其自蒸馏机制，利用模型自身的输出作为教师信号，避免了传统方法中需要离线教师生成的复杂性。这种设计使得TSD能够灵活应用于不同的模型架构。
关键设计：在实现TSD时，关键的设计包括损失函数的选择，确保早期预测与最终预测之间的距离最小化。此外，模型的网络结构也经过优化，以支持高效的并行计算和蒸馏过程。通过这些设计，TSD能够在保持模型性能的同时，显著提升推理速度。
🖼️ 关键图片
📊 实验亮点
在七个基准测试中，TSD方法显著提升了模型的推理速度和质量，推动了速度与质量的平衡，取得了与离线蒸馏相当的加速效果，展示了在低计算需求下的优越性能。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、代码生成和自动化规划等。通过提升离散扩散语言模型的推理速度，TSD能够在实时应用中提供更高效的解决方案，具有广泛的实际价值和未来影响。
📄 摘要（原文）
Diffusion language models (dLLMs) promise fast inference by generating multiple tokens in parallel, but suffer severe performance degradation when parallel decoding is pushed too aggressively. We introduce Temporal Self-Distillation (TSD), a simple on-policy method that trains dLLMs for fast inference by distilling predictions across time. Specifically, TSD distills the model's denoising distribution at earlier timesteps toward its distribution at the final timestep at which a token is committed. This encourages earlier predictions to better anticipate the model's eventual output, enabling much more aggressive parallel decoding. Because its teacher signal comes from the model itself, TSD requires no offline teacher generation and applies seamlessly to both base and post-trained policies. Across seven benchmarks in mathematics, planning, and code, TSD substantially shifts the speed--quality frontier toward the low-compute regime. TSD thus provides a simple, single-stage approach to accelerating dLLMs, achieving speedups competitive with offline distillation while avoiding a complex two-stage pipeline.