KLASS: KL-Guided Fast Inference in Masked Diffusion Models

ArXiv: 2511.05664

🎯 Pitch

KLASS accelerates masked diffusion sampling up to 2.78× while improving accuracy by unmasking multiple tokens in parallel based on a token-level KL divergence stability signal—tokens that settle early can be safely committed early. Remarkably, this KL signal separates correct from incorrect predictions better than raw confidence, enabling a simple scheduler to beat the base model's greedy decoding across math and code benchmarks without any retraining.

---

1. Executive Summary

This paper introduces KL-Adaptive Stability Sampling (KLASS) , a training-free sampling method for masked diffusion models that accelerates inference by unmasking multiple tokens in parallel—rather than one-at-a-time—based on two per-token signals: a confidence score (the model's maximum predicted probability) and a KL divergence score that tracks temporal stability of the token's predicted distribution across consecutive diffusion steps. Evaluated on reasoning benchmarks (GSM8K, MATH, HumanEval, MBPP) with LLaDA 8B Instruct and Dream 7B Instruct, KLASS achieves up to 2.78× wall-clock speedups while improving pass@1 accuracy over standard greedy Top-1 decoding—for instance, on MATH with Dream it raises accuracy from 37.97% to 43.20% while reducing sampling steps by roughly 41%—and generalizes across text, image (MMaDA), and molecular (QM9) generation without any additional model training. The method establishes that token-level KL divergence between consecutive step predictions is a strong indicator of solution correctness—correct tokens consistently exhibit lower KL than incorrect ones—but its gains are concentrated on problems within the base model's capability range, as performance on the hardest difficulty regimes remains bounded by what the underlying diffusion model can represent.

2. Context and Motivation

The Core Problem: Masked Diffusion Models Are Bottlenecked by Slow, Static Sampling

Masked diffusion models have recently emerged as a compelling alternative to autoregressive (AR) models for sequence generation. Unlike AR models that generate tokens one-by-one in a strict left-to-right order, masked diffusion models start from a fully masked sequence and iteratively "unmask" tokens over multiple refinement steps, gradually revealing a clean output. This non-autoregressive generation paradigm has demonstrated competitive results on language tasks (LLaDA [27] shows scaling laws comparable to AR models), image synthesis (MaskGIT [7]), molecular design [25, 34], and even planning algorithms [50, 51].

However, these models face a fundamental tension at inference time: the iterative refinement process that gives them their modeling power also makes them slow. The standard ancestral sampling procedure discretizes the reverse diffusion process into a fixed number of timesteps (e.g., 256 or 512) and, in its most common form, unmasks only a small number of tokens per step—often just one. This means generating a 256-token sequence requires 256 sequential forward passes through the model, each of which involves expensive matrix multiplications over the full sequence length. The result is inference latency that can be substantially higher than AR decoding, especially for short-to-medium length sequences where AR models benefit from KV-caching across tokens.

This slowness is not merely an inconvenience—it constrains the practical applicability of masked diffusion approaches in latency-sensitive deployments (interactive assistants, real-time code completion, on-device generation) and makes them less economically viable for high-throughput batch inference compared to AR alternatives. If diffusion models are to fulfill their promise as a general-purpose generative paradigm, their inference speed must improve without sacrificing the sample quality that makes them attractive in the first place.

Compounding the speed problem is a second, subtler issue: existing sampling strategies are static—they apply the same unmasking rule to every token at every timestep regardless of context. For instance, the standard Top-1 sampler always unmasks exactly one token per step (the one with highest predicted probability). A Top-2 sampler always unmasks two. A confidence-threshold sampler unmaskes any token whose predicted probability exceeds some fixed threshold (e.g., 0.9). These approaches treat all tokens as interchangeable units of generation progress, ignoring the reality that at any given diffusion step, some tokens are already highly reliable and could be safely unmasked early, while others are still uncertain and should remain masked longer so they can benefit from additional context revealed by surrounding tokens. A static schedule that forces uniform treatment across all tokens wastes computation on tokens that don't need it and risks prematurely unmasking tokens that do.

This is the gap KLASS aims to fill: adaptive, token-level decision-making about when to unmask, using signals that the model already computes during its forward pass, without requiring additional training or auxiliary models.

Why This Problem Matters: Practical Deployment and a Missing Piece in Diffusion Theory

The importance of solving this problem extends in two directions—one practical, one conceptual.

On the practical side, the deployment economics of language models increasingly depend on inference efficiency. The prevailing paradigm—train the largest model you can afford, then serve it with AR decoding—has well-understood cost characteristics. If diffusion models are to compete in this landscape, they need inference strategies that are not merely "not terrible" but genuinely competitive with AR decoding in wall-clock time. The paper reports that KLASS achieves up to 2.78× speedups over standard Top-1 decoding (Table 8), which begins to close this gap. For applications like code generation (where diffusion models show particular promise due to their ability to refine entire programs holistically) or mathematical reasoning (where step-by-step verification benefits from global coherence), reducing inference latency by half or more can make the difference between a research curiosity and a deployable system.

Moreover, the training-free nature of the approach matters enormously in practice. Methods that require fine-tuning the base model (e.g., distillation approaches like those of Deschenaux and Gulcehre [11] or Hayakawa et al. [15]) are expensive, require maintaining separate model artifacts, and may not transfer across model versions. Methods that rely on external "planner" models to guide token ordering [21, 24, 29] introduce additional computational overhead, memory footprint, and distribution-alignment challenges. A method that works with any pre-trained masked diffusion model, with negligible computational overhead (the paper reports <1.57% memory and <0.21% latency overhead per step; Table 6), and requires no model modifications, is immediately adoptable by practitioners using existing checkpoints.

On the conceptual side, this work addresses a deeper question about how masked diffusion models build confidence during generation. The clean mathematical formulation of the forward and reverse processes in discrete diffusion is well-established—the absorbing-state Markov chain described in Eq. 1–2, the simplified NELBO objective in Eq. 3–4. But the dynamics of how a trained model resolves uncertainty across a sequence during the reverse process remain poorly characterized. When does a token's predicted distribution stabilize? How does the resolution of surrounding context affect individual token predictions? Can we read off a "readiness signal" from the model's own internal computations that tells us when it's safe to commit to a prediction?

The paper's central empirical finding—that the KL divergence between a token's predicted distribution at consecutive timesteps is a strong and consistent indicator of correctness (Figure 1b)—provides a partial answer. Correct tokens exhibit systematically lower KL divergence than incorrect ones across multiple models (LLaDA and DREAM) and multiple reasoning benchmarks (GSM8K, MATH, HumanEval, MBPP). This is not an obvious result: one might have expected that the model's confidence (the maximum predicted probability) would be the primary, or even sufficient, signal. But Figure 1a shows a concrete counterexample where a high-confidence token (conf = 0.9241) is incorrect, while the correct token exhibits lower KL (0.0193 vs. 0.4517). This suggests that confidence alone is insufficient—temporal stability carries independent information about correctness.

This finding connects to broader theoretical questions in generative modeling about when to "commit" to intermediate decisions during iterative refinement. In continuous diffusion, the analogous question might be: at what noise level has a particular pixel's value been effectively determined? In discrete diffusion over tokens, the discrete nature of the state space makes this question both more tractable (we can directly observe when a token's predicted distribution stops changing) and more consequential (an incorrect early unmasking cannot be incrementally corrected, since unmasked tokens are typically kept fixed throughout the remainder of generation in simplified masked diffusion models). The KLASS framework provides a principled, measurement-driven answer to this commitment timing problem.

Where Existing Approaches Fall Short

The paper identifies several categories of prior work on accelerating discrete diffusion sampling, each with specific limitations that motivate the KLASS approach.

Caching and scheduler-based approaches. The simplest acceleration strategies exploit the structure of masked diffusion to skip unnecessary computation. Sahoo et al. [34] observe that if no tokens are unmasked between two consecutive timesteps, and the model is time-invariant, the network's output can be cached and reused, avoiding a forward pass. Zheng et al. [56] propose the First-Hitting Sampler (FHS), which analytically computes the exact time each token should be unmasked, skipping intervening timesteps entirely and unmasking exactly one token per event. While these methods reduce the number of function evaluations, they maintain the one-token-at-a-time unmasking paradigm, meaning the number of sequential steps still scales with sequence length $L$ — an $O(L)$ inherent latency that cannot be further compressed. For long sequences, this remains a significant bottleneck.

Distillation-based acceleration. Several works adapt techniques from continuous diffusion acceleration to the discrete domain. Deschenaux and Gulcehre [11] use self-distillation to train student models that can generate in fewer steps, analogous to progressive distillation for continuous diffusion models [35]. Hayakawa et al. [15] distill discrete diffusion models by exploiting dimensional correlations. These approaches can achieve substantial speedups — potentially reducing steps by an order of magnitude — but they require expensive additional training, produce separate model artifacts that must be maintained and versioned, and the distilled model's quality is bounded by the teacher's. They also lock in a fixed step count (e.g., an 8-step distilled model always takes exactly 8 steps), losing the flexibility to trade off speed against quality at inference time.

Planner-guided sampling. A more recent line of work incorporates an auxiliary "planner" model that determines the order in which tokens should be unmasked. Kim et al. [21] train a separate model to optimize token ordering ("train for the worst, plan for the best"). Liu et al. [24] propose planned denoising where a planner chooses which positions to decode. Peng et al. [29] frame the unmasking order as a path planning problem. These approaches introduce substantial computational overhead — the planner is often itself a neural network requiring its own forward passes — and face a distribution alignment problem: the planner's preferred unmasking order may not match the base model's learned denoising trajectory, potentially leading to suboptimal generation quality. As the paper states, these methods "can lead to difficulty aligning the planner's distribution with the base model's learned distribution" (Section 1).

High-order numerical solvers. Ren et al. [33] develop second-order solvers for the reverse CTMC that achieve higher accuracy per step, allowing fewer total steps for equivalent quality. This is analogous to using Runge-Kutta methods instead of Euler integration in continuous diffusion. While effective, these methods improve the per-step accuracy of the reverse process rather than fundamentally changing which tokens are unmasked when. They remain within the fixed-schedule paradigm where the total number of steps is predetermined and token unmasking decisions follow a uniform rule.

Concurrent confidence-based heuristics. Several works that appeared around the same time as KLASS explore heuristics based on model certainty for multi-token unmasking. Fast-dLLM [47] uses confidence-aware decoding to unmask high-probability tokens. Dimple [54] applies confidence-based parallel decoding for multimodal diffusion models. SlowFast Sampling [45] alternates between slow (single-token) and fast (multi-token) stages based on certainty thresholds and convergence criteria. EB-Sampler [4] uses entropy bounds to determine how many tokens can be safely unmasked per step. Prophet [23] identifies answer tokens by the gap between the top two predicted probabilities.

These concurrent approaches collectively validate the intuition that model certainty can guide efficient unmasking, but the KLASS paper argues — and provides empirical evidence — that certainty alone is insufficient. The core limitation, illustrated in Figure 1a, is that a model can be highly confident about a prediction (probability 0.9241) that turns out to be incorrect. The model's probability distribution may assign high mass to a wrong answer without "knowing" it's wrong — confidence reflects calibration, not correctness. What distinguishes KLASS from these concurrent methods is the addition of a second, orthogonal signal: temporal stability as measured by KL divergence. A token that is both high-confidence and distributionally stable across timesteps is much more likely to be correct than a token that is high-confidence but still fluctuating.

The concurrent nature of many of these works (several appearing in mid-2025, around the same time as the KLASS preprint) suggests that the field was converging on the idea of certainty-guided parallel decoding for masked diffusion, but that the specific mechanism — confidence alone vs. confidence-plus-stability — remained an open design choice. KLASS provides evidence that the stability signal adds meaningful value over confidence alone.

How This Paper Positions Itself

The paper frames KLASS as occupying a specific, previously unoccupied point in the design space of diffusion samplers. The key design constraints are:

1. Training-free: No distillation, no fine-tuning, no auxiliary model training. The method works with any pre-trained masked diffusion model off-the-shelf.
2. No external planners: The unmasking decisions rely solely on signals that the base model already computes during its forward pass (the logits at each position). There is no separate planning network, no search over token orderings, no auxiliary value function.
3. Adaptive at the token level: Unlike fixed-schedule methods that predetermine how many tokens to unmask per step, KLASS adapts to the evolving certainty and stability of each individual token position. Some positions may be ready to unmask early; others may need to wait. The method respects these differences.
4. Principled signal combination: The method is not an ad hoc heuristic but rather combines two signals — confidence and KL divergence — that have complementary relationships to correctness. Section 5 provides a theoretical rationale connecting dynamic instability to incorrect tokens, grounding the approach in the structure of the conditional distributions being modeled.

The paper is explicit about what KLASS is not: it is not a replacement for the underlying diffusion model, not a distillation method, not a planner, and not a numerical solver. It is a decoding-time strategy — a decision rule for when to unmask tokens — that "harness[es] the latent potential of the base diffusion model itself" (Section 1). This positioning is important because it makes KLASS complementary to other acceleration techniques. One could, in principle, combine KLASS with a distilled model (to further reduce steps), with high-order solvers (to improve per-step quality), or with caching strategies (to skip steps where no new tokens are unmasked). The paper doesn't explore these combinations, but the modular design makes them natural extensions.

The paper also positions itself relative to the broader autoregressive-vs-diffusion debate that has characterized language model research in recent years. The fact that LLaDA demonstrates scaling laws competitive with AR models [27] suggests that diffusion is a viable alternative generative paradigm. But for this potential to be realized, inference must become practical. KLASS is presented as a step toward closing the inference efficiency gap — not the final solution, but a meaningful improvement that requires no model modifications and generalizes across modalities. The experiments span text, code, images, and molecules, explicitly demonstrating this cross-modal generality, which strengthens the claim that the approach captures something fundamental about how masked diffusion models converge to solutions rather than being a domain-specific trick.

3. Technical Approach

3.1 Reader Orientation

This is primarily a decoding algorithm design paper whose core idea is that by monitoring two per-token signals during the reverse diffusion process—the model's confidence (maximum predicted probability) and the temporal stability of its predictions (KL divergence between consecutive step distributions)—we can identify which tokens are "ready" to be safely unmasked early, enabling parallel multi-token unmasking that substantially accelerates inference without degrading quality. The paper does not modify the model architecture, retrain the base diffusion model, or introduce auxiliary planning networks; it operates purely as a decision rule layered on top of the model's existing logit outputs at each timestep.

The problem being solved is the static, one-token-at-a-time bottleneck in standard masked diffusion sampling, where a fixed number of tokens are unmasked per step (typically one in Top-1 decoding) regardless of how confident or stable the model's predictions are for individual positions. The "shape" of the solution is a gating mechanism: at each timestep, for each still-masked token position, compute two scalar scores from the model's own probability distributions (no extra forward passes needed), compare them against global thresholds, and unmask all tokens that pass both gates simultaneously. Tokens that fail either gate remain masked, allowing them to benefit from additional context as other tokens are revealed.

3.2 Big-Picture Architecture (Diagram in Words)

The KLASS system has four major components, all operating within the existing ancestral sampling loop of a pre-trained masked diffusion model:

1. Base Masked Diffusion Model ($p_\theta$, e.g., LLaDA 8B Instruct or Dream 7B Instruct) — the frozen pre-trained model that performs forward passes at each timestep $t$, outputting logits (unnormalized log-probabilities over the vocabulary) for every masked token position. This is the only neural network in the system; no additional models are added.

