Improving Sampling for Masked Diffusion Models via Information Gain

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.

Masked diffusion models generate text or images by gradually filling in missing pieces. Existing methods often choose the next piece to fill based only on what looks most certain right now, which can lead the model into worse choices later. We propose Info-Gain Sampler, a method that instead asks how much each choice will help reduce uncertainty in the remaining generation process. This makes decoding more globally planned while requiring no additional training. Across reasoning, coding, creative writing, and image generation tasks, our method produces more reliable outputs than existing sampling strategies.