Improving Sampling for Masked Diffusion Models via Information Gain

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

cs.CLFeb 20, 2026

Improving Sampling for Masked Diffusion Models via Information Gain

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

Organizations: Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University · Beihang University · College of AI, Tsinghua University

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 https://github.com/yks23/Information-Gain-Sampler.

Figures & tables

Panel of Figure 1 — Figure 1: Motivation: Analysis of decoding strategies on the one-way multiplication experiment. (a) Illustrates the contrast between the suboptimal path chosen by the greedy certainty-based sampler and the optimal path, motivating the introduction of the Info-Gain Sampler. (b) Shows the evolution of cumulative uncertainty throughout the decoding process. While the greedy sampler prioritizes decoding $c$ first (73.2%) due to its immediate high confidence, it leads to failure because of the task’s one-way nature. In contrast, the Info-Gain Sampler optimizes global uncertainty by resolving high-entropy factors $a$ or $b$ first (84.0%), ensuring a successful decoding trajectory.

Metric	Entropy	Confidence	Margin	AR
$\tilde{H}$ $\downarrow$	32.75	31.68	35.60	25.19
Acc. (%) $\uparrow$	67	73	66	90

Table 1: Quantitative results for Case Study 2.

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{Information Gain}-\text{Immediate Cost}$ in parallel to select the optimal successor state $z_{t-1}^{*}$ , and (3) executes the state transition until reaching the final sequence $z_{0}$ .

K	Sampler	GSM8K	MATH500	HumanEval	MBPP	Sudoku	Countdown	Avg.	$\tilde{H}\downarrow$
$2$	Uniform	18.7	14.3	14.3	18.6	52.8	7.6	21.1	389.2
	AR	53.8	23.0	23.7	22.2	38.2	25.8	31.1	273.1
	Entropy	55.8	27.0	26.2	23.2	76.0	23.2	38.6	247.8
	Confidence	61.9	29.4	26.8	25.2	81.6	36.2	43.5	249.2
	Margin	65.5	28.8	28.8	23.6	74.4	28.9	41.7	287.5
	KLASS	67.3	30.4	31.9	30.1	82.1	35.4	46.2	239.3

Table 2: Performance on the full-attention MDM ( Dream-7B-Instruct ). Results are reported with decoding rates

K\in\{1,2\}

, where reasoning tasks(GSM8K, MATH500, HumanEval, MBPP) use a block size of 16 and planning tasks (Sudoku, Countdown) are decoded globally. We report the accuracy for each task, along with the average accuracy (Avg.) and cumulative entropy (

\tilde{H}

) per generation.

		SDAR-8B-Chat						TraDo-8B-Instruct
K	Sampler	GSM8K	MATH500	HumanEval	MBPP	Avg.	$\tilde{H}\downarrow$	GSM8K	MATH500	HumanEval	MBPP	Avg.	$\tilde{H}\downarrow$
2	Entropy	42.2	24.4	26.2	20.6	28.4	238.6	31.9	17.0	20.7	21.8	22.8	419.5
	Confidence	47.2	36.6	24.4	20.2	32.1	204.1	36.5	39.2	19.5	22.0	29.3	334.1
	Margin	45.2	22.4	19.5	19.8	26.7	230.9	33.2	17.0	18.9	21.8	22.7	398.0
	KLASS	50.4	32.3	30.7	26.6	35.0	210.3	50.4	32.3	30.7	26.6	35.0	334.3
	LookUM	75.3	44.9	28.2	31.8	45.0	103.2	79.8	40.4	46.8	33.9	50.2	119.2

Table 3: Performance on semi-autoregressive MDMs ( SDAR-8B-Chat and TraDo-8B-Instruct ). Results are reported with a block size of 16 and token temperature

\tau_{\text{token}}=0.7

. We report the accuracy for each task, along with the average accuracy (Avg.) and cumulative entropy (

\tilde{H}

) per generation. The Info-Gain Sampler consistently achieves the strongest performance across all settings.

Method	GenEval							ImageNet-512
Method	single obj. $\uparrow$	two obj. $\uparrow$	count $\uparrow$	colors $\uparrow$	pos $\uparrow$	attr $\uparrow$	Avg $\uparrow$	IS $\uparrow$	FID $\downarrow$	sFID $\downarrow$
Uniform	94.1	66.7	38.4	78.2	19.0	28.8	54.2	49.3	46.8	123.9
Entropy	94.3	67.3	46.0	79.9	17.8	26.8	55.3	52.4	44.8	93.4
Confidence	93.8	69.7	46.3	81.9	16.0	27.0	56.0	53.3	43.3	92.0
Margin	94.0	68.7	47.3	80.1	19.0	29.0	56.3	51.9	45.2	95.3
Info-Gain	97.5	68.7	47.5	79.8	25.0	32.0	58.2	63.0	38.1	83.7

Table 4: Text-to-Image results on GenEval and ImageNet-512 with token temperature

\tau_{\text{token}}=0.4

. The Info-Gain Sampler demonstrates superior alignment and fidelity, significantly outperforming all baselines across multimodal generation metrics.

