Title: KLASS: KL-Guided Fast Inference in Masked Diffusion Models
Venue: NeurIPS 2025
Accepted: yes
Keywords: Generative Models

Efficient Inference Methods
Abstract: Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its iterative refinement process, the inference is often bottlenecked by slow and static sampling speed. To overcome this problem, we introduce `KL-Adaptive Stability Sampling' (KLASS), a fast yet effective sampling method that exploits token-level KL divergence to identify stable, high-confidence predictions. By unmasking multiple tokens in each iteration without any additional model training, our approach speeds up generation significantly while maintaining sample quality. On reasoning benchmarks, KLASS achieves up to $2.78\times$ wall-clock speedups while improving performance over standard greedy decoding, attaining state-of-the-art results among diffusion-based samplers. We further validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness as a broadly applicable sampler across different models.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors tackle the well-known latency problem of discrete/masked diffusion models by proposing KL-Adaptive Stability Sampling (KLASS). In every reverse-diffusion step the model already produces a full categorical distribution for each still-masked position. KLASS exploits this \"free\" information in two ways:\n\n1.Token-level confidence (max p).\n\n\n2. Inter-step stability (KL divergence between the current and previous logits).\n\nA position whose confidence exceeds τ and whose KL trajectory stays below ε for n recent steps is deemed stable and is unmasked immediately. If no position meets the twin criteria, KLASS falls back to unmasking the top-u most-confident tokens. Because many positions converge early, the total number of sampling steps roughly halves, giving ≈2× wall-clock speed-ups without additional training or auxiliary models.\n\nThe paper supplies:\n\n1. Theory (Theorem 4.3) arguing that low-KL tokens are (near) optimal under the true data distribution, and a cost analysis (Prop 4.4).\n\n2. Experiments on four reasoning benchmarks, unconditional text generation, ImageNet 256×256, and QM9 molecule design. Gains of +1–3 pp accuracy (Math, GSM8K, HumanEval) with ≈50 % fewer steps, lower text perplexity and higher MAUVE, FID↓ and IS↑ on images, and higher QED for molecules are reported.\n\n3. A detailed checklist and an openly released, anonymised code package.\n\nThe proposal is simple, training-free, and broadly applicable.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nThe paper addresses slow inference in masked diffusion models and proposes a delightfully simple yet powerful remedy. By harvesting information the model already computes (per-token confidence and inter-step KL divergence), KL-Adaptive Stability Sampling (KLASS) doubles generation speed without retraining or auxiliary networks. The empirical study is unusually broad: the authors test on four reasoning benchmarks, open-domain text generation, ImageNet-scale image synthesis, and QM9 molecular design, and in every case KLASS maintains or improves quality while cutting steps roughly in half. The ablation analyses are informative, showing how confidence and KL interact and why multi-token unmasking is important. Reproducibility is also taken seriously. Hyper-parameters, compute budgets, and anonymised code are all supplied which makes it straightforward for the community to verify and build upon the work. Finally, the contribution is conceptually elegant. It extracts extra mileage from the model’s own logits instead of bolting on planners or performing expensive distillation, a design that is likely to be adopted quickly by practitioners.\n\nWeaknesses\n\nDespite those merits, several issues temper the enthusiasm. The theoretical underpinning (Theorem 4.3) is sketched at a high level and leaves key quantities such as the size of the \"mild context change\" and the precise δ-approximation norm undefined, so the guarantee feels more heuristic than rigorous. The results table reports only the number of diffusion steps, not actual wall-clock latency -  without single-query timings it is hard to gauge the practical speed-up!!!! The reasoning benchmarks are presented as single scores with no variance estimates, leaving open the question of statistical significance for the modest (+1 pp) improvements on GSM8K. Parameter sensitivity is explored only for math tasks, so the assertion that one set of ε and τ works universally remains unsubstantiated, and the memory/time overhead of storing KL histories for thousand-token sequences is acknowledged only qualitatively. Finally, the draft’s presentation could benefit from polishing (minor typos were found).\n\nQUESTIONS:\nWill we see this made available via a GUI or code repo?\n\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-06-25"}

comments: SUMMARY:
This paper introduces ‘KL-Adaptive Stability Sampling’ (KLASS), a fast sampling method that exploits token-level KL divergence to identify stable, high-confidence predictions. This paper also theoretically analyzes why ‘KLASS’ improves sample quality in certain scenarios.

STRENGTHS AND WEAKNESSES:
Strengths:
1. The proposed low-KL as stabe token is rational. 
2. The theoretical analyses are presented on why ‘KLASS’ improves sample quality in certain scenarios. 
3. The experiemntal results effectively verified the effectiveness. 

Weaknesses:
1. What is the relationship between KL metric and confidence. It seems that these two metrics largely overlap on most of the predicted tokens. Given that confidence metric is a basic strategy, the relationship between these two metrics is important. 
2. Compared to the original confidence setting in MaskGIT，the proposed KL metric brings very marginal performance lift.
3. The hyperparameter selection of threshold varies across tasks and models, limiting the practicality.

QUESTIONS:
Under mask based generation, the total sampling steps are usually predefined, for example, by setting the cosine scheduler. Under the proposed threshold strategy, it seems that the generated tokens at each step is uncertain. Then, how to define the final total sampling steps.

LIMITATIONS:
Please refer to the weakness part.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents KLASS, a method for accelerating sampling from masked diffusion models, based on estimating both the confidence and KL divergence between adjacent timesteps to identify stable tokens in which to unmask. Experimental results show consistent 2x speedups while maintaining performance.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- Overall approach is both theoretically sound and empirically motivated, and is very intuitive to understand.\n- Results of KLASS seem to flatly improve methods across modalities on all tested metrics, which clearly shows their approach works in practice.\n\n\n**Weaknesses**\n- Lines 144-146: I'm not sure if this is true? The plots seem to suggest that while correct samples have *slightly* higher confidence than incorrect samples across timesteps, they seem to be well within the variance bounds for significance (though proper error bars are not reported in lieu of a scatter plot, which I think could be changed). If anything, I think this argues that internal model confidence is *not* indicative of final correctness. Additionally, the error bounds for the KL divergence plot (1b) seem even larger, and it is hard to visually tell whether the plotted smoothing line is actually a reasonable fit. The proposed method obviously works, but more could be done to give explanation as to *why* empirically other than just this plot.\n- I think the authors do not fairly recognize their limitation that while their speedup method is training-free, precisely because of that its overall acceleration ability is upper-bounded and cannot achieve the order-of-magnitude speedups that distillation methods get.\n\nQUESTIONS:\n- What is \"Halton-based\" (line 247)?\n- Is there a reason the appendix seems to just contain a completely different (assumed earlier) version of the paper?\n- Lines 146-147: This should say \"KL score for corrects tends to be significantly **lower**\", right?\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces KL-Adaptive Stability Sampling (KLASS), a novel training-free sampler designed to accelerate inference in masked diffusion models. The method dynamically unmasks multiple \"stable\" tokens in each iteration by leveraging token-level Kullback-Leibler (KL) divergence and model confidence. The authors provide theoretical justification and extensive empirical validation across various domains, including reasoning, text generation, image synthesis, and molecular generation, demonstrating significant speedups and, in many cases, improved performance.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n* Novel and Intuitive Method: It intelligently exploits the internal dynamics of the diffusion process to identify which parts of the sequence have converged, allowing for more aggressive and efficient unmasking. This contrasts with more rigid, pre-defined scheduling or methods that require external \"planner\" models.\n\n* Strong and Broad Empirical Validation: The paper's most significant strength lies in its comprehensive experimental evaluation. The authors test KLASS across a remarkably diverse set of tasks and models:\n    1. Reasoning and Code Generation (GSM8K, MATH, HumanEval, MBPP): Demonstrates not just a 2x speedup but also an improvement in accuracy, which is a compelling result.\n    2. Open-Ended Text Generation: Shows improved MAUVE scores and reduced perplexity, indicating higher-quality text generation.\n    3. Image and Molecular Generation: The method's effectiveness is further solidified by its successful application to different modalities, achieving better FID/IS scores for images and higher reward metrics for molecules.\n\n* Training-Free and Efficient: KLASS is a plug-and-play sampler that requires no additional training or architectural modifications. This is a significant practical advantage over methods that require auxiliary models, making it easy to adopt and implement with existing masked diffusion models.\n\n* Theoretical Justification: The paper provides a theoretical argument to ground the heuristic. The theorem formally links suboptimal token choices to high KL divergence under context shifts, providing a principled reason why filtering by KL stability should lead to better and more robust generation.\n\n* Thorough Ablation Studies: The authors include valuable ablation studies that analyze the interplay between confidence and KL thresholds and the benefit of multi-token unmasking over single-token strategies.\n\n**Weaknesses**\n* Hyperparameter Sensitivity: The effectiveness of KLASS appears to be dependent on several key hyperparameters: the KL threshold, the confidence threshold, and the history length. The paper provides the values used for each experiment but could benefit from a more in-depth analysis of the sensitivity to these parameters. A more detailed discussion on how to select these values for new tasks or models would enhance the practical utility of the paper.\n\n* Computational Overhead of KL Divergence: While the paper argues that the overhead of computing KL divergence is minimal compared to the cost of a full diffusion step, this overhead is not zero. In scenarios with very large vocabularies, calculating KL divergence for every token at every step could become a non-trivial computational cost. A more detailed analysis of this overhead in practice, especially in relation to vocabulary size, would be beneficial.\n\n* Comparison to Other Adaptive Sampling Methods: While the paper compares against several strong baselines, the field of accelerating diffusion models is rapidly evolving. The \"Related Works\" section could be strengthened by a more direct comparison to other contemporary adaptive sampling techniques, even if they are for continuous diffusion, to better contextualize the contribution. For instance, methods that dynamically adjust the number of sampling steps based on other criteria could be relevant points of comparison.\n\n* Were the inference time costs measured in the experiments? It would be insightful to report wall clock times.\n\nTypo on line 142 (\"reasoining\")\n\nQUESTIONS:\nSee weaknesses above.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}