Diffusion language models decode text faster with adaptive token grouping
Early-Bird Decoding: Accelerating Diffusion LLMs with Learnable Block Sizes and Parallel Sampling
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Generating text with some language models takes many steps, making it slow to produce words. The authors noticed that groups of uncertain words could be predicted together earlier than usual. They created a new method that learns how to group and predict these words in blocks of different sizes, speeding up text generation without retraining the model. This approach works as a plug-in that makes the process faster while keeping the text quality similar.
What this means in practice
- •For machine learning engineers: Speed up text generation in diffusion-based language models by implementing adaptive token grouping to reduce decoding time without changing the base model.
- •For cloud service providers: Improve throughput of large language model services by adding plug-in modules that enable earlier token decoding and parallel sampling.
Authors
Lixuan Wei, Wei Zhou, Jianwen Wu, Yipeng Shen, Meiling Wang, Haoran You
Abstract
Diffusion large language models (dLLMs) offer a promising parallel decoding paradigm as an alternative to autoregressive generation through iterative unmasking. However, dLLMs typically require many steps before token confidence reaches the decoding threshold, resulting in inefficient inference even with block-wise KV caching. To accelerate dLLM inference, we for the first time propose an "early-bird (EB)" decoding framework, motivated by the observation that tokens with similarly low entropy tend to cluster and can be jointly decoded earlier, before reaching the confidence threshold. In particular, our EB-Decode framework integrates two key enablers: (1) a learnable network that adaptively groups tokens with similar uncertainty into variable-length blocks, rather than relying on fixed block sizes; (2) a position-aware sampler that learns to unmask tokens in parallel using fewer decoding steps within predicted variable-length blocks. Both components are developed without modifying pretrained dLLM weights and can therefore be directly deployed as plug-ins during serving, with negligible training and inference overhead. Extensive experiments across three models and four benchmarks consistently validate our observation and the effectiveness of EB-Decode, achieving 3.53-18.76$\times$ higher throughput than the vanilla decoding method and up to 1.58$\times$ higher throughput over the strongest baseline, Fast-dLLM, with comparable accuracy.