Research map / Image generation / Fast diffusion sampling
KLASS: KL-Guided Fast Inference in Masked Diffusion Models
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 KLASS, a training-free sampler for masked diffusion models that unmasks multiple tokens per step when they have low token-level KL divergence between consecutive timesteps and high confidence, speeding up inference while matching or improving accuracy.
Problem
Masked diffusion models' iterative sampling is slow, and static schedules (Top-k, one token per step) are inefficient and prone to premature unmasking of incorrect tokens.
Approach
Use token-level KL divergence between consecutive denoising steps as a stability signal, combined with a confidence threshold, to select a dynamic set of stable tokens to unmask in parallel at each step, with a Top-u fallback; no extra training or external planner.
Key results
- On reasoning benchmarks, KLASS with LLaDA improves MATH accuracy from 31.4 (Top-1) to 33.8 while reducing steps to 128.6 from 256, and reaches up to 2.78x wall-clock speedup (Dream on HumanEval).
- On unconditional OpenWebText generation with MDLM, KLASS lowers LLaMA2 generative perplexity from 30.88 to 26.94 and raises MAUVE from 0.115 to 0.179 at 512 steps.
- On MMaDA ImageNet generation at 16 steps, KLASS reduces FID from 34.48 to 30.48 and raises IS from 75.72 to 93.07 over the confidence sampler.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply KLASS instead of Top-1 decoding | LLaDA 8B Instruct · MATH | accuracy: 31.4 -> 33.8 | Helped |
| Apply KLASS instead of Top-1 decoding | LLaDA 8B Instruct · GSM8K | accuracy: 75.13 -> 76.50 | Helped |
| Apply KLASS instead of Top-1 decoding | Dream 7B Instruct · GSM8K | accuracy: 79.55 -> 79.43 | Hurt |
| Apply KLASS instead of Top-1 decoding | Dream 7B Instruct · MATH | accuracy: 37.97 -> 43.20 | Helped |
| Use KL-guided adaptive unmasking instead of MDLM sampling | MDLM · OpenWebText | MAUVE: 0.115 -> 0.179 | Helped |
| Use KL-guided adaptive unmasking instead of MDLM sampling | MDLM · OpenWebText | LLaMA2 perplexity: 30.88 -> 26.94 | Helped |
| Use KLASS rather than confidence-based sampling at 16 steps | MMaDA · ImageNet validation | FID: 34.48 -> 30.48 | Helped |
| Use KL-guided sampling instead of MDLM | MDLM · QM9, QED | Reward: 0.526 -> 0.546 | Helped |
| Use KL-guided sampling instead of MDLM | MDLM · QM9, Ring count | Reward: 4.123 -> 4.258 | Helped |
| Unmask multiple stable tokens per step rather than only the highest-confidence stable token (ablation) | LLaDA 8B Instruct with KLASS criteria · MATH | accuracy: 31.2 -> 33.8 | Helped |
| Unmask multiple stable tokens per step rather than only the lowest-KL stable token (ablation) | LLaDA 8B Instruct with KLASS criteria · MATH | accuracy: 29.0 -> 33.8 | Helped |
| Unmask multiple stable tokens per step rather than only the highest-confidence stable token (ablation) | LLaDA 8B Instruct with KLASS criteria · GSM8K | accuracy: 72.86 -> 76.50 | Helped |
Surprising findings
- Confidence alone is insufficient: the KL-only sampler and confidence-only sampler each underperform the combined rule, and single-token KL/confidence variants lose accuracy relative to parallel KLASS.
- Parallel KLASS unmasking improves accuracy over sequential single-token unmasking from the same stable pool (e.g., MATH 33.8 vs 31.2 for single-confidence).
Limitations
- Optimal thresholds vary substantially between models and tasks, requiring a validation-based hyperparameter search.
- Theoretical rationale is a proposition bounding KL along context paths under strong assumptions (conditional delta-approximation), not a full guarantee of sample quality.
- Image and molecular experiments are limited in scale and comparisons (two samplers on one model; MDLM baseline only for molecules).
Experimental setup
Datasets: GSM8K MATH500 HumanEval MBPP OpenWebText ImageNet QM9
Benchmarks: GSM8K MATH500 HumanEval MBPP OpenWebText ImageNet QM9
Metrics: pass@1 accuracy sampling steps wall-clock time generative perplexity MAUVE entropy FID Inception Score QED reward ring count reward
Baselines: Top-1 greedy decoding Random unmasking Top-2 decoding Confidence threshold sampling KL threshold sampling MDLM SEDD D3PM AR MMaDA confidence sampler
Builds on: masked diffusion models MDLM LLaDA Fast-dLLM
Closest papers by meaning
- Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.90)
- Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.90)
- Encoder-Decoder Diffusion Language Models for Efficient Training and Inference (NeurIPS 2025 · similarity 0.89)
- Learn from Your Mistakes: Self-Correcting Masked Diffusion Models (NeurIPS 2026 · similarity 0.89)
- Attention-Discounted Adaptive Sampler for Masked Diffusion Language Models (NeurIPS 2026 · similarity 0.89)
- Remasking Discrete Diffusion Models with Inference-Time Scaling (NeurIPS 2025 · similarity 0.89)
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.