Looped Uniform Diffusion Language Models

Looped Uniform Diffusion Language Models. By Saiteja Utpala and three others, from MBZUAI, Sony Group Corporation, and one other institution. Published on the twenty-first of October. No code repository is linked to it.

Uniform diffusion language models, or UDLMs, generate text by repeatedly denoising sequences of ordinary vocabulary tokens, any of which can be corrupted and then revised. During training they must denoise from a context that may itself be unreliable. Masked diffusion language models, MDLMs, always condition on a clean context instead. That harder objective makes UDLMs more FLOP-intensive to train, but it yields stronger downstream reasoning. Separately, selective looping, which reapplies a few early Transformer layers recursively inside each denoising step, has been shown to improve training efficiency for MDLMs at a fixed FLOP budget. The claim here is that looping helps UDLMs too, and substantially, with the training objective and the sampler left exactly as they were.

Training efficiency is the first result. A non-looped UDLM is trained on about half a trillion tokens, and looped UDLMs reach its validation NLL with up to two point eight times fewer training FLOPs on OpenWebText. On FineWeb-Edu the saving is up to three point three times.

Generation improves as well. At a matched NFE budget, looped UDLMs reach up to nine point eight percent lower generation perplexity than the non-looped model, and most of that gain arrives at moderate loop counts.

The last result concerns how loop counts are chosen during training. Sampling smaller loop counts more often improves inference efficiency. On downstream tasks this strategy uses twelve percent fewer loops than sampling loop counts uniformly, without degrading accuracy.

The bounds matter here. Only looping is changed: the objective and the sampler are held fixed, so these results say nothing about combining looping with other modifications. The efficiency comparison is made against a non-looped UDLM, not against MDLMs, and the reference point is one trained on roughly half a trillion tokens, measured on OpenWebText and FineWeb-Edu. The savings are stated as upper bounds, up to a given factor, and the perplexity gains concentrate at moderate loop counts rather than growing without limit. The loop-count result is reported for downstream tasks against one alternative, uniform sampling.

UDLMs had a known trade: stronger reasoning, at a higher training cost. This paper shows that a technique already proven for masked models carries over to the uniform kind, and narrows that cost without touching what made UDLMs attractive in the first place.