Diffusion language models speed up token generation with parallel block decoding
One Block Ahead: Promoting Sampler State in Diffusion Language Models
Computation and Language
Summary
Making language models faster at producing words is important for many applications. The authors found a way to speed up a specific type of model, called a block-diffusion language model, by letting it work on two blocks of words at the same time instead of one after the other. They introduced a method called One Block Ahead that cleverly shares information between blocks so the model can decode multiple parts together without retraining. This approach boosted the number of tokens generated per step by over 40% and also improved speed with only a small drop in accuracy.
What this means in practice
- •For natural language processing engineers: Accelerate language model deployment by decoding multiple token blocks simultaneously to increase throughput without retraining.
- •For real-time translation developers: Improve speed of translation systems by using concurrent block decoding to deliver faster outputs on diffusion models.
Authors
Yangyi Shen, Pin-Jui Ku, Nicolai Oswald, Jean-Marie Lemercier, Yonggan Fu, Stefano Ermon, Jure Leskovec, Boris Ginsburg, Ante Jukić
Abstract
Block-diffusion language models denoise a whole block of token positions in parallel so each forward pass can emit many tokens. However, blocks are still denoised in strict sequence, a cost that within-block acceleration leaves untouched. Here, we investigate decoding blocks concurrently and propose One Block Ahead (OBA), a training-free decoding algorithm that denoises the next block alongside the current one. We demonstrate that for uniform-state diffusion LMs, where a block's denoising progress is not visible in its intermediate tokens, concurrent decoding is advantageous if the next block's tokens are promoted together with the corresponding sampler state, including the self-conditioning signal and the denoising time. We also derive a break-even criterion that relates the benefit of OBA to the relative cost of a paired-block forward pass. On DiffusionGemma, a state-of-the-art diffusion LM, OBA increases the tokens per forward pass by 42.3% (from 15.64 to 22.25) and the single-request denoising throughput by 16.3%, averaged over ten benchmarks, with a small accuracy regression.