Temperature	K	Win rate vs. baseline (%)
Temperature	K	AR	Confidence	Entropy	Margin
0.5	1	60.1	65.8	59.1	63.6
0.5	2	63.9	68.9	70.4	64.7
1.0	1	54.8	57.7	60.1	55.2
1.0	2	65.2	61.1	65.7	57.5
1.5	1	58.4	53.0	60.3	54.6

Table 5: Win rate of Info-Gain Sampler against baselines (%)

Panel of Figure 3 — 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.

Sampler	Selection Criteria	Temperature Sensitivity	Greedy Selection
Uniform ( Austin et al., 2021a )	$-$	High	$\times$
Confidence ( Chang et al., 2022 )	$\sum_{\ell\in A_{t}}\max p_{\theta}(x_{\ell}\mid z_{t})$	High	$\checkmark$
Entropy ( BenHamu et al., 2025 )	$-\sum_{\ell\in A_{t}}H^{(\ell)}(z_{t})$	High	$\checkmark$
Margin ( Kim et al., 2025a )	$\sum_{\ell\in A_{t}}(p_{\text{top1}}-p_{\text{top2}})$	High	$\checkmark$
KLASS † ( Kim et al., 2025b )	$\sum_{\ell\in A_{t}}(\max p_{\theta}(x_{\ell}\mid z_{t})+\mathbf{1}_{D_{KL}<\epsilon})$	High	$\checkmark$
PC-Sampler ( Huang et al., 2025 )	$\sum_{\ell\in A_{t}}(w_{\ell}\cdot\mathcal{C}_{t}^{(\ell)})$	Moderate	$\checkmark$

Table 6: Unified view of MDM samplers. Selection Criteria denotes the objective function optimized at each decoding step. Temperature Sensitivity indicates robustness of performance to temperature-based stochastic sampling, while Greedy Location Selection marks whether the sampler selects positions based solely on immediate certainty without considering future information gain.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Empirical distribution of $J_{\mathrm{IG}}$ values sorted from highest to lowest. The 5th percentile is $-5\times 10^{-4}$ , indicating that the utility remains close to its theoretical upper bound in practice.

Panel of Figure 7 — Figure 7: Visual results on ImageNet-512 with an extreme budget of only 5 decoding steps. The Info-Gain Sampler maintains superior structural coherence compared to baseline heuristics.

Figure 8: Evolution of cumulative entropy during image generation on the ImageNet-512 benchmark. The results are averaged over all labels using a 5-step linear schedule. The curves illustrate how Info-Gain Sampler manages global uncertainty compared to other methods throughout the decoding process.

Sampler	Collapse Step ( $\downarrow$ )
Entropy	13
Confidence	12
Margin	12
Info-Gain	8

Appendix

Table 7: Comparison of Collapse Steps on ultra-low step writing tasks. Lower values indicate better robustness in extreme acceleration scenarios.

Method	Dual-Cache Count	Memory Scaling
Info-Gain Sampler	1 shared outside-block cache	$O(M_{W}+M_{C}+N\cdot M_{A})$
Info-Gain Beam Search	$B$ divergent outside-block caches	$O(M_{W}+B\cdot M_{C}+B\cdot N\cdot M_{A})$

Appendix

Table 8: Memory efficiency comparison under shared dual-cache reuse.

Figure 9: Effect of utility threshold on cumulative entropy reduction and generation time.

Figure 10: Memory usage comparison. Info-Gain Sampler maintains low overhead via shared dual-cache reuse.

Explore similar work

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models
Jun 9, 2026Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher +1Masked Diffusion Models Parallel Decoding
University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland
DOS: Dependency-Oriented Sampler for Masked Diffusion Language Models
Mar 16, 2026Xueyu Zhou, Yangrong Hu, Jian HuangLanguage Model Decoding Diffusion Language Model Decoding
Department of Data Science and AI The Hong Kong Polytechnic University, Hong Kong, China
DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking
Mar 2, 2026Gilad Turok, Chris De Sa, Volodymyr KuleshovMasked Diffusion Models Diffusion Language Models
Department of Computer Science, Cornell University, New York, New York, USA.

Jun 9, 2026J/K move · Enter open · S save

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher, Paolo Favaro

University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland

Masked Diffusion Models Parallel Decoding Diffusion Language Model Decoding Diffusion Language Models Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration, but this parallelism is fragile: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers such as Top-

k

, Fast-dLLM, and EB-Sampler mainly control how many tokens to reveal, while often ranking candidates by token-wise scores that ignore interactions within the selected set. We propose ADAS, a training-free reranking rule that leaves the base sampler's stopping rule unchanged and greedily discounts each token-wise confidence score according to its attention to already selected positions, weighted by their prediction uncertainty. Across LLaDA-8B-Base and Dream-7B-Base on the reasoning benchmarks GSM8K and MATH500 and the code benchmarks HumanEval and MBPP, plugging ADAS into all three samplers improves low-NFE performance at matched denoiser evaluations by

9.11

and

10.46

percentage points on average, respectively, with

3.1\%

per-forward runtime overhead. Code is available at https://github.com/yusufsahin99/ADAS.

Abstract

Explore similar work

Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models

DOS: Dependency-Oriented Sampler for Masked Diffusion Language Models

DUEL: Exact Likelihood for Masked Diffusion via Deterministic Unmasking