Improving Sampling for Masked Diffusion Models via Information Gain
Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb
Abstract
Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the Info-Gain Sampler , a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.
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.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext ab3a4e1c-3876-4479-b299-67e237abd2e4Cited by top-tier papers2
- CORE: Context-Robust Remasking for Diffusion Language ModelsKevin Zhai, Sabbir Mollah, Zhenyi Wang, Mubarak ShahICML 2026 · 10 citations
- TEAM: Temporal–Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model AccelerationLINYE WEI, Zixiang Luo, Pingzhi Tang, Meng LiICML 2026 · 7 citations
Builds on18
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- The Curious Case of Neural Text DegenerationAri Holtzman, Jan Buys, Li Du, Maxwell Forbes et al.ICLR 2020 · 4,112 citations
- Structured Denoising Diffusion Models in Discrete State-SpacesJacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow et al.NeurIPS 2021 · 2,256 citations
- Large Language Diffusion ModelsShen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang et al.NeurIPS 2025 · 949 citations
- Diffusion Forcing: Next-token Prediction Meets Full-Sequence DiffusionBoyuan Chen, Diego Marti Monso, Yilun Du, Max Simchowitz et al.NeurIPS 2024 · 751 citations
Related papers
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded UnmaskingHeli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte et al.NeurIPS 2025 · 131 citations
- Empirical Analysis of Decoding Biases in Masked Diffusion ModelsPengcheng Huang, Tianming Liu, Zhenghao Liu, Yukun Yan et al.ACL 2026 · 15 citations
- Search or Accelerate: Confidence-Switched Position Beam Search for Diffusion Language ModelsMingyu Cao, Alvaro Correia, Christos Louizos, Shiwei Liu et al.ICML 2026
- 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