2. Per-Token Signal Extractor — a purely computational module (no learned parameters) that takes the logits from the current timestep $t$ and the cached logits from previous timesteps, and computes two scalar signals per masked position: (a) confidence score $\text{conf}_t^i$, the maximum softmax probability at position $i$ at time $t$, and (b) KL score $d_t^i$, the Kullback-Leibler divergence between the softmax distribution at position $i$ at time $t$ and its distribution at the previous timestep $t+1$, averaged over a short history window.

3. Stable-Token Selector — a thresholding module that applies two global hyperparameters (confidence threshold $\tau$ and KL threshold $\epsilon_{\text{KL}}$) plus a history length $n$ to determine, at each timestep, which masked tokens satisfy the "stability" condition and can be unmasked. The condition is conjunctive: a token is stable only if its confidence exceeds $\tau$ AND its KL divergence has stayed below $\epsilon_{\text{KL}}$ for all of the last $n$ consecutive timestep transitions.

4. Unmasking Executor — the final decision layer at each timestep. If the stable-token selector identifies a non-empty set $S_t$ of ready tokens, all tokens in $S_t$ are unmasked in parallel by sampling from their current predicted distributions. If $S_t$ is empty (no token meets both criteria), a fallback strategy kicks in: the top-$u$ tokens by confidence score are unmasked (where $u$ is a small fixed integer, typically 1), ensuring that generation always makes forward progress.

Information flow at each timestep $t$ (starting from $t = T$ with the fully masked sequence and proceeding to $t = 0$):

1. Forward pass: the frozen base model processes the current partially-masked sequence and outputs logits for all masked positions.
2. Signal extraction: confidence scores are computed as $\max(\text{softmax}(\text{logits}))$ at each masked position; KL scores are computed as $D_{\text{KL}}(p_t^i \| p_{t+1}^i)$ using the cached softmax from the previous timestep.
3. Stable-token selection: each masked position is checked against the two-threshold rule (Eq. 7); positions passing both gates are collected into $S_t$.
4. If $S_t$ is non-empty: all tokens in $S_t$ are unmasked (sampled from their current predicted distributions), the previous-step logits cache is updated, and the process moves to $t-1$. If $S_t$ is empty: the top-$u$ highest-confidence masked tokens are unmasked as a fallback.
5. The loop continues until all positions are unmasked (or the maximum number of timesteps is reached).

The key architectural insight is that steps 2–4 add zero additional forward passes through the neural network—they operate entirely on the logits that the base model already produces. The overhead is the storage of one previous-step softmax distribution per masked token (memory) and the computation of per-token KL divergences (time), both of which are $O(|I_m| \cdot |V|)$ where $|I_m|$ is the number of currently masked tokens and $|V|$ is the vocabulary size.

3.3 Roadmap for the Deep Dive

• First, the per-token confidence score and KL score (Definitions 4.1 and 4.2) — what exactly they measure, how they are computed from the model's logits, and the empirical motivation from Figure 1b showing that correct tokens have systematically lower KL than incorrect ones.
• Second, the stable-token selection rule (Eq. 7) — how the two scores are combined conjunctively with history-length and threshold hyperparameters to form a "stability" predicate, and why this particular combination (rather than a weighted sum or learned function) is chosen.
• Third, the complete KLASS unmasking algorithm (Eq. 8 + Algorithm 1) — how stable-token selection is integrated into the ancestral sampling loop with a fallback mechanism, the precedence rules, and the pseudocode walkthrough.
• Fourth, the theoretical rationale (Proposition 5.3) — why dynamic instability (high average KL) is guaranteed for incorrect tokens in a well-trained conditional model, providing a formal justification for using KL as a correctness signal.
• Fifth, the hyperparameter landscape and selection strategy — the three specific thresholds ($\tau$ for confidence, $\epsilon_{\text{KL}}$ for KL divergence, $n$ for history length, and $u$ for fallback unmasking count), the ranges explored, the sensitivity analysis from the ablation grid (Figure 3), and the lightweight guideline for setting these values on new models and tasks (Appendix D.1.2).
• Sixth, the computational overhead analysis — the $O(|I_m| \cdot |V|)$ scaling of the KL computation, the empirical measurements in Table 6 (<1.57% memory overhead, <0.21% latency overhead per step), and why this cost is negligible relative to the base model's forward pass.

3.4 Detailed, Sentence-Based Technical Breakdown

---

Confidence Score: Measuring Model Certainty Per Token

At each timestep $t$ in the reverse diffusion process, the base masked diffusion model processes the current partially-masked sequence and outputs logits—unnormalized scores—for every token position that is still masked. Let $p_t^i$ denote the categorical probability distribution over the vocabulary $V$ at position $i$ at timestep $t$, obtained by applying the softmax function to the model's logits:

$p_t^i = \text{softmax}(\ell_t^i)$

where $\ell_t^i \in \mathbb{R}^{|V|}$ is the vector of logits output by the model at position $i$ at time $t$, and $\text{softmax}(x)_v = \exp(x_v) / \sum_{v' \in V} \exp(x_{v'})$ for each vocabulary item $v$.

The confidence score for token position $i$ at timestep $t$ is then defined as:

$\text{conf}_t^i = \max_{v \in V} p_t^i(v)$

where $V$ is the model's vocabulary (size 126,464 for LLaDA, 152,064 for Dream), $p_t^i(v)$ is the predicted probability of vocabulary item $v$ at position $i$ at time $t$, and $\max_{v \in V}$ selects the single largest probability across all vocabulary items.

What it computes: the maximum value of the softmax distribution at the given position—the probability the model assigns to its single best guess for that token. A confidence of 0.95 means the model is putting 95% of its probability mass on one particular vocabulary item; a confidence of 0.01 (for a large vocabulary) means the distribution is essentially uniform with no clear preference for any token.

Why this form: maximum probability is the simplest scalar summary of distributional certainty that can be read directly from the model's output without any additional computation or stored state. Alternative measures of certainty exist (entropy, which captures the spread of the full distribution; margin, the gap between the top two probabilities) but confidence has the advantage of being interpretable in absolute terms: a confidence of 0.9 means "the model assigns at least 90% probability to some token," regardless of vocabulary size. The concurrent work Prophet [23] uses the top-2 margin instead, arguing it better captures the model's indecision between competing answers; KLASS uses max-probability because it corresponds most directly to the intuitive notion of "the model has made up its mind about this token."

The key limitation of confidence alone, which motivates the KL score, is that a distribution can be confident without being stable. A model that assigns 0.95 probability to token A at timestep $t$ might assign 0.95 probability to token B at timestep $t-1$—the confidence is identical, the distribution has dramatically changed, and the model is clearly not reliably committed to either prediction. Confidence measures the sharpness of the distribution but not its temporal consistency. This is the gap the KL score fills.

---

KL Score: Measuring Temporal Stability of Per-Token Predictions

While confidence measures how certain the model is at a single instant, the KL score measures how much the model's prediction for a token has changed between consecutive timesteps. The core insight is that for a token whose true value the model has effectively determined, the predicted distribution should stop changing—further refinement steps should not significantly alter the probabilities, because the model has already converged on the correct answer given the available context. For a token where the model is still uncertain, or where resolving context elsewhere in the sequence is shifting its predictions, the distribution should exhibit higher temporal variability.

Formally, the KL score for token position $i$ at timestep $t$ is defined as:

$d_t^i = D_{\text{KL}}\left(p_t^i \;\|\; p_{t+1}^i\right)$

where $D_{\text{KL}}$ is the Kullback-Leibler divergence, $p_t^i$ is the categorical distribution over the vocabulary at position $i$ at the current timestep $t$, and $p_{t+1}^i$ is the categorical distribution at position $i$ at the previous timestep $t+1$ (recall that the reverse process goes from high $t$ to low $t$, so $t+1$ is the step immediately before $t$ in chronological generation order).

The KL divergence between two discrete distributions $P$ and $Q$ is defined as:

$D_{\text{KL}}(P \;\|\; Q) = \sum_{v \in V} P(v) \log \frac{P(v)}{Q(v)}$

where the sum is taken over all vocabulary items $v$. When $P$ and $Q$ are identical, the divergence is zero; as $P$ diverges from $Q$ (placing mass on tokens where $Q$ placed little mass), the divergence grows without bound.

