循环Uniform Diffusion语言模型：训练FLOPs降低3.3倍

当扩散模型学会“循环”：训练效率的新杠杆

扩散语言模型（Diffusion Language Models）正成为文本生成领域的一条新路径。与自回归模型逐词生成不同，扩散模型通过迭代去噪来生成整段文本。其中，Uniform Diffusion Language Models（UDLMs）对序列中任意位置的词元进行加噪和去噪，训练目标更具挑战性，但也带来了更强的下游推理能力。最近一篇关于循环UDLM的研究表明，通过递归复用早期Transformer层，可以在不改变训练目标和采样器的前提下，将训练效率提升数倍。

UDLM vs. MDLM：训练更难，推理更强

UDLM在训练时需要从可能被污染的上下文中去噪，而Masked Diffusion Language Models（MDLMs）始终基于干净上下文进行预测。这种更困难的训练目标使UDLM的FLOP消耗高于MDLM，但研究明确指出，UDLM在下游推理任务上表现更强。这形成了一个有趣的权衡：更高的训练成本换来了更好的推理性能。

扩散语言模型（Diffusion Language Models） 是一类借鉴了图像生成中扩散模型思想的文本生成模型。其核心过程包括两个阶段：前向加噪过程逐步向干净文本序列中注入噪声（如随机替换或掩码词元），反向去噪过程则训练一个神经网络从带噪序列中逐步恢复出原始文本。与自回归模型从左到右逐词生成不同，扩散语言模型可以并行地迭代精炼整个序列，从而在生成时具备全局视野。根据加噪方式的不同，衍生出了多种变体，例如Masked Diffusion Language Models（MDLMs）采用掩码式加噪，而Uniform Diffusion Language Models（UDLMs）则对序列中任意位置的词元进行均匀加噪，这使得UDLM的去噪任务更难，但也可能带来更强的推理能力。

循环机制：在MDLM中已验证，UDLM中效果如何？

选择性循环（selective looping）此前在MDLM中被证明能在等FLOP预算下提升训练效率——具体做法是在每个去噪步骤中递归应用少数几个早期Transformer层。该研究将这一思路迁移到UDLM，保持训练目标和采样器不变，仅引入循环结构。

选择性循环（selective looping） 是一种模型架构技术，其灵感来源于对Transformer层计算冗余的观察。具体而言，在去噪过程的每个步骤中，模型并非使用全部Transformer层，而是递归地多次应用少数几个早期层（例如前几层）。这种设计允许模型在等FLOP预算下进行更深的计算，因为循环复用早期层比使用全部层计算成本更低。在MDLM中，选择性循环已被证明能提升训练效率，即在相同计算量下达到更低的验证损失。该技术的关键优势在于其与训练目标解耦，可以作为一种纯架构修改独立应用，而不需要改变扩散模型的训练目标或采样算法。

循环UDLM的效率提升：数据说话

训练FLOPs最多减少3.3倍

实验结果显示，循环UDLM在达到非循环UDLM相同的验证NLL时，训练FLOPs显著降低。在OpenWebText数据集上，循环UDLM仅需约非循环模型0.5万亿token训练量对应的FLOPs的1/2.8；在FineWeb-Edu上，这一比例达到1/3.3。这意味着在相同计算预算下，循环结构可以让模型看到更多数据或训练更充分。

生成困惑度降低9.8%

在匹配的NFE（Number of Function Evaluations，函数评估次数）预算下，循环UDLM的生成困惑度比非循环版本最多低9.8%。说个细节，大部分收益在中等循环次数时已经实现，继续增加循环次数带来的边际收益递减。这为实际部署提供了参考：无需追求极高的循环次数。

NFE（Number of Function Evaluations，函数评估次数） 是衡量扩散模型推理成本的一个关键指标。在扩散模型的生成过程中，每个去噪步骤都需要调用一次神经网络（即一次函数评估），因此总NFE等于去噪步骤数乘以每步中的网络前向传播次数。在循环UDLM中，由于每个去噪步骤内可能包含多次循环（即多次复用早期层），NFE会相应增加。因此，在比较不同模型的生成质量时，必须在匹配的NFE预算下进行，以确保计算成本相当。生成困惑度（Generative Perplexity）则是评估生成文本质量的常用指标，值越低表示模型对生成文本的预测越自信、质量越高。

循环次数分布：不均匀采样提升推理效率

研究还探索了不同的循环次数分布。一个有趣的发现是：训练时更频繁地采样较小的循环次数，可以在不损失下游任务准确率的前提下，将推理时的循环次数减少12%（相比均匀循环次数采样）。这意味着训练策略的微调可以直接转化为推理成本的降低。

对从业者的启示

对于关注扩散语言模型的研究者和工程师，这项工作的价值在于：循环机制是一种与训练目标解耦的架构修改，可以独立应用；训练时的循环次数分布设计是一个被低估的效率杠杆；UDLM在推理能力上的优势结合循环带来的效率提升，可能使其在实际应用中更具竞争力。

多源交叉验证

本文核心数据来自单一研究来源（rss feed），暂未发现其他独立来源的重复验证。建议读者关注后续复现研究。

总结

循环UDLM在保持训练目标和采样器不变的前提下，通过递归复用早期Transformer层，实现了训练FLOPs最多3.3倍的降低，同时在匹配NFE预算下生成困惑度降低9.8%。训练时偏向小循环次数的采样策略还能进一步减少12%的推理循环次数。这些发现为扩散语言模型的效率优化提供了新思路。