arXiv:2604.02340cs.LGcs.CL2026-04被引 1

通过智能调度,用小模型替换部分去噪步骤,加速文本生成。

Not All Denoising Steps Are Equal: Model Scheduling for Faster Masked Diffusion Language Models

  • 在去噪过程的早期和晚期用小模型替代大模型,保持生成质量。
  • 相比原方法,计算量减少17%,困惑度仅小幅上升。
  • 适合追求高效生成的开发者,尤其对资源受限场景友好。

近期的掩码扩散语言模型(MDLMs)缩小了与自回归语言模型的质量差距,但生成代价高昂:需多次完整序列去噪,且无法像自回归解码那样利用键值缓存。本文利用扩散框架的灵活性,研究模型调度策略,即在部分去噪步骤中用较小的MDLM替代完整模型。在OpenWebText和LM1B数据集上训练的模型表明,早期和晚期去噪步骤对替换更为鲁棒,而中间步骤最敏感。该策略可实现最高17%的浮点运算量(FLOPs)降低,同时在无条件与前缀条件生成下仅造成轻微困惑度损失,样本多样性保持不变。通过基于损失与小大模型间KL散度的逐步分析,以及对粗粒度步骤段的穷举搜索,一致验证了中间阶段的敏感性。结果表明,无需复杂架构调整,简单的调度规则即可显著加速采样并维持生成质量。

原文摘要 · Abstract (English)

Recent advances in masked diffusion language models (MDLMs) narrow the quality gap to autoregressive LMs, but their sampling remains expensive because generation requires many full-sequence denoising passes with a large Transformer and, unlike autoregressive decoding, cannot benefit from KV caching. In this work, we exploit the flexibility of the diffusion framework and study model scheduling, where a smaller MDLM replaces the full model at a subset of denoising steps. Across models trained on OpenWebText and LM1B, we show that early and late denoising steps are substantially more robust to such replacement than middle steps, enabling up to a 17% reduction in FLOPs with only modest degradation in generative perplexity under both unconditional and prefix-conditional generation, while preserving sample diversity. We support these findings with a step-importance analysis based on loss and KL divergence between small and large models across timesteps, as well as an exhaustive search over coarse step segments, both of which identify the middle of the diffusion trajectory as most sensitive consistently across datasets. Our results suggest that simple, architecture-agnostic scheduling rules can significantly accelerate MDLM sampling while largely preserving generation quality.

扩散模型文本生成加速推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。