Looping technique speeds training and improves efficiency in diffusion language models
Looped Uniform Diffusion Language Models
Computation and LanguageMachine Learning
Summary
Training some advanced language models that fix mistakes in generated text can be very slow because they have to learn from possibly messy contexts. The authors studied a method called looping, where certain parts of the model are repeated multiple times during training steps. They found that looping helped these models learn faster and generate better text with fewer computing resources. They also discovered ways to choose how often to loop during training that make the models work better and more efficiently when actually writing sentences.
What this means in practice
- •For natural language processing engineers: Train diffusion-based language models faster and more efficiently by using looped early transformer layers without changing the model’s output quality.
- •For cloud ai service operators: Reduce the computational cost of serving and training large diffusion-based language models while maintaining or improving text generation quality.$Commercial implications: Enables more cost-effective deployment of advanced language models in commercial AI cloud platforms by lowering necessary compute resources.
Authors
Saiteja Utpala, Chieh-Hsin Lai, Yuki Mitsufuji, Praneeth Vepakomma
Abstract
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.