💬 NLP

Improving Sampling for Masked Diffusion Models via Information Gain

This paper introduces the Info-Gain Sampler, a principled decoding framework for Masked Diffusion Models that leverages their non-causal nature to balance immediate uncertainty with future information gain, thereby significantly outperforming existing greedy heuristics across diverse tasks like reasoning, coding, and creative writing.

Original authors: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

Published 2026-03-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a giant, complex jigsaw puzzle, but you have a magical assistant (the AI) who can see the whole picture at once, not just the piece right in front of them. This is how Masked Diffusion Models (MDMs) work. Unlike traditional AI that writes a story one word at a time from left to right (like reading a book), MDMs can look at the whole blank canvas and fill in pieces in any order they want.

However, there's a catch: How do you decide which piece to fill in first?

The Problem: The "Easy First" Trap

Currently, most AI samplers act like a person who is afraid of making mistakes. Their strategy is simple: "Pick the piece that looks the easiest and most obvious right now."

The paper calls this Greedy Certainty.

The Analogy: Imagine you are solving a math problem: $2 \times 3 = ?$ $2 \times 3 = ?$ .
- The "Easy First" AI sees the answer ($6$) is very obvious. It fills that in immediately because it feels 100% confident.
- But then it gets stuck. It filled in the answer before it figured out the numbers ($2$ and $3$). Now it has to guess the factors based on the answer, which is incredibly hard and prone to errors.
- Result: The AI gets the easy part right but fails the whole puzzle because it didn't plan ahead. It's like building a house by putting the roof on first because it looks nice, only to realize you have no walls to support it.

The Solution: The "Information Gain" Sampler

The authors propose a new strategy called the Info-Gain Sampler. Instead of just asking, "What is the easiest piece to fill in?" it asks, "Which piece, if I fill it in, will help me solve the rest of the puzzle the fastest?"

The Analogy: Think of it like a detective solving a mystery.
- Old Way (Greedy): The detective sees a suspect who looks guilty and arrests them immediately. But later, they realize that suspect was innocent, and the real clue was missed.
- New Way (Info-Gain): The detective looks at the whole crime scene. They realize that finding the missing weapon (a hard clue) will instantly reveal who the killer is. Even though finding the weapon is harder than just looking at the suspect, doing so gives them the most Information Gain. It clears up the confusion for the entire case.

How It Works (The Magic Trick)

The paper highlights a special superpower of MDMs: They can see the future.

Because MDMs look at the whole sentence or image at once (bidirectional), they can simulate: "If I fill in this specific word now, how much clearer will the rest of the sentence become?"

The Old Way: "I'm 90% sure this word is 'cat'. Let's write 'cat'." (Ignores that 'cat' might make the next sentence make no sense).
The New Way: "I'm only 60% sure this word is 'cat', but if I write it, the uncertainty of the entire sentence drops by 50%. Let's write 'cat'."

The new sampler balances Immediate Confidence (how sure we are right now) with Future Clarity (how much this helps us later).

Why Does This Matter?

The researchers tested this on many difficult tasks:

Math & Logic: It solved more complex equations by figuring out the hard parts first, rather than getting stuck on the easy answers.
Coding: It wrote better code by planning the structure before filling in the details.
Creative Writing: It wrote more coherent stories that didn't ramble or contradict themselves.
Images: It generated pictures where the objects were in the right places and had the right colors.

The Bottom Line

The paper introduces a smarter way for AI to "think" before it "acts." Instead of rushing to fill in the obvious blanks, the Info-Gain Sampler takes a step back, looks at the big picture, and chooses the move that unlocks the most clarity for the rest of the task.

It's the difference between a student who memorizes answers to pass a test (Greedy) and a student who understands the underlying concepts to solve any problem (Info-Gain). The result? Smarter, more accurate, and more creative AI.

1. Problem Statement

Masked Diffusion Models (MDMs) have emerged as a powerful alternative to autoregressive models (ARMs) for discrete sequence generation, offering flexibility through bidirectional attention and non-causal decoding paths. However, a critical training-inference mismatch exists:

Training: MDMs are trained under random masking patterns, learning to predict any token given any context.
Inference: Generation requires a multi-step, order-sensitive decoding process where the sampler must decide which masked positions to reveal next and what tokens to assign.

The Core Limitation: Existing samplers rely on greedy, local certainty heuristics (e.g., selecting positions with the highest confidence or lowest entropy).

Myopia: These methods prioritize immediate local certainty without considering the downstream impact of current decisions on future steps.
Failure Cases: As demonstrated in the paper's "One-way Multiplication" and "Binary Judgment" case studies, greedy samplers often choose paths that look locally certain (e.g., decoding a binary answer token before the reasoning steps) but lead to high cumulative uncertainty and logical errors later in the generation.
Untapped Potential: Unlike ARMs, where evaluating future impacts is computationally prohibitive due to the "next-token bottleneck," MDMs' bidirectional nature theoretically allows for efficient evaluation of how a decoding decision reshapes uncertainty across the entire remaining sequence.

