KLASS: KL-Guided Fast Inference in Masked Diffusion Models
Authors: Seo Hyun Kim
Sunwoo Hong
Hojung Jung
Youngrok Park
Se-Young Yun
Venue: NeurIPS 2025 spotlight ()
Keywords: Generative Models
Efficient Inference Methods
Primary area: deep_learning
TL;DR: We propose KLASS, a fast KL-guided sampling method for masked diffusion models that improves accuracy while cutting inference time by over 2x.

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.

KLASS: KL-Guided Fast Inference in Masked Diffusion Models
Seo Hyun Kiml
Sunwoo Hong
Hojung Jung' KAIST AI
Youngrok Park' Se-Young Yun'
(shkimsally, sunwoo, Jeaass.
Abstract
Masked diffusion models have demonstrated competitive results on various tasks including language generation. However, due to its iterative refinement process, the inference often bottlenecked by slow and static sampling speed. To overcome this problem, we introduce 'KI "KL-Adaptive Stability Sampling' (KLASS), a fast yet ef- fective sampling method that exploits token-level KI. divergence to identify stable, high-confidence predictions. By unmasking multiple tokens in each iteration with- out any additional model training, our approach speeds up generation significantly while maintaining sample quality. On reasoning benchmarks. KLASS achieves up to 2.78x wall-clock speedups while improving performance over standard greedy decoding, attaining ttate-of-the-ar results among diffusion-based samplers. We fur- ther validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness broadly applicable sampler across differ- ent models. Our code available at https: //gi thub. com/ com/shkim0) 16/KLASS.
1 Introduction
Masked diffusion models [1, 28, 34, 38] have attracted growing attention for their ability to model joint distribution of sequences by iteratively refining samples from partially masked sequences to clean data, achieving competitive performance complex language tasks [27], image generation [7]. biological sequences [25, 34], and planning algorithms [50, 51].
Despite recent successes, these models are often restricted by slow and static sampling strategics such Top-k or stochastic sampling, where only limited number of high-confidence tokens unmasked at each step. As a result, the generation process can become inefficient and prone to local suboptimalities, thus constraining the practical applicability of masked diffusion approaches.
Several works investigate efficient samplers by caching the logits no tokens are unmasked at the specific timestep [34] or design specific scheduler to unmask one token at time [56]. Another natural solution might be to incorporate an additional "planner" or auxiliary distribution to guide sampling [48, 55]. However, doing typically incurs substantial computational overhead, increases inference latency, and can lead to difficulty aligning the planner's distribution with the base model's leamned distribution. Instead, our goal is to develop lightweight yet effective sampling method that remains within the model's 's own capabilities, yielding speedups in generation while simultaneously improving or maintaining overall accuracy.
To address these challenges, we propose 'KL-Adaptive Stability Sampling' (KLASS), an adaptive sampling strategy that leverages the diffusion model's own feedback to guide unmasking. Unlike previous approaches that rely on fixed schedules (i.e., a predetermined number of tokens unmasked at each timestep), our method adapts to the evolving confidence of the model during generation. We accelerate inference by identifying stable tokens as low-risk candidates for early unmasking. To
"Equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Top-k Confidence
LIADA Conrect Incorrect
DREAM Correct Incorrect
herefoore the number of cars that drave through in the first 15 minutes snfn vinn 20 20 10hm
Therafore, the number cars that drove through the traffi jam in the first: 15 minutes Melboxed Po..
KLDverreence
KLDivergence
0.08
0.08
KLASS
0.06
0.06
Therofore, the number care drove through
in the first 15 minutes isfn n 20 20 Din 10
Thersfore, the number of cars that drove through the traffi jam in the first 15 minutes NeYboredi
0.04
cont 8.7587, 0.0193
MATH HumanEval MBPP
MATH HumanE val
(a) Case study comparing Top-k confi- (b) Average KL divergence of tokens unmasking for correct and dence and KLASS solutions.
incorrect predictions on LLaDA and DREAM.
Figure I: KL divergence as a strong indicator of solution correctness. (a) The Top-k method selects an incorrect solution despite high confidence, whereas KLASS identifies the correct solution, which exhibits a significantly lower KL divergence. (b) KL divergence distributions for the LLaDA and DREAM models show that correct predictions consistently have lower KL divergence than incorrect ones across all datasets.
quantify stability, we track the token-level Kullback-Leibler (KL) divergence between conditional distributions at consecutive timesteps. Tokens are unmasked when their distributions remain similar (KL below threshold) and predicted with high confidence (probability exceeding confidence threshold). This dynamic allocation of unmasking tokens results in significant acceleration of generation speed while maintaining sample quality by avoiding premature suboptimal token unmasking without additional model training or extra memory burden.
We empirically validate our method on challenging reasoning benchmarks, including GSM8K, MATH, HumanEval, and MBPP. We show that applying KLASS with large-scale masked diffusion models not only halves the number of sampling steps compared to standard greedy or Top-k decoding [19], but also achieves higher accuracy, achieving state-of-the-a results compared to other diffusion samplers. Figure la presents comparison between solutions generated by Top-k confidence and KLASS sampler. KLASS successfully identifies the correct token with lower KL, whereas Top-k confidence tends to unmask incorrect tokens even with higher confidence. Furthermore, text generation also proves the effectiveness our method which results in reduced perplexity while maintaining entropy, thereby mitigating the inefficiencies inherent in conventional sampling. We further show that our sampler works in other modalities, including images and molecules.
experiment plain
Overall, our proposed sampler for masked diffusion models is both simple and practical, harnessing the latent potential of the base diffusion model itself, rather than relying on complex external planners. By strategically identifying stable tokens at each iteration, the algorithm accelerates generation and fosters more robust coverage of viable token candidates. We believe this work provides a practical and scalable way for large-scale masked diffusion models, particularly where reliable and efficient generation is essential, such as complex reasoning tasks.
We summarize our main contributions below:
We propose KLASS, training-free ffre sampler that leverages the model's internal dynamics in terms of token level KL divergence and confidence without requiring external planners.
We achieve up to 2.78 x faster sampling by more than halving the number of diffusion steps through parallel unmasking of stable tokens.
We provide comprehensive empirical validation, showing improved quality on reasoning bench- marks across math and code generation, text generation, image synthesis, and molecular generation.
2 Related Works
Discrete diffusion models D3PM [1] investigate how forward and backward processes can be constructed in discrete state spaces which is analogous to the continuous diffusion models [18, 40]. [6] leverage continuous time Markov chain (CTMC) theory to formulate the borward-backward process of discrete diffusion models with providing negative ELBO in continuous time limit as an objective. Following the success of denoising score matching [41], Lou et al. [25], Meng et al. [26] suggest

---

discrete score matching loss by defining Stein score in discrete space. Ou et al. [28], Sahoo et al. [34], Shi al. [38] further shows that simplified version of masked diffusion model can significantly boost the performance of diffusion models closing the performance gap with AR models in language domains. Recently, L.LaDA [27] demonstrates scaling law of discrete diffusion models in language domain and further shows reasoning abilities.
Discrete diffusion samplers Generating a text from language diffusion models involves iteratively refining sequence from a noisy or masked state. Ancestral Sampling [25, 34] starts from fully masked sequence and iteratively applies the learned reverse denoising process over a series of discrete timesteps to produce a clean sequence. SUBS parametrization [34] of the reverse step dictates how model predictions are used to unmask tokens, often by ensuring that already revealed tokens remain unchanged. To improve sample quality, ReMDM [43] adopts remasking strategies, where some newly predicted tokens are reset to a mask based on confidence or timestep.
Accelerated Sampling of Discrete diffusion models The iterative nature of ancestral sampling can result in high latency due to the large number of sequential steps. Consequently, much research has focused on reducing the number of function evaluations (NFEs) in diffusion models. Deschenaux and Gulcehre [11], Hayakawa et al. [15] leverage distillation methods to train the model with reduced NFEs in analogous to fast sampling of continuous diffusion models [35, 41, 533. Ren et al. [33] improve discrete diffusion solvers by considering second-order numerical solver in CTMC framework. Zheng et al. [56] propose First-Hitting Sampler (FHS) to skip the unnecessary timesteps and unmask one token at a time. Most of the existing samplers of masked diffusion models, however, resort to additional training rely other models (i.e., planners) to choose unmasking tokens each timestep 121, 24, 29]. This could help avoiding suboptimal token selection but with considerable computational overhead and may fail to be aligned with the model's intrinsic capability.
Recent training-free strategies for accelerating masked diffusion language models have emerged concurrently, with several works exploring heuristics based on model certainty to guide this process. Fast-dLLM [47] and Dimple [54] use confidence-awarc decoding. SlowFast Sampling [45] alternates decoding stages based on certainty, convergence, and position principles, EB- EB-Sampler [4] unmasks multiple tokens based on entropy bounds, and Prophet [23] uses the Top-2 confidence gap. While these concurrent approaches validate the utility of heuristics largely based on certainty, empirically demonstrate that this signal alone is insufficient. To ensure tokens are not unmasked prematurely, we propose novel method that utilizes KL divergence to identify stable tokens for parallel decoding.
3 Preliminaries
3.1 Masked diffusion models
In masked diffusion models, requires additional mask index m for each tokens and forward process is defined by following absorbing process [1]:
(z) = Cat(2, xx + (1 - -x.))))))
where at is predefined schedule, monotonically decreasing in t. Then one can analytically obtain posterior distribution as:
Cat (t)
if 21 * m,
g(z, Zz, x)
Cat
Zz: (1- a, (m.
if z m.
1-ac
The goal of the masked diffusion model is to learn this reverse process by parameterizing the posterior (Eq. 2) by a neural network with po (z, (z:) == q(2, Iz,, Mo(z,, t)).
In simplified masked diffusion models [28, 34, 38), learning objective can be simplified by parame- terizing the models to focus on estimating only masked tokens while maintaining unmasked tokens throughout the generation.
Then the learning objective to minimize Negative ELBO (NELBO) whose continuous form the following:
3

---

Confidence max (lubaa)
KL: Dau. Ciluaahh Iulluss)
Satiafied
6-7
lotal
lakes
Confidence: mux( anlyy 0.928
0.9
lakes
Mask
Both Satisfled
lotal
[akes]]
bolls
Confidence: max( M
0.9
Mlaall:
Token probabiity distribution,
Masked,
Unmasked
Fixed
KL: DEL C.ntthh Jnn
0.01
Unmask
Figure 2: Illustration of parallel decoding with KLASS. Tokens are unmasked when they meet the two criteria: high predictive confidence and stable probability distribution. Stability is measured by a low KL divergence between consecutive steps (illustrated with history length of for simplicity). On the right it shows the sampling process for position 245: it remains masked due to low confidence or high KL score, and unmasked when both conditions are satisfied.
2o E. Rx (x { 0a 1-a4 az
16x.m log pno (Ze, ()]
Here, Go denotes data distribution and af is the derivative noise schedule a in time. In this continuous time framework, [34] further proves that above objective is invariant of noise schedule aq. The above can be generalized sequence-level of token length L modeling as follows.
[/'
oi - Ex (2 x)
((L)
Σ x(1) log u
1 (2),t)
dt.
(4)
1x2 =m
3.2 Inference via Ancestral Sampling
At inference, we discretize t E [0, 1] into times (tr > >... ... > t1 R 0}, initializing X [mask]] "
We then sample backward:
X po Xt._
xx.): T.
In simplified MDM, unmasked tokens remain fixed and masked tokens are drawn from the model's prediction. After T steps, we obtain a complete sequence Xto- We provide additional analysis of other sampling strategies in Appendix C.
4 Method
4.1
Defining Confidence Score and KL Score
KLASS aims to identify which tokens are stable enough to be unmasked at each step of the inference process, which we index by timesteps
T...., 1. To guide this selection, we introduce key metrics: a confidence score to measure the model's certainty on given token and a KL score to measure the temporal consistency of its predictions.
Definition 4.1. (Confidence score) Denoting Pt as the categorical distribution predicted by the diffusion model at timestep for token position i, we define the confidence conf; to be the largest value of the probability function among vocabulary space V (v E VV:
conf: max p(()).
Intuitively, higher confidence score indicates the model more certain about estimating the current token, which increases the chance that the model's estimate for that token is correct.

---

Definition 4.2. (KL score) We define KL score d; of the token position i at timestep as the Kullback-Leibler divergence between previous estimates and current estimates of the given token:
d =
DKL (p' P(+1),
(6)
where we denote ph: Pi+1 be the probability distribution of the model estimates of token index at time and at time + 1, respectively.
KI. score should be low only when the model's s estimate is consistent throughout the reverse diffusion process, which implies the estimated token more reliable.
To empirically demonstrate how KL score behaves in practical scenario, we first generate samples for variety of math and programming reasoning benchmarks. As shown in Figure 1b, correct samples consistently exhibit significantly lower KL scores than incorrect ones, for all models and datasets. This observation motivates our use of KL scores guiding signal in the sampling algorithm of masked diffusion models, which we formally introduce in the next section.
4.2 KLASS: KL-Adaptive Stability Sampling
We introduce KLI-Adaptiv Stability Sampling' (KLASS), novel sampling algorithm for masked diffusion models. As illustrated Figure 2, KLASS leverages confidence score and KI. score during the unmasking process of the masked diffusion models (Eq. 2), by selectively choosing unmasking tokens that have low KL score and high confidence score.
Stable-token selection. To effectively set the standard using both KL. and confidence score, we propose stable-token selection in the following way: Given history length n, KL threshold éKL. and a confidence threshold т, we select the set of stable tokens at step as,
S.
vk (1, (1 n)
DKL 1l P1++)
CKL A conf >> >3}
all recent KL's below threshold
high confidence
Unmasking rule.
KLASS adaptively chooses which tokens to unmask given timestep with above defined stable index (Eq. 7). At each diffusion step I, we apply
unmask token at position i,
St: S, 0,
8
otherwise, unmask the Top-u positions by conf: e
where u a fixed fallback unmasking count. provide pseudocode of our algorithm with further analysis in Appendix B.
Theoretical Rationale
We provide theoretical perspective why using KL divergence can improve sample quality. We show that, for a well-trained model, a token that is predicted as incorrect at the current step cannot remain uniformly stable the context is progressively resolved.
Definition 5.1. For each context (iinstantiation of variables outside Xj). let C(c) be the nonempty set of task-correct conditionals. Let C := =t:: u(- c) E C(c) Vel. We say pa is conditional 3-approximation the task if
inf sup TV(po( c), T(. [c) < 6.
REC
Definition 5.2. Fix i. Let x, be optimal under x(. (²)) at near-optimal context ct. Let x! + xx be suboptimal. Assume true margin > 0 al et: <xx] c*) 2 T c*) + Assume the model currently prefers 2t at cM by margin 3 2 O: po(x! cM) 2 (x cm) + B.
Proposition 5.3. Suppose po is a conditional 8-approximation of म. For any context path eM - CM-11 en (changing only variables outside X,) ending at = c*, let P. == po (t) and : (3 + 2.) Then
M-1
2
Po)
Σ
(14) 2
M
M
Proof. The proof is in Appendix A.
5

---

Table Performance and sampling steps on reasoning benchmarks for different diffusion samplers.
MATH
GSM8K
HumanEval
MBPP
Method
Parallel
Acct
Steps
Acct
Steps
Acct
Steps..
Acct
Steps.
LLaDA
Top-1
31.4
256
75.13
256
39.63
256
46.69
256
Random
X
26.2
256
67.10
256
20.21
256
29.18
256
Top-2
29.6
128
72.40
128
33.54
128
37.74
128
Confidence
31.6
96.46
75.21
74.35
37.80
54.41
47.08
85,20
KL divergence
32.6
172.21
74.52
155.88
40.24
111.93
45.53
150.47
KLASS (ours)
33.8
128.62
76.50
98.57
40.85
91.98
47.86
119.59
Dream
Top-1
37.97
256
79.55
256
58.53
256
63.81
256
Random
X
18.73
256
37.35
256
18.09
256
28.14
256
Top-2
33.60
128
71.69
128
42.88
128
47.08
128
Confidence
41.80
95.10
73.67
74.81
50.00
52.47
57.59
72.49
KI. divergence
41.27
162.49
76.70
150.02
59.35
73.94
62.65
108. 15
KLASS (ours)
43.20
149.72
79.43
155.67
59.35
74.88
64.59
111.24
Remarks. A token that is wrong at but correct at c* must be dynamically unstable somewhere along the path: its average per-step KL is bounded away from In essence, incorrect tokens cannot remain dynamically stable. Accordingly, KLASS delays unmasking until tokens exhibit dynamic stability thereby improving generation quality.
Experiments
To show effectiveness our proposed sampler, we conduct experiments on multiple benchmarks including reasoning benchmarks with large scale models in Section 6.1, text generation in Section 6.2, along with other modalities including images in Section 6.3 and molecules Section 6.4. We also present ablation studies in Section 6.5 and analyze computational overhead in Section 6.6.
6.1 Reasoning tasks
Experimental setup We evaluate four reasoning benchmarks: GSM8K [10] and MATH500 [16] for math, and HumanEval [9] and ABPP-sanitized [2] for code synthesis. We use two instruction- tuned models, LLaDA 8B Instruct [27] and Dream 7B Instruct [52]. For both models we set the generation length to 256 tokens, with LLaDA using a block size 64. The generation temperature set to 0 for LLaDA and 0.2 for Dream. We report both the number of sampling steps and the pass@ accuracy. The maximum inference timestep is set to 256. In KLASS, we compute per-token KL divergence over history length of 2, and apply KL thresholds ranging from 0.001 to 0.01 and confidence thresholds from 0.5 to 0.9. Full configuration details and a lightweight guideline for hyperparameter selection provided in Appendix D.1.2.
Baselines We compare KLASS against baselines across two categories. The first is sequential unmasking (single-token), which includes: (i) Top- sampling, selecting the highest-confidence token at each step [7]; and (ii) random sampling [1]. The second category is parallel unmasking, which accelerates generation by revealing multiple tokens per step: (iii) Top-2 sampling, decoding the two cigfest-confidence tokens per step halve the total number of steps; (iv) seffiddnncethreehhe sampling. unmasking all tokens with a predicted probability over 0.9; and (v) KL-threshold sampling. unmasking all tokens with KL. divergence under 0.001, using history length
in KLASS.
As shown in Table 1, KLASS consistently improves accuracy across most tasks compared to the standard greedy decoding (Top-1) baseline. It demonstrates robust generalization for both LLaDA and Dream models across math and code synthesis benchmarks. Beyond accuracy, KLASS
Results

---

Table Generative perplexity, MAUVE, and entropy on unconditional text generation sampled with
512 steps.
Method
MAUVE1
LLaMA2
LLaMA3
GPT-2
Entropy
Data
1.0
7.0
9.4
14.8
5.44
AR
0.855
10.97
15.12
12.07
5.21
SEDD
0.037
53.09
109.60
105.40
5.62
D3PM
0.022
41.82
72.85
76.70
5.40
MDLM
0.115
30.88
54.15
51.78
5.46
KLASS (Ours)
0.179
26.94
49.19
45.50
5.43
Table Generative FID and IS on MMMDA with different step sizes.
Table 4:
Molecular generation results on the QM9 dataset conditioned different molecular properties.
Method
Steps
FID +
IS f
Method
Property
Reward f NFEs |
Confidence
16
34.48
75.72
MDLM
QED
0.526
32.0
KLASS (ours)
16
30.48
93.07
KLASS (ours)
QED
0.546
18.8
Confidence
32
36.45
72.40
MDLM
Ring count
4.123
32.0
KLASS (ours)
32
32.00
89.17
KLASS (ours)
Ring count
4.258
24.4
is also highly efficient. It reduces sampling steps by 40-70% relative the full 256-step schedule, yiclding wall-clock speedups of up to 2.78x (Appendix D.1.3). Unlike other acceleration strategies such as halving steps with confidence-basee Top-2 method, which degrades accuracy, KLASS improves accuracy with fewer steps overall. KLASS achieves a superior balance between speed and accuracy compared to methods that rely single threshold for either confidence KI score alone. This proves that the effectiveness of KLASS comes from its novel approach of combining token confidence with KL-divergence trajectories.
6.2 Text generation
Experimental setup We evaluate KLASS on Masked Diffusion Language Model (MDLM) [34] pre-trained the WetTtt corpus [13]. As baselines, include (i) the original autoregressive sampler, (ii) SEDD [25], and (iii) two variants of MDLM: one parameterized with SUBS (the standard 512-step sampler) and one parameterized with D3PM [3] (the absorbing variant). For all diffusion- based methods, we generate 1,000 sequences of length 1,024 tokens under a fixed 512-step schedule, with nucleus Top-p filtering at 0.9, history length = 2, KL threshold сKL le and confidence threshold т = 0.57.
Evaluation We report generative perplexity by exponentiating the average token-level negative log- likelihood under three oracle models: LLaMA2 (7B) [42], LLaMA3 (8B) [14], and GPT-2 [30]. We measure Shannon entropy of the predicted token distributions and compute MAUVE by comparing our 1,000 generated samples to 1,000 held-out segments from the Teb Baseline (*Data) results given from the corresponding literatures [43, 48).
Results Table 2 shows that KLASS substantially improves generative quality over existing discrete diffusion samplers. Our method higher MAUVE and lower perplexity across all oracle models while maintaining comparable entropy. These results highlight that stability-aware multi-token unmasking guided by KLASS leads to more coherent and fluent text generation, all without any additional model training. We provide experimental details in Appendix D.2.
6.3 Image generation
Experimental setup evaluate KLASS on the MMaDA (Multimodal Large Diffusion Language Models) [49], multimodal diffusion foundation model. We compare two samplers: (i) the standard confidence-baseee sampler used by MMaDA, and (ii) our proposed KLASS. For each method,

---

none
25.4
29.2
30.6
30.8
31.6
none
22.2
25.0
25.0
36.8
41.8
40
32
0.02
29.2
33.2
38.6
39.4
40.2
KL Threshaid
0.02
27.8
31.0
31.4
31.6
30.8
Accuract (%)
KL Threshold
Accuracy (%))
0.015
32.4
33.4
38.4
39.2
40.4
35
30
0.015
30.4
33.2
33.0
31.0
30.8
0.01
31.0
29.6
37.8
39.2
42.0
30
0.01
31.6
33.8
32.4
30.8
30.6
28
0.005
32.6
35.6
37.8
40.6
43.2
0.005
32.6
31.4
31.6
30.8
30.2
26
0.001
25
0.5
0.7
0.8
0.5
0.6
0.7
0.9
Confidence Threshold
Confidence Threshold
(a) LLaDA
(b) Dream
Figure KL Effect Across Confidence Levels MATH.
generate 10,000 images conditioned labels drawn uniformly from the 1,000 ImageNet classes, using 16 and 32 step decoding schedules. KLASS is configured with history length = 1, KI. divergence threshold 'LL 0.3, and confidence threshold 0.1.
Evaluation We assess sample fidelity using two widely adopted metrics. First, we compute Fréchet Inception Distance (FID) [17] between our 10,000 generated samples and the ImageNet validation set, using the official Inception v3 implementation. Second, we measure Inception Score (IS) [36] on same samples with the standard protocol.
Results Table 3 shows that KLASS improves image quality MMaDA the standard confidence- based sampler. Across both decoding schedules, KLASS yielding lower FID and higher IS. The trend holds under the same decoding schedules and faimess controls, indicating that KLASS improves fidelity and class-consistency without modifying the backbone or adding auxiliary guidance. We provide experimental details in Appendix D.3.
6.4 Molecular generation
Experimental setup We use QM9 [31], which contains molecules with up to nine heavy atoms, represented in SMILES [46]. For models we follow the training recipe of [37] train seperate models conditioned on dtrugllieneee (QED) [5] and number of rings using classifier-fred training of masked diffusion models.
Evaluation We test KLASS conditional generation of small molecules using CFG guidance. Specifically, we aim to generate molecules with higher score QED or maximizing ring counts while fixing the CFG strength for fair comparison. We generate 1,024 samples for each task and provide average value of number of function evaluation (NFEs). Further details of the experimental setups provided in Appendix D.4.
The result shows that KLASS effectively reduces the total sampling steps while maintaining target reward in the conditional generation scenario for both target reward (QED and Ring count). We provide further experimental results in this setup in Appendix D.4.
Results
6.5 Ablation Studies
Effect of confidence and KL score thresholds Our evaluation of different confidence and KL thresholds on the MATH dataset reveals that combining both is essential for optimal performance. As shown in Figure 3. applying levels compared to relying on confidence threshold alone ('none' row). This synergistic relationship further substantiated by Table 1, which demonstrates that using a single criterion leads to notable reduction in accuracy.
KL threshold consistently enhances accuracy across all confidence

---

Table Comparison of single-token and paral- lel unmasking strategies under KLASS criteria.
Table Memory and computational overhead of KL divergence per decoding step.
Unmasking
MATH
GSM8K
Model
Memory (MB)
Time (s)
Acc f
Steps
Acc 1 Steps
Overhead
Total
Overhead
Total
Single (conf)
31.2
256
72.86
256
Single (KL)
29.0
256
73.46
256
LLaDA
247
18,702
0.000255 0.1218
Parallel
33.8
128.6
76.50
98.57
Dream
296
18,875
0.000177
0.1275
While the optimal hyperparameter settings vary significantly between models, each model's perfor- mance remains stable and robust around its unique optimal point. For example, LLLDDA performs best with lower confidence threshold, whereas Dream requires higher one achieve maximum accuracy. In both cases, however, accuracy does not degrade sharply near these values, indicating low sensitivity to minor hyperparameter adjustments. more detailed sensitivity analysis, featuring additional tasks and finer-grained grid of thresholds, is provided in Appendix D.5.1.
We evaluate whether unmasking multiple tokens per step improves LLaDA's performance. Using KLASS. which selects tokens based on fixed thresholds, we compare parallel multi-token unmasking two sequential variants. These variants unmask only single token from the same stable pool satisfying the KLASS criteria: 'Single (conf)' unmasks the one with the highest confidence and 'Single (KL)) unmasks the one with the lowest KL score.
Effect of unmasking multiple tokens
As shown in Table 5, parallel sampling of KLASS boosts both accuracy and efficiency. On MATH, improves accuracy by up to 4.8 points while cutting sampling steps by nearly 50%. Similar trends hold on GSM8K. These results suggest that LLaDA benefits from unmasking multiple stable tokens in parallel, leading faster and even more accurate reasoning.
6.6 Analysis on Computational Overhead
The overhead of KL computation is negligible, as is a lightweight post-processing step on existing
logits that requires additional forward pass. For the of masked tokens I,,, = {i 2 m), compute the KL score d; = Dx (p) (p and cache the prior distribution. This yields a combined computational and memory overhead of O( -VV),, linear cost that is negligible compared to the expensive matrix multiplications and multi-gigabyte footprint of the main diffusion step. Table empirically supports this conclusion. We measure the overhead for LLaDA and Dream, with vocab sizes of 126,464 and 152,064, respectively, using a generation length of 256. The results show memory overheads below 1.57% of total memory and latency overheads below 0.21 per decoding step, confirming that KL computation adds only minimal cost.
7 Conclusion
We proposed KL.-Adaptive Stability Sampling (KLASS), an efficient and adaptive sampling method for masked diffusion models that leverages token-level KI. divergence and model confidence to guide the unmasking process. KLASS substantially reduces the number of sampling steps while main- taining or improving accuracy, achieving state-of-the-art performance benchmarks. Our approach simple, requires no additional training, and generalizes well across multiple modalities, making it practical solution for faster and more reliable generation in masked diffusion models.
math and code reasoning
For future work, one could extend this approach to discrete diffusion models with alternative noise schedules, such as the uniform or marginal prior [1]. Another direction is to evaluate the proposed sampler with larger models as they become available. We also discuss the broader impact and limitations of our work in Appendix G.

---

Acknowledgments
This work was supported by Institute of Information & Communications Technology Planning Evaluation (IITP) grant funded by the Korea government (MSIT) (RS-2024-00457882, AI Research Hub Project, 10%) and 1RS-2022-1120311, Development of Goal-Oriented Reinforcement Learning Techniques for Contact-Rich Robotic Manipulation of Everyday Objects, 90%).
References
[1] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in neural information processing systems, 34:17981-17999 2021.
[2] Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint WrXiv:2108.07733 2021.
[3] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces, 2023. URL https://arxiv. org/ /ass///// 03006.
[4] Heli Ben-Hamu, Itai Gat, Daniel Severo, Niklas Nolte, and Brian Karrer. Accelerated sampling from masked diffusion models via entropy bounded unmasking. arXiv preprint 2rx2v:2505.24857, 2025.
[5] G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):90-98, 2012.
[6] Andrew Camphell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-28279, 2022.
[7] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 11315-11325, 2022.
[8] Hongrui Chen and Lexing Ying. Convergence analysis of discrete diffusion model: Exact implementation through uniformization. arXiv preprint arXiv:2402.08800 2024.
[9] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint urXiv:2107.03374 2021.
[10] Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichhro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
[11] Justin Deschenaux and Caglar Gulcehre. Beyond autoregression: Fast 1lms via self-distillation through time. arXiv preprint arXiv:2410.21033 2024.
[12] Daniel T Gillespic. A general method for numerically simulating the stochastic time evolution of coupled chemical reactions. Journal of computational physies, 22(4):403-434, 1976.
[13] Aaron Gokaslan and Vanya Cohen. Open WebTeet corpus. http://Skylloon007, wwt..ue..........iio.o DpenWebTextCorpue 2019. Accessed: 2025-05-16
[14] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ah- mad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, Amy Yang. Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis. Damien Allonsius,
10

