Diffusion language models speed up response by learning token reveal timing
Learning When to Commit: Schedule Distillation for Fast Diffusion Language Models
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
Generating text with diffusion language models can be slow because they reveal parts of the answer bit by bit. The authors found a way to teach these models when to reveal certain words so that it helps guess the rest faster. Their method focuses on revealing words that unlock easier predictions for others, speeding up the process without losing quality. They tested this on math problems and code generation, showing big improvements in speed while keeping accuracy.
What this means in practice
- •For natural language processing engineers: Speed up text generation in AI systems for faster interaction without harming output quality by learning optimal token reveal order.
- •For automated coding tool developers: Produce code generation outputs more quickly by integrating learned token scheduling that reduces decoding rounds.
Authors
Aidar Myrzakhan, Zhiqiang Shen
Abstract
Diffusion large language models (dLLMs) can generate multiple tokens in parallel, but often require many decoding rounds to complete a response. Efficient generation depends not only on which tokens are easy to predict, but also on how committing to them helps predict the remaining positions. A highly predictable token may provide little useful context, while revealing a less predictable token can help several other positions become predictable sooner. We propose dSpark, a new schedule distillation-based framework for learning when to commit tokens for faster diffusion language generation. To achieve this, we construct unmasking schedules that prioritize tokens whose reveal makes other positions easier to predict, reducing uncertainty earlier in decoding. We then distill these schedules into the model to enable faster parallel sampling, with a KL-divergence regularizer to help preserve generation quality. We extensively evaluate our approach across mathematical reasoning and code generation tasks. On GSM8K, our method achieves remarkable 10.12$\times$ and 5.12$\times$ speedups over the various baselines on LLaDA-8B-Instruct and Dream-7B-Instruct, respectively, while maintaining superior task performance. Our code is available at https://github.com/VILA-Lab/dSpark.