Discrete average generator speeds up diffusion language models
Acceleration of Diffusion Language Model through Discrete Average Generator
Artificial IntelligenceMachine Learning
Summary
Generating text with diffusion language models can be slow and inefficient when using only a few steps. The authors developed a new method called the Discrete Average Generator that speeds up the process without losing accuracy. This approach uses mathematical principles from continuous-time Markov chains to better approximate how text changes over time. Their method shows improved performance on text and image data, making generation much faster.
What this means in practice
- •For language model engineers: Generate text with diffusion models much faster while maintaining accuracy in language applications.
- •For image generation teams: Produce images efficiently using diffusion methods with faster sampling, demonstrated on ImageNet datasets.
Authors
Yidong Ouyang, Zhengyan Wan, Themis Haris, Tian Tan, Liqian Peng, Henry Li, Ziqian Lin, Jianhang Chen, Maryam Karimzadehgan, Alec Go, George Michailidis
Abstract
Discrete diffusion models and flow matching have emerged as powerful frameworks for generative modeling over discrete state spaces, yet efficient few-step generation remains a fundamental challenge. In this work, we introduce the Discrete Average Generator, a principled extension of MeanFlow to Continuous-Time Markov Chains (CTMCs). Analogously to how MeanFlow defines an average velocity field over a time interval in continuous spaces, we define an average generator as the normalized increment of the transition kernel over a time interval. We show that this average generator satisfies a self-consistency identity, which provides the foundation for our training objective. We further develop training strategies that align with the standard training paradigm of diffusion language models while keeping the resulting objective tractable. When projected onto per-coordinate marginals, the self-consistency identity admits a closed-form expression, enabling efficient training and inference. In Potts model simulations, our objective reduces the total variation distance of the $K$-step sampler by up to 67%. On OpenWebText, our method achieves the lowest generative perplexity among the evaluated methods for 8 to 64 sampling steps while enabling a $16\times$ acceleration, and achieves comparable performance to existing methods on ImageNet.