Improving Sampling for Masked Diffusion Models via Information Gain
Mirrored from arXiv — NLP / Computation & Language for archival readability. Support the source by reading on the original site.
Computer Science > Computation and Language
Title:Improving Sampling for Masked Diffusion Models via Information Gain
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 this https URL.
| Comments: | this https URL Accepted by ICML2026 Accepted by ICML2026 |
| Subjects: | Computation and Language (cs.CL) |
| Cite as: | arXiv:2602.18176 [cs.CL] |
| (or arXiv:2602.18176v3 [cs.CL] for this version) | |
| https://doi.org/10.48550/arXiv.2602.18176
arXiv-issued DOI via DataCite
|
Submission history
From: Kaisen Yang [view email][v1] Fri, 20 Feb 2026 12:26:03 UTC (13,673 KB)
[v2] Wed, 18 Mar 2026 07:32:59 UTC (13,674 KB)
[v3] Fri, 22 May 2026 13:29:19 UTC (13,674 KB)
Access Paper:
- View PDF
- HTML (experimental)
- TeX Source
References & Citations
Bookmark
Bibliographic and Citation Tools
Code, Data and Media Associated with this Article
Demos
Recommenders and Search Tools
arXivLabs: experimental projects with community collaborators
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
More from arXiv — NLP / Computation & Language
-
Evaluating Large Language Models in a Complex Hidden Role Game
May 25
-
A Survey of Text and Speech Resources for Hausa and Fongbe: Availability, Quality, and Gaps for NLP Development
May 25
-
Query-Adaptive Semantic Chunking for Retrieval-Augmented Generation: A Dynamic Strategy with Contextual Window Expansion
May 25
-
Knowledge Distillation for Low-Resource Open-source Text-to-SQL Model
May 25
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.