arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2610.35804cs.CLcs.AIcs.LG
收藏

学习何时提交:用于快速扩散语言模型的调度蒸馏

Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models

MBZUAI(穆罕默德·本·扎耶德人工智能大学)

Aidar Myrzakhan, Zhiqiang Shen

AI总结:

针对扩散语言模型解码轮次多的问题,提出基于调度蒸馏的 dSpark 框架,构建优先揭示关键 token 的去掩蔽调度,在 GSM8K 上实现 LLaDA-8B-Instruct 10.12×、Dream-7B-Instruct 5.12× 加速且性能更优。

摘要:

扩散大语言模型(dLLMs)可并行生成多个 token,但通常需要多轮解码才能完成响应。高效生成不仅取决于哪些 token 易于预测,还取决于提交这些 token 对预测剩余位置的帮助程度:高度可预测的 token 可能提供的有用上下文很少,而揭示一个较难预测的 token 可帮助其他多个位置更快变得可预测。我们提出 dSpark,一种基于调度蒸馏的新型框架,用于学习在扩散语言生成中何时提交 token。为实现这一点,我们构建了去掩蔽调度,优先选择那些其揭示能让其他位置更易预测的 token,以在解码早期减少不确定性。随后,我们将这些调度蒸馏到模型中,以实现更快的并行采样,并使用 KL 散度正则化器帮助保持生成质量。我们在数学推理和代码生成任务上对我们的方法进行了广泛评估:在 GSM8K 上,我们的方法在 LLaDA-8B-Instruct 和 Dream-7B-Instruct 上分别实现了显著的 10.12× 和 5.12× 加速,同时保持了更优的任务性能。我们的代码可在此 https URL 获取。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