---

Daniel Song, Danielle Pintz, Danny Livshits, Danny Wyatt, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Carcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Francisco Guzmán, Frank Zhang. Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Govind Thattai, Graeme Nail, Gregoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, llyann Zarov, Imanol Arricta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jack Zhang. Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park. Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Va- suden Alwala, Karthik Prasad, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Kushal Lakhotia, Lauren Rantala-Yean Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Maria Tsimpoukelli, Mathew Oldham. Mathieu Rita, Maya Pavlova, Melanie Kam- badur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji. Ning Zhang, Olivier Duchenne, Onur Çelebi. Patrick AIrassy, Pengchuan Zhang, Pengwei Li, Petar Vasic, Peter Weng. Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong, Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Raileanu, Rohan Maheswari, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Ronnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang. Saghar Hos- seini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra, Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher. Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta. Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginic Do. Vish Vogeti, Vítor Albiero, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet. Xiaodong Wang. Xiaofang Wang. Xiaoqing Ellen Tan, Xide Xia, Xinfeng Xie, Xuchao Jia, Xuewei Wang, Yaelle Goldschlag. Yashesh Gaur, Yasmine Babaei, Yi Wen, Yiwen Song. Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aayushi Srivastava, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Amos Teo, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Dong, Annie Franco, Anuj Goyal, Aparajita Saraf. Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan. Beau James, Ben Maurer, Benjamin Leonhardi. Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Ce Liu, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Cynthia Gao, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang, Due Le, Dustin Holland, Edward Dowling, Eissa Jamil, Elaine Montgomery, Eleonora Presani, Emily Hahn, Emily Wood, Eric-Tuan Le, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun. Felix Kreuk, Feng Tian, Filippos Kokkinos, Firat Ozgenel, Francesco Caggioni. Frank Kanayet, Frank Seide, Gabricla Medina Florez, Gabriclla Schwarz, Gada Badecr, Georgia Swee, Gil Halper. Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hakan Inan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Hongyuan Zhan, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Ilias Leontiadis, Irina-Elena Veliche, Itai Gat, Jake Weissman, James Geboski, James Kohli, Janice Lam, Japhet Asher, -aat-Ste Gaya, Jeff Mareus, Jeff Tang, Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul, Jessica Zhong, Jian Jin, Jingyi Yang. Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjie Wang, Kai Wu, Kam Hou U, Karan Saxena, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kiran Jagadeesh,
11

