Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Sampling for Masked Diffusion Models via Information Gain

Kaisen Yang, Jayden Teoh, Kaicheng Yang, Yitong Zhang, Alex Lamb

TL;DR

The Info-Gain Sampler is proposed, a principled decoding framework that balances immediate uncertainty with information gain over future masked tokens, and demonstrates that Info-Gain Sampler consistently outperforms existing samplers for MDMs.

Abstract

Masked Diffusion Models (MDMs) offer greater flexibility in decoding order than autoregressive models but require careful planning to achieve high-quality generation. Existing samplers typically adopt greedy heuristics, prioritizing positions with the highest local certainty to decode at each step. Through failure case analysis, we identify a fundamental limitation of this approach: it neglects the downstream impact of current decoding choices on subsequent steps and fails to minimize cumulative uncertainty. In particular, these methods do not fully exploit the non-causal nature of MDMs, which enables evaluating how a decoding decision reshapes token probabilities/uncertainty across all remaining masked positions. To bridge this gap, we propose the Info-Gain Sampler, a principled decoding framework that balances immediate uncertainty with information gain over future masked tokens. Extensive evaluations across diverse architectures and tasks (reasoning, coding, creative writing, and image generation) demonstrate that Info-Gain Sampler consistently outperforms existing samplers for MDMs. For instance, it achieves a 3.6% improvement in average accuracy on reasoning tasks and a 63.1% win-rate in creative writing. Notably, on reasoning tasks it reduces cumulative uncertainty from 78.4 to 48.6, outperforming the best baseline by a large margin. The code will be available at https://github.com/yks23/Information-Gain-Sampler.

Improving Sampling for Masked Diffusion Models via Information Gain

TL;DR

The Info-Gain Sampler is proposed, a principled decoding framework that balances immediate uncertainty with information gain over future masked tokens, and demonstrates that Info-Gain Sampler consistently outperforms existing samplers for MDMs.

Abstract

Masked Diffusion Models (MDMs) offer greater flexibility in decoding order than autoregressive models but require careful planning to achieve high-quality generation. Existing samplers typically adopt greedy heuristics, prioritizing positions with the highest local certainty to decode at each step. Through failure case analysis, we identify a fundamental limitation of this approach: it neglects the downstream impact of current decoding choices on subsequent steps and fails to minimize cumulative uncertainty. In particular, these methods do not fully exploit the non-causal nature of MDMs, which enables evaluating how a decoding decision reshapes token probabilities/uncertainty across all remaining masked positions. To bridge this gap, we propose the Info-Gain Sampler, a principled decoding framework that balances immediate uncertainty with information gain over future masked tokens. Extensive evaluations across diverse architectures and tasks (reasoning, coding, creative writing, and image generation) demonstrate that Info-Gain Sampler consistently outperforms existing samplers for MDMs. For instance, it achieves a 3.6% improvement in average accuracy on reasoning tasks and a 63.1% win-rate in creative writing. Notably, on reasoning tasks it reduces cumulative uncertainty from 78.4 to 48.6, outperforming the best baseline by a large margin. The code will be available at https://github.com/yks23/Information-Gain-Sampler.
Paper Structure (37 sections, 1 theorem, 25 equations, 12 figures, 8 tables)

This paper contains 37 sections, 1 theorem, 25 equations, 12 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Masked Diffusion Models (MDMs)
  4. Method
  5. Motivation
  6. Information Gain Sampler
  7. Objective of Info-Gain Sampler
  8. Implementation of Info-Gain Sampler
  9. Efficient Implementation of Info-Gain Sampler
  10. Why is Info-Gain effective
  11. Experiments
  12. Experimental Setup
  13. Results and Analysis
  14. Ablation Study
  15. Limitation and Future Work
  16. ...and 22 more sections

Key Result

Proposition 3.1

For any state $z_t$ and position set $A_t \subseteq \mathcal{M}_t$,

Figures (12)

  • Figure 2: The Info-Gain Sampler workflow. Starting from state $z_{T_0}$, the sampler iteratively: (1) samples candidate actions, (2) evaluates $J_{\text{IG}} = \text{Immediate Cost} - \text{Information Gain}$ in parallel to select the optimal successor state $z_{t-1}^*$, and (3) executes the state transition until reaching the final sequence $z_0$.
  • Figure 3: Analysis of Cumulative Entropy. (a) Cumulative entropy trajectories for the Entropy baseline and Info-Gain Sampler on a synthetic set of 100 simple arithmetic problems that can be answered within a short window. We use global decoding with a fixed length of 64 tokens. (b) Correlation between average accuracy and average cumulative entropy across various sampling configurations.
  • Figure 4: Impact of different beam sizes on the MATH-500 dataset. Specifically: Beam Size = 1 is a special case equivalent to the Info-Gain Sampler; Beam Size = Expansion Budget is equivalent to the Best-of-$N$ (BoN) baseline; and Intermediate Values represent a look-ahead beam search algorithm using Info-Gain as the pruning heuristic.
  • Figure 5: Temperature Sensitivity. Cumulative trajectory uncertainty under varying position and token temperatures on the 100 simple arithmetic problems, evaluated using global decoding with a fixed length of 64 tokens.
  • Figure 6: Empirical distribution of $J_{IG}$ values sorted from highest to lowest. The 5th percentile is $-5 \times 10^{-4}$, indicating the bound is rarely violated in practice.
  • ...and 7 more figures

Theorems & Definitions (2)

  • Proposition 3.1: Upper bound on expected information gain
  • proof