Operationally, at each timestep $t$, the KLASS algorithm computes $d_t^i$ for each currently-masked position $i$ using the softmax distribution from the current forward pass ($p_t^i$) and the cached softmax distribution from the previous forward pass ($p_{t+1}^i$). The previous-step distribution must have been stored at time $t+1$; this is the primary memory cost of the method. Since $p_t^i$ and $p_{t+1}^i$ are both already computed (the former for the current confidence check and for the eventual unmasking decision, the latter from the previous iteration's forward pass), the KL computation introduces no additional forward passes—only the element-wise sum over the vocabulary.

Why KL divergence specifically: The paper chooses KL divergence over other distributional distance measures for two reasons. First, KL divergence is asymmetric—$D_{\text{KL}}(p_t^i \| p_{t+1}^i)$ measures how much the current distribution diverges from the previous distribution, which matches the causal direction of the reverse process: we want to know whether the model's prediction at time $t$ is consistent with where it was at time $t+1$. An alternative symmetric measure like Jensen-Shannon divergence would treat changes in either direction symmetrically, which is less interpretable for tracking convergence. Second, KL divergence has a direct connection to information theory: the KL divergence represents the number of extra nats (or bits, if using log base 2) needed to encode samples from $p_t^i$ using a code optimized for $p_{t+1}^i$. A high KL means the current distribution carries substantially new information relative to the previous step—the model has "changed its mind." A low KL means the current distribution is informationally redundant—the model is maintaining its prediction.

Why not just use the change in the predicted token identity? One might ask: why not simply check whether the argmax token at position $i$ has changed since the previous step? This would collapse the entire distribution into a binary signal (same token or different token), losing information about the magnitude of the change. A token whose predicted probability shifts from 0.91 to 0.89 for the same argmax token is essentially stable; a token that flips from 0.51-0.49 between two competing tokens has fundamentally unstable support. The binary change signal would flag the latter but not the former; the KL divergence captures both the identity change and the probability mass shift in a continuous, graded way.

---

Empirical Motivation: KL Divergence Separates Correct from Incorrect Tokens

Figure 1b in the paper provides the key empirical justification for using KL as a correctness signal. The experiment procedure (inferred from the figure and Section 4.1 text):

1. For each of four reasoning benchmarks (GSM8K, MATH, HumanEval, MBPP), generate complete solutions using the standard diffusion process with LLaDA and Dream models.
2. For each token in each generated solution, compute the average KL divergence across the timesteps during which that token was unmasked (i.e., the mean $d_t^i$ for all $t$ where position $i$ transitions from masked to unmasked).
3. Partition tokens into two groups: those that belong to correct solutions (final answer matches ground truth) and those that belong to incorrect solutions.
4. Plot the distribution of average KL scores for each group.

The result: across all four benchmarks and both models, the KL divergence distribution for correct tokens is consistently and substantially lower than for incorrect tokens. The box plots in Figure 1b show median KL values for correct tokens clustering around 0.04–0.06, while incorrect tokens show medians around 0.08–0.10, with significant separation between the distributions.

Why this matters for algorithm design: This separation means that KL divergence carries independent information about correctness that is not captured by confidence alone. The case study in Figure 1a illustrates this concretely. In a mathematical reasoning problem (from GSM8K), the model generates an incorrect arithmetic step: "25 - 20 = 10" has high confidence (0.9241) but a high KL divergence (0.4517). The correct alternative "25 - 20 = 5" has a lower KL divergence (0.0193). The Top-k confidence sampler selects the incorrect token because it has higher confidence; KLASS selects the correct token because the low KL flag overrides the confidence difference.

This case study demonstrates the complementary nature of the two signals. Confidence captures the model's current certainty—a snapshot. KL captures the model's convergence trajectory—a dynamic property. A token can be confidently wrong (the model is sure but hasn't converged), and a token can be uncertain but correct (the model has converged to the right answer with modest probability). KLASS's conjunctive condition—both high confidence AND low KL—filters for tokens that are both certain AND converged, which is a stricter and more reliable correctness indicator than either signal alone.

---

Stable-Token Selection: The Conjunctive Threshold Rule

The stable-token selection rule (Eq. 7 in the paper) combines the confidence score and KL score into a binary predicate that determines whether position $i$ is ready to be unmasked at timestep $t$. The rule has three hyperparameters: a history length $n$, a KL threshold $\epsilon_{\text{KL}}$, and a confidence threshold $\tau$.

Formally, the set of stable tokens at step $t$ is:

$S_t = \left\{ i \;\middle|\; \underbrace{\forall k \in \{1, \ldots, n\} \; D_{\text{KL}}\left(p_{t+k-1}^i \;\|\; p_{t+k}^i\right) < \epsilon_{\text{KL}}}_{\text{all recent KL values below threshold}} \;\wedge\; \underbrace{\text{conf}_t^i > \tau}_{\text{high confidence}} \right\}$

where:

• $i$ indexes token positions in the sequence,
• $n$ is the history length (how many consecutive past timestep transitions must have low KL),
• $k$ indexes over the $n$ most recent timestep transitions (from $t+1 \to t$, $t+2 \to t+1$, ..., $t+n \to t+n-1$),
• $p_{t+k-1}^i$ and $p_{t+k}^i$ are the predicted distributions at position $i$ at timesteps $t+k-1$ and $t+k$ respectively,
• $\epsilon_{\text{KL}}$ is the KL divergence threshold (must be strictly below this value),
• $\text{conf}_t^i$ is the confidence score at position $i$ at the current timestep $t$, and
• $\tau$ is the confidence threshold (must strictly exceed this value).

What it computes: a set of token indices. A token is included only if it passes two simultaneous checks: (1) its predicted distribution has been "stable" for the last $n$ transitions (all $n$ KL divergences are below $\epsilon_{\text{KL}}$), and (2) its current predicted distribution is "confident" (the maximum probability exceeds $\tau$). Both conditions must hold; a token that has low KL but low confidence (the model is consistently uncertain) stays masked, as does a token with high confidence but high KL (the model just became certain but was changing recently).

Why history length $n > 1$: A single low-KL transition ($n = 1$) only tells us that the distribution didn't change between the last two timesteps. But the model might have converged to a stable-but-wrong prediction two steps ago and maintained it since—the KL would be low, suggesting stability, but the token would still be incorrect. Requiring $n > 1$ (the paper uses $n = 2$ for all main experiments) demands evidence of stability over a longer window, which filters out tokens that only recently settled and may not have truly converged. The ablation in Table 16 examines $n \in \{1, 2, 3\}$ and generally finds $n = 2$ to be optimal—$n = 1$ unmaskes too aggressively (premature unmasking), while $n = 3$ is too conservative (many tokens never stabilize enough to unmask before the process ends, reducing the speedup benefit).

Why conjunctive rather than a weighted combination: An alternative design would compute a single scalar "readiness score" as a weighted combination of confidence and KL (e.g., $\text{conf}_t^i - \lambda \cdot \text{KL}_t^i$), then threshold on that combined score. The paper argues that confidence and KL carry qualitatively different information that should not be traded off linearly. A token with confidence 0.99 and KL 10 (rapidly changing despite high certainty) is fundamentally different from a token with confidence 0.5 and KL 0.001 (stable but uncertain). A linear combination would allow the high confidence to "compensate" for the high KL, masking the instability; the conjunctive rule requires both signals to independently pass their thresholds, ensuring neither type of unreliability can slip through.

Threshold ranges explored: The paper sweeps confidence thresholds $\tau \in \{0.5, 0.6, 0.7, 0.8, 0.9\}$ and KL thresholds $\epsilon_{\text{KL}} \in \{0.001, 0.005, 0.01, 0.015, 0.02\}$ (see Figure 3 and Table 7). The specific optimal values vary by model and task—LLaDA on MATH uses $(\tau = 0.6, \epsilon_{\text{KL}} = 0.01)$, while Dream on MATH uses $(\tau = 0.9, \epsilon_{\text{KL}} = 0.005)$—but the sensitivity analysis in Appendix D.5.1 demonstrates that performance is robust around the optimal point: nearby threshold values produce similar accuracy, and the degradation is gradual rather than catastrophic.

Connection to the theory (Section 5): Proposition 5.3 provides formal justification for the KL condition: for a well-trained model that closely approximates the true conditional distributions of the data, a token that is predicted incorrectly at the current step cannot remain uniformly stable as additional context is resolved. Specifically, if the model currently prefers an incorrect token $x^\dagger$ at context $c_M$ (by some margin $\beta$), but the correct token $x^\star$ is preferred at the optimal context $c^\star$ (by some margin $\gamma$), then along any path from $c_M$ to $c^\star$ the average per-step KL divergence is bounded below by $2\Delta^2/M^2$, where $\Delta = \frac{1}{2}(\beta + \gamma - 2\delta)_+$ and $\delta$ is the model's approximation error. In plain language: an incorrect prediction cannot stay stable—as the surrounding context gets resolved, the model's distribution for that position must shift, producing non-zero KL divergence that will eventually exceed any fixed threshold $\epsilon_{\text{KL}}$. Correct predictions, by contrast, can remain stable because they are already consistent with the optimal conditional distribution. This theoretical result justifies delaying unmasking until KL drops below threshold: it ensures we wait until the token's prediction has "settled" on the correct answer, which is the only state that can remain persistently stable.

---

The Complete KLASS Unmasking Algorithm

The stable-token selection rule (Eq. 7) defines which tokens are ready, but it doesn't fully specify what happens at each timestep—there are edge cases to handle, most critically: what if no token satisfies both criteria?

The complete unmasking rule at each timestep $t$ (Eq. 8) is:

$x_t^i = \begin{cases} \text{unmask token at position } i, & i \in S_t \\ \text{otherwise, unmask the top-} u \text{ positions by } \text{conf}_t^i, & S_t = \emptyset \end{cases}$

where:

• $x_t^i$ denotes the state of token position $i$ at timestep $t$ (either a vocabulary token or the mask token),
• $S_t$ is the set of stable tokens as defined in Eq. 7,
• $u$ is a small fixed integer (typically $u = 1$) serving as the fallback unmasking count,
• $\text{conf}_t^i$ is the confidence score used to rank tokens when no stable tokens exist.

What happens operationally, step by step (Algorithm 1 in Appendix B):

1. Initialization: The sequence $x$ is initialized to all mask tokens: $x \leftarrow [\text{MASK}]^L$ where $L$ is the generation length (256 for all reasoning experiments). A buffer $\text{P}_{\text{prev}}$ is initialized to store the previous-step softmax distributions (initially zero or a placeholder). A circular buffer $\text{KL}_{\text{buf}}$ of size $n$ (history length) is initialized for each masked position to track recent KL values.

2. Timestep loop: For each diffusion step $t$ from $T$ down to $1$ (where $T = 256$ in the main experiments):

   a. Forward pass: Call the base model $M(x)$ on the current partially-masked sequence $x$ to obtain logits $\ell$ for all masked positions. Compute the softmax distributions $P = \text{softmax}(\ell)$.

   b. Signal computation: For each masked position $i$:

   • Confidence: $c^i = \max_v P^i(v)$
   • KL divergence: $\delta^i = D_{\text{KL}}(P^i \;\|\; P_{\text{prev}}^i)$, requiring the cached $P_{\text{prev}}^i$ from the previous timestep

   c. History update: Roll the KL buffer for each position (shift old values, make room for the new $\delta^i$) and store the current $\delta^i$ at the front. Cache $P$ as $P_{\text{prev}}$ for the next timestep.

   d. Stable-token identification: For each masked position $i$, check:

   • stable_kl: Are all $n$ entries in $\text{KL}_{\text{buf}}^i$ strictly less than $\epsilon_{\text{KL}}$? (The buffer must be fully populated; if fewer than $n$ KL values have been recorded, the condition cannot be satisfied.)
   • high_conf: Is $c^i > \tau$?
   • is_masked: Is position $i$ currently masked?
   • ready[i] = stable_kl AND high_conf AND is_masked

   e. Unmasking decision:

   • If any ready[i] is True: For all positions where ready[i] is True, sample a token from the current distribution $P^i$ (or take the argmax, depending on temperature) and replace the mask token at position $i$ with the sampled token.
   • If no ready[i] is True (fallback): Compute a score for each masked position as c^i \cdot \text{is_masked}[i], select the top-$u$ positions by this score, and unmask those $u$ tokens (sampling from their current predicted distributions). This ensures exactly $u$ tokens are unmasked when the stability criteria identify none.

   f. Advance timestep: $t \leftarrow t-1$, repeat from 2a until $t = 0$ or all positions are unmasked.

3. Output: Return the fully unmasked sequence $x$.

The fallback mechanism—why $u$ is needed: The stable-token selection rule can, in principle, return an empty set at any timestep. This could happen early in generation when context is too sparse for any token to have converged, or late in generation if the remaining masked tokens are intrinsically uncertain and never stabilize. If the algorithm waited for stabilization without a fallback, it could stall indefinitely—no new tokens would be unmasked, context wouldn't change, and distributions would remain frozen, creating a deadlock. The top-$u$ fallback guarantees progress: at each step, at least $u$ tokens (typically 1) are always unmasked, ensuring the sequence is completed in at most $L$ steps. The value $u = 1$ is used in all main experiments, matching the single-token unmasking rate of standard Top-1 decoding and providing a conservative floor: KLASS never does worse than Top-1 in terms of step count, and typically does much better because it unmaskes many more tokens when stability allows.

Why apply the KL history check only to masked tokens: Tokens that have already been unmasked in previous timesteps are kept fixed—the simplified masked diffusion framework (Section 3.1, Eq. 2) guarantees that unmasked tokens never revert to the mask state. So their distributions are no longer relevant for unmasking decisions. The KL history buffer only needs to be maintained for currently-masked positions, which shrinks as generation progresses. This means the computational and memory overhead of the KL computation is front-loaded (largest early in generation when many tokens are masked) and decreases as the sequence fills in.

Temperature and sampling: For the LLaDA experiments, temperature is set to 0 (deterministic argmax), so "unmasking" simply sets the token to the argmax of the current predicted distribution. For Dream experiments, temperature is set to 0.2, so tokens are sampled stochastically from the softmax distribution with temperature scaling. The KLASS algorithm itself is agnostic to the sampling strategy—it only affects the final token selection within the unmasking step, not the stability gate decisions. The confidence and KL scores are always computed on the full softmax distributions, not on the temperature-scaled or top-p filtered versions.

Handling the initial few timesteps: The history-length condition requires $n$ prior KL values to evaluate. At the very beginning of generation (timesteps $T, T-1, \ldots, T-n+1$), the KL buffer is not yet fully populated, so the stable_kl condition cannot be satisfied. During these early steps, the algorithm relies entirely on the fallback mechanism—unmasking the top-$u$ tokens by confidence. This is explicitly designed: early in the reverse process, the model has almost no context (the sequence is mostly or entirely masked), and no token can be expected to have converged. The fallback ensures forward progress while the model builds up enough context for meaningful stability measurements.

---

Theoretical Rationale: Why Unstable Tokens Tend to Be Wrong

Section 5 of the paper provides a theoretical argument connecting dynamic instability (high average KL) to incorrectness. The argument is not a tight performance bound but rather a consistency check: it shows that for a sufficiently well-trained model, an incorrectly predicted token cannot remain stable as context is resolved, which justifies using KL divergence as a filter for correctness. The proof proceeds by contradiction: assume an incorrect token is stable; then the model's distribution cannot be changing; but as context improves, the model must eventually shift toward the correct conditional distribution; contradiction.

The formal framework:

Definition 5.1 (Conditional $\delta$-approximation): A model $p_\theta$ is a conditional $\delta$-approximation to the task if there exists some set of "task-correct" conditional distributions $\mathcal{C}$ (the distributions that would be produced by a perfect model) such that for any context $c$ (the values of all tokens outside position $i$), the total variation distance between $p_\theta(\cdot \mid c)$ and some $\pi(\cdot \mid c) \in \mathcal{C}(c)$ is at most $\delta$. In operational terms: the model's predictions are always within TV distance $\delta$ of some correct conditional distribution, where $\delta$ quantifies the model's imperfection.

Definition 5.2 (Incorrect prediction with margin): Fix a token position $i$. Let $x^\star$ be the correct token under the optimal context $c^\star$ (the fully-revealed sequence), and let $x^\dagger$ be an alternative incorrect token. Assume that under the optimal context, the correct model prefers $x^\star$ by at least a margin $\gamma > 0$: $\pi(x^\star \mid c^\star) \geq \pi(x^\dagger \mid c^\star) + \gamma$. However, at the current (partially-masked, suboptimal) context $c_M$, the model $p_\theta$ prefers the incorrect token $x^\dagger$ by at least a margin $\beta \geq 0$: $p_\theta(x^\dagger \mid c_M) \geq p_\theta(x^\star \mid c_M) + \beta$.

Proposition 5.3 (Instability lower bound): Under these conditions, for any path of contexts $c_M \to c_{M-1} \to \cdots \to c_0$ that progressively reveals more tokens (eventually reaching the optimal context $c_0 = c^\star$), the total variation distance between the model's distribution at $c_M$ and at $c^\star$ is at least $\Delta = \frac{1}{2}(\beta + \gamma - 2\delta)_+$, where $(y)_+ = \max(y, 0)$. Moreover, the average per-step KL divergence along this path is bounded below:

$\frac{1}{M} \sum_{t=0}^{M-1} \text{KL}\left(P_t \;\|\; P_{t+1}\right) \geq \frac{2\Delta^2}{M^2}$

where $P_t = p_\theta(\cdot \mid c_t)$, $M$ is the number of context transitions along the path, and $\Delta$ is as defined above.

What this means operationally: If the model is currently wrong about token $i$ (it prefers $x^\dagger$ over $x^\star$ at the current context), then as the remaining masked tokens in the sequence are progressively revealed (moving from $c_M$ to $c^\star$), the model's distribution at position $i$ must change substantially—the total variation distance between the current distribution and the final distribution is at least $\Delta$, and the average KL divergence per step is at least $2\Delta^2/M^2$. The quantity $\Delta$ is larger when (1) the model is more confidently wrong at the current step (larger $\beta$), (2) the true margin favoring the correct answer is larger (larger $\gamma$), and (3) the model is a better approximator (smaller $\delta$).

Why this justifies the KL threshold: Since the lower bound on average per-step KL is strictly positive when $\beta + \gamma > 2\delta$, an incorrect prediction forces non-zero KL divergence at some point along the path to full context. A token that remains below $\epsilon_{\text{KL}}$ for an extended history window cannot be incorrect (for sufficiently small $\epsilon_{\text{KL}}$, relative to the bound), because if it were incorrect, the required distributional shift would eventually push the KL above any small threshold. This is not a guarantee that every correct token will have low KL—some correct tokens might also change as context refines—but it guarantees that the set of tokens with persistently low KL is a subset of mostly correct tokens, with the error rate depending on $\epsilon_{\text{KL}}$, $\delta$, and the distribution of margins $\gamma$.

The key assumptions and their limitations: The theory assumes the model is a $\delta$-approximation to the true conditionals—that there exists some "correct" distribution that the model approximates within TV distance $\delta$. For models trained on finite data with finite capacity, $\delta$ may be non-negligible, and if $\delta > (\beta + \gamma)/2$, the bound becomes vacuous ($\Delta = 0$). The theory also assumes a path that ends at the optimal context $c^\star$, but the actual reverse diffusion process follows a stochastic path determined by the model's own predictions, not a prescribed context-revelation order. Tokens may be unmasked in any order, and the context path may not reach $c^\star$ at all if the model makes errors along the way. So the theoretical result is best understood as an existence argument: for a sufficiently well-trained model, instability is a necessary condition for incorrectness, making stability a useful (though not perfect) filter.

---

Hyperparameter Configuration and Selection Strategy

KLASS introduces four hyperparameters that control the tradeoff between generation speed (how many tokens are unmasked in parallel) and generation quality (how reliably those early-unmasked tokens are correct):

1. Confidence threshold $\tau$ (range explored: 0.5–0.95): The minimum probability the model must assign to its best-guess token for that position to be considered "confident." Higher $\tau$ means stricter filtering—fewer tokens qualify as stable, so fewer are unmasked per step, reducing speedup but improving reliability. Optimal values range from 0.6 (LLaDA on MATH) to 0.9 (Dream on MATH, LLaDA on HumanEval). The wide variation across models and tasks reflects differences in model calibration: LLaDA is generally less calibrated (assigns lower probabilities to correct tokens), requiring a lower threshold to achieve meaningful parallel unmasking, while Dream is better calibrated and can use a higher threshold without choking off too many tokens.

2. KL divergence threshold $\epsilon_{\text{KL}}$ (range explored: 0.001–0.02): The maximum KL divergence allowed for a token's distribution to be considered "stable" at any single transition. Lower $\epsilon_{\text{KL}}$ means stricter stability requirements—only tokens whose distributions are nearly frozen qualify. Optimal values range from 0.001 (Dream on several tasks) to 0.015 (LLaDA on GSM8K). A natural unit interpretation: $\epsilon_{\text{KL}} = 0.01$ means the distribution can change by roughly 0.01 nats, which corresponds to a very small shift in probability mass (e.g., a token going from 95% to 94% while another goes from 5% to 6%).

3. History length $n$ (range explored: 1–3): The number of consecutive timestep transitions that must all have KL below $\epsilon_{\text{KL}}$. Higher $n$ demands longer evidence of stability. The ablation in Table 16 shows $n = 2$ is optimal across most settings—$n = 1$ produces higher speedups but lower accuracy (premature unmasking of tokens that haven't truly converged), while $n = 3$ produces lower speedups with marginal or no accuracy improvement (many tokens never stabilize for three consecutive steps, so the fallback fires more often).

4. Fallback unmasking count $u$ (value used: 1 in all main experiments): The number of tokens to unmask via confidence ranking when no tokens satisfy the stability criteria. Setting $u = 1$ ensures progress is never slower than standard Top-1 decoding. Larger $u$ values would accelerate the fallback case but risk unmaskeding unstable tokens; the paper doesn't explore $u > 1$, treating this as a safety valve rather than a tunable speedup lever.

The lightweight hyperparameter selection guideline (Appendix D.1.2) describes a three-step procedure requiring only about 100 validation examples:

1. Initial KL threshold estimation: Run the diffusion model on a few validation examples and record the distribution of per-token KL values across all timesteps. The initial $\epsilon_{\text{KL}}$ estimate is set to approximately the 25th–50th percentile of observed KL values—low enough to be selective, high enough that some tokens will pass.

2. Confidence threshold sweep: With $\epsilon_{\text{KL}}$ fixed at the initial estimate, sweep $\tau$ from 0.9 downward (0.9, 0.8, 0.7, 0.6, 0.5) and select the value that yields the best tradeoff between accuracy and decoding step count on the validation set.

3. KL threshold refinement: With $\tau$ fixed at the best value from step 2, refine $\epsilon_{\text{KL}}$ through a finer-grained search around the initial estimate (e.g., ±50% in small increments). Select the value that maximizes accuracy.

The paper reports the final configurations in Table 7. Notable patterns: LLaDA uses lower confidence thresholds (0.6–0.9) and higher KL thresholds (0.01–0.015) than Dream (0.8–0.9 and 0.001–0.005, respectively), suggesting LLaDA's predictions are less calibrated but more stable, while Dream's are better calibrated but more volatile. The MATH benchmark consistently requires different thresholds than code benchmarks (HumanEval, MBPP)—math reasoning involves longer chains of dependent steps, where premature unmasking of an intermediate calculation can cascade into downstream errors, favoring stricter thresholds.

Sensitivity analysis (Appendix D.5.1, Tables 11–14): The grid search results demonstrate that KLASS performance is robust around the optimal point. For LLaDA on HumanEval (Table 11), configurations near the chosen ($\tau = 0.9, \epsilon_{\text{KL}} = 0.01$) all achieve accuracy within 1.2 percentage points of the 40.85% optimum while still using 90–103 steps (vs. 256 for Top-1). For Dream on MATH (Table 14 in the appendix, actually the MBPP table), accuracies range from 62.65% to 65.37% across a broad region of the hyperparameter space, always exceeding the Top-1 baseline of 63.81% while using 96–113 steps. The absence of sharp cliffs in the accuracy surface means practitioners don't need precise hyperparameter tuning—approximate values from a small validation set suffice.

---

Computational Overhead: Memory and Latency

The KLASS algorithm adds computation to each diffusion step, but the paper argues—and provides empirical measurements (Section 6.6, Table 6) to support—that this overhead is negligible relative to the base model's forward pass.

Operation count: At each timestep, for each currently-masked token position $i$, the algorithm must:

• Compute $p_t^i = \text{softmax}(\ell_t^i)$ — but this is already done as part of the confidence computation and the eventual unmasking step, so it's not additional work.
• Compute $D_{\text{KL}}(p_t^i \| p_{t+1}^i) = \sum_{v \in V} p_t^i(v) \log(p_t^i(v) / p_{t+1}^i(v))$ — this is a sum over the vocabulary size $|V|$, requiring one log, one division, and one multiplication per vocabulary item, plus accumulation. This is $O(|I_m| \cdot |V|)$ where $|I_m|$ is the number of currently masked tokens.
• Maintain the circular buffer of $n$ past KL values per masked position — this is $O(|I_m| \cdot n)$ space and $O(|I_m|)$ time per step for the roll operation.

Why this is negligible: The base model's forward pass at each timestep involves attention computations that scale as $O(L^2 \cdot d_{\text{model}})$ (where $L$ is sequence length and $d_{\text{model}}$ is hidden dimension) and feedforward computations that scale as $O(L \cdot d_{\text{model}} \cdot d_{\text{ff}})$. For LLaDA 8B, $d_{\text{model}}$ is approximately 4096 and $d_{\text{ff}}$ is approximately 14336, while $|V| = 126,464$. The KL computation is $O(|I_m| \cdot |V|) \leq O(L \cdot |V|)$, and $L \cdot |V| = 256 \times 126,464 \approx 3.2 \times 10^7$ floating-point operations (a direct sum), while the attention alone is $O(L^2 \cdot d_{\text{model}}) = 256^2 \times 4096 \approx 2.7 \times 10^8$ operations with additional multiplicative factors from multi-head attention and the quadratic attention complexity. The KL computation is one to two orders of magnitude cheaper than the forward pass.

Empirical measurements (Table 6): The paper reports measurements on a single NVIDIA RTX A5000 GPU with generation length 256:

The memory overhead (247–296 MB) comes from caching the previous-step softmax distributions for all masked positions, which requires storing $|I_m| \times |V|$ floating-point values (at 4 bytes per float, $256 \times 152,064 \times 4 \approx 156$ MB for Dream, plus buffer overhead). This is <1.6% of the total GPU memory occupied by the 7–8B parameter models. The time overhead per decoding step (0.00018–0.00026 seconds) is <0.21% of the per-step forward pass time (0.12–0.13 seconds).

Key practical implication: These overhead measurements mean that the wall-clock speedups reported in Table 8 (1.32× to 2.78× over Top-1) are realized speedups, not theoretical step-count reductions that get partially eaten by per-step overhead. If KLASS reduces the average number of steps from 256 to, say, 100, the total KL overhead across those 100 steps is approximately $100 \times 0.00026 = 0.026$ seconds for LLaDA—essentially invisible compared to the $100 \times 0.1218 = 12.18$ seconds for the forward passes. The speedup comes almost entirely from the reduction in the number of expensive forward passes, not from any clever caching or skipping of computation.

Where the overhead could become non-negligible: For models with much larger vocabularies (e.g., multilingual models with vocabularies of 250K+ tokens) or very long sequences (e.g., 2048 tokens), the KL computation cost grows linearly in both dimensions, while the attention cost grows quadratically in sequence length. For sufficiently long sequences, the quadratic attention dominates, and the KL overhead remains negligible. For very large vocabularies and short sequences, the linear $|V|$ factor could become a larger fraction of total cost, but for typical language model configurations (vocabularies of 50K–150K, sequences of 128–2048 tokens), the overhead remains in the low single-digit percentages.

4. Key Insights and Innovations

Innovation 1: Temporal Stability as an Orthogonal Correctness Signal Independent of Confidence

The dominant assumption across concurrent and prior work on accelerating masked diffusion sampling is that model certainty — typically measured as the maximum predicted probability, entropy, or top-2 probability margin — is the right (and perhaps sufficient) signal for deciding when to unmask a token. Fast-dLLM [47] uses confidence thresholds. Dimple [54] applies confidence-based parallel decoding. Prophet [23] uses the Top-2 confidence gap. SlowFast Sampling [45] alternates stages based on certainty. These methods share a core premise: if the model is confident about a token, it's probably correct, so unmask it early.

KLASS makes a fundamentally different diagnostic move: it argues that confidence is necessary but insufficient, and that a second, orthogonal signal — temporal stability of the predicted distribution — carries independent information about correctness that confidence alone misses. This is not an incremental refinement of the confidence-thresholding idea; it's a conceptual shift from a static view of model state (what does the distribution look like right now?) to a dynamic view (how has the distribution been evolving?).

The distinction matters because the two signals have different failure modes. A model can be confidently wrong — as shown in Figure 1a, where "25 - 20 = 10" gets confidence 0.9241 but is incorrect. A static confidence check has no way to detect this: the distribution is sharp, the model is sure, and the sampler commits to the error. A dynamic stability check, however, can catch it: the incorrect token exhibits KL divergence of 0.4517, indicating that the model's distribution for that position has been shifting substantially as surrounding context resolved, even though it momentarily landed on a high-confidence wrong answer. The correct token ("25 - 20 = 5") shows KL of only 0.0193 — the model converged to it and stayed there.

The empirical evidence for this separation is Figure 1b, which shows that across four reasoning benchmarks and two different model families (LLaDA and Dream), correct tokens consistently exhibit significantly lower KL divergence than incorrect tokens. This is not a subtle statistical difference needing careful measurement — the distributions are visibly separated in the box plots. The finding is robust across math reasoning (GSM8K, MATH) and code generation (HumanEval, MBPP), suggesting it reflects something fundamental about how masked diffusion models converge to solutions, not a domain-specific quirk.

What elevates this from a useful observation to a genuine conceptual contribution is the theoretical grounding in Section 5 (Proposition 5.3). The argument formalizes why an incorrect token cannot remain dynamically stable: if the model currently prefers a wrong answer but a correct answer exists under the true conditional distribution, then as masked context gets resolved, the model's predictions must shift toward the correct conditional — producing non-zero KL divergence along the way. A correct token, by contrast, can remain stable because it's already aligned with the optimal conditional. This transforms the empirical pattern ("KL is lower for correct tokens") into a principled justification ("KL must eventually rise for incorrect tokens in a sufficiently well-trained model"). The theory is not tight enough to provide guarantees — it depends on the model being a good approximator (small δ) and the correct answer having a clear margin (large γ) — but it provides a conceptual framework that explains why the stability signal works, not just that it works.

The intellectual shift here parallels developments in other areas of machine learning where researchers realized that confidence alone is an unreliable guide. In active learning, uncertainty sampling was supplemented with diversity criteria because confident predictions could still be wrong in unexplored regions. In Bayesian deep learning, epistemic uncertainty (what the model doesn't know) was distinguished from aleatoric uncertainty (inherent noise), because high confidence doesn't imply high knowledge. KLASS applies a similar insight to the temporal dimension of diffusion: confidence is a snapshot; stability is a trajectory; and you need both to reliably identify tokens that are ready to be committed.

Comparison to prior work: Prior approaches using confidence alone implicitly assume that the model's maximum probability is well-calibrated to correctness. The KLASS results demonstrate this assumption is violated in practice — Figure 3 shows that adding a KL threshold consistently improves accuracy over confidence alone across all confidence levels tested. This is a diagnostic finding with implications beyond KLASS itself: it suggests that future work on diffusion sampling should incorporate temporal dynamics, and that evaluating samplers solely on their ability to identify high-confidence tokens may be measuring the wrong thing.

Innovation 2: The Conjunctive Gating Architecture as a Minimalist Alternative to Learned Token Ordering

The paper makes a deliberate architectural choice that represents a philosophical stance about how token unmasking decisions should be made. Rather than training a separate planner network to learn optimal token ordering [21, 24, 29], or distilling a fast student model [11, 15], or designing complex numerical integration schemes [33], KLASS implements unmasking decisions as a thresholded conjunction of two hand-designed, zero-parameter signal extractors operating on the base model's existing outputs.

What makes this distinctive is not the specific thresholds chosen — those are hyperparameters anyone could tune — but the demonstration that such a simple mechanism is competitive with or superior to far more complex alternatives. Table 1 shows KLASS matching or exceeding the accuracy of planner-based and distillation-based approaches (by proxy, since those require additional training and aren't directly compared) while requiring no model modifications and adding negligible overhead (Table 6: <1.6% memory, <0.21% latency per step). The concurrent confidence-based parallel decoders [47, 54, 45, 4, 23] share the training-free philosophy but use only a single signal (confidence or entropy or margin); KLASS shows that adding a second, orthogonal signal (stability) yields consistent improvements over single-signal approaches (the "none" row vs. KL-threshold rows in Figure 3).

This is a design philosophy contribution as much as a technical one. The field of diffusion sampling acceleration has been pulled in two directions: toward increasingly sophisticated training-based methods (distillation, planning) that promise large speedups at the cost of complexity, and toward simple heuristics (confidence thresholding) that are easy to implement but leave performance on the table. KLASS stakes out a third position: multi-signal, zero-training decision rules that capture enough of the relevant dynamics to be practically useful without crossing the complexity threshold where retraining or auxiliary models become necessary.

The significance of this position extends beyond the specific KLASS algorithm. It suggests a research program: what other model-internal signals could be extracted at zero training cost and conjoined to improve generation? Candidates might include: attention entropy (is the model attending diffusely or sharply to context?), gradient norms (how sensitive is the prediction to small input perturbations?), or representation-space distances (how far has the token's hidden representation moved?). KLASS provides a template — identify a signal that carries independent correctness information, threshold it, conjoin it with existing signals, measure the overhead — that other researchers can follow.

The fallback mechanism as a safety guarantee is a subtle but important part of this architecture. The top-$u$ fallback (Eq. 8, with $u = 1$) ensures that KLASS can never be slower than standard Top-1 decoding in terms of step count: in the worst case where no tokens ever stabilize, exactly one token is unmasked per step, exactly matching the Top-1 baseline. This means adopting KLASS is a no-regret decision for practitioners: you cannot accidentally make generation slower by setting thresholds poorly; you can only fail to achieve the available speedup. This is a deliberately conservative design choice that prioritizes robustness over maximal acceleration, and it contrasts with learned planners, which could in principle make actively bad token ordering decisions that degrade quality below the baseline. The paper doesn't dwell on this point, but it's an important practical property that reduces the barrier to adoption.

Comparison to prior work: Trainer planners [21, 24] require designing, training, and maintaining a separate model — a significant engineering investment that may not transfer across base model versions. Distillation approaches [11, 15] require expensive training runs and produce artifacts that are locked to a specific step count. KLASS works with any pre-trained masked diffusion model immediately, with a hyperparameter search that requires only ~100 validation examples and negligible compute (Appendix D.1.2). This is a pragmatic advantage that changes the adoption calculus for practitioners. It's not that KLASS is theoretically more elegant than planner-based approaches — it's that KLASS achieves most of the benefit at a fraction of the implementation cost, which is a meaningful type of innovation in applied machine learning.

Innovation 3: Cross-Modal Generality as Evidence for a Fundamental Diffusion Dynamic

The paper demonstrates KLASS on four different modalities — text reasoning (GSM8K, MATH), code generation (HumanEval, MBPP), unconditional text generation (OpenWebText with MDLM), image generation (MMaDA on ImageNet), and molecular generation (QM9) — with consistent improvements across all of them. On reasoning, accuracy improves while steps decrease (Table 1). On unconditional text, MAUVE increases from 0.115 to 0.179 and perplexity decreases across all three oracle models (Table 2). On images, FID improves from 34.48 to 30.48 at 16 steps and from 36.45 to 32.00 at 32 steps (Table 3). On molecules, property reward is maintained or improved while reducing function evaluations by 25–41% (Table 4).

The paper treats this cross-modal validation as a strength (which it is), but doesn't fully articulate why it's intellectually significant. The reason is this: if KLASS worked only on math reasoning, one might attribute its success to domain-specific properties of mathematical text — perhaps mathematical tokens have more deterministic relationships, or the model's uncertainty is better calibrated for arithmetic. The fact that it works on images (where "tokens" are visual patches in a VQ-VAE codebook) and molecules (where "tokens" are SMILES string characters with chemical constraints) suggests that the relationship between temporal stability and correctness is a property of the masked diffusion process itself, not of any particular data domain.

This has implications for how we understand masked diffusion models as a generative paradigm. The standard theoretical framework (Section 3, Eqs. 1–4) characterizes the forward and reverse processes in terms of transition probabilities and ELBO objectives. It says nothing about how individual tokens converge during the reverse process — whether they stabilize gradually, abruptly, monotonically, or with oscillations. The KLASS results provide empirical evidence that there is a systematic, cross-domain dynamic: tokens that eventually end up correct tend to stabilize earlier and more smoothly than tokens that end up incorrect, which tend to shift around as context resolves. This dynamic appears to be a universal property of masked diffusion generation when the model is well-trained, regardless of whether the tokens represent words, code, image patches, or molecular substructures.

This universality claim is strengthened by the fact that the two base models tested for reasoning (LLaDA and Dream) have different architectures, training procedures, and calibration properties — yet both show the same KL-separation pattern in Figure 1b, and both benefit from KLASS in Table 1. The image model (MMaDA) is a multimodal diffusion foundation model, architecturally distinct from the language-focused MDLM used for text generation. The molecular model is a domain-specific DiT trained on SMILES representations. The consistent benefit across this diversity of architectures and data types argues against the explanation that KLASS is exploiting some idiosyncratic property of a particular model's training.

Why this matters for future research: If temporal stability is a universal property of masked diffusion convergence, then (a) it can be studied theoretically — what properties of the training objective and architecture produce this dynamic? — and (b) it can be exploited systematically — future sampler designs can assume that stability signals will be available and informative across domains, without needing to re-validate the basic phenomenon on each new application. The paper provides initial theoretical grounding for (a) in Proposition 5.3, but the result is limited to the conditional δ-approximation framework and doesn't fully characterize when and how strongly the stability-correctness relationship holds. The empirical cross-modal validation fills some of this gap, establishing the phenomenon as sufficiently robust to build on.

Innovation 4: KL Divergence Between Consecutive Predictions as a Practical, Low-Overhead Convergence Diagnostic

The paper's choice of KL divergence as the measure of temporal stability is simultaneously obvious (in retrospect) and not the only option. Other measures of distributional change could have been used: total variation distance, Jensen-Shannon divergence, Hellinger distance, or even simpler metrics like whether the argmax token changed identity. The paper doesn't provide an extensive ablation comparing KL to these alternatives (a notable absence), but the choice of KL is justified by properties that make it particularly well-suited for this application.

First, KL divergence naturally captures the magnitude of distributional change in an information-theoretically principled way. A token whose predicted probability shifts from 0.51 to 0.49 for the same argmax (a small movement) produces a small KL; a token that flips from 0.9 on token A to 0.9 on token B (a large movement) produces a large KL. Total variation distance would collapse these into a binary same/different argmax check, losing the graded signal. Jensen-Shannon divergence is symmetric and bounded, which is nice for theory but doesn't match the causal direction of the reverse process (we care about how the current distribution relates to the previous one, not a symmetric average). The asymmetry of $D_{\text{KL}}(p_t \| p_{t+1})$ — measuring how surprising the current distribution is given the previous one — aligns with the semantics of convergence: we want to know if the model has "changed its mind" since the last step.

Second, KL divergence is computation-friendly in this specific context. Per-token KL requires only the current and previous softmax distributions, both of which are already computed (the former for confidence scoring and unmasking, the latter cached from the previous forward pass). The computation is a simple sum over the vocabulary — no log-determinants, no matrix inversions, no eigendecompositions. The empirical overhead (Table 6) confirms this: <0.00026 seconds per step and <300 MB of memory for models with 126K–152K vocabularies. This computational lightness is essential to the paper's training-free, zero-overhead positioning.

Third, KL divergence has a natural threshold interpretation that other measures lack. A KL of 0.01 nats means the current distribution requires about 0.01 extra nats to encode relative to the previous distribution — a tiny amount of new information. A KL of 0.5 nats means substantially new information has arrived. This interpretability helps practitioners set the $\epsilon_{\text{KL}}$ threshold without needing extensive calibration: values around 0.001–0.02 represent very small distributional shifts, consistent with the intuition that "stable" means "hardly changing at all." The paper's recommended initialization procedure (Appendix D.1.2, step 1) exploits this interpretability by examining the empirical distribution of KL values during decoding to set a reasonable starting threshold.

What's missing from the innovation: The paper does not rigorously compare KL divergence against alternative stability measures (total variation, JS divergence, cosine similarity of logit vectors, change in argmax identity). Table 1 includes a "KL divergence" baseline (using only KL threshold without confidence), but this compares KL to confidence, not KL to other stability measures. The ablation in Figure 3 shows that adding KL to confidence helps, but doesn't test whether some other stability measure would help even more. This is a limitation — the claim that KL divergence is the right stability measure is empirically supported for the specific hyperparameter ranges tested but not systematically proven against alternatives. Future work could strengthen this by showing that KL outperforms other distributional distances as a stability signal for this application.

Despite this limitation, the paper's introduction of KL divergence as a practical, principled convergence diagnostic for discrete diffusion is a meaningful contribution. Prior work on diffusion dynamics — both continuous and discrete — has not emphasized per-token KL as an observable signal that can guide decoding decisions. The paper demonstrates that this signal is (a) cheap to compute, (b) empirically informative about correctness, and (c) theoretically motivated. This opens up KL divergence as a tool for other diffusion analysis tasks beyond sampling acceleration — for instance, diagnosing when and where a model is uncertain during generation, or detecting out-of-distribution inputs where no token ever stabilizes.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The primary evaluation uses four reasoning benchmarks: GSM8K [10] and MATH500 [16] for mathematical reasoning, and HumanEval [9] and MBPP-sanitized [2] for code synthesis. MATH500 refers to the 500-question test split from Lightman et al. (2022). For unconditional text generation, the paper uses 1,000 held-out segments from the OpenWebText corpus [13], comparing 1,000 generated sequences of length 1,024 against them. Image generation uses 10,000 class-conditional generations evaluated against the 50K ImageNet validation set using Inception-v3 features. Molecular generation uses QM9 [31], which contains molecules with up to nine heavy atoms represented as SMILES strings [46], conditioned on drug-likeness (QED) [5] and ring count.

• Base model(s). Reasoning experiments use two instruction-tuned masked diffusion models: LLaDA 8B Instruct [27] and Dream 7B Instruct [52]. LLaDA 8B is described as demonstrating diffusion scaling laws competitive with AR models; Dream 7B is a separate diffusion LM trained by a different group. Text generation uses MDLM [34] pre-trained on OpenWebText. Image generation uses MMaDA [49], a multimodal diffusion foundation model. Molecular generation uses a custom DiT-based diffusion transformer (12 blocks, hidden dimension 768) trained on QM9 following the recipe of Schiff et al. [37] with classifier-free guidance. The diversity of model families, scales, and architectures across these experiments supports the paper's claim of cross-modal and cross-architecture generality.

• Metrics. For reasoning: pass@1 accuracy (percentage of problems where the final generated answer matches the ground truth, using the grading function from Lightman et al. (2022)). For reasoning, the paper also reports average number of sampling steps as an efficiency metric. Wall-clock time per sample (seconds) and speedup factor are reported in Table 8. For text generation: MAUVE (comparing 1,000 generated sequences to 1,000 held-out OpenWebText segments), generative perplexity (exponentiated average negative log-likelihood under three oracle models: LLaMA2 7B [42], LLaMA3 8B [14], and GPT-2 [30]), and Shannon entropy of predicted token distributions. For image generation: Fréchet Inception Distance (FID) [17] and Inception Score (IS) [36]. For molecules: property reward (target QED or ring count value achieved) and average number of function evaluations (NFEs).

• Baselines. The reasoning experiments (Table 1) compare against five baselines organized into two categories. Sequential unmasking (single token per step): (i) Top-1 — selects the highest-confidence token at each step [7], the standard greedy decoding baseline for masked diffusion; (ii) Random — unmasks tokens in a purely random order [1], representing a lower bound on structured unmasking. Parallel unmasking (multiple tokens per step): (iii) Top-2 — decodes the two highest-confidence tokens per step, halving the total number of steps; (iv) Confidence-threshold — unmasks all tokens with predicted probability exceeding 0.9; (v) KL-threshold — unmasks all tokens with KL divergence below 0.001 using history length n = 2. Appendix E, Table 18 additionally reports comparisons against Top-k Margin [21] (unmasks the token with the largest probability gap between top two predictions) and Entropy (unmasks tokens ranked by negative Shannon entropy, prioritizing lower uncertainty). For text generation (Table 2), baselines include the original AR (autoregressive) sampler, SEDD [25], MDLM with SUBS parameterization (standard 512-step), and D3PM [3] (the absorbing variant of MDLM). For image generation (Table 3), the baseline is the standard confidence-based sampler used by MMaDA. For molecular generation (Table 4), the baseline is MDLM with standard 32-step sampling.

• Generation budget / compute accounting. The primary unit of generation cost is the number of sampling steps (also called function evaluations or NFEs). The maximum inference timestep is set to 256 for LLaDA and Dream on reasoning benchmarks, 512 for MDLM on text generation, 16 or 32 for MMaDA on images, and 32 for molecular MDLM. All methods use the same base model and the same maximum timestep budget; the differences lie in how many tokens are unmasked per step (and therefore how many steps are actually executed). For reasoning experiments, generation length is fixed at 256 tokens, with LLaDA using a block size of 64 for semi-autoregressive sampling. Temperature is set to 0 (deterministic) for LLaDA and 0.2 for Dream in the main experiments, with an additional temperature-0 ablation for Dream in Appendix D.5.3 (Table 17). For text generation, nucleus top-p filtering at p = 0.9 is applied. The paper accounts for KL computation overhead separately (Section 6.6, Table 6) and demonstrates it is negligible (<0.21% latency per step, <1.57% memory), so step count reduction translates directly to wall-clock speedup. All sampling experiments are conducted on a single NVIDIA RTX A5000 GPU (reasoning and images), with text generation on an RTX A6000 and molecular experiments on an RTX 3090.

• Cross-validation / statistical protocol. For the deterministic LLaDA experiments (temperature 0), results are single-run and deterministic. For Dream experiments with temperature 0.2, the paper reports mean and standard deviation over three runs with different random seeds (Table 9, Appendix D.1.4). For text generation, each method is repeated with three fixed random seeds and results are reported as mean ± one standard deviation over replicates (Table 10). Hyperparameter selection (confidence threshold τ, KL threshold ε_KL) uses a lightweight three-step search on a small validation set of approximately 100 examples per task (Appendix D.1.2), separate from the test set evaluation. The final thresholds per model and dataset are reported in Table 7.

Main Quantitative Results

Reasoning Tasks (Table 1, Table 8, Tables 11–15)

The headline result for reasoning benchmarks is that KLASS simultaneously improves accuracy and reduces sampling steps compared to the standard Top-1 baseline, and it achieves a superior accuracy-efficiency tradeoff compared to all other baselines tested.

LLaDA 8B Instruct results (Table 1): Across four benchmarks, KLASS improves accuracy while using 50–64% fewer steps than the 256-step Top-1 baseline:

Compared to the parallel-unmasking baselines that use only a single signal, KLASS consistently outperforms both confidence-only and KL-only approaches. For instance, on MATH, confidence-threshold (τ = 0.9) achieves 31.6% at 96.5 steps; KL-threshold (ε_KL = 0.001) achieves 32.6% at 172.2 steps; KLASS achieves 33.8% at 128.6 steps — better accuracy than either and a balanced step count between the two. On GSM8K, confidence-threshold achieves 75.21% at 74.4 steps; KLASS achieves 76.50% at 98.6 steps — better accuracy at a moderate step increase. The Top-2 baseline (always two tokens per step, fixed 128 steps) uniformly degrades accuracy compared to Top-1 (e.g., 29.6% vs. 31.4% on MATH, 33.54% vs. 39.63% on HumanEval), confirming that naive fixed-rate parallel unmasking hurts quality. KLASS's adaptive per-token gating avoids this degradation while still achieving substantial step reduction.

Dream 7B Instruct results (Table 1, Table 9 for error bars): The pattern is similar, with even larger absolute accuracy gains:

Notably, on MATH the accuracy gain is +5.23 percentage points — a substantial improvement — while still reducing steps by over 40%. The GSM8K result is essentially tied with Top-1 in accuracy (within one standard deviation) but with ~39% fewer steps. On HumanEval, KLASS achieves the highest accuracy (59.35%) with only 74.9 steps — more than 70% reduction — meaning generation is more than 3× faster by step count while maintaining or slightly improving accuracy. Table 8 confirms this translates to real wall-clock speedup: Dream on HumanEval achieves 2.78× speedup (11.52s vs. 32.01s for Top-1).

Comparison among baselines: The confidence-only and KL-only baselines in Table 1 demonstrate the value of the conjunctive approach. Confidence-threshold (conf > 0.9) is fast but sacrifices accuracy: on Dream MATH, it achieves 41.80% at 95.1 steps vs. KLASS's 43.20% at 149.7 steps — KLASS trades some speedup for meaningful accuracy gain. KL-threshold (KL < 0.001) achieves better accuracy than confidence alone (41.27% at 162.5 steps on Dream MATH) but at the cost of many more steps — KLASS achieves higher accuracy with fewer steps than KL-only, suggesting the confidence gate prunes tokens that are stable but uncertain, which KL-only would prematurely unmask. The appendix comparison to Top-k Margin and Entropy (Table 18) reinforces this: on LLaDA MATH, Top-k Margin achieves 32.0% at 256 steps (worse accuracy than Top-1, same cost) while KLASS achieves 33.8% at 128.6 steps; Entropy achieves 34.6% at 256 steps (slightly better accuracy than KLASS but at 2× the steps).

Wall-clock speedups (Table 8): Across all benchmarks and models, KLASS achieves speedups of 1.32× to 2.78× over Top-1:

The speedup variation across benchmarks and models reflects differences in how many tokens stabilize early. LLaDA achieves higher speedups than Dream on GSM8K and MATH despite having lower absolute accuracy, suggesting LLaDA's predictions stabilize more quickly (consistent with its lower optimal confidence threshold of 0.6 vs. Dream's 0.9). The largest speedup (2.78× on Dream HumanEval) corresponds to the benchmark with the largest step reduction (70.7%), as expected given the negligible per-step overhead.

Hyperparameter sensitivity (Appendix D.5.1, Tables 11–15): The grid search results demonstrate robustness around the selected thresholds. For LLaDA on HumanEval (Table 11), configurations with τ ∈ {0.8, 0.85, 0.9} and ε_KL ∈ {0.005, 0.01, 0.015} all achieve accuracy within roughly 1.8 percentage points of the optimum (40.85%) while using 83–103 steps. For Dream on MBPP (Table 14), the range of achievable accuracies spans 62.65–65.37% across a broad sweep (τ ∈ {0.8–0.95}, ε_KL ∈ {0.0005–0.005}) — always at or above the Top-1 baseline of 63.81% while using 96–114 steps. There are no catastrophic cliffs where a slightly suboptimal threshold causes performance collapse; the accuracy surface is relatively flat near the optimum. For molecular generation (Table 15), QED reward varies between 0.515 and 0.546 across τ ∈ {0.96–0.999} and ε_KL ∈ {0.0005–0.01}, with the optimum at (τ = 0.98, ε_KL = 0.001) achieving 0.546 at 18.8 steps vs. the MDLM baseline of 0.526 at 32.0 steps.

Temperature effects (Appendix D.5.3, Table 17): For Dream with temperature 0 (fully deterministic), the gains are even more dramatic: MATH accuracy improves from 25.80% (Top-1) to 33.80% (KLASS), a gain of +8.00 percentage points with steps reduced from 256 to 121. GSM8K improves from 41.70% to 47.92% (+6.22), HumanEval from 29.27% to 37.19% (+7.92), and MBPP from 33.46% to 40.86% (+7.40), all with step reductions of 48–79%. This suggests KLASS is particularly beneficial in deterministic settings where the model's convergence dynamics are cleaner and stability signals are more informative.

---

Text Generation (Table 2, Table 10)

The text generation experiments evaluate KLASS on unconditional generation with MDLM at a fixed 512-step schedule, comparing against AR and several discrete diffusion baselines.

Headline result: KLASS achieves the highest MAUVE (0.179) among all diffusion-based methods and substantially lower perplexity under all three oracle models compared to the standard MDLM sampler (SUBS), while maintaining comparable entropy to the data distribution (5.43 vs. 5.44 for the held-out data).

Detailed comparisons (Table 10, which reports mean ± std over three seeds):

KLASS improves MAUVE by ~56% over MDLM (0.179 vs. 0.115), reduces LLaMA2 perplexity by ~13% (26.94 vs. 30.88), LLaMA3 perplexity by ~9% (49.19 vs. 54.15), and GPT-2 perplexity by ~12% (45.50 vs. 51.78). Entropy is nearly identical to the data (5.43 vs. 5.44), indicating the improvement in coherence does not come at the cost of diversity collapse — a common failure mode where a sampler produces fluent but repetitive text that scores well on perplexity but poorly on distribution-level metrics like MAUVE. KLASS maintains high MAUVE while reducing perplexity, suggesting genuine quality improvement rather than mode collapse.

The gap between KLASS and AR (MAUVE 0.179 vs. 0.855, LLaMA2 ppl 26.94 vs. 10.97) remains substantial — diffusion models are not yet matching autoregressive quality on unconditional text — but KLASS narrows this gap significantly relative to the standard MDLM sampler. Notably, SEDD and D3PM perform extremely poorly on this task (MAUVE 0.022–0.037, perplexity 42–110), serving primarily as lower bounds.

The qualitative examples in Appendix F (Figures 4–5) illustrate the type of improvement. The baseline MDLM sample (Figure 4) on the topic of "SolarCity" begins coherently but quickly degenerates into repetition ("SolarCitySolarCity"), nonsensical phrases ("informs of capacity"), and eventually a spam link. The KLASS sample (Figure 5) on "urban sprawl" maintains topical consistency and journalistic style throughout the full 1024-token sequence, with plausible citations and multi-paragraph structure. These examples are cherry-picked (only one pair is shown) but align with the quantitative metrics.

---

Image Generation (Table 3)

The image generation experiments evaluate KLASS on MMaDA, a multimodal diffusion foundation model, under two decoding step budgets (16 and 32 steps) using 10,000 class-conditional ImageNet generations. KLASS is configured with history length n = 1, ε_KL = 0.3, τ = 0.1.

Headline result: At both step budgets, KLASS yields lower FID (better fidelity) and higher IS (better class-consistency) than the standard confidence-based sampler used by MMaDA.

At 16 steps, KLASS reduces FID by 4.00 (34.48 → 30.48, an 11.6% relative improvement) and increases IS by 17.35 (75.72 → 93.07, a 22.9% improvement). At 32 steps, FID improves by 4.45 (36.45 → 32.00, 12.2%) and IS by 16.77 (72.40 → 89.17, 23.2%). The gains are consistent across both step budgets, suggesting KLASS provides benefit independent of the total number of decoding steps.

A curious pattern: the confidence-based sampler actually performs worse at 32 steps than at 16 steps (FID 36.45 vs. 34.48, IS 72.40 vs. 75.72). KLASS at 32 steps also shows slightly higher FID than at 16 steps (32.00 vs. 30.48), though IS remains high. The paper doesn't discuss this non-monotonicity. Possible explanations include: the confidence-based sampler may over-commit to early predictions at the finer-grained 32-step schedule, or the specific 16-step discretization may happen to align better with the model's convergence dynamics. KLASS's stability gating appears to partially mitigate this effect, as the degradation from 16 to 32 steps is smaller for KLASS (FID +1.52) than for confidence (+1.97).

The hyperparameter settings for image generation (ε_KL = 0.3, τ = 0.1) are notably different from those used for language (ε_KL ≈ 0.001–0.02, τ ≈ 0.6–0.9). The much higher KL threshold (0.3 vs. 0.01) suggests that image token distributions are inherently more volatile during diffusion than language token distributions — image patches in a VQ-VAE codebook may shift more dramatically as global structure resolves. The low confidence threshold (0.1) means the confidence gate is nearly inactive, letting the KL gate do most of the filtering. This configuration difference across modalities supports the paper's claim that KLASS adapts to domain-specific convergence dynamics, though it also means practitioners must re-tune thresholds per modality rather than using one universal setting.

---

Molecular Generation (Table 4)

The molecular generation experiments evaluate KLASS on conditional generation from QM9 using classifier-free guidance, targeting two molecular properties: drug-likeness (QED) and ring count. Each method generates 1,024 samples.

Headline result: KLASS maintains or improves target property reward while reducing the average number of function evaluations by 24–41%.

For QED, KLASS improves the reward from 0.526 to 0.546 (+3.8%) while reducing NFEs from 32.0 to 18.8 (−41.3%). For ring count, KLASS improves from 4.123 to 4.258 (+3.3%) while reducing NFEs to 24.4 (−23.8%). These are modest absolute improvements in property score, but the key result is that KLASS achieves them with significantly fewer sampling steps — the speedup doesn't come at the cost of quality, and in fact slightly improves it.

The hyperparameter settings for molecules (Table 15 shows the sweep; the chosen values are ε_KL = 0.001, τ = 0.98) are again different from other modalities, with a very high confidence threshold (0.98) and a very low KL threshold (0.001) — both quite strict. This suggests molecular SMILES generation has particularly clean convergence dynamics where the model either knows a token with very high certainty or doesn't, with little middle ground, and where stable tokens are almost always correct. The strict thresholds filter aggressively, but enough tokens pass to achieve meaningful parallel unmasking.

Table 15 provides the hyperparameter sensitivity for QED. The reward varies from 0.515 to 0.546 across the sweep — a range of about 6% of the mean — while NFEs stay in a relatively narrow band of 18.2–18.8. The optimum (0.546) is at a slightly sub-maximal NFE (18.78 vs. the minimum of 18.22), suggesting a small speed-accuracy tradeoff even in this well-behaved domain.

---

Ablation Studies and Robustness Checks

Effect of confidence and KL score thresholds (Figure 3, Tables 11–15): Applying a KL threshold consistently improves accuracy across all confidence levels compared to relying on confidence alone (the 'none' row in Figure 3). For LLaDA on MATH, the best accuracy (33.8%) occurs at (τ = 0.6, ε_KL = 0.01), while the confidence-only baseline ('none' row at τ = 0.6) achieves only 25.4% — the KL gate alone accounts for a +8.4 percentage point gain. For Dream on MATH, the best accuracy (43.2%) occurs at (τ = 0.9, ε_KL = 0.005), with the confidence-only baseline at τ = 0.9 achieving 41.8% — a more modest +1.4 point gain. The different magnitudes reflect model-specific dynamics: LLaDA is less calibrated (low confidence thresholds work best), so the KL gate provides more additional signal; Dream is better calibrated, so confidence alone is more informative. The heatmaps in Figure 3 show that accuracy varies smoothly across the grid — there are no sharp discontinuities where a small threshold change causes accuracy collapse. The full sensitivity tables in Appendix D.5.1 (Tables 11–15) confirm this smoothness across additional benchmarks and models.

Effect of unmasking multiple tokens vs. single token (Table 5): Using KLASS's stability criteria but unmasking only one token per step (selecting either the highest-confidence stable token or the lowest-KL stable token) reduces both accuracy and efficiency compared to parallel multi-token unmasking. On MATH, single-token (conf) achieves 31.2% at 256 steps (no speedup), single-token (KL) achieves 29.0% at 256 steps, while parallel KLASS achieves 33.8% at 128.6 steps — higher accuracy with roughly half the steps. On GSM8K, parallel KLASS achieves 76.50% at 98.6 steps vs. single (conf) at 72.86% at 256 steps and single (KL) at 73.46% at 256 steps. This ablation demonstrates that the benefit of KLASS is not just about identifying stable tokens — it's specifically about unmasking multiple stable tokens simultaneously. The model benefits from resolving several converged positions at once, likely because it creates richer context for the remaining masked tokens in the next step. Single-token unmasking from the stable set is actually worse than standard Top-1 (31.2% and 29.0% vs. 31.4% for Top-1 on MATH), perhaps because waiting for stability delays all unmasking without the compensating benefit of parallel context revelation.

Effect of history length n (Table 16): The history length ablation on MATH explores n ∈ {1, 2, 3} across multiple threshold configurations for both LLaDA and Dream. For LLaDA at the optimal configuration (τ = 0.6, ε_KL = 0.010): n = 2 achieves 33.8% at 128.6 steps; n = 1 achieves only 32.2% at 77.1 steps (faster but less accurate — premature unmasking); n = 3 achieves 32.2% at 153.6 steps (comparable accuracy to n = 1 but slower — overly conservative). For Dream at the optimal (τ = 0.9, ε_KL = 0.005): n = 2 achieves 43.2% at 149.7 steps; n = 1 achieves 41.0% at 126.4 steps; n = 3 achieves 40.2% at 165.4 steps. The pattern is consistent: n = 1 sacrifices accuracy for speed, n = 3 sacrifices speed without accuracy gain, n = 2 balances both. At stricter confidence thresholds (τ = 0.9 for LLaDA), history length has less impact — accuracy varies only from 30.6 to 31.4 across n values — because the confidence gate is already filtering aggressively, leaving fewer tokens for the history check to discriminate among.

Effect of temperature (Table 17): For Dream, using temperature 0 (deterministic argmax) instead of 0.2 produces substantially larger absolute accuracy gains from KLASS across all benchmarks: MATH +8.00 (vs. +5.10 at temp 0.2), GSM8K +6.22 (vs. +0.69), HumanEval +7.92 (vs. +1.23), MBPP +7.40 (vs. +0.78). Step reductions are also larger: 106–203 steps saved at temp 0 vs. 74–182 at temp 0.2. This suggests that stochastic sampling introduces noise that partially obscures the stability signal — with non-zero temperature, even converged tokens show some distributional fluctuation due to sampling stochasticity, making it harder to distinguish genuinely unstable tokens from sampling noise. In deterministic mode, the KL signal is cleaner, and KLASS can unmask more aggressively without quality loss.

### Critical Assessment

The experimental evaluation provides broad evidence that KLASS works across multiple modalities, models, and tasks, but there are important qualifications and gaps that affect how confidently the paper's central claims can be accepted.

Claim: KLASS achieves up to 2.78× wall-clock speedups while improving accuracy over standard greedy decoding. This claim is supported with qualifications by the data in Tables 1 and 8. The 2.78× figure is specific to Dream on HumanEval (11.52s vs. 32.01s) — a real measurement, but also the single best-case scenario across all benchmarks and models. More typical speedups are in the 1.3–2.5× range, and the accuracy improvements are modest (+0.8 to +5.2 points on most tasks). On GSM8K with Dream, accuracy is essentially flat (−0.12 points, within noise), so the benefit there is purely speed. The claim of "improving accuracy" holds for 7 of 8 benchmark-model pairs but not universally.

The more fundamental qualification is that Top-1 is a weak baseline for speed comparisons — it artificially forces one token per step regardless of whether the model is ready to unmask more. A fairer baseline would be Top-2 or a simple confidence-threshold method at the same accuracy level, which Table 1 shows can already achieve most of the speedup (e.g., confidence-threshold on LLaDA MATH achieves 31.6% at 96.5 steps vs. KLASS's 33.8% at 128.6 steps — KLASS trades some speed for accuracy). The paper's comparison to single-signal baselines (confidence-only, KL-only) makes the case for the value of combining signals, but the absolute speedup over the most naive baseline (Top-1) somewhat overstates the practical improvement over what a practitioner might already be using.

Claim: KL divergence is a strong indicator of solution correctness, and combining it with confidence yields better results than confidence alone. This claim is well-supported by Figure 1b (the KL separation between correct and incorrect tokens), Figure 3 (KL threshold consistently improving over confidence-only), and Table 1 (KLASS outperforming confidence-threshold and KL-threshold baselines). The evidence is consistent across two model families and four benchmarks, and the theoretical grounding in Proposition 5.3 provides a conceptual framework for why this should hold. The robustness is not perfect — the magnitude of improvement from adding KL varies substantially (large for LLaDA, modest for Dream) — but the direction is consistent.

The missing ablation is a comparison against other stability measures. Would total variation distance, Jensen-Shannon divergence, or the change in argmax identity perform equally well or better? The paper doesn't test this, so we cannot conclude that KL divergence is uniquely suited to this role, only that it works. Given how central the KL score is to the method, this is a notable gap. The paper's theoretical argument for KL specifically (asymmetric, information-theoretic) is persuasive but not empirically validated against alternatives.

Claim: KLASS generalizes across modalities (text, images, molecules). This claim is supported in breadth but not in depth. The paper reports results on four modalities, which is unusually broad for a sampling method paper. However, the experimental depth varies dramatically. Reasoning benchmarks receive detailed hyperparameter sweeps, multiple baselines, ablation studies, and wall-clock measurements (Tables 1, 5, 8, 11–14, 16–17; Figures 3, 1b). Text generation receives a single set of metrics with error bars (Table 10) but no step-count comparison to baselines at fixed quality levels, no hyperparameter sensitivity, and no ablation of the KLASS components specifically for this task. Image generation reports FID and IS at two step budgets (Table 3) but provides no hyperparameter sweep, no comparison to other parallel-decoding methods, and no qualitative examples. Molecular generation reports only mean reward and NFEs (Table 4) with a sensitivity table (Table 15) but no statistical error bars and a small sample (1,024 molecules). The cross-modal claim is directionally supported but the evidence is thinner than for the reasoning experiments.

A specific concern: the hyperparameters used for different modalities are wildly different (ε_KL: 0.001–0.02 for language, 0.3 for images, 0.001 for molecules; τ: 0.6–0.9 for language, 0.1 for images, 0.98 for molecules). This is partially expected — different data types have different convergence dynamics — but it also means that deploying KLASS on a new modality requires non-trivial hyperparameter search. The lightweight guideline in Appendix D.1.2 requires ~100 validation examples for language; it's unclear whether this transfers to images or molecules, or how much tuning was required to find the reported configurations.

Claim: The overhead of KL computation is negligible. This claim is well-supported by Table 6, which reports <1.57% memory overhead and <0.21% latency overhead per step for the two language models tested. The theoretical $O(|I_m| \cdot |V|)$ scaling is also correctly analyzed. However, the measurement is done only for LLaDA and Dream (vocabularies 126K and 152K) with generation length 256. For image models with different codebook sizes and sequence lengths, or for models with much larger vocabularies, the overhead could differ. The paper doesn't report overhead measurements for MMaDA or the molecular DiT.

Missing experiments that would strengthen the evaluation:

• Comparison to concurrent parallel-decoding methods: Fast-dLLM [47], Dimple [54], SlowFast [45], EB-Sampler [4], and Prophet [23] are mentioned in the related work but never directly compared in experiments. This is understandable given the concurrent timing, but it means we cannot quantify KLASS's advantage over the broader class of certainty-based heuristics — we can only compare to the generic confidence-threshold and KL-threshold baselines that the paper implements itself.

• Comparison to distillation-based acceleration: Distillation methods [11, 15] can reduce steps by an order of magnitude (e.g., to 8–32 steps) whereas KLASS reduces steps by 40–70%. A direct comparison would clarify whether the training-free advantage compensates for the more modest speedup, but this comparison is absent.

• Evaluation on harder benchmarks: All reasoning experiments use MATH500, GSM8K, HumanEval, and MBPP, which are standard but not the most challenging contemporary benchmarks. The paper acknowledges (Appendix G.1) that larger diffusion models are not yet available, limiting evaluation on more difficult tasks like agentic reasoning or very long-form generation.

• Statistical rigor for image and molecule experiments: The image FID and IS numbers are reported without error bars, making it impossible to assess whether the 4-point FID improvement is statistically significant or within noise. The molecular results lack any indication of variance. For reasoning, error bars are provided only for Dream (three runs) but not for LLaDA (deterministic, so variance is zero by design).

• Scaling analysis: All experiments use models of a single scale (8B for LLaDA, 7B for Dream). We don't know whether KLASS's benefits increase, decrease, or stay constant as model scale grows. The dynamics of convergence — how quickly token distributions stabilize — might change with model capacity, and a scaling study would be informative for predicting KLASS's utility on future larger diffusion models.

• The difficulty-dependent analysis that characterizes the reference paper's Summary (e.g., "easy questions benefit from this, hard questions don't") is entirely absent from the KLASS paper. There is no breakdown of where the accuracy gains come from — are they concentrated on easy problems where the model would likely get the right answer anyway? On medium problems where stability signals are most informative? On hard problems where no method helps? This type of analysis would substantially strengthen the claim that KLASS is improving genuine reasoning quality rather than just accelerating convergence on already-solved problems. The broader implication: we cannot tell whether KLASS is overcoming specific failure modes of standard sampling or merely making the model faster at producing the same distribution of correct and incorrect answers.

Bottom line: The experiments convincingly demonstrate that KLASS provides a meaningful speed-accuracy improvement over standard single-token and confidence-threshold baselines for masked diffusion language models on reasoning benchmarks, with preliminary evidence of cross-modal generality. The evaluation is thorough within its scope but leaves open questions about scaling behavior, comparison to the broader landscape of acceleration methods, statistical significance in non-language domains, and whether the gains represent improved reasoning or merely accelerated convergence. The paper's central claim — that combining confidence and temporal stability signals yields better parallel decoding than confidence alone — is well-supported by the data presented.

6. Limitations and Trade-offs

Limited Scalability to Larger Models and Harder Tasks

The assumption or constraint. All reasoning experiments use models at the 7–8B parameter scale (LLaDA 8B Instruct and Dream 7B Instruct). The paper does not evaluate KLASS on larger diffusion models, agentic reasoning tasks, or benchmarks more challenging than MATH and HumanEval. The authors acknowledge this explicitly in Appendix G.1:

"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."

The consequence. We have no evidence about whether the stability-correctness relationship that KLASS exploits (Figure 1b) persists at larger model scales. It is possible that larger models converge differently—their token distributions might stabilize more quickly (making KLASS's stability filter less discriminative, since many tokens pass simultaneously), or more slowly (reducing the speedup benefit), or with different calibration properties that change the optimal confidence thresholds. The absence of scaling analysis also means we cannot predict whether the 1.3–2.8× speedups observed at 7–8B scale would grow, shrink, or stay constant as model capacity increases. For practitioners evaluating whether to adopt KLASS in production pipelines that will eventually use larger models, this uncertainty is material.

What evidence exists in the paper. The paper provides no evidence at all on model scaling—no experiments with LLaDA or Dream variants at different sizes, no analysis of how KL divergence statistics or optimal thresholds change with model capacity. The cross-model comparison (LLaDA 8B vs. Dream 7B) shows qualitatively similar KL-separation patterns (Figure 1b) but very different optimal hyperparameters (τ = 0.6 vs. 0.9 for MATH; Table 7), hinting that model-specific properties matter strongly, which does not inspire confidence that findings at one scale transfer cleanly to another. The limitation is acknowledged but not addressed experimentally.

Mitigation status. The paper identifies this as future work in Section 7: "one could extend this approach... evaluate the proposed sampler with larger models as they become available." The mitigation is entirely deferred—no experiments are run, no scaling trends are extrapolated, and no guidance is provided for practitioners anticipating larger model deployments.

---

Difficulty Estimation Cost Is Not Accounted For (Hyperparameter Search)

The assumption or constraint. KLASS introduces four hyperparameters (τ, ε_KL, n, u) that must be tuned per model and task. The paper provides a "lightweight guideline" (Appendix D.1.2) requiring approximately 100 validation examples per benchmark, and a three-step procedure: (1) estimate an initial KL threshold by inspecting KL value distributions during decoding, (2) sweep confidence thresholds from 0.9 downward, and (3) refine the KL threshold through finer-grained search. The paper states this is "efficient, requiring only a small number of validation samples and negligible computation relative to training."

The consequence. This framing understates the practical cost. Step 1 alone requires running the full diffusion decoding process on some number of validation examples just to observe KL statistics—each decoding run involves up to 256 forward passes through an 8B-parameter model, which is computationally non-trivial. The grid search that produced Figure 3 sweeps 5 confidence values and 5–6 KL values (25–30 configurations), and the sensitivity tables in Appendix D.5.1 demonstrate that while performance is robust near the optimum, finding that neighborhood may require evaluating a substantial portion of the grid. For a practitioner deploying KLASS on a new model or task, the total hyperparameter search cost—measured in GPU-hours spent on validation decoding—could equal or exceed the compute saved by KLASS's speedup on the actual test set, depending on test-set size.

This is analogous to the difficulty estimation problem in the reference paper's compute-optimal test-time scaling framework: the cost of estimating per-prompt difficulty (2048 samples) was excluded from the headline efficiency numbers. Here, the cost of hyperparameter tuning is excluded from the speedup calculations. The paper reports wall-clock speedups of 1.32–2.78× after hyperparameters are selected, but does not amortize the tuning cost.

What evidence exists in the paper. The paper does not report the total GPU-hours required for hyperparameter tuning, nor does it provide a cost-benefit analysis comparing tuning cost to inference savings at various deployment scales. The sensitivity tables (Tables 11–15) show that performance surfaces are relatively flat, which suggests that coarse tuning may suffice—but this is an observation, not a quantified claim about minimum tuning budget. Table 7 lists the final hyperparameters per task without indicating how many validation decoding runs were needed to find them.

Mitigation status. The paper partially mitigates this by demonstrating that performance is robust to small hyperparameter perturbations (Tables 11–15), which reduces the precision required in tuning. The three-step procedure is a concrete, actionable guideline. However, the total compute cost of executing that procedure is never measured or reported, and there is no guidance for practitioners on how to balance tuning budget against deployment scale. The paper does not address the question: "How many test-set inferences do I need to run before the speedup pays back the tuning cost?"

---

Single Benchmark Family for Core Reasoning Claims

The assumption or constraint. The primary quantitative claims about KLASS—accuracy improvements of +1.2 to +5.2 percentage points with 40–70% step reduction—are based entirely on four reasoning benchmarks: GSM8K, MATH500, HumanEval, and MBPP. These are all structured-reasoning tasks with well-defined correct answers and grading functions. The paper evaluates text generation (OpenWebText), image generation (ImageNet via MMaDA), and molecular generation (QM9) in secondary experiments, but these evaluations are less thorough: text generation reports metrics at a single step budget (512) with no hyperparameter sweep; image generation reports only FID and IS at two budgets; molecular generation reports only mean reward and NFEs on 1,024 samples.

The consequence. We cannot assess whether KLASS's gains on reasoning benchmarks generalize to other important language tasks—summarization, translation, dialogue, instruction-following, long-form question answering—where correctness is less crisply defined and model convergence dynamics may differ. The reasoning benchmarks share structural properties (multi-step logical deduction, arithmetic computation, algorithmic problem-solving) that may make the stability signal particularly informative: in math problems, intermediate calculation steps have deterministic correct answers that the model either converges to or doesn't. In more open-ended generation tasks, tokens may have multiple acceptable values (stylistic choices, paraphrases), and the KL score may be less informative about output quality—a stable distribution over equally-valid completions is not the same as a correct completion.

The image and molecule experiments provide preliminary evidence of cross-modal generality, but the depth is insufficient to claim robustness. For images, we don't know whether the FID improvement is statistically significant (no error bars), whether KLASS helps across a range of hyperparameters (only one configuration is reported), or whether the benefit transfers to other image generation models or tasks. For molecules, the sample size (1,024) is small and the property improvements are modest (+3.8% QED, +3.3% ring count).

What evidence exists in the paper. The reasoning experiments are thoroughly evaluated (Tables 1, 5, 8, 11–14; Figure 3) and provide the strongest evidence for KLASS. The text generation experiment (Table 10) shows improvements in MAUVE and perplexity but evaluates only one model (MDLM) at one step budget with one set of hyperparameters—there is no ablation of KLASS components for this task, no comparison to confidence-only or KL-only baselines, and no step-count reduction analysis. The cross-modal experiments are best characterized as existence proofs (KLASS can work outside reasoning) rather than rigorous evaluations.

Mitigation status. The paper explicitly positions cross-modal generality as a contribution ("We further validate KLASS across diverse domains, including text, image, and molecular generation, showing its effectiveness as a broadly applicable sampler"; Section 1), but does not claim the same level of rigor for these domains. The limitation is not discussed in the limitations section (Appendix G), which focuses on model scale and hyperparameter cost. The absence of deeper cross-task evaluation within the language domain (beyond reasoning) is not acknowledged.

---

The Theory Provides a Necessary Condition, Not a Sufficient One — No Guarantees Against Premature Unmasking

The assumption or constraint. Proposition 5.3 proves a lower bound on average per-step KL divergence for tokens that are predicted incorrectly at the current context but would be predicted correctly at the optimal context. The result shows that an incorrect token cannot remain stably predicted—it must exhibit KL divergence of at least 2Δ²/M² on average along the path to full context. The paper uses this to justify the KL threshold: "Accordingly, KLASS delays unmasking until tokens exhibit dynamic stability thereby improving generation quality."

The consequence. The theorem establishes that dynamic instability is a necessary condition for incorrectness, which implies that stability is a necessary condition for correctness (the contrapositive). However, it does not establish that stability is a sufficient condition for correctness. A token can be stably wrong if the model has settled on an incorrect prediction and the remaining masked context does not provide enough signal to dislodge it. Section 5 never claims sufficiency, but the algorithm behaves as if stability-plus-confidence implies readiness—it unmasks any token that satisfies both gates. The theory provides no guarantee about the error rate among unmasked tokens, which depends on the model's approximation quality (δ), the margins of correct answers (γ), and the specific context-revelation path taken during generation.

This matters practically because KLASS can still prematurely unmask incorrect tokens that happen to be both confident and temporarily stable. The case study in Figure 1a shows one where it works—the incorrect token has high KL and gets filtered—but no analysis quantifies how often KLASS fails to filter incorrect tokens. The overall accuracy improvements in Table 1 show that, on average, KLASS does better than baselines at avoiding such errors, but they provide no per-token error analysis or breakdown of where the remaining errors come from.

What evidence exists in the paper. The accuracy improvements in Table 1 are the primary evidence that KLASS's stability gate reduces premature unmasking relative to confidence alone. However, there is no direct measurement of false-positive rate among stable tokens—what fraction of tokens that pass both gates are actually incorrect? This could be measured by comparing KLASS's per-token accuracy to the baseline's per-token accuracy, but the paper reports only sequence-level pass@1. The theory (Section 5) explicitly acknowledges the dependence on model quality: the bound becomes vacuous when δ > (β + γ)/2, meaning the result provides no guarantee for poorly-trained or poorly-calibrated models. But the empirical experiments don't test how KLASS degrades as model quality varies (e.g., by evaluating intermediate checkpoints or models of different capacities).

Mitigation status. The paper does not address the sufficiency gap. The theoretical section is presented as a "rationale" rather than a guarantee, which is appropriate given its assumptions, but the gap between "instability implies incorrectness" and "stability implies correctness" is not discussed. This is a fundamental limitation of the approach: KLASS cannot detect stably-wrong tokens, and the paper provides no analysis of how often these occur in practice or how they affect downstream task performance.

---

Sequential vs. Parallel Latency Tradeoff Is Not Analyzed

The assumption or constraint. The paper measures efficiency in terms of number of sampling steps and wall-clock time per sample (Table 8), where wall-clock time is total end-to-end generation time on a single GPU. This metric conflates two dimensions of efficiency: total FLOPs (which step count approximates) and latency (which depends on the serial dependency structure). KLASS reduces the number of sequential forward passes (steps), which reduces latency. But it does not fundamentally change the fact that each step depends on the previous step's output—the reverse diffusion process remains inherently sequential.

The consequence. For deployment scenarios where throughput (samples per second across a batch) matters more than latency (time per individual sample), KLASS's step reduction translates directly to throughput improvement—fewer forward passes per sample means more samples can be processed per GPU-hour. This is the scenario the paper's measurements capture. However, for latency-critical applications (interactive assistants, real-time code completion, on-device generation), the absolute latency floor is determined by the per-step forward pass time multiplied by the number of sequential steps. KLASS reduces the number of steps from ~256 to ~100–150, which is a meaningful latency improvement (1.3–2.8×). But compared to autoregressive models with KV-caching—where generating token N costs only O(1) additional computation on top of token N−1—the per-step cost structure is fundamentally different. An AR model generating 256 tokens requires 256 sequential forward passes, but each one processes only the new token with cached attention states. A diffusion model at each of its ~100–150 steps processes the entire 256-token sequence through full attention. The paper doesn't compare diffusion+KLASS latency to AR latency for equivalent-quality outputs, which is the comparison a practitioner deciding between paradigms would need.

The paper does not discuss batching or throughput at all. KLASS's per-step overhead (KL computation) scales with the number of masked tokens, which is batch-size-invariant—but the paper doesn't measure how KLASS's relative speedup changes under batched inference, where GPU utilization and memory bandwidth constraints may differ from the single-sample regime.

What evidence exists in the paper. Table 8 provides per-sample wall-clock times (e.g., 11.52s for Dream on HumanEval with KLASS), which is a latency measurement. But there is no comparison to AR latency on the same benchmarks, no throughput measurement (samples/second at various batch sizes), and no analysis of how the speedup ratio changes with batch size. The text generation experiment (Table 10) uses a fixed 512-step schedule for all diffusion methods—KLASS is capped such that it cannot reduce steps below the fixed schedule, so no speedup is measured; this experiment evaluates quality only, not efficiency.

Mitigation status. The paper does not address the latency-vs-throughput distinction, does not compare to AR wall-clock time, and does not evaluate KLASS under batched inference. This is not acknowledged as a limitation. For practitioners evaluating whether to adopt diffusion models with KLASS instead of AR models for latency-sensitive applications, this missing analysis is a significant gap.

7. Implications and Future Directions

How This Work Changes the Landscape

KLASS introduces a diagnostic shift rather than a paradigm shift in how the field approaches inference-time acceleration for discrete diffusion models. Before this work, the dominant research thrust for speeding up masked diffusion decoding fell into roughly two camps: (1) training-based acceleration — distill the base model into a student that requires fewer steps [11, 15], or train an auxiliary planner to optimize token ordering [21, 24, 29] — and (2) heuristic single-signal parallel decoding — unmask all tokens whose predicted probability exceeds a threshold [47, 54], whose entropy is low [4], or whose top-2 probability gap is large [23]. KLASS argues, and provides empirical evidence (Figure 3, Table 1), that the second category leaves performance on the table because confidence alone is an insufficient signal: a model can be confidently wrong, and a static confidence check has no mechanism to detect this.

The conceptual shift is from a snapshot-based view of model state to a trajectory-based view. Rather than asking "is the model certain about this token right now?" (confidence), KLASS also asks "has the model's prediction for this token stopped changing?" (KL divergence across consecutive timesteps). The conjunction of these two signals — certainty plus temporal stability — defines a stricter readiness criterion that empirically yields better accuracy-efficiency tradeoffs (Table 1, Figure 3). This is not a radical reconceptualization of diffusion models, but it reframes the decoding acceleration problem from "how do we identify high-certainty tokens?" to "how do we identify tokens whose distributions have converged?" — a question that naturally invites investigation of other convergence signals beyond KL divergence.

The work also contributes a principled diagnostic tool to the growing toolkit for understanding diffusion model internals. The finding that per-token KL divergence between consecutive steps cleanly separates correct from incorrect tokens (Figure 1b) — across two model families and four benchmarks — is not an acceleration technique; it is a discovery about how masked diffusion models converge during the reverse process. Correct predictions tend to stabilize early and stay stable; incorrect predictions tend to shift around as additional context is resolved. This dynamic was not previously characterized, and it opens up KL divergence as a general-purpose probe for analyzing diffusion model behavior beyond sampling acceleration: diagnosing where and when a model is uncertain during generation, detecting out-of-distribution inputs where no tokens stabilize, or understanding how architectural choices affect convergence speed.

In terms of reconciling prior contradictions, KLASS helps explain a tension in the concurrent literature. Several works published around the same time (mid-2025) explored confidence-based parallel decoding [47, 54, 45, 4, 23], and collectively they demonstrated that model certainty can guide efficient unmasking. But they also showed inconsistent results: confidence thresholds that work well for one model or task sometimes fail for others, and the relationship between confidence and correctness varies substantially across models (as Table 7 shows, LLaDA's optimal τ is 0.6 while Dream's is 0.9 — a large gap). KLASS resolves part of this tension by showing that confidence alone is underdetermined: a high-confidence token might be correct (converged) or incorrect (model is sure but hasn't finished reasoning), and you cannot distinguish these cases without a temporal signal. Adding the KL stability check filters out confidently-wrong tokens, making the parallel-decoding approach more robust across models and potentially explaining why confidence-only approaches see model-dependent performance.

Research directions that become more attractive:

• Temporal dynamics analysis for diffusion models. The KL-separation phenomenon (Figure 1b) suggests a research program around characterizing how individual tokens converge during the reverse process. Do different architectural choices (attention patterns, positional encodings, noise schedules) produce different convergence trajectories? Can we predict which tokens will be hard to converge from early-timestep signals? This direction was not obvious before KLASS because the field lacked a simple, cheap-to-compute convergence diagnostic.

• Multi-signal zero-training decoding rules. KLASS demonstrates that combining two hand-designed signals (confidence + KL) with a conjunctive threshold rule outperforms either signal alone. This invites exploration of additional signals extractable at zero training cost: attention entropy (is the model attending to a narrow or broad set of context tokens?), gradient norms with respect to token embeddings (how sensitive is the prediction to small perturbations?), or representation-space distances (how far has the token's hidden state moved between steps?). The template — identify a signal that carries independent correctness information, threshold it, conjoin it with existing signals, measure overhead — is directly replicable.

• Verifier-guided sampling for discrete diffusion. The paper's finding that stability signals correctness (Proposition 5.3, Figure 1b) parallels the role of process reward models (PRMs) in the reference paper's analysis of autoregressive LLM inference. In that work, a PRM scores intermediate reasoning steps and guides search or revision strategies. In discrete diffusion, the KL score functions as a model-internal verifier — it estimates, without any training, whether the model has converged on the correct answer for each token position. This analogy suggests that more sophisticated verifier architectures (trained on diffusion trajectories) could further improve parallel decoding decisions, analogous to how trained PRMs improved over simple confidence scoring in the reference paper.

Research directions that become less attractive:

• Pure confidence-thresholding as a research contribution. Given KLASS's demonstration that adding a stability signal consistently improves over confidence alone (Table 1, Figure 3), future work that proposes a new parallel-decoding heuristic based solely on model certainty — without addressing the confidently-wrong failure mode — faces a higher evidentiary bar. It must either show that its novel certainty measure substantially outperforms KL divergence as a correctness signal, or it must explicitly situate itself as a simpler alternative that deliberately trades accuracy for ease of implementation.

• Token-ordering planners that ignore temporal dynamics. The planner-based approaches [21, 24, 29] are not rendered obsolete — they address the more general problem of optimal token ordering, while KLASS addresses only safe ordering — but KLASS's results suggest that any planner that fails to incorporate temporal stability information is leaving a free signal on the table. A planner that optimizes token ordering while also using KL-like stability checks as a constraint (e.g., "don't plan to unmask this token early if its distribution hasn't converged") could outperform either approach alone.

---

Follow-Up Research This Work Enables

1. Systematic comparison of distributional distance measures as stability signals for discrete diffusion decoding.

The paper chooses KL divergence as its stability measure and provides a theoretical rationale (Proposition 5.3) and empirical evidence (Figure 1b, Table 1) that it works. But it never compares KL divergence against alternative measures of distributional change: total variation distance $D_{\text{TV}}(P, Q) = \frac{1}{2}\sum_v |P(v) - Q(v)|$, Jensen-Shannon divergence $D_{\text{JS}}(P \| Q) = \frac{1}{2}D_{\text{KL}}(P \| M) + \frac{1}{2}D_{\text{KL}}(Q \| M)$ with $M = (P+Q)/2$, Hellinger distance, cosine similarity of logit vectors, or the simple binary signal of whether the argmax token identity changed. A strong follow-up would replicate Figure 3 on MATH with LLaDA and Dream, replacing the KL score with each of these alternatives, sweeping their thresholds equivalently, and measuring (a) whether they produce similar accuracy-speedup tradeoffs, (b) whether they exhibit the same separation between correct and incorrect tokens (analogous to Figure 1b), and (c) whether their computational overhead differs meaningfully. This would strengthen or qualify the paper's implicit claim that KL divergence is particularly well-suited by establishing whether any reasonable stability measure works roughly as well, or whether KL's asymmetry and information-theoretic properties give it a genuine edge. The experiment is straightforward — it reuses the paper's existing evaluation infrastructure and requires only a new metric implementation at step 2b of Algorithm 1.

2. Difficulty-conditioned analysis of where KLASS accuracy gains originate.

The most conspicuous gap in the paper's evaluation is the absence of any breakdown by problem difficulty — in contrast to the reference paper, which carefully separates easy, medium, and hard problems and shows test-time compute strategies work primarily on easy-to-medium difficulty tiers. For KLASS, we do not know whether the +2.4 to +5.2 percentage point accuracy improvements on reasoning benchmarks (Table 1) come from (a) correcting errors that standard sampling makes on problems the model would otherwise get right (reducing "sloppy" mistakes), (b) solving genuinely harder problems that standard sampling cannot handle (extending capability), or (c) a mixture. A follow-up study would bin MATH500 and GSM8K problems by the base model's pass@1 rate under Top-1 decoding (analogous to the oracle difficulty bins in the reference paper), then measure KLASS's accuracy improvement separately per bin. If gains concentrate in easy-medium bins where the model already has non-trivial pass@1, then KLASS is primarily a reliability improvement — it helps the model not screw up. If gains extend to the hardest bin, then KLASS is actually extending capability by enabling more effective exploration during the reverse process. This distinction matters for practitioners: a reliability improvement helps on routine tasks; a capability extension helps on novel problems. The experiment requires no new model training or infrastructure — only a post-hoc stratification of the existing evaluation results, ideally with validation-set selection of KLASS hyperparameters per difficulty bin (to check whether optimal thresholds vary with difficulty, as the reference paper found for search strategies).

3. Combining KLASS with distillation or high-order numerical solvers to push toward single-digit step counts.

The paper explicitly positions KLASS as complementary to training-based acceleration methods (Section 1: "requires no additional training"), but never demonstrates this complementarity. A natural next step is to apply KLASS on top of a distilled model. For instance, take the 512-step MDLM text generation model, distill it to a 32-step student following the recipe of Deschenaux and Gulcehre [11] or Hayakawa et al. [15], then apply KLASS during inference from the 32-step student. The research question is: does the stability signal remain informative when the model has been trained to produce good outputs in many fewer steps? It is possible that distilling to very low step counts forces the model to make larger per-step jumps, reducing the informativeness of per-step KL (the distribution might change more abruptly, making KL less discriminative between correct-convergence and random-fluctuation). Alternatively, it is possible that KLASS's stability gating helps the distilled model avoid over-committing to early errors, partially recovering the quality loss that typically accompanies aggressive distillation. The experiment would measure accuracy at matched step counts with and without KLASS on the distilled student, compared to the undistilled teacher with KLASS. A similar integration could be tested with the high-order solvers of Ren et al. [33] — use a θ-RK-2 solver to improve per-step accuracy, then apply KLASS to determine which tokens to unmask at each of the (fewer) steps.

4. Detecting and mitigating verifier over-optimization in masked diffusion through KL trajectory analysis.

The reference paper identifies verifier over-optimization as the primary bottleneck for test-time compute scaling in autoregressive models: aggressive search against a learned verifier eventually finds solutions that score highly but are incorrect. Masked diffusion models with KLASS face an analogous risk: the stability gate could, in principle, be "over-optimized" if the model learns to produce artificially stable distributions that are wrong — for instance, by confidently predicting a common-but-incorrect token early and maintaining that prediction rigidly. A diagnostic follow-up would measure, on a benchmark with known difficult edge cases (e.g., MATH problems where LLaDA's pass@1 is near zero), whether KLASS's unmasked tokens exhibit lower error rates than confidence-only unmasking at matched step counts, or whether KLASS sometimes commits more confidently to wrong answers that happen to be stable. If over-optimization against the stability signal occurs, the per-token KL values for incorrect KLASS-unmasked tokens would be low (the gate let them through) despite being wrong — measuring this false-negative rate would quantify the gap between the theoretical guarantee (instability implies incorrectness, but stability does not imply correctness) and practical performance. Mitigations could include dynamic threshold adjustment (tightening ε_KL when the model's predictions appear to be coherent but potentially wrong, as detected by external consistency checks) or ensembling stability signals across multiple parallel generation trajectories.

5. KL divergence as a lightweight difficulty estimator for routing and compute allocation in diffusion-based systems.

The KLASS finding that per-token KL divergence separates correct from incorrect tokens (Figure 1b) suggests a use case beyond decoding acceleration: difficulty estimation. A prompt that produces consistently high per-token KL divergence throughout the reverse process — even after many steps — indicates that the model is struggling to converge, analogous to the "hard problem" regime in the reference paper's analysis. A prompt that produces rapid convergence (low KL early) is easy. This signal could be used to implement adaptive compute allocation for diffusion models: start decoding with a small step budget; if average per-token KL remains high after a fixed proportion of steps, allocate additional steps (restart or extend the reverse process); if KL drops quickly, terminate early. This would be the diffusion analog of the reference paper's difficulty-conditioned test-time compute allocation, but using the model's own stability signal rather than a separately trained difficulty estimator. A strong follow-up would measure the correlation between early-timestep average KL and final answer correctness across a held-out set of prompts, then implement a simple threshold-based routing policy (e.g., "if average KL after 64 steps exceeds X, re-run with 256 steps") and measure whether this recovers accuracy that would be lost by always using a fixed 64-step budget, while still providing a net speedup over always using 256 steps. The experiment piggybacks on KLASS's existing infrastructure and would connect the diffusion acceleration literature to the broader test-time compute allocation framework.

6. Extension to discrete diffusion models with non-absorbing noise schedules (uniform, discretized Gaussian, marginal prior).

The paper evaluates KLASS exclusively on masked (absorbing-state) diffusion models, where tokens are either fully masked or fully revealed, and unmasked tokens stay fixed. The theoretical framework (Section 3, Eqs. 1–4) and the ancestral sampling procedure (Section 3.2) are specific to this absorbing parameterization. However, discrete diffusion encompasses a broader family of noise schedules, including the uniform prior (tokens transition to a uniform distribution over the vocabulary rather than to a mask state) and the discretized Gaussian / multinomial diffusion of D3PM [1, 3], where tokens can transition to any other token in the forward process and the reverse process involves more complex state transitions. In non-absorbing diffusion, tokens are never fully "masked" — they exist in some intermediate state between clean and noisy — and the concept of "unmasking" a token does not directly apply. The research question is whether KL divergence between consecutive-step distributions remains informative about prediction correctness in these settings, and how the KLASS gating mechanism would need to be adapted. An experiment would implement KLASS on top of D3PM or SEDD [25] (which use non-absorbing parameterizations) on a text generation or sequence modeling task, measuring whether per-token KL divergence still separates correct from incorrect predictions and whether a stability gate improves quality-efficiency tradeoffs. Negative results (KL is not informative in non-absorbing diffusion) would clarify the scope of the KLASS approach and suggest whether the observed dynamics are specific to the absorbing parameterization or general to iterative refinement processes. Positive results would substantially broaden the method's applicability.

---

Practical Applications and Downstream Use Cases

1. Cost-efficient batch inference for language diffusion models in reasoning-heavy applications.

Organizations deploying LLaDA or Dream for mathematical reasoning (e.g., automated grading, math tutoring systems, quantitative problem-solving APIs) currently face a choice: use the standard 256-step Top-1 sampler and pay the full compute cost, or use a faster but lower-accuracy parallel method. KLASS provides a third option with both higher accuracy and lower cost. From Table 8, deploying KLASS instead of Top-1 on Dream for MATH reduces per-sample inference time from 30.76 seconds to 23.31 seconds (a 1.32× speedup) while increasing accuracy from 37.97% to 43.20% (+5.23 points). On LLaDA for GSM8K, the speedup is 2.34× (37.04s → 15.86s) with accuracy improving from 75.13% to 76.50%. For a batch inference pipeline processing 100,000 GSM8K problems, this reduces GPU-hours from approximately 37,040,000 seconds (≈10,300 GPU-hours) to 15,860,000 seconds (≈4,400 GPU-hours) — a savings of roughly 5,900 GPU-hours — while simultaneously producing more correct answers. The cost savings alone justify the modest hyperparameter tuning investment (Appendix D.1.2: ~100 validation examples, negligible relative to batch size), and the accuracy improvement is a bonus that would require substantially larger models to achieve through scaling alone. The implementation barrier is minimal: KLASS requires no model retraining, no architectural changes, and adds <1.57% memory overhead (Table 6), meaning existing deployment infrastructure works without modification.

2. Real-time code completion with diffusion models at competitive latency.

A practical barrier to deploying diffusion models for code generation in interactive settings (IDE plugins, notebook assistants) is that 256 sequential forward passes at ~0.12 seconds each produces ~30 seconds of latency — unacceptable for real-time use. AR models with KV-caching achieve much lower per-token latency for code completion because the amortized cost per new token is small. KLASS reduces Dream's code generation latency on HumanEval from 32.01 seconds to 11.52 seconds (2.78× speedup; Table 8) while matching or slightly improving accuracy (59.35% vs. 58.53% for Top-1). While 11.5 seconds is still too slow for truly interactive use (sub-second response times are expected), it brings diffusion models into the range where they could be deployed for batch code generation (generating multiple candidate completions offline for a developer to review) or asynchronous code suggestion (background generation while the developer continues typing). This opens up a deployment mode — diffusion-based code generation as a background process producing diverse, globally-coherent completions — that was previously impractical due to latency. Further combining KLASS with distillation (follow-up direction 3 above) could push latency into the 2–5 second range, at which point real-time interactive use becomes plausible.

3. Quality-improved unconditional text generation from off-the-shelf MDLM checkpoints without retraining.

The OpenWebText experiment (Table 10) demonstrates that KLASS improves generation quality on a standard MDLM checkpoint without any model modification: MAUVE increases from 0.115 to 0.179 (+56%), and generative perplexity under LLaMA2 drops from 30.88 to 26.94 (−13%). This is immediately actionable for practitioners using MDLM-based text generators who care about output coherence and fluency. The improvement comes at no additional training cost and with only the hyperparameter tuning described in Appendix D.1.2. A deployment scenario: a company using an MDLM checkpoint fine-tuned on domain-specific text (e.g., legal documents, medical reports) currently accepts the base MDLM sampler's quality. Swapping in KLASS as the sampler yields better outputs on the same model — higher coherence (MAUVE), lower perplexity under external LLM evaluators — without the cost and risk of fine-tuning a new model. This is a zero-cost quality upgrade for existing diffusion model deployments.

4. Accelerated molecular generation for virtual screening pipelines.

In drug discovery, generating large libraries of candidate molecules with desired properties (drug-likeness, target binding affinity) and then filtering them is a standard workflow. The molecular generation experiment (Table 4) shows KLASS reducing the average number of function evaluations from 32.0 to 18.8 (−41%) for QED-conditioned generation, while slightly improving the achieved drug-likeness score (0.546 vs. 0.526). For a virtual screening pipeline generating 10 million candidate molecules, this reduces the total NFEs from 320 million to 188 million — a savings of 132 million forward passes through the DiT model, translating to roughly 41% reduction in GPU-hours. The quality improvement (+3.8% QED) means the generated molecules are slightly more drug-like on average, potentially increasing the hit rate in downstream filtering. The implementation is straightforward: KLASS operates on the trained model's outputs with no architectural changes, and the hyperparameter sweep in Table 15 shows performance is stable across a range of threshold values near the optimum (τ = 0.96–0.999, ε_KL = 0.0005–0.01 all achieve comparable QED), so coarse tuning suffices. This is a drop-in efficiency improvement for any molecular generation pipeline built on masked diffusion — no changes to training, data, or downstream filtering code required.