论文精选

LoopMDM:循环层提升掩码扩散语言模型效率与性能

Looped Diffusion Language Models

精选理由

做扩散语言模型或高效 Transformer 架构的开发者值得关注——LoopMDM 用简单循环层技巧同时省训练算力、提推理性能,直接可复现。

AI 摘要

本文提出 LoopMDM(Looped Masked Diffusion Model),通过在掩码扩散模型的早期-中间层选择性循环,显著提升训练效率和模型性能。训练时循环层产生深度缩放效果而不增加参数,推理时可变循环次数实现灵活计算缩放。在多个预训练语料上,LoopMDM 匹配同尺寸 MDM 性能但节省高达 3.3 倍训练 FLOPs,在 GSM8K 等推理基准上提升最多 8.5 分,甚至超越更深层非循环模型。注意力分析表明,循环促进了掩码位置间的交互。代码和权重将开源。

原文 · arXiv cs.LG

Looped Diffusion Language Models

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.