PaperlayerView exact revision on arXiv

Source-linked AI summary

Improving Sampling for Masked Diffusion Models via Information Gain

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

arXiv:2602.18176v3cs.CL

TL;DR

Existing MDM samplers often make greedy, locally certain choices despite order-sensitive inference and downstream effects on uncertainty. Info-Gain Sampler balances immediate uncertainty with information gained over remaining masked positions; across reasoning, coding, planning, and image generation, it reports consistent performance improvements.

  • Problem

    MDM inference is order-sensitive, yet existing samplers largely use local certainty heuristics without accounting for downstream effects.

    Source: §1 · p. 1

  • Method

    Info-Gain Sampler uses MDM bidirectionality to balance immediate certainty with information gained over remaining masked positions.

    Source: §1 · p. 1

  • Results

    Info-Gain Sampler consistently improves performance across reasoning, code generation, planning, and image generation, with reasoning accuracy gains of 2.9–11.6 percentage points.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The sampler reduces cumulative uncertainty and is compatible with full-attention and semi-autoregressive architectures.

    Source: §7 · p. 9

    The search process still incurs higher computational cost than greedy decoding.

    Source: §5 · p. 8

Abstract

from arXiv · show

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.

1. Introduction

MDM decoding is order-sensitive, but existing samplers often choose locally certain tokens without considering effects on later uncertainty. Info-Gain Sampler uses bidirectional information gain to guide decoding and outperforms sampler baselines across several task domains.

  • Random-mask training contrasts with order-sensitive inference, making generation quality dependent on how the sampler selects tokens to reveal.Source: §1 · p. 1
  • Greedy certainty heuristics can favor syntactically confident but semantically suboptimal tokens, propagating errors and compromising generation quality.Source: §1 · p. 1
  • MDM bidirectionality allows assessment of how a decoding decision affects uncertainty across all remaining masked positions.Source: §1 · p. 1
  • Info-Gain Sampler balances immediate certainty with information gained over remaining masks, while also proposing computationally efficient implementations.Source: §1 · p. 1
  • Info-Gain consistently outperforms sampler baselines across math, coding, planning, writing, and image generation tasks.Source: §1 · p. 2

2. Preliminary

MDMs learn to reverse a masking process, predicting original tokens at masked positions. Their samplers choose which positions to unmask and what tokens to place there.

2.1. Masked Diffusion Models (MDMs)

  • The forward process progressively masks coordinates until the sequence is fully masked at time T.Source: §2.1 · p. 2
  • During reverse generation, a denoising network predicts each masked position’s original-token distribution, trained with weighted cross-entropy over masked positions.Source: §2.1 · p. 2
  • At each reverse step, a sampler chooses positions to unmask and assigns tokens; certainty-based methods prioritize positions with higher confidence scores.Source: §2.1 · pp. 2–3

3. Method

Greedy certainty-based decoding can choose locally easy tokens that leave difficult decisions unresolved, while MDM bidirectionality permits efficient estimates of downstream uncertainty. Info-Gain Sampler evaluates actions by balancing their immediate uncertainty cost against information gained over remaining masks.

3.1. Motivation

  • In one-way multiplication, greedy product-first decoding chooses the 73.2% path and leaves high residual uncertainty, whereas resolving factors first is the optimal strategy.The task generates decimal factors a and b and their binary product c; the factors-first path is associated with 84.0% in Figure 1.Source: §3.1 · p. 3
  • In binary judgment, greedy samplers can commit to a low-uncertainty Yes/No answer before reasoning is resolved; the AR baseline achieves higher accuracy and lower cumulative uncertainty.Table 1 reports the Case Study 2 comparison.Source: §3.1 · p. 3, Table 1 · p. 3
  • MDMs’ bidirectional attention can estimate how an action changes uncertainty across remaining masks in one forward pass, avoiding iterative look-ahead computations.Source: §3.1 · pp. 3–4

3.2. Information Gain Sampler

  • Info-Gain Sampler uses MDM bidirectionality to balance the immediate uncertainty cost of a decoding decision against expected information gain over remaining masks.Source: §3.2 · p. 4