---

Kun Huang, Kunal Chawla, Kyle Huang, Lailin Chen, Lakshya Garg. Lavender A, Leandro Silva, Lee Bell, Lei Zhang. Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Miao Liu, Michael L.. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang. Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal, Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikhil Mehta, Nikolay Pavlovich Laptev, Ning Dong, Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj. Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Rangaprabhu Parthasarathy, Raymond L.i, Rehekkah Hogan, Robin Battey, Rocky Wang, Russ Howes, Ruty Rinott, Sachin Mehta, Sachin Siby, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan, Saurabh Mahajan, Saurabh Verma, Seiji Yamamoto. Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng. Shenghao Lin, Shengxin Cindy Zha, Shishir Patil, Shiva Shankar, Shugiang Zhang, Shuqiang Zhang, Sinong Wang. Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala, Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Summer Deng. Sungmin Cho, Sunny Virk. Suraj Subramanian, Sy Choudbury, Sydney Goldman. Tal Remez, Tamar Glaser, Tamara Best, Thilo Koehler, Thomas Robinson, Tianhe Li. Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vlad Ioneseu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li, Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Con- stable, Xiaocheng Tang, Xiaojian Wu, Xiaolan Wang. Xilun Wu, Xinbo Gao, Yaniv Kleinman, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yu Zhao, Yuchen Hao, Yundi Qian, Yunlu Li, Yuzi He, Zach Rait, Zachary DeVito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang. Zhiwei Zhao, and Zhiyu Ma. The llama 3 herd of models, 2024. URL https: /ttp///axvvvv////i//ee//..a /ttp//aww...g.aa/////////aaaaa/a 220.. 21783.
[15] Satoshi Hayakawa, Yuhta Takida, Masaaki Imaizumi, Hiromi Wakaki, and Yuki Mitsufuji. Distil- lation of discrete diffusion through dimensional correlations. arXiv preprint 24110:2410.08709, 2024.
[16] Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprini arXiv:2 103.03874 2021.
[17] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge to a local nash equilibrium, 2018. URL https://arxiv wrgg/bs////roooo.oo/..//.e abs/ 7706008500.
[18) Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 14.6840-68511 2020.
[19] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural 2019.
text degeneration. arXiv preprint
[20] Tero Karras, Miia Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35: 26565-26577, 2022.
[21] Jaeyeon Kim, Kulin Shah, Vasilis Kontonis, Sham Kakade, and Sitan Chen. Train for the worst, plan for the best: Understanding token ordering in masked diffusions. arXiv preprint arXiv:2502.066: 06766000..666. 2025.
[22] Bin Lei, Yi Zhang, Shan Zuo, Ali Payani, and Caiwen Ding. Macm: Utilizing a multi-agent system for condition mining in solving complex mathematical problems. arXiv preprint 47Xi::2404.04735 2024.
[23] Pengxiang Li, Yefan Zhou, Dilxat Muhtar, Lu Yin, Shilin Yan, Li Shen, Yi Liang, Soroush Vosoughi, and Shiwei Liu. Diffusion language models know the answer before decoding. arXiv preprint a50iv:2508.1 19982, 2025.
12

---

[24] Sulin Liu, Juno Nam, Andrew Campbell, Hannes Stärk, Yilun Xu, Tommi Jaakkola, and Rafael cómezz-Bombane Think while you generate: Discrete diffusion with planned denoising. arXiv preprini arXiv:2410.06264, 2024.
[25] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.1663 2023.
[26] Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Advances in Newral Information Processing Systems, 35:34532-34433 34545, 2022.
[27] Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992, 2025.
[28] Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprini urXiv:2406.037766 2024.
[29] Fred Zhangzhi Peng, Zachary Bezemek, Sawan Patel, Jarrid Rector- Rector-Brooks, Sherwood Yao, Alexander Tong, and Pranam Chatterjee. Path planning for masked diffusion model sampling. arXiv preprint urXiv:2502.03544 2025.
[30] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models unsupervised multitask learners. 2019.
[31] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[32] Yinuo Ren, Haoxuan Chen, Grant M Rotskoff, and Lexing Ying. How discrete and continuous diffusion meet: Comprehensive analysis of discrete diffusion models via stochastic integral framework. arXiv preprint arXiv:2410.03601, 2024.
[33] Yinuo Ren, Haoxuan Chen, Yuchen Zhu, Wei Guo. Yongxin Chen, Grant M Rotskoff, Molei Tao, and Lexing Ying. Fast solvers for discrete diffusion models: Theory and applications of high-order algorithms. arXiv preprint arXiv:2502.0023. 2025.
[34] Subham Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. Advances in Neural Information Processing Systems, 10016-13016 2024.
[35] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022.
[36] Tim Salimans, Ian Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
[37] Yair Schiff, Subham Sekhar Sahoo, Hao Phung, Guanghan Wang. Sam Boshar, Hugo Dalla- torre, Bernardo de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. arXiv preprint arXiv:2412.10119,, 2024.
[38] Jiaxin Shi, Kehang Han, Zhe Wang, Amaud Doucet, and Michalis Titsias. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131-100166 2024.
[39] Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36:8634 8652, 2023.
[40] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning, pages 2256-2265 pmlr, 2015.
13

---

