论文

并非所有去噪步骤都是相同的:更快的掩蔽扩散语言模型的模型调度

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

模型推理推理加速

摘要

掩码扩散语言模型 (MDLM) 的最新进展缩小了与自回归 LM 的质量差距,但它们的采样仍然昂贵,因为生成需要使用大 Transformer 进行许多全序列去噪传递,并且与自回归解码不同,无法从 KV 缓存中受益。在这项工作中,我们利用扩散框架的灵活性并研究模型调度,其中较小的 MDLM 在去噪步骤的子集中替换完整模型。在 OpenWebText 和 LM1B 上训练的模型中,我们发现早期和晚期去噪步骤对于这种替换来说比中间步骤更加稳健,使得 FLOP 减少了 17%,并且在无条件和前缀条件生成下生成困惑度仅略有下降,同时保留了样本多样性。我们通过基于跨时间步长的小模型和大模型之间的损失和 KL 散度的步长重要性分析以及对粗步长段的详尽搜索来支持这些发现,这两者都将扩散轨迹的中间确定为跨数据集一致最敏感的。我们的结果表明,简单的、与架构无关的调度规则可以显着加速 MDLM 采样,同时在很大程度上保持生成质量。