3.2.1. OBJECTIVE OF INFO-GAIN SAMPLER

  • The objective combines immediate cost, the summed marginal entropy of selected tokens, with information gain, the reduction in uncertainty over remaining masks.Source: §3.2.1 · p. 4

3.2.2. IMPLEMENTATION OF INFO-GAIN SAMPLER

  • At each step, the sampler proposes candidate actions, evaluates their objective in parallel, and executes the highest-scoring action until the sequence is complete.Source: Figure 2 · p. 4, §3.2.2 · p. 4

3.2.3. EFFICIENT IMPLEMENTATION OF INFO-GAIN SAMPLER

  • Parallel candidate evaluation, active-block information-gain computation, shared caches, and a high-confidence bypass reduce inference latency relative to unrestricted evaluation.The bypass directly fixes sufficiently high-confidence positions, capped by the step’s action-set size.Source: §3.2.3 · p. 5

3.2.4. WHY IS INFO-GAIN EFFECTIVE

  • Because state uncertainty can signal deviation from coherent, fluent regions, the information-gain term discourages actions that increase uncertainty and helps maintain coherence and fluency, including at high temperatures.Source: §3.2.4 · p. 5

4. Experiments

The experiments evaluate Info-Gain across reasoning, coding, planning, text-to-image generation, and creative writing, then analyze cumulative uncertainty, beam variants, and temperature sensitivity. Across these settings, it improves over sampling baselines, while ablations examine its uncertainty and compute trade-offs.

4.1. Experimental Setup

  • Experiments cover fully attentive and semi-autoregressive reasoning, text-to-image generation, and creative writing using multiple model families and task-specific baselines.Reasoning evaluations include math, code, and planning benchmarks; image generation uses GenEval and ImageNet-512, and writing uses AlpacaEval.Source: §4.1 · p. 5

4.2. Results and Analysis

  • Info-Gain exceeds the best-performing full-attention baselines in average accuracy by 3.1 percentage points at K = 2 and 2.9 points at K = 1.The comparison is on Dream-7B-Instruct; cumulative entropy reaches 65.9% (K = 2) and 62.0% (K = 1) of PC-Sampler's level.Source: Table 2 · p. 6, §4.2 · pp. 6–7
  • Info-Gain leads the best baseline in average accuracy by 11.6 points on SDAR-8B-Chat and 11.3 points on TraDo-8B-Instruct at K = 1.These semi-autoregressive evaluations use token temperature τtoken = 0.7 and block size 16.Source: Table 3 · p. 6, §4.2 · p. 7
  • Info-Gain records a 58.2 average GenEval score versus 56.3 for Margin, and improves ImageNet-512 FID from 43.3 to 38.1 and IS from 53.3 to 63.0.The reported GenEval positional and attribute scores are 25.0 versus 19.0 and 32.0 versus 29.0, respectively.Source: Table 4 · p. 6, §4.2 · p. 7
  • Info-Gain maintains a 62.8% average creative-writing win rate across settings and baselines, reaching 80.3% against Entropy at τtoken = 1.5.The evaluations use length-controlled win rates; the paper reports that increased stochasticity often degrades coherence.Source: Table 5 · p. 7, §4.2 · p. 7

4.3. Ablation Study

  • On 100 short arithmetic problems, Info-Gain's cumulative entropy trajectory stabilizes earlier than Entropy's, and cumulative entropy correlates negatively with accuracy (Pearson’s r = −0.70).The figure uses global decoding with a fixed length of 64 tokens; the authors interpret cumulative entropy as a proxy for decoding quality.Source: Figure 3 · p. 7, §4.3 · p. 7
  • Under a fixed computational budget, Info-Gain variants outperform Best-of-N; B = 1 is near the Pareto frontier, while larger beams add memory overhead for marginal uncertainty reduction.The beam-size comparison treats B = 1 as Info-Gain and the expansion-budget beam as Best-of-N.Source: Figure 4 · p. 8, §4.3 · p. 8
  • Info-Gain maintains stable, low trajectory uncertainty across temperature scales, whereas other baselines are temperature-sensitive and show decoding instability.The authors report preserved diversity and competitive creative-writing win rates, distinguishing lower cumulative entropy from mode collapse.Source: Figure 5 · p. 8, §4.3 · p. 8

5. Limitation and Future Work

