LUGS: Latent-aware Guidance for Efficient Unmasking in Diffusion Large Language Models
Nuanqiao Shan, Kairong Han, Xinpeng Dong, Kun Kuang
Abstract
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.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on20
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Let's Verify Step by StepHunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards et al.ICLR 2024 · 3,045 citations
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
Related papers
- Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language ModelsMingyu Cao, Alvaro Correia, Christos Louizos, Shiwei Liu et al.ICML 2026
- Lookahead Unmasking Elicits Reliable Decoding in Diffusion Language ModelsSanghyun Lee, Seungryong Kim, Jongho Park, Dongmin ParkICML 2026 · 13 citations
- Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked DiffusionsJaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham M. Kakade et al.ICML 2025
- Reward-Weighted Sampling: Enhancing Non-Autoregressive Characteristics in Masked Diffusion LLMsDaehoon Gwak, Minseo Jung, Junwoo Park, Minho Park et al.EMNLP 2025
- LaDiR: Latent Diffusion Enhances LLMs for Text ReasoningHaoqiang Kang, Yizhe Zhang, Nikki Lijing Kuang, Nicklas Majamaki et al.ICLR 2026 · 25 citations