2. Methodology: The Info-Gain Sampler

The authors propose the Info-Gain Sampler, a principled decoding framework that balances immediate decoding costs with the expected information gain for the remaining sequence.

Key Concepts

State Uncertainty ( $H(z_t)$ ): Defined as the average marginal entropy over all currently masked positions in state $z_t$ .
Information Gain ($IG$): The reduction in state uncertainty induced by a specific decoding action $a_t$ .
$IG(a_t; z_t) = H(z_t) - H(z_{t-1})$
Objective Function: The sampler optimizes a score $J_{IG}$ that balances the Immediate Cost (uncertainty of the tokens being decoded) against the Information Gain (reduction in uncertainty for remaining tokens):
$J_{IG}(a_t | z_t) = \underbrace{IG(a_t; z_t)}_{\text{Information Gain}} - \underbrace{C(a_t | z_t)}_{\text{Immediate Cost}}$
Where $C(a_t | z_t)$ is the sum of marginal entropies of the selected tokens.

Algorithm Workflow

Action Sampling: Instead of deterministically picking the top-k certain tokens, the sampler generates a diverse candidate set $\mathcal{C}$ $C$ of $N$ $N$ actions. This involves:
- Token Sampling: Drawing candidate tokens with temperature scaling.
- Position Sampling: Selecting positions based on certainty scores.
Parallel Evaluation: For each candidate action in $\mathcal{C}$ , the model performs a single batched forward pass to compute the resulting state uncertainty ( $H(z_{t-1})$ ) for all candidates simultaneously. This leverages the bidirectional architecture to estimate information gain efficiently without expensive tree searches (unlike ARMs).
Transition: The action maximizing $J_{IG}$ is selected, the state is updated, and the process repeats until the sequence is complete.

Efficiency Optimizations

Block-wise Computation: Restricts information gain computation to the current active block to enable KV-caching.
High-Confidence Bypass: If a token's probability exceeds a threshold $\gamma$ , it is directly fixed into the action set, skipping the evaluation step for high-confidence positions to reduce latency.

3. Key Contributions

Identification of Limitations: Empirically demonstrated that greedy certainty-based samplers fail to minimize cumulative uncertainty, often leading to error propagation in complex reasoning tasks.
Novel Decoding Framework: Introduced the Info-Gain Sampler, which utilizes the unique bidirectional nature of MDMs to perform efficient, global planning by balancing immediate costs with future information gain.
Comprehensive Evaluation: Validated the method across diverse architectures (Full-Attention and Semi-Autoregressive) and tasks (Reasoning, Coding, Planning, Creative Writing, and Image Generation), showing consistent superiority over existing baselines.

4. Experimental Results

The Info-Gain Sampler was evaluated against baselines like Uniform, Entropy, Confidence, Margin, KLASS, PC-Sampler, and a concurrent Lookahead method (LookUM).

Reasoning Tasks (Math & Code):
- On the Dream-7B model, Info-Gain achieved a 3.6% improvement in average accuracy over the best greedy baseline.
- It significantly reduced Cumulative Entropy (a proxy for decoding quality), dropping from 78.4 (best greedy baseline) to 48.6 in specific reasoning settings.
- Outperformed the concurrent LookUM method on Code and Planning tasks, which require high token-level precision.
Creative Writing:
- Achieved a 63.1% win-rate against baselines in length-controlled evaluations.
- Demonstrated superior robustness to temperature scaling; at high temperatures ( $\tau=1.5$ ), it maintained a 80.3% win-rate against the Entropy baseline, balancing creativity with coherence.
Image Generation (Text-to-Image):
- On GenEval, achieved the highest average score (58.2 vs. 56.3 for the next best).
- On ImageNet-512, significantly improved FID (38.1 vs. 43.3) and Inception Score (63.0 vs. 53.3).
Efficiency: The method incurred only a marginal increase in generation time (+24%) and GPU memory (+20%) compared to greedy baselines, thanks to parallel evaluation and caching optimizations.

5. Significance and Impact

Bridging Local and Global Planning: The paper provides a principled solution to the "myopia" of current MDM samplers, proving that non-autoregressive models can effectively perform global planning without the computational cost associated with ARMs.
Data Manifold Preservation: The authors argue that the Info-Gain Sampler effectively keeps the decoding trajectory close to the training data manifold. By penalizing actions that lead to high future uncertainty (deviations from logical coherence), it prevents the model from entering "hallucination" states.
Generalizability: The framework is training-free and compatible with various MDM architectures (Full-Attention, Semi-AR, Block Diffusion), making it a versatile upgrade for the current generation of discrete diffusion models.
New Metric: The paper highlights Cumulative Entropy as a reliable, negative-correlated proxy for generation accuracy, offering a new lens for evaluating decoding strategies.

In summary, the Info-Gain Sampler represents a significant step forward in making Masked Diffusion Models more reliable for complex, multi-step reasoning tasks by shifting the decoding paradigm from "what is most certain now" to "what provides the most information for the future."