Improving Sampling for Masked Diffusion Models via Information Gain
On the map: Diffusion models → Discrete and text diffusion
Abstract (arXiv)
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.
Key facts. Cited by 6 papers in the map (data through 7 Oct 2026). Placed in Discrete and text diffusion, one of 14 directions in Diffusion models. Builds on 35 of its 45 references that are in the map. First posted on arXiv 20 Feb 2026.
Builds on
35 of 45 references.
Closest ancestors · 14
Same direction, newest first
- 2026TEAM: Temporal-Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration9 Feb
- 2025WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference28 Dec
- 2025KLASS: KL-Guided Fast Inference in Masked Diffusion Models7 Nov
- 2025SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation7 Oct
- 2025Revolutionizing Reinforcement Learning Framework for Diffusion Large Language Models8 Sep
- 2025Dream 7B: Diffusion Large Language Models21 Aug
- 2025Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking30 May
- 2025Dimple: Discrete Diffusion Multimodal Large Language Model with Parallel Decoding22 May
- 2025Large Language Diffusion Models14 Feb
- 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions10 Feb
- 2024Scaling Diffusion Language Models via Adaptation from Autoregressive Models23 Oct
- 2024Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning18 Oct
2 more
Other references · 21
- 2021High-Resolution Image Synthesis with Latent Diffusion Modelscited by 18,694
- 2021Training Verifiers to Solve Math Word Problemscited by 9,743
- 2021Evaluating Large Language Models Trained on Codecited by 8,053
- 2021Measuring Mathematical Problem Solving With the MATH Datasetcited by 5,203
- 2021Program Synthesis with Large Language Modelscited by 3,509
- 2016Improved Techniques for Training GANscited by 2,938
- 2019The Curious Case of Neural Text Degenerationcited by 2,881
- 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,665
- 2016Google's Neural Machine Translation System: Bridging the Gap between Human and Machine Translationcited by 1,471
- 2018Hierarchical Neural Story Generationcited by 1,182
- 2022MaskGIT: Masked Generative Image Transformercited by 1,082
- 2023GenEval: An Object-Focused Framework for Evaluating Text-to-Image Alignmentcited by 946
9 more
- 2024Length-Controlled AlpacaEval: A Simple Way to Debias Automatic Evaluatorscited by 865
- 2024Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusioncited by 634
- 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 404
- 2025MMaDA: Multimodal Large Diffusion Language Modelscited by 241
- 2017Beam Search Strategies for Neural Machine Translationcited by 232
- 2024Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputscited by 148
- 2025Fast-dLLM v2: Efficient Block-Diffusion LLMcited by 128
- 2025Next-Latent Prediction Transformers Learn Compact World Modelscited by 29
- 2025To Backtrack or Not to Backtrack: When Sequential Search Limits Model Reasoningcited by 24
Built on by
Citing papers per year: 2026: 6.
6 papers.
- 2026Visual-Redundancy-Controlled Parallel Decoding for Diffusion-Based Multimodal Large Language Modelscited by 1
- 2026VoidPadding: Let [VOID] Handle Padding in Masked Diffusion Language Models so that [EOS] Can Focus on Semantic Terminationcited by 1
- 2026Looped Diffusion Language Modelscited by 0
- 2026Decoding in Order-Agnostic Language Models: Chain-Rule Deviation and Uniform Spreadingcited by 0
- 2026Visual Information-Guided Parallel Decoding for Diffusion Multimodal Large Language Modelscited by 0
- 2026Entropy Can Flow, or It Can Guide. Be Entropy. LEDFlow: Introducing Entropy-guided Generation Order into Uniform Discrete Flowcited by 0
- 2026Entropy Can Flow, or It Can Guide. Be Entropy. LEDFlow: Introducing Entropy-guided Generation Order into Uniform Discrete Flow20 Sep 2026
- 2026Visual Information-Guided Parallel Decoding for Diffusion Multimodal Large Language Models27 Aug 2026
- 2026VoidPadding: Let [VOID] Handle Padding in Masked Diffusion Language Models so that [EOS] Can Focus on Semantic Termination16 Jun 2026
- 2026Decoding in Order-Agnostic Language Models: Chain-Rule Deviation and Uniform Spreading31 May 2026
- 2026Looped Diffusion Language Models25 May 2026
- 2026Visual-Redundancy-Controlled Parallel Decoding for Diffusion-Based Multimodal Large Language Models25 May 2026
Similar, not cited
Close in wording, no citation either way
- 2026Generation Order and Parallel Decoding in Masked Diffusion Models: An Information-Theoretic Perspective
- 2025Optimizing Decoding Paths in Masked Diffusion Models by Quantifying Uncertainty
- 2026Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
- 2025Guided Star-Shaped Masked Diffusion
- 2025Lookahead Unmasking Elicits Accurate Decoding in Diffusion Language Models
- 2025Reward-Weighted Sampling: Enhancing Non-Autoregressive Characteristics in Masked Diffusion LLMs
- 2026DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking
- 2025Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking