Source-linked AI summary
Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking
Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, Brian Karrer
TL;DR
Efficient sampling from masked diffusion models is underexplored because independently unmasking multiple tokens can harm generation quality. The paper introduces EB-Sampler, an entropy-bounded adaptive sampler that dynamically unmaskes multiple tokens using existing models, yielding 2-3x speed-ups without performance loss.
Problem
Efficient MDM sampling remains limited because independently unmasking more tokens can quickly degrade generation quality.
Method
EB-Sampler adaptively chooses which and how many tokens to unmask using an entropy bound that approximately limits their dependence, as a drop-in replacement requiring no further training.
Results
2-3x speed-ups without loss in performance were obtained on coding and math reasoning benchmarks, with validation on maze navigation and Sudoku.
Takeaways & Limitations
EB-Sampler improves the compute-versus-performance trade-off for masked diffusion sampling across math, code, and reasoning evaluations.
Takeaways & Limitations
Future work is needed to learn parameterized adaptive samplers and extend EB-Sampler to revisit previously unmasked tokens.
Abstract
from arXiv · showhide
Recent masked diffusion models (MDMs) have shown competitive performance compared to autoregressive models (ARMs) for language modeling. While most literature has focused on performance enhancing sampling procedures, efficient sampling from MDMs has been scarcely explored. We make the observation that often a given sequence of partially masked tokens determines the values of multiple unknown tokens deterministically, meaning that a single prediction of a masked model holds additional information unused by standard sampling procedures. Based on this observation, we introduce EB-Sampler, a simple drop-in replacement for existing samplers, utilizing an Entropy Bounded unmasking procedure that dynamically unmasks multiple tokens in one function evaluation with predefined approximate error tolerance. We formulate the EB-Sampler as part of a broad family of adaptive samplers for which we provide an error analysis that motivates our algorithmic choices. EB-Sampler accelerates sampling from current state of the art MDMs by roughly 2-3x on standard coding and math reasoning benchmarks without loss in performance. We also validate the same procedure works well on smaller reasoning tasks including maze navigation and Sudoku, tasks ARMs often struggle with.
1 Introduction
Masked diffusion models offer competitive language modeling, but efficient sampling remains difficult because independently unmasking multiple tokens degrades generation quality. EB-Sampler adaptively selects and jointly bounds multi-token unmasking, achieving faster sampling without performance loss.
- MDMs have shown competitive performance with similarly sized autoregressive models on code, text, and mathematical reasoning benchmarks.
- Efficient MDM sampling is challenging because full attention prevents key-value caching, while independently unmasking more tokens quickly degrades generation quality.
- EB-Sampler adaptively decides which tokens to unmask and how many to unmask using an entropy bound that approximately limits dependence among unmasked tokens.
- EB-Sampler is a drop-in replacement that applies directly to existing masked diffusion models without further training.
- 2-3x speed-ups on standard coding and math reasoning benchmarks occur without loss in performance, with additional validation on maze navigation and Sudoku.
- The adaptive sampler family is supported by an error decomposition separating model error from joint dependence error.
2 Preliminaries
Masked diffusion models operate on discrete token sequences by iteratively replacing mask tokens with vocabulary tokens. Their factorized conditional predictions can support sequential unmasking, while EB-Sampler illustrates adaptive multi-token unmasking.
- A masked diffusion model generates fixed-length discrete sequences by iteratively replacing mask tokens with vocabulary tokens until all positions are unmasked.
- The model learns factorized conditional predictions for masked tokens given any set of unmasked tokens, ideally matching clean-data conditionals.
- Sequential sampling can unmask one token at a time, while EB-Sampler’s illustrated step selects tokens and determines the number to unmask using prediction entropies.
3 Known challenges of MDM sampling
MDM sampling faces local model error from imperfect conditionals and joint dependence error from independently unmasking dependent tokens. Greedy model-guided orders can reduce total model error, while fixed-size multi-token unmasking trades efficiency for increasing dependence error.
- Order of unmasking matters: For optimal factorized conditionals, unmasking order does not change the model distribution, but finite capacity and data create order-dependent model error.The paper formalizes nonzero local KL divergence between data and model conditionals.
- Order of unmasking matters: Greedy samplers outperform random unmasking, indicating that model-guided orders with lower total model error exist.Figure 2 supports confidence, entropy, and margin criteria as local model-error proxies.
- Two sources of error: MDM sampling errors comprise local model error from incorrect conditionals and joint dependence error from independently sampling non-independent tokens.The latter is identified with joint mutual information and can grow when multiple tokens are unmasked together.
- Sampling efficiency: Standard language MDM samplers predict one token per function evaluation, limiting efficiency because full attention prevents ARM-style key-value caching.This motivates using predictions for multiple masked tokens per evaluation.
- Sampling efficiency: Top-k samplers unmask a fixed k tokens independently, with larger k improving efficiency but increasing joint dependence error and degrading performance.The paper reports this degradation for k ∈{1, 2, 4, 8, 16} across error proxies.
4 Entropy Bounded (EB) Sampler
EB-Sampler replaces fixed-size Top-k unmasking with adaptive multi-token selection. It uses model-error ordering and an entropy bound to vary the number of tokens unmasked according to estimated dependence.
- Evaluation: Figure 3 evaluates the efficiency-accuracy tradeoff of Top-k sampling on MBPP using NFE.The caption identifies the plotted comparison but does not state a specific outcome.
- EB-Sampler: EB-Sampler is a direct replacement for Top-k samplers motivated by controlling both local model error and joint dependence error.Its design combines low-error token selection with bounded dependence during parallel unmasking.
- EB-Sampler: EB-Sampler first sorts masked tokens by an error proxy, then selects a subset whose entropy approximately bounds the introduced joint dependence error.Low-error tokens are prioritized before the entropy-based subset selection.
- Implementation: EB-Sampler requires only a minimal PyTorch change relative to common Top-k implementations.The procedure remains compatible with adjustments such as temperature and unsupervised classifier-free guidance.
- Adaptive unmasking: Unlike Top-k sampling, EB-Sampler determines k dynamically from the entropy bound rather than fixing the number of tokens per step.It unmasks more tokens when predicted dependence is low and fewer when dependence may be high.
5 Adaptive unmasking samplers
The paper embeds EB-Sampler in a general family of adaptive samplers represented by ordered partitions of token indices. An error decomposition separates model error from joint dependence error and motivates controlling the latter after selecting low-error tokens.
- Adaptive sampler family: An ordered partition z=(z1,z2,...,zd) represents varying-length unmasking steps, with each zi containing token indices unmasked at step i.Partition validity requires disjoint subsets drawn from the full index set, with empty subsets allowed under the stated convention.
- Adaptive sampler family: The adaptive-sampler family produces a joint distribution over sequence state and unmasking partition using approximate clean-data conditionals from a pretrained model.The distribution over partitions determines which token indices are unmasked at each step.
- Expressiveness: This family includes fixed-size Top-k samplers, random unmasking, and procedures that dynamically determine how many tokens to unmask.Examples include Top-k margin, entropy, and confidence samplers.
- Error decomposition: The sampling error decomposes into model error from incorrect conditionals and joint dependence error from treating dependent tokens as independent.The second component is identified with joint mutual information.
- Choosing the sampler: After identifying low-model-error tokens, the sampler can control overall error by selecting subsets with bounded joint dependence.EB-Sampler applies the model-error proxies from the earlier section to implement this strategy.
6 Experiments
The experiments evaluate EB-Sampler’s efficiency–accuracy tradeoff on code, math, maze, and Sudoku tasks. EB-Sampler improves sampling efficiency while preserving performance across these settings.
- Experimental setup: The experiments compare EB-Sampler with Top-k samplers using confidence, entropy, and margin error proxies across language and logic tasks.The language evaluation uses LLaDa 8B Base and Dream 7B Base on HumanEval, MBPP, GSM8K, and Math; maze and Sudoku use small discrete diffusion models.
- Code and math reasoning: 2-4x speed-ups are obtained versus the Top-1 sampler at the same accuracy across datasets and error proxies.NFE is measured as the average number of function evaluations until task-specific generate_until logic is satisfied.
- Efficiency measurement: 21.19 function evaluations achieve the same MBPP performance as comparison schemes, with tokens generated at 2.4 per step.The mean answer length is around 50 tokens, while semi-autoregressive block generation requires 64.59 NFE in this test case.
- Maze navigation: EB-Sampler preserves most maze accuracy until a sharp drop below 5 NFEs, whereas Top-k baselines decline more steeply around 10 NFEs.Maze accuracy is defined as the fraction of validation mazes fully solved.
7 Related Work
Related work has improved discrete diffusion sampling performance and explored efficiency through solvers, schedules, distillation, and speculative methods. EB-Sampler differs by adaptively controlling multi-token unmasking for masked diffusion models.
- Performant sampling for discrete diffusion: Prior work improves MDM sampling through token-order planning, remasking, predictor-corrector iterations, and forward-backward sampling.These methods primarily address which tokens to unmask or whether to revisit masked states.
- Efficient sampling for discrete diffusion: Efficiency-oriented approaches include higher-order solvers, globally optimized non-adaptive schedules, MaskGIT unmasking, and distillation into one-step generators.EB-Sampler determines the number of tokens per step adaptively rather than using a fixed schedule.
- Speculative decoding: Speculative decoding accelerates language models with a draft model, but full-attention MDMs cannot directly apply it because target sequence probabilities are expensive to compute.EB-Sampler is complementary and could also speed up an MDM used as a speculative draft model.
8 Conclusions and Future Work
The paper concludes that EB-Sampler is a theoretically grounded, training-free adaptive sampler for masked discrete diffusion and flow models. It improves the compute–performance tradeoff, while future work could learn adaptive policies or revisit previously unmasked tokens.
- Conclusions: EB-Sampler adaptively controls which and how many tokens to sample using an interpretable entropy bound and serves as a drop-in replacement.The method is evaluated on math, code, and reasoning benchmarks with contemporary diffusion models.
- Conclusions: 2-3x speed-ups are achieved without loss of performance, with EB-Sampler outperforming existing samplers on the compute-vs-performance Pareto frontier.The reported conclusion covers the paper’s evaluated masked diffusion and flow-model sampling setting.
- Future work: Future work could learn a parameterized adaptive sampler, optimize the KL bound with respect to ϕ, or incorporate revisiting past unmasked tokens.These directions are presented as possible extensions for improving efficient and performant sampling.
A.1 KL divergence error decomposition
The KL-divergence decomposition separates sampling error into joint dependence error and model error. Under deterministic sampling schemes, the relevant KL quantities coincide, so optimizing the divergence directly optimizes likelihood.
- Joint dependence error: The decomposition identifies joint dependence error as the KL divergence between joint and factorized conditional distributions within the unmasking subset.This term measures dependence lost when multiple tokens are sampled independently.
- Model error: Model error is the sum of KL divergences between factorized conditionals, which are the conditionals pθ is trained to learn.The derivation obtains this term after marginalizing variables absent from the expectation.
- KL equality: The two KL divergences are equal only when qϕ(z|x) = pϕ(z|x), a condition occurring under special ϕ.Thus equality is not generally available for arbitrary sampling procedures.
- Deterministic sampling: For deterministic ϕ, the partition is determined by x, making the posterior distributions point masses and equating KL optimization with likelihood optimization.This deterministic case is relevant because many schemes in the main text are deterministic or nearly so.
- EB-Sampler procedure: EB-Sampler initializes masked-token predictions, sorts tokens by an error proxy, and adds tokens while an entropy bound remains below γ.The procedure samples selected tokens and stops adding tokens when the bound is exceeded.
C.1.1 Datasets
The experiments evaluate EB-Sampler on coding and mathematical reasoning benchmarks using LLaDa 8B and Dream 7B under controlled sequence-length and generation-length settings. Runtime measurements use the same 8×H100 computational environment, with evaluation-specific handling for Dream 7B on HumanEval.
- Datasets: The benchmark suite includes HumanEval, MBPP, GSM8K, and Math, spanning code generation and mathematical reasoning tasks.The tasks use 0-shot, 4-shot, or 8-shot prompting as specified for each benchmark.
- Models and setup: EB-Sampler is evaluated on the state-of-the-art MDMs LLaDa 8B and Dream 7B.Their maximal sequence lengths are 4096 and 2048, respectively.
- Generation setup: Generation pads prompts with mask tokens and permits unmasking only within the designated max_gen_len range.Generation stops after all tokens in that range are unmasked.
- Efficiency measurement: All benchmarks were run on 8×H100, while runtime comparisons report confidence and entropy proxies and note extra sorting cost for margin.LLaDa 8B also runs longer than Dream 7B because its maximal sequence length is twice as large.
- Accuracy evaluation: Dream 7B HumanEval outputs required special handling because raw outputs produced around an 8% performance drop relative to reported results.The evaluation investigated answer-template formatting as a possible source of the discrepancy.
D Additional experiments - code and math reasoning
Additional experiments compare pass@1 accuracy with full-sequence function evaluations on code and math reasoning tasks. Margin-based results are included alongside confidence and entropy results, but typically show inferior full-NFE accuracy.
- Full-sequence NFE: Figures 9 and 10 plot pass@1 accuracy against full max_gen_len NFE for code and math reasoning tasks, respectively.These plots examine accuracy as a function of full-sequence sampling cost.
- Margin error proxy: Margin error proxy usually yields inferior accuracy at full NFE, so its curves are omitted from the main-body plots for readability.The appendix presents margin results together with confidence and entropy results.
D.2 Measuring efficiency of MDMs
The paper distinguishes apparent efficiency from answer-generation efficiency by using generate_until stopping logic and effective tokens per step. These measurements motivate semi-autoregressive generation when samplers continue unmasking tokens after the answer has ended.
- Measuring efficiency: Full-sequence generation can overstate efficiency gains because tokens beyond the evaluated answer are still generated.The paper therefore proposes generate_until-based measurements and semi-autoregressive variants.
- Effective tokens per step: Effective Tokens/Step equals mean_answer_len divided by mean_NFE_to_condition.mean_NFE_to_condition counts evaluations needed to generate through the answer markers and unmask preceding tokens.
- Generate-until evaluation: Figure 11 reports pass@1 accuracy against NFE using generate_until logic for code and math reasoning tasks.This evaluates sampling cost up to the stopping condition rather than the full padded sequence.
- Post-stopping tokens: At γ = 0 or Top1, effective speed is often below 1 because models unmask tokens unused in the final answer or after the stopping phrase.This observation motivates semi-autoregressive generation to reduce post-stopping unmasking.
- Semi-autoregressive generation: Semi-autoregressive generation is especially relevant for MBPP, where many tokens are generated after the stopping phrase before all earlier tokens are unmasked.The paper reports an MBPP ablation with Dream 7B and LLaDa 8B.
- Semi-autoregressive generation: On GSM8K, semi-autoregressive generation changes efficiency little because effective tokens per step are around 1 and post-stopping generation is less likely.This comparison is reported for both Dream 7B and LLaDa 8B.