Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models for language modeling, yet the effective design of transformer architectures for MDMs remains underexplored. In this paper, we show that selectively looping the early-middle transformer layers significantly improves both training efficiency and model performance in MDMs. We call this approach LoopMDM(Looped Masked Diffusion Model), which brings two key benefits: looping layers at training-time yields a depth-scaling effect without adding parameters, while varying the number of loops at inference-time enables flexible compute scaling. Despite the simplicity, the results are striking: across multiple pre-training corpora, LoopMDM matches the performance of same-size MDMs with up to 3.3 fewer training FLOPs, while its final performance outperforms them on various reasoning benchmarks, including up to 8.5 points on GSM8K. It even surpasses deeper non-looped MDMs trained with comparable per-step compute, indicating that selective looping is more effective than naive depth scaling. Furthermore, LoopMDM can scale inference-time compute by increasing the number of loops. Adaptively adjusting the number of loops throughout the sampling process further yields additional gains in compute efficiency while maintaining performance. Lastly, with attention analysis, we provide evidence that looping is effective in MDMs by promoting interactions among masked positions. Our code and weights will be publicly released.

翻译：掩码扩散模型（MDMs）已成为自回归模型在语言建模中的有力替代方案，然而针对MDMs的Transformer架构有效设计仍鲜有探索。本文证明，选择性循环使用早期-中间Transformer层可显著提升MDMs的训练效率与模型性能。我们将此方法命名为LoopMDM（循环掩码扩散模型），其带来两项关键优势：训练时循环层数可产生深度缩放效应且不增加参数，而推理时调整循环次数则能实现灵活的计算量缩放。尽管方法简洁，但效果显著：在多个预训练语料库上，LoopMDM以最高3.3倍更少的训练FLOPs实现了与同规模MDMs相当的性能，并在各类推理基准测试中超越后者，包括在GSM8K上最高提升8.5个百分点。该方法甚至超越了以可比单步计算量训练的更深层非循环MDMs，表明选择性循环比朴素深度缩放更有效。此外，LoopMDM可通过增加循环次数扩展推理阶段计算量。在采样过程中自适应调整循环次数，可在保持性能的同时进一步提升计算效率。最后，通过注意力分析，我们证明了循环通过促进掩码位置间的交互有效提升MDMs性能。我们的代码和权重将公开发布。