Learned unmasking speeds up token generation in diffusion language models
What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
Artificial Intelligence
Summary
Generating text quickly and accurately is tricky for some language models that guess many words at once. The authors studied how different ways of deciding which guessed words to keep can improve speed and quality. They found that being too relaxed about which words to accept slows things down, and that combining the model's confidence with internal info works best. Using these insights, they built a smaller, faster system that makes good predictions with fewer checks and runs well on different tests.
What this means in practice
- •For language model engineers: Design faster decoding algorithms for diffusion models by optimizing which token predictions to commit early, balancing speed and accuracy.
- •For ai service operators: Improve throughput and reduce compute costs when serving diffusion language models by implementing learned unmasking policies with trajectory-based targets.
Authors
Jinya Sakurai, Xun Xu, Jaehong Yoon
Abstract
Masked diffusion language models generate tokens in parallel, but their efficiency depends on which predictions are committed. Learned unmasking policies improve the speed--accuracy trade-off, yet differ in supervision, inputs, and capacity, obscuring which choices matter. We view supervised policies as distilling a decoding procedure through a teacher and examine independently what the teacher labels as committable, on which states the policy imitates it, what the policy observes, and how much the number of function evaluations (NFE) must decrease to offset policy overhead. We find that labeling predictions matching the final output is overly permissive at low NFE, and hidden states are weak alone but can complement confidence. About a 10% NFE reduction suffices to offset policy overhead for policies up to 7.5M parameters across tested settings. We therefore train a policy on trajectory-contraction targets, which accept only commitments staying on a reference trajectory, using confidence and hidden states. On LLaDA-8B-Instruct, the resulting 856K-parameter policy achieves the best accuracy--NFE frontier and highest throughput on GSM8K, is competitive at moderate targets on MATH-500 and HumanEval, and largely transfers to Dream-7B-Instruct.