LUGS: Latent-aware Guidance for Efficient Unmasking in Diffusion Large Language Models

Diffusion Language Models (DLMs) have emerged as a flexible alternative to autoregressive (AR) models. They can decode tokens in any order, but the generation quality critically depends on the decoding strategy. Existing approaches predominantly rely on local heuristics, such as confidence or entropy, which may fail to capture sequence-level dependencies and the semantics in the context. To solve this problem, we propose Latent-aware Unmasking Guidance Search (LUGS), a novel decoding framework that leverages the model's internal hidden states to guide the unmasking process. By incorporating latent-aware scores to compensate for the limitations of local heuristics such as confidence or entropy, LUGS improves the model's performance. Extensive experiments on various downstream tasks demonstrate that our approach consistently outperforms existing baselines on LLaDA-8B-instruct and LLaDA-1.5 models. In Science and Reason tasks, LUGS improved performance by more than 1\% on both base models. And LUGS obtains an average improvement of 3.5\% in code generation. Remarkably, LUGS outperforms the beam search baseline by more than 5\% on average using LLaDA-8B-Instruct on code tasks. These results highlight the potential of latent-aware guidance for advancing controllable and high-quality generation.

Traditional text models generate words strictly from left to right. A newer approach called diffusion models allows for more flexibility by filling in words in any order, similar to solving a crossword puzzle. However, these models are often short-sighted because they rely on simple word confidence to decide what to reveal next. This leads to early mistakes that cannot be easily corrected, ultimately damaging the coherence and quality of the final text. We introduce LUGS, a framework that provides a more strategic guide for the text generation process. Instead of just picking the most obvious words, LUGS uses the deep internal signals of the AI to understand the semantic importance of each position. We trained this module by first conducting a thorough and time-consuming search for the best possible generation paths. We then distilled the knowledge from these high-quality paths into a lightweight component that can efficiently steer the model toward better decisions during use. Our evaluations demonstrate that this approach makes AI generation far more reliable in technical domains. Specifically, LUGS provides significant improvements in computer programming and mathematical reasoning tasks. By enabling the model to see beyond simple word probability, our research paves the way for higher-quality and more controllable AI assistants in complex scientific fields.