少步无限掩模扩散 蒸馏

掩蔽扩散模型 (MDM) 已成为语言建模中自回归模型的一种有前景的替代方案，在简单而有效的框架内提供并行解码和双向上下文处理的优势。具体来说，它们在屏蔽标记和数据之间的明确区别是其简单框架和有效条件生成的基础。然而，由于同时词元更新引起的分解错误，MDM 通常需要多次采样迭代。我们观察到存在因式分解误差的理论下限，标准 MDM 由于使用确定性单状态掩码而无法降低该下限。在本文中，我们提出了无限掩模扩散模型（IMDM），它引入了随机无限状态掩模来减轻理论界限，同时直接继承 MDM 的优点，包括与预训练权重的兼容性。我们凭经验证明，由于分解误差界限，即使在简单的合成任务中，MDM 也无法执行几步生成，而 IMDM 可以为同一任务找到有效的解决方案。最后，当配备适当的 蒸馏 方法时，IMDM 在 LM1B 和 OpenWebText 上以小步数超越了现有的少步 蒸馏 方法。代码可在 https://Ugness.github.io/official_imdm 获取。