The authors identify higher computational cost than greedy decoding and reliance on local uncertainty for action-candidate generation as areas for further refinement.

5. Limitation and Future Work

  • The authors identify further refinement as an avenue for improving the sampler.The specific limitations and future directions include inference cost and action-candidate sampling.Source: §5 · p. 8

More Efficient Implementation

  • The search process incurs higher computational cost than greedy decoding, despite parallel evaluation and acceleration techniques intended to minimize latency.Suggested directions include more efficient lookahead, adaptive branching, and hardware-level optimization.Source: §5 · p. 8

Refinement of Action Sampler

  • The action sampler uses local uncertainty to generate candidates; more sophisticated strategies could improve candidate-set diversity and quality.The authors describe the current heuristic as robust and yielding high-quality plans.Source: §5 · p. 8

6. Related Work

Related work situates MDMs as bidirectional alternatives to autoregressive models and describes a sampling problem in which decoding position matters as well as token choice. Existing MDM samplers largely use greedy local criteria, while the paper’s comparison framework distinguishes samplers by selection criteria and other properties.

Samplers

  • The sampler comparison organizes methods by selection criteria, temperature sensitivity, and whether location selection is greedy.The table caption identifies these as comparison dimensions, not comparative outcomes.Source: Table 6 · p. 9

Masked Diffusion Models

  • MDMs support bidirectional reasoning, with research spanning models trained from scratch, adapted pretrained autoregressive models, and semi-autoregressive architectures.The related work also notes optimizations and models advancing complex reasoning and long-text generation.Source: §6 · pp. 8–9

Samplers

  • Unlike autoregressive sampling, MDM samplers choose both what token to decode and where, making early choices consequential and local uncertainty criteria insufficient.Existing samplers typically select positions greedily using entropy or margin, with some adding calibration or stability refinements.Source: §6 · p. 9

7. Conclusion

The conclusion presents Info-Gain as a training-free sampler that uses bidirectional attention to balance immediate certainty with future uncertainty reduction. The appendix formalizes baseline samplers through their position-selection rules.

7. Conclusion

  • 2.9–11.6 percentage points in average reasoning accuracy across full-attention and semi-autoregressive settings, alongside a 62.8% average win rate in creative writing and 1.9 percentage points on GenEval.The conclusion reports consistent improvements across reasoning, code generation, planning, and image-generation tasks.Source: §7 · p. 9
  • Info-Gain is described as training-free, compatible with full-attention and semi-autoregressive architectures, and associated with reduced cumulative uncertainty.The conclusion frames it as a bridge from local heuristics to global planning for non-autoregressive generation.Source: §7 · p. 9

A. Formal Definitions of Baseline Samplers

  • Uniform, confidence, entropy, and margin baselines select masked positions randomly or according to local confidence, uncertainty, or top-token margin.The associated descriptions respectively specify random selection, top-1 certainty, minimum predictive uncertainty, and the gap between the two most likely tokens.Source: A. Formal Definitions of Baseline… · p. 12
  • KLASS prioritizes positions satisfying a KL-divergence threshold, ranks them by confidence, and fills any remaining decoding budget from other positions.The experiments use a threshold of 5 × 10^-4 and adapt KLASS to decode a fixed number of tokens per step.Source: A. Formal Definitions of Baseline… · p. 12
  • Block Diffusion restricts candidate positions to the active diffusion block, which slides over time according to a scheduling rule.The schedule typically traverses positions sequentially or randomly.Source: A. Formal Definitions of Baseline… · p. 12
  • LookUM generates multiple candidate unmasking trajectories and selects one by sequence-level certainty, using negative entropy as the paper’s base metric.The comparison omits LookUM’s SMC and NIS designs because they require task-specific parameter tuning.Source: A. Formal Definitions of Baseline… · p. 13

B. Detailed Hyperparameter Settings

The experiments use NVIDIA A800 GPUs and specify task-dependent sampling temperatures, candidate counts, and decoding schedules. Settings vary across reasoning and coding, image generation, and creative writing.

B. Detailed Hyperparameter Settings

C. Theoretical Analysis of the Info-Gain Objective

The analysis defines an action’s utility as information gained over remaining masked positions minus the immediate uncertainty cost. It bounds expected utility above by zero and reports that observed values are usually close to that bound.

