循环扩散语言模型

掩蔽扩散模型 (MDM) 已成为语言建模自回归模型的一种有前途的替代方案，但 MDM 的 Transformer 架构的有效设计仍未得到充分探索。在本文中，我们表明，有选择地循环早中 Transformer 层可显着提高 MDM 中的训练效率和模型性能。我们将这种方法称为 LoopMDM（循环掩模扩散模型），它带来了两个主要好处：训练时的循环层无需添加参数即可产生深度缩放效果，而推理时改变循环数量可以实现灵活的计算缩放。尽管很简单，但结果却是惊人的：在多个 预训练 语料库中，LoopMDM 的性能与相同大小的 MDM 相当，但训练 FLOP 减少了 3.3 次，而其最终性能在各种推理基准上都优于它们，包括在 GSM8K 上高达 8.5 分。它甚至超过了使用可比较的每步计算训练的更深的非循环 MDM，这表明选择性循环比简单的深度缩放更有效。此外，LoopMDM 可以通过增加循环数量来扩展推理时间计算。在整个采样过程中自适应调整循环数量可以进一步提高计算效率，同时保持性能。最后，通过注意力分析，我们提供了证据，证明循环通过促进屏蔽位置之间的交互在 MDM 中是有效的。我们的代码和权重将公开发布。