[41] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprini arXiv:201 1. 13456, 2020.
[42]
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshomn, Saghar Hosseini. Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiao- qing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL https: / /arxiv. wttg 09288.
[43] Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking discrete diffusion models with inference-time scaling, arXiv preprint arXiv:2503.003077 2025.
[44] Xingyao Wang, Yangyi Chen, Lifan Yuan, Yizhe Zhang, Yunzhu Li, Hao Peng, and Heng Ji. Executable code actions elicit better 1lm agents. In Forty-first International Conference on Machine Learning, 2024.
[45] Qingyan Wei, Yaojie Zhang, Zhiyuan Liu, Dongrui Liu, and Linfeng Zhang. Accelerating diffusion large language models with slowfast: The three golden principles. arXiv preprint arXiv:2506. 10848, 2025.
[46] David Weininger. Smiles, chemical language and information system. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 28 (1):31-36, 1988.
[47] Chengyue Wu, Hao Zhang, Shuchen Xue, Zhijian Liu, Shizhe Diao, Ligeng Zhu, Ping Luo, Song Han, and Enze Xie. Fast-dllm: Training- free acceleration of diffusion Ilm by enabling kv cache and parallel decoding. arXiv preprint 25Xi::5505.22666 2025.
[48] Minkai Xu, Tomas Geffner, Karsten Kreis, Weili Nie, Yilun Xu, Jure Leskovec, Stefano Ermon, and Arash Vahdat. Energy-based diffusion language models for text generation. arXiv preprint arXiv:2410.21 1335, 2024.
[49] Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. Mmada: Multimodal large diffusion language models, 2025. URL https://arxiv orgg abs/ 2505. 15809.
[50] Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning. arXiv preprint arXiv:2410. 14157, 2024.
[51] Jiacheng Ye, Zhenyu Wu, Jiahui Gao, Zhiyong Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Implicit search via discrete diffusion: A study on chess. arXiv preprint drXiv:2502.19800 2025.
[52] Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b, 2025. URL https::/hkunl.. gi wtw.bb.oo/loo///e/ 2025/// dream.
[53] Tianwei Yin, Michaël Gharbi, Richard Zhang, Eli Shechtman. Fredo Durand, William T Freeman, and Taesung Park. One-step diffusion with distribution matching distillation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6613 (613,,663,00 2024.
[54] Runpeng Yu, Xinyin Ma, and Xinchao Wang. Dimple: Discrete diffusion multimodal large language model with parallel decoding. arXiv preprint arXhv:2505.16990, 2025.
14

---

[55] Yixiu Zhao, Jiaxin Shi, Feng Chen, Shaul Druckmann, Lester Mackey, and Scott Linderman. Informed correctors for discrete diffusion models. arXiv preprint arXiv:2407.21245, 2024.
[56] Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Y Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint arXiv:2409.02908, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction in Section clearly state the motivation, pro- posed method (KLASS), and main contributions, including faster sampling and improved performance. These claims are supported by the theoretical and empirical results presented in the paper.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper includes dedicated discussion of limitations in Appendix G.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
16

---

Answer: [Yes]
Justification: The paper clearly states the assumptions for its theoretical result and provides complete and correct proof in Appendix A. The result is discussed in the method section to support the core intuition behind KL-based sampling.
Guidelines:
The answer NA means that the paper does not include theoretical results, All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
The paper provides sufficient information to reproduce the main experimental results in Section 6. More details are provided in Appendix D.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
17

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We include our anonymized source code and detailed reproduction instructions in the supplementary material to enable faithful reproduction of all main experimental results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
sublic/guides/CoeeSuuee www.uiiiiiooooliiioooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurIPS code and data submission guidelines (https: / httppsscccppbll// wwwubbic/guu seuudee//oeeeS/ siimisssionn ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper specifies necessary training and evaluation details in the main experimental section and provides full hyperparameters, data splits, and implementation choices Appendix D, ensuring the experimental setup is transparent and reproducible.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We provide the experimental settings and code used in KLASS to ensure reproducibilityy Some KLASS experiments are deterministic, yielding identical results in the same settings. For the stochastic experiments, we report statistical measures such as standard deviation the Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
18

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
The paper provides detailed information about compute resources in Ap- pendix including GPU types, runtime per experiment, and memory usage, allowing others to estimate and reproduce the computational costs involved in the reported results.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported didn't 'tttt make it into the paper).
paper (e.g., preliminary or failed experiments that
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https:
Answer: [Yes]
Justification: We have carefully reviewed the NeurIPS Code of Ethics and ensured that our work is conducted accordance with its principles.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
19

---

The paper discusses both potential positive and negative societal impacts in Appendix G including the benefits of faster, more efficient generation for real-world applications and the risks of misuse in sensitive domains, along with possible mitigation strategies.
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not involve the release of pretrained models, scraped datasets, or other high-risk artifacts that would require special safeguards against misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (c.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All used assets are properly credited with license and usage terms stated in Appendix D.
Guidelines:
20

---

The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paper thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA]
The paper does not introduce new assets such as datasets or pretrained models; it proposes new sampling algorithm applied to existing models and datasets.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: The paper does not involve any crowdsourcing tasks or experiments with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
21

---

Answer: [NA]
Justification: The research does not involve human subjects and therefore does not require IRB approval equivalent review.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification:
The core methodology does not involve the use large language models as original or non-standard components; KLASS is novel sampling algorithm applied to masked diffusion models.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our L.M policy (ht /ttpss//////ipp.eeeeeeeee /tteerrppp..cccoofeeeeeee/////// for what should or should not be deseribed.
22

---

Table of contents
Introduction
2 Related Works
3 Preliminaries
3.1 Masked diffusion models
3.2 Inference via Ancestral Sampling
4
Method
4.1
Defining Confidence Score and KI. Score
4.2
KLASS: KL-Adaptive Stability Sampling
5
Theoretical Rationale
Experiments
6.1
Reasoning tasks
6.2
Text generation
6.3
Image generation
6.4
Molecular generation
8
6.5
Ablation Studies
6.6
Analysis on Computational Overhead
7
Conclusion
Theoretical Proofs
25
Pseudo Code
25
C Further Prior Works
25
Experiment details and additional results
27
D.1
Reasoning tasks
27
D.1.1
Experiment details
................................
27
D.1.2
Hyperparameter Selection Guideline
................................
27
D.1.3
Wall-clock time comparison
28
D.1.4
Statistical significance of Dream 7B Instruct results
.............
28
D.2
Text generation
28
D.2.1
Experiment details
28
D.3
Image generation
29
D.3.1
Experiment details
29
D.4
Molecules
29
D.4.1
Experiment details
29
23

---

D.5 Ablations
30
D.5.1
Hyperparameter Sensitivity
30
D.5.2
Effect of history length
30
D.5.3
Effect of temperature
30
Comparison other diffusion samplers
30
E,1
Performance on reasoning tasks
30
F Examples of generated samples
32
Limitation & Broader impact
33
G.1 Limitations
33
G.2 Broader impact
33
24

---

A Theoretical Proofs
In this section, we provide the proof for Section 5.
Definition A.1. For each context (instantiation of variables outside X2). let C(c) be the nonempty set of cosk-correct conditionals. Let == -1=: nu(. c) E C(c) Vel. We say po is conditional 8 5-approximation the task if
inf sup TV(po(- 1 (c), T(. (o)) s 6. REC
Definition A.2. Fix Let xt be optimal under x(: c*) at near-optimal context c*. Let x + x be suboptimal. Assume a true margin y > 0 at c²: <x;; c*) 2 < c*) + y. Assume the model currently prefers z! at cM by margin 8 2 0: po(x) cM) 2 pe (2 cm) + 3.
Proposition A.3. Suppose pe is conditional 5-approximation of т. For any context path CM - CM-1 + co (changing only variables outside X₂) ending at co = c*, let Pt pe (e) and
= (3 + 28)+. Then
1 M
2
TV(PM, Po) 2 A.
Σ M
KL(P.
2
M²
Proof. Let f = 1(2; = x11 - 1f2; 2;} so |f S 1. Then
2TV Po) 2 ER. P,, [f] (pb(ear) p)((m)) (p((c")-p(("" pb(c)))
By the margin assumptions (Definition A.2) and the conditional 8 a-appoximmtiion (Definition A.1), Ph(CM) (cM) > and pp(c*) (c() 2 7 25. Hence 2 TT (PM, Po) 2 + 26, which implies TV(PM. Po) 2 A.
By the triangle inequality in total variation,
TV(P=I, P.) TV(PM, Po) A.
Let T. == TV(P++1 P.). Pinsker's inequality gives KL(P2|(P(42) 2T, for each (since (V(P+,, = (P P.)). Averaging and applying Kauchy-Schwarz,
M Σ
(L(()||PH))) 2
2
M-1
ΣT
2
(E T.)
2 2
M
M
M
M
t=
B Pseudo Code
We provide the pseudo code for our KLASS algorithm, implementing the unmasking rule described in Section 4.2.
C Further Prior Works
In this section, we review several approaches for sampling from discrete diffusion models described in prior work.
Ancestral sampling Generation proceeds by discretizing the reverse diffusion time-interval [o, 1] into
= to < f < ...
ty 1.
To sample a sequence of length L, one initializes
(T)
|mask]] L
25

---

Algorithm 1: KL-Adaptive Stability Sampling (KLASS)
Input: model M, total steps T, sequence length L, confidence threshold KL threshold c, history window H, fallback count
Output: Generated sequence x
Initialize x <- [MASKL.]] P... Prev +- 0, KLbuf +-
for t + 1 to do
L + Moxologgsss
- softmax(()
c f max(P) Get per-token confidence scores
5 - DRL(P PPeev)))) // Get per-token KL scores
// Update KL history buffer
KLbuf roll(KLbuf,, shift -1)
KLbuff; H] < 6
Prwe f P
// Identify stable tokens
stable_kl + Y(KLbuf < <) // Check all history high_ is_ sa.maseed isMask(x)
11
(c>
12
13
ready < stable_kl high_conf A is is_masked if any(ready) then
14
unmask_at_ at.iiiieet_idicexxx.... P, ready)
16
// Fallback:
unmask top-u tokens
17
scores c- ii_mmsked // Zero out unmasked tokens U +- topk_indices(scores.. u)
unmask_at_ famnask_a__nnices(x: P,
Return
and for i TTT I,..., draws each coordinate independently:
(1-1)
((((() po (ze 111)
* mask, l -= L.
28
(6) mask,
Because unmasked tokens remain unchanged, if at step i no new tokens are decoded then 21:L (() Zi. and-when the denoiser μl is time-invariant-its output at 6 can be reused at 41-1, skipping that network evaluation [28, 34]. Models whose μ4 depends on & (e.g. SEDD [25]) must recompute at every 4; and cannot exploit this caching [25, 34, 38).
(1-1)
Exact simulation via uniformization: one bounds the time-varying generator (t) by X, samples
Exact simulation of the reverse CTMC in absorbing masked diffusion is achieved
M Poisson (XT),
{ iid (O,
and at each 7; transitions from state x to y with probability (T.)/ preserving the exact law of the reversed path [8]. While unbiased, as the chain nears absorption the number of proposals-and hence cost-can grow large.
Alternatively, the first-hitting sampler of Zheng et al. [56] draws each unmasking time without discretization error: when tokens remain masked, one samples
T a a-
u. 1/n a(Tw))). U, Unif(0,
then un- un-masks exactly one token (chosen uniformly among the n) according to the model's conditional distribution. This procedure is unbiased but incurs O(L) sequential events, making runtime scale with sequence length [56].
r-leaping Tau-leaping discretizes the reverse CTMC into fixed intervals of length т, holding all jump rates constant and applying transitions in parallel. Let R2 (xx x²) be the learned rate from r to I' at time Over i F, one draws
K
Poisson RV (x,xx')) (x,
26

---

and updates
Σ Ku-tz (ex -x.) x'zx;
It-r = +
Under mild regularity, the global weak error scales as O(T), recovering Gillespie's exact algorithm when r 0 [12]. Ren et al. derive a KL divergence bound
DEL. ( (Paata) < O(+ + O(MT) + O(E cT/ log D),
showing modest suffices even in high dimensions [32]. In practice, coordinates with multiple jumps are rejected to enforee categorical integrity negligible event under suitable rates), trading (1) network evaluations per leap for an ((T) discretization bias.
To improve first-order 7-leaping, multi-stage integrators achieve higher local accuracy. Ren et al. [33] introduce two-stage schemes with second-order convergence in KL. The 6-RK-2 method computes an intermediate state
High-order samplers
y Yt + Σ
Poisson (oisson(pe(v)) Ot)),
then updates
Σ v Poisson( (1--)(11( ((()) mi(v)) St,
Dt Vt +
WE
where oe is the intensity at y² The t-tapezzoiall variant replaces the second stage with Poisson 14 02 pc) 0) At),
U1- yy
Σ
FED
These reduce the discretization error O(A/T)), enabling 3-5x fewer evaluations for comparable fidelity [33], and draw on high-order continuous schemes [20].
with 01 20(100 02
+02
28(1-8)
D Experiment details and additional results
D.1 Reasoning tasks
D.1.1 Experiment details
We conduct our experiments using LLaDA 8B Instruct [27] and Dream 7B Instruct [52], which are masked diffusion models capable of advanced reasoning. Across all settings, we fix the generation length to 256 tokens. For LLaDA 8B Instruct, which supports semi- ammi-autoregressvvee sampling, we set the block size to for all experiments. Sampling temperature set for LLaDA 8B Instruct, and to 0.2 for Dream 7B Instruct. Additional experiments using Dream 7B Instruct with temperature of are reported in Appendix D.5.3. All sampling experiments are conducted on single NVIDIA RTX A5000 GPU.
D.1.2 Hyperparameter Selection Guideline
For hyperparameter selection, we adopt a lightweight three-step search procedure using small validation set of around 100 examples.
Initial KI Threshold Estimation: We obtain an initial estimate of the KI. threshold by inspecting the distribution of KI. values during decoding.
2. Confidence Threshold Search: With the KL threshold fixed, we sweep confidence values
(0.9 down to 0.6) to identify the best trade-off between accuracy and decoding speed.
3. KL Threshold Refinement: With the confidence threshold fixed, we refine the KL threshold through finer-graine geernegai search around the initial estimate.
This procedure is efficient, requiring only small number of validation samples and negligible computation relative to training. The final confidence and KI. threshold configurations for each dataset and model are summarized in Table 7.
27

---

Table Experiment threshold configurations for KLASS
Model
MATH
GSMSK
HvaanEval
MBPP
Conf
KL
Conf
KL
Conf
KL
Conf
KL
LLaDA
0.6
0.010
0.6
0,015
0.9
0.010
0.7
0.010
Dream
0.9
0.005
0.9
0.001
0.8
0.001
0.9
0.001
Table 8: Wall-clock time per sample for Top-1 and KLASS decoding.
Model
Dataset
Accuracy
Time (s)
Speedup
Top-1
KLASS
Top-1
KLASS
GSM8K
75.13
76.50
37.04
15.86
2.34
MATH
31.40
33.80
38.40
21.41
1.79
LLaDA
HumanEval
39.63
40.85
39.54
16.04
2.47
MBPP
46.69
47.86
39.12
20.68
1.89
GSM8K
79.75
80.44
29.66
22.26
1,33
MATH
38.00
43.20
30.76
23.31
1.32x
Dream
HumanEval
58.53
59.76
32.01
11.52
2.78x
MBPP
63.81
64.59
31.89
17.65
1.81x
D.1.3 Wall-clock time comparison
We compare the actual wall-clock time of the KLASS and Top-k samplers on the MATH dataset in Table 8. Compared the Top-k sampler with 256 steps (i.e., unmasking one token per step), KLASS reduces generation time by 47.4% for LLaDA and by 16. 1% for Dream).
When comparing KLASS with the Top-k sampler using similar number of steps (128 for LLaDA and 149 for Dream), the Top-k method achieves slightly lower generation times. However, this speed gain comes the cost of reduced accuracy. KLASS not only maintains comparable runtime but also improves accuracy, demonstrating its efficiency and effectiveness.
D.1.4 Statistical significance Dream 7B Instruct results
In Table 9, we report the mean and standard deviation over three runs for all methods using Dream 7B Instruct with temperature of 0.2. Since the experiments with LLaDA 8B Instruct were conducted with temperature of O, the results are deterministic, so we only report statistics for Dream.
D.2
Text generation
D.2.1
Experiment details
We evaluate KL-Adaptive Stability Sampling (KLASS) on Masked Diffusson Language Model (MDLM) [34] pre-trained the OpenWebText corpus [13]. As baselines, include (i) the original autoregressive sampler (one-token-at-a-iin
unmasking). (ii) SEDD [25], and (iii) two MDLM variants: the standard 512-step sampler and the "absorb" variant [3]. For all diffusion-based methods, we generate 1, 1,000 sequences of length 1,024 tokens under fixed 512-step schedule, applying nucleus (top-p) filtering at p = 0.9, history length
2, a KL divergence threshold 6KL 1e and a confidence threshold r = 0.57. To ensure fair comparison under this fixed step count, we cap the maximum number of tokens accepted by the thresholds each step. In the fallback case where tokens do not pass this criterion, we revert to the original MDLM. We report generative perplexity by exponentiating the average token-level negative log- likelihke under three oracle models (LLaMA2 7B, LLaMA3 8B, and GPT-2), measure Shannon entropy of the predicted token distributions, and compute MAUVE by comparing our 1,000 generated samples to 1,000 held-out segments from OpenWebTest All runs were executed on single NVIDIA RTX A6000 GPU. To quantify run-to-run variability, each method was repeated with three fixed random seeds, and we report mean one standard deviation over these replicates.
28

---

Table Mean and standard deviation for each sampler across three runs (mean std).
MATH
GSM8K
Sampler
Acc (%)
Step
Acc (%)
Step
Top-1
37.97 = 0.12
256.00 0.00
79.55 0.14
256.00 + 0.00
Random
18.73 + 1.61
37.35 0.53
256.00 0.00
Top-2
33.60 0.16
71.69 0.00
128.00 + 0.00
conf > 0.9
41.80 0.00
9..00
0.15
74.81 - 0.08
KT 0.001
41.27 0.09
162.49
1.14
150.02 + 0.32
KLASS (ours)
43.20 + 0.00
149.72 0.00
7 + 0.72
155.67 0.41
(a) MATH GSMSK
HumanEval
MBPP
Sampler
Acc (%)
Step
Acc (%)
Step
Top-1
58.53 + 0.00
256.00 + 0.00
63.81 0.00
256.00 + 0.00
Random
18.09 + 2.51
256.00=0.00
28.14 0.91
256.00 +
Top-2
42.88 = 0.29
1.000
+7.08
128.00
conf > 0.9
50.00 = 0.00
52.47 0.00
0.00
72.49 0.00
KL 0.001
59.35 + 0.29
73.94 + 0.57
62.65 0.00
108.15 0.00
KLASS (ours)
59.35 + 0.29
74.88 + 0.74
64.59 0.00
111.24 + 0.00
(b) HumanEval MBPP
Table 10: Generative perplexity, MAUVE and entropy on unconditional text generation. Here, 'D3PM' denotes an MDLM that is parameterized using D3PM.
Method
MAUVET
LLaMA2
LLaMA3
GPT2 1
Entropy 1
Data
1.000
7.00
9.40
14.80
5.44
AR
0.855 0.033
10.97 0.10
15.12 + 0.18
12.07 + 0.12
5.21 0.02
SEDD
0.037 0.012
53.09 0.24
109.60 + 0.79
105.40
0.67
5.62 0.00
D3PM
0.022 + 0.006
41.82 + 5.88
72.85 + 12.69
76.70
5.40 0.00
MDLM
0.115 + 0.033
30.88 + 0.20
54.15 + 0.27 0.27
51.78
5.46 0.00
KLASS
0.179 + 0.041
26.94 0.24
49.19 + 0.40
45.50
5.43 0.01
D.3
Image generation
D.3.1
Experiment details
We evaluate KLASS the MMaDA [49]. We compare two samplers-confideer
and KLASS.
For each sampler, we generate 10,000 class conditiona images with uniformly sampled ImageNet labels under 16-step 32-step decoding budget. For KLASS, fix the hyperparameters history length = 1, KL divergence threshold EKL = 0.3, and confidence threshold т = 0.1. For a fair comparison with fixed step count, restrict the per-step reveal count allowed by the thresholds. FID is computed against the 50k ImageNet validation set using official Inception- v3 features, and IS follows the standard protocol. All gengeeeinn runs use single NVIDIA RTX A5000.
D.4 Molecules
D.4.1 Experiment details
We mainly follow the experimental settings in [37] for the experiment. For training, we train independent models for two target conditions: QED and ring count. We use diffusion step size 32, taking 25,000 gradient steps. Small size diffusion transformer (DIT) which is composed of 12 DIT
29

---

blocks, hidden dimension of 768 is utilized for the architecture. We train the model with classifier-free guidance (CFG) training with dropout condition probability of 0.1. We generate samples with CFG strength y = 1 for the experiment. Reported values in Table 4 for KLASS use K = 0.001 and r 0.98 for both QED and Ring Count experiments. We utilize a single RTX 3090 GPU for both training and the inference.
D.5 Ablations
D.5.1
Hyperparameter Sensitivity
To address sensitivity, performed grid search across diverse models and tasks (LLaDA/Dream for reasoning, MDLM for molecular), demonstrating the robustness of KLASS.
Our findings show that configurations near the selected optimum consistently yield high accuracy while significantly reducing sampling steps. For instance, on HumanEval with LLaDA (Table 11), settings near (confidence = 0.9, KL = 0.01) maintain or improve upon the baseline accuracy of 39.63% while using far fewer than 256 steps. Similar trends are observed for other settings (Tables 12, 13, 14, 15). Slightly different hyperparameter settings can sometimes outperform the main reported configuration, indicating both robustness and potential for further tuning.
D.5.2 Effect of history length
We evaluate the effect of varying the KL score history length in KLASS across different KL divergence thresholds and confidence thresholds the MATH dataset. Results are reported in Table 16 for both LLaDA and Dream models.
For LLaDA, history length of 2 offers the best balance of accuracy and efficiency, particularly KL threshold of 0.015 and a confidence threshold of 0.6. At the stricter threshold of 0.9, history length has less impact, suggesting that more relaxed confidence threshold allows more informative token candidates to be considered for unmasking. For Dream, the highest accuracy achieved with history length 2, EKL 0.005, and 0.9. At lower confidence of 0.6, overall accuracy decreases, and longer history helps stabilize token predictions. summary, history length 2 is optimal across most settings, providing improved accuracy with moderate computational cost. Therefore, use history length 2 for all reasoning tasks.
D.5.3 Effect of temperature
Table 17 shows that KLASS consistently improves over a deterministic Top-1 sampler across tasks and temperature settings. Gains are especially strong temperature where KLASS boosts accuracy by 6.22 to 8.00 percentage points and reduces steps by up to 79%. At temperature 0.2, still provides solid improvements, with accuracy gains of 0.69 5.10 points and step reductions of 39% to 71% 71%. These results highlight KLASS's s ability to accelcrate reasoning and improve accuracy, with even greater boosts more deterministic settings.
E
Comparison to other diffusion samplers
E.1 Performance on reasoning tasks
To recap the baselines used the main experiment (Table 1), we consider:
Top-k: Tokens are generated by selecting the one with the highest confidence [7].
Random: Tokens are generated in purely random order [1].
Table 18 reports results for two additional samplers:
Top-k Margin: Unmasks the token with the largest probability margin between the highest and second-highest confidence [21]]
Tokens are ranked by their negative Shannon entropy, prioritizing those with lower uncertainty (i.c., higher confidence) in the model's prediction.
Entropy:
30

---

Table 11: Hyperparameter sensitivity on HumanEval with LLaDA. We report values as Accuracy (Steps). Conf = 0.9 and KL = 0.01 are chosen for KLASS. The baseline accuracy is 39.63% with 256
steps.
KL = 0.015
KL 0.01
KL = 0.005
Conf 0.95
39.63 (65.19)
40.24 (99.29)
40.24 103.14)
Conf = 0.9
40.24 (89.29)
40.85 (91.98)
40.24 (96.75)
Conf = 0.85
39.63 (84.98)
39.63 (88.78)
40.24 (94.48)
Conf = 0.8
39.02 (83.14)
39.02 (87.21)
40.85 (94.01)
Table 12: Hyperparameter sensitivity on MBPP with LLaDA. We report values as Accuracy (Steps).
Conf 0.7 and KI. = 0.01 chosen for KLASS. The baseline accuracy is 48.64% with 256 steps.
KI. = 0.015
=I. 0.01
KI. 0.005
Conf = 0.8
49.42 (122.37)
49.42 (134.52)
49.42 (134.52)
Conf = 0.75
49.03 (118.61)
49.42 (123.45)
48.25 (132.11)
Conf 0.7
48.25 (116.24)
49.03 (127.81)
49.03 (130.67)
Conf = 0.65
46.30 (113.22)
48.64 (118.54)
49.42 (128.99)
Table 13: Hyperparameter sensitivity on HumanEval with Dream. We report values as Accuracy (Steps). Conf = 0.8 and KL = 0.001 are chosen for KLASS. The baseline accuracy is 58.53% with 256 steps.
KL 0.005
= 0.003
KL 0.001
KL 0.0005
Conf = 0.9
54.87 (64.78)
56. 16.10 (67.41)
57.93 (74.86)
61.59 (79.57)
Conf = 0.85
57.32 (59.39)
55.49 (65.31)
59.15 (73.38)
60.98 (79,43)
Conf = 0.8
51.22 (58.09)
54.27 (62.82)
59.76 (73.73)
60.96 (79.51)
Conf = 0.75
48.78 (55.21)
53.05 (62.61)
59.15 (73.41)
60.37 (79.37)
Table 14: Hyperparameter sensitivity on MBPP with Dream. We report values as Accuracy (Steps).
Conf 0.9 and KL 0.001 are chosen for KLASS. The baseline accuracy 63.81 % with 256 steps.
KI. 0.005
KI. 0.003
KI. = 0.001
KI. 0.0005
Conf = 0.95
65.37 (108.93)
65.37 (110.03)
64.59 (112.56)
64.59 (113.14)
Conf = 0.9
62.65 (103.99)
64.20 (107.09)
64.59 (111.24)
64.59 (112.65)
Conf 0.85
64.20 (101.34)
63.81 (105.43)
65.37 (112.54)
64.59 (112.76)
Conf = 0.8
63.81 (96.02)
63.04 (103.22)
65.37 (112.52)
64.59 (1 (2.82)
Table 15: Hyperparameter sensitivity on Molecule QED (MDLM). Conf= 0.98 and KL 0.001 are chosen for KLASS. The baseline QED is 0.526 with 32 steps. We report values QED (Steps).
KL = 0.01
KL 0.005
KL = 0.001
KL 0.0005
Conf 0,999
0.527 (18.61)
0.526 (18.61)
0.524 (18.58)
0.538 (18.45)
Conf 0.99
0.515 (18.22)
0.521 (18.33)
0.543 (18.69)
0.537 (18.66)
Conf 0.98
0.531 (18.43)
0.517 (18.38)
0.546 (18.78)
0.534 (18.82)
Conf 0.96
0.529 (18.43)
0.535 (18.51)
0.534 (18.63)
0.543 (18.80)
KLASS consistently outperforms Top-k Margin and Entropy-based methods with fewer sampling steps. On LLaDA, it achieves top results on GSM8K and HumanEval in less than half the usual iterations. While Entropy yields the highest accuracy on MATH and MBPP with 256 steps, its performance drops sharply with fewer steps. In contrast, KLASS maintains high accuracy at lower computational cost. On Dream, also improves MATH and GSM8K accuracy while reducing steps, demonstrating more efficient and effective sampling.
31

---

Table 16: Ablation results on the history length, across KL thresholds and models, grouped by confidence thresholds.
LLaDA
Dream
Conf
KL
History Length
Ace (%)
Steps
Acc (%)
Steps
32.2
77.13
37.6
123.73
0.010
2
33.8
128.62
39.6
165.68
32.2
153.60
41.8
182.93
30.4
89.04
34.8
83.11
0.6
0.015
33.2
121.11
35.6
119.33
3
33.2
146.19
36.6
146.40
31.0
74.42
34.4
73.66
0.020
31.0
117.01
29.6
104.19
3
30.6
140.56
34.6
128.46
31.4
128.76
42.2
141.36
0.010
30.6
152.75
41.6
169.45
30.8
170.66
42.0
183.51
31.0
127.05
41.0
126.42
0.9
0.015
2
30.8
149.72
43.2
149.72
31.0
167.11
40.2
165.37
30.8
125.85
40.4
120.35
0.020
2
30.8
148.92
42.0
142.75
31.2
164.75
40.4
158.01
Table 17: Effect of temperature on KLASS gains over Top-1 sampler with Dream.
Method
MATH
GSM8K
HumanEval
MBPP
Acc
Steps
Acc
Steps
Acc
Steps
Acc
Steps
Temperature 0.2
Top-1 KLASS
38.10
256
79.75
256
58.53
256
63.81
256
43.20+s.10
150-106
80.44 +0.69
156 100
59.76-1.23
74 182
64.59 +0.78
60.59
111-145
Temperature
Top-1 KLASS
25,80
256
41.70
256
29.27
256
33.46
256
33.80-8.00
121-135
47.92
106-150
37.19
53 -03
40.86+7.40
76-180
F
Examples of generated samples
We present a qualitative comparison generated samples to demonstrate KLASS'S improvements. Table 19 illustrates mathematical reasoning task. The Top-1 baseline fails due a simple arithmetic eиror in the first step (f(-2) =8 -4), while the random sampling disregards the function's 's structure. contrast, KLASS correctly computes each intermediate step and combines them reach the correct answer, 3 highlighting its improved reasoning reliability. We also compare long-form text coherence. Figure 4 shows an MDLM baseline sample that begins on-topic ('SolarCity') but quickly degenerates. It exhibits severe generation artifacts such as repetition ("SolarCitySolarCity nonsensical phrases ('informs of capacity'). The sample's coherence completely breaks down, ending with an unrelated spam link. Conversely, Figure 5 shows KLASS sample on 'urban sprawl,' This text relatively coherent and topically consistent from start to finish, maintaining a plausible journalistic style.
2
14
32

---

Table 18: Performance and sampling steps on reasoning benchmarks for Top-k Margin and Entropy samplers.
MATH
GSM8K
HvmanEval
MBPP
Model
Method
Acct
Stepl
Accf
Stepl
Acct
Step
Acct
Step.
Top-k Margin
32.0
256
74.14
256
39.63
256
47.86
256
Top-k Margin
31.4
128
74.45
128
30.48
128
40.08
128
LLaDA
Entropy
34.6
256
75.43
256
35.97
256
51.75
256
Entropy
32.6
128
73.01
128
25.60
128
40.08
128
KLASS (ours)
33.8
128.62
76.50
98.57
40.85
91.98
47.86
119.59
Top-k Margin
39.4
256
79,45
256
58.53
256
63.81
256
Top-k Margin
32.4
128
71.49
128
43.29
128
46.69
128
Dream
Entropy
39.4
256
79.45
256
58.53
256
63.81
256
Entropy
32.6
128
71.49
128
43.29
128
46.69
128
KLASS (ours)
43.2
149.72
80.44
156.24
59.76
73.73
64.59
111.24
G Limitation & Broader impact
G.1 Limitations
Further scalability larger models In the absence of larger-size discrete diffusion models compared to AR models, our method cannot be evaluated on the more challenging benchmarks such as in agentic systems of LLMs [22, 39, 44].
Hyperparameter Search Cost Although we provide lightweight guideline in Appendix D.1.2, KLASS introduces a minimal search cost for its additional hyperparameters compared simpler baselines like Top-1 sampling. While our analysis in Appendix D.5.1 shows performance is robust around the optimal points, further tuning could still be performed to find the best possible settings.
G.2 Broader impact
This work focuses on how to improve discrete diffusion sampler. However, improving diffusion samplers, if misused, could facilitate the creation of harmful or misleading text, images, or even novel molecules for malicious purposes.
To mitigate these risks, improvements in sampling techniques should be accompanied by strong emphasis on AI safety research and ethical guidelines. This includes developing robust mechanisms for detecting and preventing the generation of harmful content and fostering public discourse on the socictal implications of ad adnnneed generative models.
33

---

Table 19: Generated solutions with Dream on MATH.
Question
If f()
3r-2 what is the value of f(-2) + f(-1) f(0)?
x-2
Express your answer as a common fraction.
Answer
14
Top-1
First, let's find f(-2): f(-2) 3( 3(-2)-2
-2-2
Next, let's find f(-1): f(-1)
3(-1)-2
-5
-1-2
-3
Now, let's find f(0): f(0) = 3(0)-2
-2
= 1
0-2
-2
Finally, add: ( f0) -4 +
4
12
The final answer is J..
Random
f(0) -
=
3(-2) 1.
0-2
So (--1) 33
The final answer is 3.
KLASS
First, let's find f(-2): f(-2) 3(-2)-2
-8
22-2
-4
Next, let's find f(-1): f(-1) 3( 3(-1)-2
-5
Now, let's find /(0): f(0) 3(0)-2
-2 1
0-2
Finally, add: f(-2) + f(-1) + f(0) 2 +2+3 +
The final answer is 4
34

---

Figure 4: MDLM Generated Sample (512 tokens)
<endoftext>City\ \u20195 providing up to times New York'u2019s current capacity as energy efficient storage is on the way, according to new reporting StarTheSunCity reports.An Awwnnnllnnnaityyy..o.......a says a site in Los Angeles is designed to replace NYC\u2819s current capacity in the emergency room. As SolarCity Krieger confirmed, facility solarcitysolarcity. An\nThe company plans adopt the city\u2019s own capacity taking advantage that the city is constantly being built to low energy standards.
Solar City facility in San Diego will be
Richmond
planned for
The NYTChronicle reports that SolarCity is aiming to provide capacity increase by 10 to 16 percent of capacity, greater than New York. Even the non-Manhattan capacity increase is about percent. Antnstarhentatttccncc says it will soon capacity' Au2019 including some AnAnThe indoor solar center will built to pay solar customers a $3.40 price per day, and based according to the budget for olarrittsanoolssoorrttyyyz22
out of space for an increasing number of 'informs of
not dictated by capacity requirements
facility opens in 2017, and the City Center, located in just one block from Microsoft Place, seats 25 million New Yorkers. SolarCity spokesman Barry said, \u201cIt is clear that the city represents the best leaders in energy efficiency. In fact, providing so much more energy for people is a vital Leadership position for consumers u2019 transportation, health, and finances over the next decade. An nclick to Download The New York Charger and look for True New York Charger to feel free download or for support <endoftext>
Figure 5: KLASS Generated Sample (512 tokens)
</endoftext] |> an order of magni tude smaller than rise in the number of people living in other cities. An\nAs Aities.An\nA central to realignment in how planners view the contribution of each individual population and location to different cities. An\nThe University researchers measured growth in densities by factor of two and space and pollution by 31 per cent in the period of urban sprawl more broadly. A nnTThere is, of course, period of growth in the number of inhabitants relocated public space by quarter to two per cent.
result, urban sprawl is now
Again, at this volume, it has tended to grow unevenly in other cities. of low urban sprawl areas
nThe Berlin Centre for Urban Research has around the borders in Brussels, Belgium..nknLiihhnee AnAnLichtner and co-authors of the study, published in UrbanLabs West, say places with more intense urban sprawl need infrastructure to expand public space.Annnd a result, their boundaries are denser, that means more land for construction and more roads for the new developers want to build. Public libraries and classrooms also go up dramatically, supposedly making urban sprawl severe and steeper. steeper. AnAnLichtner urban sprawl are powerful in this respect. They say much of the impact of the study is ultimately centred the means to thrive. \"If sprawl is the city, it might be infrastructure that allows for housing in sprawl, creating productive concentration between the wealthier and the poor that stand for new develapment. \" they said. < <ndoftexxt]
co-authors say the foundations for
the urban sprawl itself, that offers
35


Reviews:
[{"review_id": "TvTN2qQYb0", "reviewer": "Reviewer_hkNH", "summary": "This 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.", "questions": "See weaknesses above.", "limitations": "yes", "rating": 5, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**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\")", "was_revised": true, "final_justification": "updating score from 4 to 5 upon reading the rebuttal, but do note my low confidence in this topic.", "rebuttal": "We sincerely appreciate your insightful review and are grateful for your acknowledgement that our work:\n- Proposes a novel and intuitve method\n- Provides strong and broad empricial validation across diverse domains, demonstrating both significant speedups and peformance improvements.\n- Offeres a practical, training-free and plug-and-play sampler.\nWe have carefully considered your concerns and provide our detailed responses below.\n---\n\n### [W1-1] Hyperparameter Sensitivity\nWe clarify below how we select hyperparameters in practice and provide additional evidence that KLASS is robust to variations around the chosen values.\n\n**1. Hyperparameter Selection Guideline**\n\nFor hyperparameter selection process, we follow a lightweight three-step search procedure using a small validation set, around 100 examples.\n\n- **Step 1. Initial KL Threshold Estimation**:\nWe estimate a reasonable KL threshold by inspecting the distribution of KL values during decoding.\n- **Step 2. Confidence Threshold Search**:\nWith a fixed KL threshold, we sweep confidence values (e.g., 0.9 down to 0.6) to find the best balance of accuracy and decoding speed.\n- **Step 3. KL Threshold Refinement**:\nWith a fixed confidence threshold, we perform a finer search for the optimal KL value around the initial estimate.\n\nThis process is fast, requiring only a small number of samples and negligible computation compared to training.\n\n**2. Hyperparameter Sensitivity Grid Search**\n\nTo address sensitivity, we performed a grid search across diverse models and tasks (LLaDA/Dream for reasoning, MDLM for molecular), demonstrating KLASS's robustness.\n\nOur findings show that configurations near the selected optimum consistently yield high accuracy with significantly reduced sampling steps. For instance, on HumanEval with LLaDA, settings near (confidence=0.9, KL=0.01) maintain or improve upon the baseline accuracy of **39.63%** while using far fewer than its 256 steps. Moreover, these searches often reveal that slightly different settings can outperform our main reported configuration—such as (confidence=0.95, KL=0.005) on MBPP with Dream—indicating both robustness and potential for further tuning.\n\n**HumanEval (LLaDA)** : Conf = 0.9, KL = 0.01\n\\ **Baseline Performance**: 39.63 (NFE=256)\n|**Conf**|KL = 0.015|**KL = 0.01**|KL = 0.005|\n|---|---|---|---|\n|0.95|39.63 (65.19)|40.24 (99.29)|40.24 (103.14)|\n|**0.9**|40.24 (89.29)|**40.85 (91.98)**|40.24 (96.75)|\n|0.85|39.63 (84.98)|39.63 (88.78)|40.24 (94.48)|\n|0.8|39.02 (83.14)|39.02 (87.21)|40.85 (94.01)|\n\n**MBPP (LLaDA)**: Conf = 0.7, KL = 0.01\n\\ **Baseline Performance**: 48.64 (NFE=256)\n|**Conf**|KL = 0.015|**KL = 0.01**|KL = 0.005|\n|---|---|---|---|\n|0.8|49.42 (122.37)|49.42 (134.52)|49.42 (134.52)|\n|0.75|49.03 (118.61)|49.42 (123.45)|48.25 (132.11)|\n|**0.7**|48.25 (116.24)|**49.03 (127.81)**|49.03 (130.67)|\n|0.65|46.30 (113.22)|48.64 (118.54)|49.42 (128.99)|\n\n**HumanEval (Dream)**: Conf = 0.8, KL = 0.001\n\\ **Baseline Performance**: 58.53 (NFE=256)\n|**Conf**|KL = 0.005|KL = 0.003|**KL = 0.001**|KL = 0.0005|\n|---|---|---|---|---|\n|0.9|54.87 (64.78)|56.10 (67.41)|57.93 (74.86)|61.59 (79.57)|\n|0.85|57.32 (59.39)|55.49 (65.31)|59.15 (73.38)|60.98 (79.43)|\n|**0.8**|51.22 (58.09)|54.27 (62.82)|**59.76 (73.73)**|60.96 (79.51)|\n|0.75|48.78 (55.21)|53.05 (62.61)|59.15 (73.41)|60.37 (79.37)|\n\n**MBPP (Dream)**: Conf = 0.9, KL = 0.001\n\\ **Baseline Performance**: 63.81 (NFE=256)\n|**Conf**|KL = 0.005|KL = 0.003|**KL = 0.001**|KL = 0.0005|\n|---|---|---|---|---|\n|0.95|65.37 (108.93)|65.37 (110.03)|64.59 (112.56)|64.59 (113.14)|\n|**0.9**|62.65 (103.99)|64.20 (107.09)|**64.59 (111.24)**|64.59 (112.65)|\n|0.85|64.20 (101.34)|63.81 (105.43)|65.37 (112.54)|64.59 (112.76)|\n|0.8|63.81 (96.02)|63.04 (103.22)|65.37 (112.52)|64.59 (112.82)|\n\n*Each cell: Accuracy (Steps)*\n\n**Molecule (QED)**: Conf = 0.98, KL = 0.001\n\\ **Baseline Performance**: 0.526 (NFE=32)\n|**Conf**|KL = 0.01|KL = 0.005|**KL = 0.001**|KL = 0.0005|\n|---|---|---|---|---|\n|0.999|0.527 (18.607)|0.526 (18.611)|0.524 (18.580)| 0.538 (18.452)|\n|0.99|0.515 (18.223)|0.521 (18.325)|0.543 (18.689)| 0.537 (18.664)|\n|**0.98**|0.531 (18.434)|0.517 (18.377)|**0.546** (18.782)| 0.534 (18.823)|\n|0.96|0.529 (18.431)|0.535 (18.511)| 0.534 (18.634)| 0.543 (18.801)|\n\n*Each cell: QED (Steps)*\n\nWe will include the results for all datasets in our final version of the manuscript.\n\n### [W1-2] Computational Overhead of KL Divergence\nThe overhead is minimal because it is a lightweight post-processing step on existing logits, requiring no additional forward pass. For the set of masked tokens $I_m = \\lbrace i \\mid z^i_t = m \\rbrace$, the process computes the KL score $d^i_t = D_{KL}(p^i_t || p^i_{t-1})$ and caches the prior distribution. This yields a combined computational $\\Delta C$ and memory overhead of $O(|I_m| \\cdot |V|)$. This linear cost is negligible, as the simple vector operations are insignificant compared to the expensive matrix multiplications and multi-gigabyte footprint of the main diffusion step.\n\nEmpirically, our wall-clock time measurements validate this conclusion, showing **memory overheads under 1.57%** and **latency overheads under 0.21%** per decoding step. Additional results for sequence lengths up to 1,024 are provided in [W4-4].\n\n|Model|Vocab Size|Length|Memory Overhead (MB)|Total Memory (MB)|Time Overhead (s)|Total Time per Step (s)|\n|---|---|---|---|---|---|---|\n|**LLaDA**|126,464|256|247|18,702|0.000255|0.1218|\n|**Dream**|152,064|256|296|18,875|0.000177|0.1275|\n\nWe will include this analysis in our revised manuscript.\n\n### [W1-3] Comparison to Other Adaptive Sampling Methods\nWe thank the reviewer for the valuable suggestion. We agree that a broader discussion of adaptive sampling methods, including those for continuous diffusion, will better contextualize our work. We will incorporate a full comparative analysis, summarized below, into our final manuscript.\n\nKLASS operates dynamically, creating a sample-specific schedule on-the-fly guided by an internal measure of temporal stability (low KL divergence). This self-planning approach contrasts with methods that use fixed, pre-computed schedules from theoretical error bounds (e.g., JYS [1], DUS [2]) or rely on heavy external models for path guidance and error-correction remasking (P2 [3], DDPD [4]). Furthermore, its stability criterion advances beyond simpler heuristics like low confidence (ReMDM [5]) or high entropy (Diffusion-EAGS [6]) used for single-token correction, as it allows KLASS to finalize multiple converged tokens at once to efficiently shorten the generation path.\n\nFurthermore, this step-reduction strategy is orthogonal and complementary to techniques that reduce per-step computational cost. Methods like caching (FreeCache [7], dKV-Cache [8]), which make each denoising step cheaper, can be combined with KLASS to compound acceleration.\n\nWhile KLASS shares an adaptive spirit with solvers for continuous diffusion, their mechanisms are fundamentally different. Continuous solvers like DDIM [9] and DPM-Solver [10] reframe the reverse process as an ODE, using numerical methods to modulate a global sampling trajectory's step size based on local error, enabling 10-20 step generation. In contrast, KLASS operates in the discrete domain by adapting the inference lifespan of individual tokens—a mechanism tailored for token-level convergence.\n\nIn summary, KLASS is a lightweight, training-free sampler that uses a unique temporal stability criterion to dynamically shorten inference, setting it apart from methods that rely on fixed schedules, external planners, simpler heuristics, or orthogonal computational optimizations. \n\n### [W1-4] Wall Clock Time\nYes, we have reported wall-clock latency for MATH in the appendix (Table 8). We have now newly measured and included results for all datasets:\n\n|**Model**|**Dataset**|**Top-1 Acc**|**KLASS Acc**|**Top-1 Time (s)**|**KLASS Time (s)**|**Speedup**|\n|---|---|---|---|---|---|---|\n|**LLaDA**|GSM8K|75.13|76.50|37.04|15.86|**2.34x**|\n|**LLaDA**|MATH|31.40|33.80|38.40|21.41|**1.79x**|\n|**LLaDA**|Humaneval|39.63|40.85|39.54|16.04|**2.47x**|\n|**LLaDA**|MBPP|46.69|47.86|39.12|20.68|**1.89x**|\n|**Dream**|GSM8K|79.75|80.44|29.66|22.26|**1.33x**|\n|**Dream**|MATH|38.00|43.20|30.76|23.31|**1.32x**|\n|**Dream**|Humaneval|58.53|59.76|32.01|11.52|**2.78x**|\n|**Dream**|MBPP|63.81|64.59|31.89|17.65|**1.81x**|\n\nThese results validate that KLASS achieves substantial speedups with comparable or improved accuracy across diverse reasoning tasks. Notably, we observe over 2× speedup on GSM8K for LLaDA, and over 2× speedup on HumanEval for both models, confirming that our method provides meaningful efficiency gains in practical settings.\n\n**[Typo]** Thank you for pointing it out. We will thoroughly proofread and refine the final version.\n\n---\n### **References**\n[1] Park et al. \"Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models.\" ICLR2025.\n\n[2] Luxembourg et al. \"Plan for Speed: Dilated Scheduling for Masked Diffusion Language Models.\" arXiv:2506.19037 (2025).\n\n[3] Peng et al. \"Path Planning for Masked Diffusion Model Sampling.\" arXiv:2502.03540 (2025).\n\n[4] Liu et al. \"THINK WHILE YOU GENERATE: DISCRETE DIFFUSION WITH PLANNED DENOISING.\" ICLR2025.\n\n[5] Wang et al. \"Remasking Discrete Diffusion Models with Inference-Time Scaling.\" arXiv:2503.00307 (2025).\n\n[6] Koh et al. \"PLM-Based Discrete Diffusion Language Models with Entropy-Adaptive Gibbs Sampling.\" arXiv:2411.06438 (2024).\n\n[7] Hu et al. \"Accelerating Diffusion Language Model Inference via Efficient KV Caching and Guided Diffusion.\" arXiv:2505.21467 (2025).\n\n[8] Ma et al. \"dKV-Cache: The Cache for Diffusion Language Models.\" arXiv:2505.15781 (2025).\n\n[9]Song et al. \"Denoising Diffusion Implicit Models.\" ICLR2021.\n\n[10]Lu et al. \"DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps.\" NeurIPS2022."}, {"review_id": "Sa9HFQL24w", "reviewer": "Reviewer_SCro", "summary": "This 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.", "questions": "- 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?", "limitations": "See weaknesses.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "Overall, I am now confident in my score of accept. The paper's contribution is strong and straightforward, with clear reasoning and empirical results across the board. In the rebuttal, the author's addressed my only real concern with additional experiments and theoretical insight, and thus I think the paper is deserving of acceptance at NeurIPS.", "rebuttal": "We are deeply grateful for your thoughtful and constructive review, which has been invaluable in improving our work. We were particularly encouraged by your acknowledgment of:\n\n- The theoretical soundness, empirical motivation, and intuitive nature of our approach.\n\n- Its consistent performane improvements across all testeed modalites and metrics, demonstrating practical effectiveness.\n\nWe have carefully considered the concerns you raised and provide our detailed responses below.\n\n---\n\n### [W2-1] Empirical Justification of the Stability Criterion\n\nWe appreciate the reviewer’s insightful feedback regarding the interpretation of the plots and the relationship between internal model confidence, KL divergence, and correctness.\n\nTo address this, we conducted a newly measured statistical analysis across four reasoning benchmarks, comparing average confidence and KL divergence between correct and incorrect samples for both LLaDA and Dream. As detailed in the tables below, correct samples consistently exhibit higher confidence and lower KL divergence compared to incorrect samples across all tasks and models. This consistent pattern across multiple benchmarks and both models provides empirical evidence that internal confidence and KL divergence serve as reliable indicators of correctness throughout the decoding process.\n\n**LLaDA: Confidence**\n| |GSM8K|MATH|HumanEval|MBPP|\n|---|---|---|---|---|\n|Correct|0.9626 ± 0.0104|0.9512 ± 0.0179|0.9756 ± 0.0136|0.9517 ± 0.0123|\n|Incorrect|0.9497 ± 0.0147|0.9253 ± 0.0263|0.9667 ± 0.0155|0.9450 ± 0.0201|\n\n**LLaDA: KL Divergence**\n| |GSM8K|MATH|HumanEval|MBPP|\n|---|---|---|---|---|\n|Correct|0.0620 ± 0.0161|0.0773 ± 0.0261|0.0449 ± 0.0226|0.0821 ± 0.0234|\n|Incorrect|0.0812 ± 0.0236|0.1091 ± 0.0376|0.0537 ± 0.0235|0.0928 ± 0.0319|\n\n**Dream: Confidence**\n| |GSM8K|MATH|HumanEval|MBPP|\n|---|---|---|---|---|\n|Correct|0.9417 ± 0.0176|0.9309 ± 0.0215|0.9813 ± 0.0138|0.9469 ± 0.0185|\n|Incorrect|0.9224 ± 0.0224|0.9218 ± 0.0248|0.9724 ± 0.0133|0.9204 ± 0.0335|\n\n**Dream: KL Divergence**\n| |GSM8K|MATH|HumanEval|MBPP|\n|---|---|---|---|---|\n|Correct|0.0743 ± 0.0215|0.0873 ± 0.0279|0.0254 ± 0.0154|0.0612 ± 0.0231|\n|Incorrect|0.0988 ± 0.0287|0.0933 ± 0.0319|0.0383 ± 0.0188|0.0891 ± 0.0394|\n\nWe will include comprehensive results across all datasets in the final manuscript to further substantiate these findings.\n\nFurthermore, to provide a theoretical explanation for *why* these empirical patterns are reliable indicators of correctness, we have developed a new information-theoretic framework.\n\nWe formalize the \"unmasking risk\" of a token $i$ using Conditional Mutual Information (CMI), $I_t(i) = I(X_i; Z \\mid V_t)$, which measures the token's dependency on other unknown parts of the sequence. Our core theoretical result proves that the criteria used by KLASS (low cumulative KL divergence and high confidence) are sufficient to guarantee that this CMI is bounded by a provably value, $\\Gamma_t(i) \\le 0.6 \\, \\text{bits}$.\n\nIn essence, this means a token's distributional stability—which we show empirically in the tables above—is a strong indicator of its informational independence. A token that is informationally independent is \"safe\" to unmask and highly likely to be correct, as its value does not depend on other unknowns. This provides the formal link between our empirical observations and the success of our method.\n\nFor the full derivation and detailed discussion of this new framework, we kindly refer the reviewer to our response to Reviewer 5Uk3 [w4-1].\n\n### [W2-2] Comparison with Distillation-Based Methods\nWe agree that distillation can achieve greater speedups. However, we argue that our training-free method, KLASS, is a complementary and orthogonal approach, not a competing one.\n\nA distilled model still relies on an iterative sampling process, and KLASS, as a universal sampler, is designed to optimize this remaining step. This strategy of stacking optimizations has been shown to be effective; for example, Di4C [1] further accelerated the already-distilled SDTT [2] model.\n\nThe two methods address different aspects of optimization. Distillation aims to create a new, faster model artifact through a costly, model-specific training procedure. In contrast, KLASS is a training-free sampler that optimizes the inference process for any given model. This makes KLASS a low-cost, universal, and plug-and-play tool that can be applied immediately for efficiency gains, complementing the more intensive approach of distillation.\n\n**[Q2-1]What is \"Halton-based\" (line 247)?**\n\nThis refers to the experimental setup from the paper \"Halton Scheduler for Masked Generative Image Transformer\" [3] For our image generation task, we adapted the publicly released code for their confidence-based MaskGIT implementation to serve as our baseline.\n\n**[Q2-2]Is there a reason the appendix seems to just contain a completely different (assumed earlier) version of the paper?**\n\nThank you for pointing it out. The slight differences you may have noted are intended to improve overall readability for the reviewers.\n\n**[Q2-3]Lines 146-147: This should say \"KL score for corrects tends to be significantly lower\", right?**\n\nThank you for the careful reading and the helpful suggestion. You are correct. Our intended meaning is that the KL score for correct tokens is significantly lower, as shown in Figure 1b. We will revise the text to make this clearer in the final version.\n\n---\n### **References**\n[1] Hayakawa et al. \"Distillation of Discrete Diffusion through Dimensional Correlations.\" ICML2025.\n\n[2] Deschenaux et al. \"Beyond Autoregression: Fast LLMs via Self-Distillation Through Time.\" ICLR2025.\n\n[3] Besnier et al. \"Halton Scheduler for Masked Generative Image Transformer\" ICLR2025."}, {"review_id": "JonwGOzimS", "reviewer": "Reviewer_yKia", "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.", "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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The proposed low-KL as stabe token is rational. \n2. The theoretical analyses are presented on why ‘KLASS’ improves sample quality in certain scenarios. \n3. The experiemntal results effectively verified the effectiveness. \n\nWeaknesses:\n1. 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. \n2. Compared to the original confidence setting in MaskGIT，the proposed KL metric brings very marginal performance lift.\n3. The hyperparameter selection of threshold varies across tasks and models, limiting the practicality.", "was_revised": true, "final_justification": "Most of my concerns are addressed. I lean towards raise my score.", "rebuttal": "We are grateful for your detailed and constructive feedback. We were pleased to see you recognize that our work:\n\n- Rests on a rational principle supported by theoretical analysis.\n\n- Validates its effectiveness with comprehensive experimental results.\n\nWe have carefully considered your comments and provide our detailed responses below.\n\n---\n\n### [W3-1] Relationship between KL metric and confidence\nWe agree their relationship is critical. Below, we clarify its relationship.\n\n**1. Correlation between confidence and KL metric**\n\nTo better understand the relationship between confidence and KL divergence, we measured the Pearson correlation coefficient between the two metric. We conducted this analysis by recording the confidence and KL divergence of the decoded tokens. As shown in the table below, the correlation between these two metrics is relatively low across all datasets. This indicates that the confidence and KL metrics capture different aspects of model behavior, thus making them complementary rather than overlapping.\n| Model     | GSM8K  | MATH   | HumanEval | MBPP   |\n| --------- | ------ | ------ | --------- | ------ |\n| **LLaDA** | -0.092 | -0.174 | -0.213    | -0.187 |\n| **Dream** | -0.247 | -0.253 | -0.306    | -0.297 |\n\n**2. Impact of using only one metric**\n\n As shown in Table 1 of our manuscript, when we rely on either confidence or KL divergence alone, the model's performance drops compared to when both metrics are utilized together. This supports the argument that while each metric has its own utility, their combined use yields a more robust decision-making process.\n\n**3. Semantic Distinction**\n\nOur empirical and theoretical analyses confirm that confidence and KL divergence are complementary, not redundant. Their Pearson correlation is consistently low across all datasets, and as we show in response **[W2-1]**, correctly generated tokens exhibit both higher confidence and lower KL divergence. This motivates their combined use, as they capture fundamentally different aspects of the model's prediction process.\n\nConfidence provides a static snapshot of certainty at a single timestep ($p_\\theta(x_{t, \\ell} | z_t)$). It measures the peak of a probability distribution, but it contains no information about the iterative process that led to that state. Relying on confidence alone can lead to errors from premature commitments to transient spikes of certainty.\n\nDivergence, by contrast, measures the dynamics of the model's predictions over time by comparing consecutive distributions ($p_\\theta(\\cdot | z_t)$ and $p_\\theta(\\cdot | z_{t-1})$). A low KL score is a direct signal that the model's iterative refinement has concluded for that token—its thought process has converged and stabilized. It is a measure of temporal consistency.\n\nKLASS requires both a confident state and a stable flow. This two-factor check allows us to fully leverage the diffusion model's iterative refinement capability, ensuring we unmask tokens only when the model's prediction is both certain and has verifiably ceased to fluctuate.\n\n### [W3-2] Peformance gap with MaskGIT setting\nThank you for the constructive feedback. While FID score improvement seems marginal compared to the baseline method, significance of KLASS is particularly evident in the Inception Score (IS) [1]. Specifically, KLASS achieved a **9.2-point** increase in IS over the strong baseline (**43% relative gain**) which indicates substantial improvement in sample quality and diversity [2]. By waiting for a prediction to become stable (low KL-divergence) rather than exploiting only confidence value, KLASS enables generating complex classes where confidence-based samplers might often fail, which is indicated by the high IS score. Furthermore, we would like to point out that the reported value has further room for improvement as we conducted only naive hyper-parameter search for the experiment. We will include this discussion in our final manuscript. \n\n### [W3-3] Hyperparameter selection of threshold\nWe fully agree that the practicality of our method hinges on a systematic hyperparameter selection process. To this end, all our experiments were guided by a lightweight, three-step guideline we developed for this purpose. Below, we outline this practical guideline—which requires only a small validation set—and demonstrate its low sensitivity using Dream on MBPP as a concrete example. We kindly ask you to refer **[W1-1]** for additional results on hyperparameter sensitivity.\n\n**Step 1. KL Threshold Estimation**: \nRun baseline decoding on the validation set and collect KL values. Set the initial KL threshold to the median of this distribution (e.g. $\\approx$ 0.01 for reasoning tasks).\n\n**Step 2. Confidence Threshold Search**: \nFix the KL threshold from Step 1. Sweep confidence from 0.9 down to 0.6, selecting the one that maximizes accuracy while maintaining efficiency. For Dream on MBPP, fixing KL = 0.01 yields:\n\n| **Conf** | Acc (%) | NFE    |\n| -------- | ------- | ------ |\n| 0.9      | **67**  | 100.93 |\n| 0.8      | 66      | 89.36  |\n| 0.7      | 64      | 75.03  |\n| 0.6      | 62      | 62.60  |\n\nBoth 0.9 and 0.95 tie in accuracy, but confidence = 0.9 uses fewer evaluations. We therefore select conf = 0.9.\n\n\n**Step 3. KL Threshold Refinement**:\nFix confidence at the value from Step 2 and perform a finer KL sweep around the initial estimate. Results with fixing conf = 0.9 are:\n\n| **KL** | Acc (%) | NFE    |\n| ------ | ------- | ------ |\n| 0.015  | 65      | 98.98  |\n| 0.01   | 67      | 100.93 |\n| 0.005  | 67      | 105.05 |\n| 0.001  | **68**  | 109.47 |\n\nPeak accuracy at KL = 0.001 confirms both the initial estimate and its optimal refinement, giving the final configuration (conf=0.9, kl=0.001).\n\nBy optimizing confidence and KL in this order, we converge on strong configurations with far fewer evaluations than a full grid search. Moreover, since thresholded decoding reduces per-sample latency, the search process is fast and practical. \n\nUsing the selected configuration, we achieve 64.59% accuracy with 111.24 steps—outperforming the top-confidence baseline (63.81% accuracy with 256 steps). Notably, performance remains stable across nearby settings, demonstrating robustness and low sensitivity to hyperparameter choices. The table below shows results on the full dataset:\n\n**Dream on MBPP**\n\n| **Conf** | KL = 0.005     | KL = 0.003     | **KL = 0.001**     | KL = 0.0005    |\n| -------- | -------------- | -------------- | ------------------ | -------------- |\n| 0.95     | 65.37 (108.93) | 65.37 (110.03) | 64.59 (112.56)     | 64.59 (113.14) |\n| **0.9**  | 62.65 (103.99) | 64.20 (107.09) | **64.59 (111.24)** | 64.59 (112.65) |\n| 0.85     | 64.20 (101.34) | 63.81 (105.43) | 65.37 (112.54)     | 64.59 (112.76) |\n| 0.8      | 63.81 (96.02)  | 63.04 (103.22) | 65.37 (112.52)     | 64.59 (112.82) |\n\n*Each cell: Accuracy (Steps)*\n\nAccuracy stays within 1% of the selected configuration across variations, with potential for even further tuning. This stability, combined with minimal search cost, highlights the practicality and generalizability of our guideline across tasks and models.\n\n\n### [Q3-1] Final total sampling steps\nThank you for raising insightful question. KLASS operates within a predefined maximum number of sampling steps but shortens the actual steps used by unmasking a variable number of tokens per step based on the KLASS criterion. If no tokens meet the criterion, the top-k most confident tokens are selected to satisfy a minimum number of unmaskings per step—determined to ensure full sequence completion within the step budget. This strategy allows KLASS to often finish earlier while never exceeding the predefined maximum.\n\nAdditionally, KLASS can be configured to match a fixed number of steps exactly by selecting the required number of tokens from those satisfying the KLASS criterion. To enforce a predictable lower bound on the total sampling steps, we may guide the process by the analytically computable expectation of token transitions from the predefined reverse process. At each step, we decode all tokens that meet the KLASS criterion and then fill any deficit with top-k decoding to meet the target number for that step.\n\nWe have added further results demonstrating that our method also performs well under a shorter predefined budget of 128 steps. While respecting this maximum, KLASS often completes generation in fewer steps and achieves higher accuracy than the top-k confidence baseline.\n\n|**Model**|**Method**|**GSM8K**|**MATH**|**HumanEval**|**MBPP**|\n|---|---|---|---|---|---|\n|**LLaDA**|Top-k|72.40 (128)|29.6 (128)|33.54 (128)|37.74 (128)|\n|**LLaDA**|KLASS|**74.60** (82.50)|**32.2** (92.61)|**34.76** (79.04)|**40.08** (79.96)|\n|**Dream**|Top-k|72.18 (128)|33.6 (128)|42.68 (128)|47.47 (128)|\n|**Dream**|KLASS|**74.75** (118.92)|**37.6** (113.43)|**54.27** (68.03)|**63.04** (98.32)|\n\n*Each cell: Accuracy (Steps)*\n\nWe will include above discussion and results in our final manuscript to strengthen the paper.\n\n---\n### **References**\n[1] Brock et al. \"Large Scale GAN Training for High Fidelity Natural Image Synthesis.\" ICLR2019.\n\n[2] Zhou et al. \"Activation Maximization Generative Adversarial Nets.\" ICLR2018."}, {"review_id": "xeT2gM05xw", "reviewer": "Reviewer_5Uk3", "summary": "The 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.", "questions": "Will we see this made available via a GUI or code repo?", "limitations": "None", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strengths\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).", "was_revised": true, "final_justification": "Excellent quality rebuttal by the authors. This is strong work, and sampling improvements help the whole field. We need more papers like this.", "rebuttal": "We sincerely thank you for your comprehensive and positive comments on our manuscript. We are especially grateful for your recognition of our key contributions:\n\n- The conceptual elegance and practical power of our simple, training-free method.\n- The strength and breadth of our empirical study, demonstrating consistent quality improvements and significant speed-ups across diverse domains.\n- Its reproducibility, ensured by our provision of code and detailed experimental settings.\n\nWe have carefully considered your suggestions for improvement and provide our detailed responses below.\n\n---\n### [W4-1] Theoretical Guarantee\n\nWe thank the reviewer for their helpful comments on Theorem 4.3. We agree that its primary role was to build intuition. Motivated by your feedback, we have developed a new, rigorous information-theoretic framework that formally proves the effectiveness of KLASS. This now serves as the primary theoretical foundation in our revised manuscript.\n\nThe primary goal of the original Theorem 4.3 was to provide a high-level intuition for why KLASS is effective. Our intended definitions were as follows:\n\n* **\"Mild context change\"**: We intended this to mean a change between two contexts, $x_{-i}$ and $x'_{-i}$, where the Hamming distance is bounded by a small integer $k$. This is empirically grounded, as standard samplers like Top-1 or Top-2 correspond to k=1 or k=2. Our own adaptive method, KLASS, also results in a small average k (typically 2-4 tokens per step), confirming the changes are indeed mild.\n\n* **\"δ-approximation norm\"**: This refers to the distance between our model's distribution, $p_\\theta$, and an ideal target distribution, $\\pi$, as measured by KL Divergence. Formally, the assumption is that the model's approximation error is bounded such that $D_{KL}(\\pi \\| p_\\theta) \\le \\delta$.\n\nFormalizing these concepts as you suggested indeed makes the original theorem's structure more rigorous. Taking your valuable feedback a step further, we were motivated to develop an entirely new proof that provides a more direct and quantitatively verifiable guarantee, which we believe better solidifies the theoretical foundations of our work.\n\n\n#### **Provable Framework for Token Unmasking Safety**\n\nTo complement this high-level intuition with a rigorous proof, we have developed a new, self-contained theoretical framework built on standard principles of information theory. This new information-theoretic approach directly analyzes the risk of unmasking a stable token. It provides a **provable, quantitative upper bound** on this risk, thereby offering a formal and rigorous complement to the high-level intuition provided by our original theorem.\n\nTo make our new framework clear, we first define the key terms used in our theorem:\n* **Risk Measure (CMI), $I_t(i)$**: We quantify the risk of unmasking token $i$ at step $t$ using the Conditional Mutual Information, $I_t(i) = I(X_i; Z \\mid V_t)$, which measures its dependency on all other masked tokens $Z$. A low CMI means the token is \"safe\" to unmask.\n* **Cone Entry Step, $t_\\star(i)$**: This is the first step where token $i$ becomes 'stable' by satisfying both KL stability ($K_{t,n}(i) \\leq \\varepsilon$) and confidence ($\\max_v p_t(X_i=v) \\geq \\tau$) thresholds.\n* **Post-Stability KL Tail, $\\delta_t(i)$**: This term aggregates the small, confidence-weighted KL divergences ($\\Delta$) of a token *after* it has entered the stability cone:\n    $$\n    \\delta_t(i) = \\sum_{k=t_\\star(i)}^{t-1} [1-c_{k+1}(i)]\\Delta_{k+1}(i)\n    $$\n* **Initial Entropy, $H_{t_\\star(i)}(X_i)$**: This is the uncertainty (entropy) of token $i$ at the precise moment it becomes stable. Due to the high confidence requirement for entering the cone, this value is provably small.\n\n**Theorem (Provable Bound for Token Safety)**\nLet $i$ be a token that is inside the stability cone at step $t$, and let $t_\\star(i)$ be the step at which it first entered the cone. The mutual information $I_t(i)$ is bounded by:\n$$\nI_t(i) \\leq \\Gamma_t(i) := \\delta_t(i) + H_{t_\\star(i)}(X_i)\n$$where $\\delta_t(i)$ is the small(empirically lower than 1bits), accumulated KL divergence after stabilization and $H_{t_\\star(i)}(X_i)$ is the token's entropy at the moment it entered the cone.\n\n**Corollaries: Pairwise and Block MI Bounds**\nFor any set of stable tokens $\\mathcal{S}_t = \\{ i_1, \\dots, i_m \\}$ unmasked in parallel at step $t$, the total dependency within the block is bounded by the sum of individual risks:\n\n$$I(X_{i_1}; \\dots; X_{i_m} \\mid V_t) \\le \\sum_{k=1}^{m} \\Gamma_t(i_k)$$\n\n#### **How the New Theory Justifies KLASS**\nOur new theory provides a rigorous justification for KLASS's core strategy of safe, parallel unmasking. We formalize the dependency of a candidate token on other unknowns using Conditional Mutual Information $I_t(i)$ and prove that KLASS's criteria guarantee this value is bounded by a tiny amount $\\Gamma_t(i)$. Empirically, we measured this bound to be as low as 0.6 bits on GSM8K, confirming that a token's distributional stability is a strong indicator of its informational independence.\n\nThis guarantee extends to the multi-token case. Our corollaries show that the total dependency within an entire block of stable tokens also remains minimal, formally validating why KLASS's parallel decoding is not just fast, but also safe.\n\n### [W4-2] Lack of Wall-Clock Latency\nWe provide detailed wall-clock latency comparisons between standard top-1 confidence decoding and KLASS across multiple datasets and models. As shown below, KLASS achieves consistent speedups-up to 2.78× faster—while maintaining or improving accuracy. This highlights its practical benefits for reducing decoding time without sacrificing performance.\n\n| **Model** | **Dataset** | **Top-1 Acc** | **KLASS Acc** | **Top-1 Time (s)** | **KLASS Time (s)** | **Speedup** |\n| --- | --- | --- | --- | --- | --- | --- |\n| **LLaDA** | GSM8K | 75.13 | 76.50 | 37.04 | 15.86 | **2.34x** |\n| **LLaDA** | MATH | 31.40 | 33.80 | 38.40 | 21.41 | **1.79x** |\n| **LLaDA** | Humaneval | 39.63 | 40.85 | 39.54 | 16.04 | **2.47x** |\n| **LLaDA** | MBPP | 46.69 | 47.86 | 39.12 | 20.68 | **1.89x** |\n| **Dream** | GSM8K | 79.75 | 80.44 | 29.66 | 22.26 | **1.33x** |\n| **Dream** | MATH | 38.00 | 43.20 | 30.76 | 23.31 | **1.32x** |\n| **Dream** | Humaneval | 58.53 | 59.76 | 32.01 | 11.52 | **2.78x** |\n| **Dream** | MBPP | 63.81 | 64.59 | 31.89 | 17.65 | **1.81x** |\n\n### [W4-3] Uncertainty in GSM8K Improvement Significance\nWhile LLaDA results are deterministic, Dream involves sampling with temperature = 0.2, following the default settings used in prior work. Accordingly, we report variance estimates for Dream, computed over three independent runs for each method, across all four reasoning benchmarks in Table 9 of the appendix.\n\nOn GSM8K, we acknowledge that the accuracy improvement is modest and not statistically significant. However, KLASS achieves a substantial reduction in decoding steps (e.g., from 256 to 155.67), offering meaningful efficiency gains. Additionally, the accuracy improvements and step reduction on other tasks are statistically significant and consistent, reinforcing the overall effectiveness of KLASS across diverse reasoning benchmarks.\n\n### [W4-4] Unsubstantiated Universal Parameter Sensitivity\nWe thank the reviewer for raising this point. Our results indicate that performance is **robust to hyperparameter choices**, remaining stable across different settings. Full grid search results are provided in **\\[W1-1]**.\n\nWhile optimal hyperparameters may vary slightly across tasks, we also provide a **practical hyperparameter search guideline** requiring only 100 validation samples:\n\n1. **KL Initialization**: Set the KL threshold to the median per-token KL divergence on a small validation set.\n2. **Confidence Sweep**: With KL fixed, sweep the confidence threshold to balance accuracy and decoding cost.\n3. **KL Refinement**: With confidence fixed, fine-tune the KL threshold around the initial value for further gains.\n\nIn **[W1-1]**, we report a local sensitivity analysis guided by this procedure. Across all tasks, nearby configurations match or fall within 1% of the best accuracy while maintaining similar decoding efficiency. This confirms that our method is robust and exhibits low sensitivity to hyperparameter choices.\n\n### [W4-5] Qualitative KL History Overhead\nWe thank the reviewer for raising the concern about memory and time overhead for longer sequences. In practice, this overhead is very small, as KLASS is a lightweight post-processing step over existing logits and requires no additional forward pass.\nFor masked tokens $I_m = \\lbrace i \\mid z^i_t = m \\rbrace$, KLASS computes KL scores $d^i_t = D_{\\mathrm{KL}}(p^i_t \\,\\|\\, p^i_{t-1})$ and caches the prior distributions, resulting in a compute and memory cost of $O(|I_m| \\cdot |V|)$. These vector operations are negligible compared to the expensive matrix multiplications and memory usage of the diffusion process.\n\nEmpirically, even for sequences up to 1,024 tokens, our wall-clock measurements show **<5% memory overhead** and **<0.4% time overhead** per decoding step, as shown below:\n\n|Model|Vocab Size|Length|Memory Overhead (MB)|Total Memory (MB)|Time Overhead (s)|Total Time per Step (s)|\n|---|---|---|---|---|---|---|\n|**LLaDA**|126,464|256|247|18,702|0.000255|0.1218|\n|**LLaDA**|126,464|512|494|19,996|0.000673|0.1889|\n|**LLaDA**|126,464|1024|988|22,589|0.00108|0.3114|\n|**Dream**|152,064|256|296|18,875|0.000177|0.1275|\n|**Dream**|152,064|512|592|20,801|0.000272|0.1852|\n|**Dream**|152,064|1024|1184|24,663|0.000515|0.2968|\n\n### [W4-6] Requiring Presentation Refinement\nWe will thoroughly proofread and polish the final version, addressing all minor typographical errors and enhancing overall clarity.\n\n### [Q4-1] Code Releasing\n\nYes, we will make the code publicly available in an open-source repository to encourage further research and use by the community."}]

Author rebuttal:

Meta-review: This paper introduces KL-Adaptive Stability Sampling (KLASS), a simple yet powerful method to accelerate inference in masked diffusion models by leveraging token-level KL divergence and confidence to identify stable tokens for early unmasking. Reviewers praised the method as conceptually elegant, intuitive, and practically impactful, noting its broad empirical validation across reasoning, text, image, and molecule generation tasks. Concerns about hyperparameter sensitivity, computational overhead, and comparisons with related work were thoroughly addressed in the rebuttal with new analyses, statistical validation, and expanded experiments, including strong results on large-scale models. The value is further enhanced by the training-free, plug-and-play nature of KLASS, its solid theoretical grounding, and its consistent speedups of up to ~2–3× with maintained or improved quality. With multiple reviewers raising their scores after rebuttal (including one to a strong accept), the consensus is highly positive.

Overall, KLASS is a rare example of a simple, training-free method that yields consistent 2× speedups across diverse masked diffusion tasks, with no quality loss and negligible overhead. Its practicality and broad applicability make it highly likely to see real-world adoption and inspire follow-up work. I recommend it as a spotlight paper.
Decision: Accept (spotlight)