Research map / Image generation / Fast diffusion sampling
Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes EB-Sampler, an adaptive sampler for masked diffusion models that chooses both which and how many tokens to unmask per step by bounding the summed entropy of the selected tokens, giving 2-3x speedups without loss of accuracy on code, math, and puzzle tasks.
Problem
Efficient sampling from masked diffusion models (MDMs) for language is underexplored; unmasking multiple tokens per step with fixed k degrades quality due to joint dependence error.
Approach
Adaptively choose the number of tokens to unmask per step by ordering masked tokens by an error proxy (confidence/entropy/margin) and unmasking the largest prefix whose summed entropy minus maximum entropy stays under a threshold gamma, which upper-bounds joint dependence error from a KL decomposition.
Key results
- On Dream 7B MBPP, EB-sampler (entropy, gamma=0.1) reaches 58% pass@1 vs 58.8% for Top-1 with NFE 25.49 vs 101.71 under generate_until (3.99x speed-up).
- On LLaDA 8B GSM8K, EB-sampler (confidence, gamma=0.1) reaches 72.17% vs 71.79% for Top-1 with 2.08x speed-up under generate_until logic.
- On MBPP with semi-AR block generation, EB-sampler needs about 21.19 NFE for Dream 7B to match Top-1 at 64.59 NFE, roughly 3x speed-up.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use entropy-bounded adaptive unmasking at γ=0.001 instead of Top-1 | Dream 7B · MBPP | NFE: 512 -> 174.57 | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.001 instead of Top-1 | Dream 7B · MBPP | pass@1: 58.8% -> 59.2% | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | Dream 7B · MBPP, generate_until + semi-AR (block_len=64) | NFE: 64.59 -> 21.19 | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | LLaDa 8B · MBPP | NFE: 224.93 -> 64.01 | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | LLaDa 8B · MBPP | pass@1: 39.6% -> 39.2% | Hurt |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | Dream 7B · GSM8K | pass@1: 74.30% -> 74.83% | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | Dream 7B · GSM8K | NFE: 97.44 -> 35.94 | Helped |
| Use entropy-bounded adaptive unmasking at γ=0.1 instead of Top-1 | LLaDa 8B · GSM8K | pass@1: 71.79% -> 72.17% | Helped |
Surprising findings
- Under generate_until logic, Top-1 sampling of Dream 7B on MBPP yields effective tokens/step below 1, because the model unmasks tokens after the stopping phrase.
- Temperature sampling was found unhelpful for LLaDA, so zero-temperature sampling was used throughout.
Limitations
- Speed-up measurements depend on generation configuration; the authors refrain from claiming the larger 6x gain on MBPP.
- Does not revisit or remask previously unmasked tokens.
- Logic puzzle experiments use small models trained from scratch and EB-sampler drops off sharply below about 5 NFE on mazes.
Experimental setup
Datasets: HumanEval MBPP GSM8K MATH 48K generated 10x10 DFS mazes 48K generated 9x9 Sudoku puzzles
Benchmarks: HumanEval MBPP GSM8K MATH 10x10 maze navigation Sudoku
Metrics: pass@1 accuracy NFE speed-up effective tokens/step
Baselines: Top-k confidence Top-k entropy Top-k margin Top-1 sampling
Builds on: LLaDA Dream 7B MaskGIT top-k confidence/entropy/margin samplers
Closest papers by meaning
- Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling (ICLR 2025 · similarity 0.92)
- Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.92)
- Scaling up Masked Diffusion Models on Text (ICLR 2025 · similarity 0.92)
- Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions (ICML 2025 · similarity 0.92)
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.92)
- Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.