通过循环早期中间层,让扩散语言模型训练更高效、推理更灵活。
Looped Diffusion Language Models
- 选择性循环早期中层Transformer,不增参数实现深度扩展。
- 训练少3.3倍FLOPs,GSM8K上性能最高提升8.5分。
- 支持推理时动态调整计算量,适合资源受限场景。
掩码扩散模型(MDMs)作为语言建模的新范式,但其Transformer架构设计仍不充分。本文提出LoopMDM,通过在训练时有选择地循环早期与中期的Transformer层,显著提升训练效率与模型性能。该方法在不增加参数的前提下,实现深度扩展效应;在推理时通过调整循环次数,灵活控制计算量。在多个预训练语料上,LoopMDM以最多减少3.3倍训练FLOPs达到与同等规模MDMs相当的性能,且在各类推理基准测试中表现更优,如在GSM8K上最高提升8.5分。即使对比计算量相近的更深非循环模型,其表现也更优,表明选择性循环优于简单堆叠。此外,通过增加循环次数可弹性扩展推理计算。注意力分析进一步显示,循环有助于增强被掩码位置间的交互。代码与权重将公开。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。