Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models
作者: Aidar Myrzakhan, Zhiqiang Shen
分类: cs.CL, cs.AI, cs.LG
发布日期: 2026-10-26
备注: Code is available at https://github.com/VILA-Lab/dSpark
🔗 代码/项目: GITHUB
💡 一句话要点
提出dSpark框架以加速扩散语言模型的生成效率
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 扩散语言模型 调度蒸馏 生成效率 自然语言处理 数学推理 代码生成
📋 核心要点
- 现有的扩散语言模型在生成过程中需要多轮解码,导致生成效率低下。
- 论文提出dSpark框架,通过调度蒸馏学习何时承诺标记,从而加速生成过程。
- 在GSM8K数据集上,dSpark在多个基线模型上实现了显著的速度提升,同时保持了高性能。
📝 摘要(中文)
扩散大型语言模型(dLLMs)能够并行生成多个标记,但通常需要多轮解码才能完成响应。生成效率不仅依赖于易于预测的标记,还与承诺这些标记如何帮助预测剩余位置有关。我们提出了dSpark,一个基于调度蒸馏的新框架,用于学习何时承诺标记以加速扩散语言生成。通过构建优先揭示那些能使其他位置更易预测的标记的解掩码调度,我们能够在解码早期减少不确定性。我们在数学推理和代码生成任务上进行了广泛评估,结果显示在GSM8K上,我们的方法在LLaDA-8B-Instruct和Dream-7B-Instruct上分别实现了10.12倍和5.12倍的速度提升,同时保持了优越的任务性能。
🔬 方法详解
问题定义:本论文旨在解决扩散大型语言模型在生成过程中多轮解码导致的效率低下问题。现有方法在选择生成标记时未能有效考虑标记之间的相互影响,导致生成速度缓慢。
核心思路:论文的核心思路是通过调度蒸馏技术,学习何时承诺生成标记,以便更早地减少解码过程中的不确定性。具体而言,优先揭示那些能使其他位置更易预测的标记,从而加速整体生成过程。
技术框架:整体架构包括两个主要模块:一是构建解掩码调度,二是将这些调度蒸馏到模型中以实现快速并行采样。调度设计中引入了KL散度正则化,以保持生成质量。
关键创新:最重要的技术创新在于调度蒸馏的引入,使得模型能够学习在何时承诺标记,从而优化生成效率。这一方法与传统的逐步解码方法有本质区别,后者未能充分利用标记之间的关系。
关键设计:在参数设置上,调度优先级是基于标记的可预测性设计的。损失函数中引入了KL散度正则化项,以确保生成质量不受影响。此外,网络结构经过优化,以支持快速并行采样。
🖼️ 关键图片
📊 实验亮点
在GSM8K数据集上,dSpark方法在LLaDA-8B-Instruct和Dream-7B-Instruct模型上分别实现了10.12倍和5.12倍的速度提升,显著优于各类基线模型,同时保持了高水平的任务性能,展示了其在实际应用中的有效性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的对话系统、代码生成以及任何需要高效文本生成的场景。通过加速生成过程,dSpark能够在实时应用中提供更好的用户体验,具有重要的实际价值和未来影响。
📄 摘要(原文)
Diffusion large language models (dLLMs) can generate multiple tokens in parallel, but often require many decoding rounds to complete a response. Efficient generation depends not only on which tokens are easy to predict, but also on how committing to them helps predict the remaining positions. A highly predictable token may provide little useful context, while revealing a less predictable token can help several other positions become predictable sooner. We propose dSpark, a new schedule distillation-based framework for learning when to commit tokens for faster diffusion language generation. To achieve this, we construct unmasking schedules that prioritize tokens whose reveal makes other positions easier to predict, reducing uncertainty earlier in decoding. We then distill these schedules into the model to enable faster parallel sampling, with a KL-divergence regularizer to help preserve generation quality. We extensively evaluate our approach across mathematical reasoning and code generation tasks. On GSM8K, our method achieves remarkable 10.12$\times$ and 5.12$\times$ speedups over the various baselines on LLaDA-8B-Instruct and Dream-7B-Instruct, respectively, while maintaining superior task performance. Our code is available at https://github.com/VILA-Lab/dSpark.