论文:循环结构让均匀扩散语言模型训练 FLOPs 最多减少 3.3 倍
Looped Uniform Diffusion Language Models
扩散语言模型的训练成本问题有了新解法,UDLM 加上循环层结构后省了近三分之一算力,做生成模型研究的可以看看。
这篇 arXiv 论文研究了在均匀扩散语言模型(UDLM)中引入 selective looping,即在每个去噪步骤内递归复用少数早期 Transformer 层。实验显示,looped UDLM 在 OpenWebText 上以最多 2.8 倍更少的训练 FLOPs 达到非循环模型在约 0.5 万亿 token 上训练的验证 NLL,在 FineWeb-Edu 上该数字为 3.3 倍。在相同 NFE 预算下,生成困惑度最多降低 9.8%。训练时更频繁采样较小的循环次数,还能在不损失准确率的情况下将推理循环次数减少 12%。
Looped Uniform Diffusion Language Models
Uniform diffusion language models (UDLMs) generate text by iteratively denoising sequences of ordinary vocabulary tokens, any of which may be corrupted and subsequently revised. During training, UDLMs must denoise from a context that may itself be unreliable, whereas masked diffusion language models (MDLMs) always condition on a clean context. This challenging training objective makes UDLMs more FLOP-intensive to train than MDLMs, but it yields stronger downstream reasoning performance. Recently, selective looping, which recursively applies a few early Transformer layers within each denoising step, has been shown to improve training efficiency in MDLMs under an iso-FLOP budget. We investigate looping in UDLMs while keeping the training objective and sampler unchanged. We find that looping significantly improves training efficiency: looped UDLMs reach the validation NLL of a non-looped UDLM trained on about 0.5 trillion tokens with up to 2.8x fewer training FLOPs on OpenWebText and 3.3x fewer on FineWeb-Edu. We also find that, at a matched NFE budget, looped UDLMs reach up to 9.8% lower generation perplexity than the non-looped UDLM, with most gains achieved at moderate loop counts. Finally, we study different loop-count distributions and find that sampling smaller loop counts more frequently during training improves inference efficiency: on downstream tasks, this strategy uses 12% fewer loops than uniform loop-count sampling without degrading accuracy.