Lookahead Unmasking Elicits Reliable Decoding in Diffusion Language Models

Masked Diffusion Models (MDMs) as language models generate by iteratively unmasking tokens, yet their performance crucially depends on the inference-time order of unmasking. Conventional methods such as confidence-based sampling are short-sighted, focusing on local optimization which neglects test-time computation and allows early decoding errors to cascade. We propose Lookahead Unmasking (LookUM), which addresses these concerns by guiding sampling path with a verifier over alternative unmasking orders, without requiring an external reward model. Our framework couples (i) a path generator that proposes paths by sampling from pools of unmasking sets with (ii) a verifier that computes the uncertainty of the proposed paths and performs importance sampling to subsequently select the final paths. Erroneous unmasking inflates sequence-level uncertainty, and our method exploits this to avoid error-prone trajectories. We validate our framework across six benchmarks, such as mathematics, planning, and coding, and demonstrate consistent performance improvements. LookUM requires only two to three paths to achieve peak performance. LLaDA with LookUM matches the performance of RL-tuned LLaDA 1.5 and yields additional gains when applied to LLaDA 1.5, suggesting complementarity with reinforcement learning.

Masked diffusion models are a new type of language model that write text by starting with many blanks and filling them in step by step. A key challenge is deciding which blanks to fill first. If the model fills in a difficult word too early and gets it wrong, that mistake can affect the rest of the answer. We propose Lookahead Unmasking, or LookUM, a method that helps the model choose a better order for filling in these blanks. Instead of trusting only the model’s immediate confidence, LookUM tries a small number of possible filling orders and checks which one is likely to lead to a more reliable final answer. It does this using the model’s own uncertainty, so it does not need an extra reward model or human feedback signal. Our experiments show that LookUM improves performance across tasks involving math, planning, and coding. It usually needs to compare only two or three possible paths to get its best results. The method also works well with models that have already been improved using reinforcement learning, suggesting that LookUM can be used as a complementary way to make masked diffusion language models more accurate and reliable.