Utility definition

  • The Info-Gain utility compares the reduction in overall state uncertainty with the immediate uncertainty of the selected positions.The action selects positions from the masked set, and applying it produces the next state with those positions removed from the masked set.Source: C. Theoretical Analysis of the… · pp. 13–14

Theoretical bound

  • The expected information gain for any fixed position set is bounded by the immediate uncertainty cost, so expected Info-Gain utility is at most zero.The derivation considers cases α ≤ β and α > β.Source: C. Theoretical Analysis of the… · p. 14

Practical implications

  • More than 95% of observed GSM8K values satisfy JIG ≥−5 × 10−4, placing realized actions close to the theoretical upper bound of zero.The bound is equivalent to a nonnegative expected commitment cost.Source: C. Theoretical Analysis of the… · p. 14, Figure 6 · p. 15
  • The authors describe information gain as sensitive to how assignments affect uncertainty in remaining masked positions, discouraging poorly conditioned or inconsistent states.They characterize this as an implicit correction mechanism that can recover from suboptimal early decisions.Source: C. Theoretical Analysis of the… · p. 14

D. Pseudocode for Info-Gain Sampler and Info-Gain Beam Search

The appendix provides pseudocode for both the Info-Gain Sampler and Info-Gain Beam Search. It separates the sampler’s procedure from the beam-search procedure.

D. Pseudocode for Info-Gain Sampler and Info-Gain Beam Search

D.1. Info-Gain Sampler

  • The Info-Gain Sampler pseudocode includes candidate action sampling, parallel evaluation, and a state transition within its decoding loop.The listing initializes all positions as masked and performs an initial forward pass.Source: §D.1 · p. 15

D.2. Info-Gain Beam Search

  • The appendix provides a separate pseudocode listing for Info-Gain Beam Search.The listing includes beam expansion, action sampling, candidate updates, and construction of the next beam.Source: §D.2 · pp. 15–16

E. Experiments in Extremely Low-Step Generation Scenarios

The section tests generation under extremely limited decoding-step budgets in image generation and writing. Info-Gain is reported to produce more coherent outputs than greedy sampling in these constrained settings.

E. Experiments in Extremely Low-Step Generation Scenarios

  • The Dream Model writing experiment tests 50 prompts at a fixed length of 80 tokens and records the minimum step budget before fewer than half of samples show severe repetitions or grammatical errors.The section reports that most baselines fail to produce meaningful, coherent text in this ultra-low-step regime.Source: E. Experiments in Extremely Low-Step… · p. 16
  • Info-Gain produces more coherent generations than greedy sampling when decoding steps are extremely limited.The ImageNet-512 MMaDa experiment uses a linear schedule and only 5 decoding steps; the figure reports better visual quality and structural coherence.Source: E. Experiments in Extremely Low-Step… · p. 16, Figure 7 · p. 17

F.1. Memory Overhead Analysis

The sampler shares a dual cache across candidate actions, while beam search maintains separate caches for divergent trajectories. The computation analysis describes batching, cache reuse, and a confidence-based bypass as ways to limit overhead.

F.1. Memory Overhead Analysis

  • The sampler’s added memory for more candidates is dominated by transient activations, while beam-search memory scales with both beam size B and candidate count N.The sampler shares its outside-block dual cache; beam search maintains B separate dual caches and evaluates N actions for each beam element.Source: §F.1 · p. 17
  • Info-Gain Sampler has 24% memory overhead at N = 8, while Info-Gain Beam Search shows a substantial memory surge.Measurements use TraDo-8B-Instruct with 512 tokens and block size 16; the sampler shares one outside-block dual cache, whereas beam search maintains divergent caches.Source: §F.1 · pp. 17–18, Table 8 · p. 18, Figure 10 · p. 19

F.2. Computation Overhead Analysis

  • Batched evaluation, dual-cache reuse, and high-confidence bypass reduce the sampler’s practical computation overhead below the naive N× estimate.The threshold mechanism is reported to achieve near-optimal entropy reduction with minimal additional time while maintaining high decoding quality.Source: §F.2 · pp. 18–19, Figure 9 · p. 19

Related papers

Loading 2602.18176v3…