PaperlayerView exact revision on arXiv

Source-linked AI summary

KLASS: KL-Guided Fast Inference in Masked Diffusion Models

Seo Hyun Kim, Sunwoo Hong, Hojung Jung, Youngrok Park, Se-Young Yun

arXiv:2511.05664v2cs.LG

TL;DR

Masked diffusion models face slow inference because iterative refinement uses static sampling strategies that unmask few tokens at a time. KLASS uses token-level confidence and KL divergence to adaptively unmask stable tokens in parallel, achieving faster generation with maintained or improved quality across reasoning and other modalities.

  • Problem

    Masked diffusion inference is slowed by iterative, static sampling strategies that unmask limited numbers of tokens and can produce local suboptimalities.

  • Method

    KLASS is a training-free sampler that combines high token confidence with low temporal KL divergence to guide parallel unmasking without external planners.

  • Results

    KLASS reduces sampling steps while maintaining or improving accuracy, achieving state-of-the-art math and code reasoning performance and generalizing across multiple modalities.

  • Takeaways & Limitations

    KLASS provides a simple, practical sampler for faster and more reliable masked-diffusion generation without additional training.

  • Takeaways & Limitations

    Future work should evaluate KLASS with alternative noise schedules and larger models.

Abstract

from arXiv · show

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

1 Introduction

KLASS addresses slow, static masked-diffusion sampling by using model-internal confidence and token-level KL divergence to unmask stable tokens in parallel. It reduces sampling steps while maintaining or improving quality across reasoning and other generation tasks.

  • Masked diffusion models iteratively refine partially masked sequences and have shown competitive performance in language, image, biological-sequence, and planning tasks.
  • Fixed schedules and limited-token unmasking make conventional sampling inefficient, while external planners add computational overhead and alignment difficulties.
  • KLASS identifies low-risk tokens for early unmasking using high confidence and low KL divergence between consecutive timestep distributions.This adaptive strategy avoids premature or suboptimal unmasking without extra model training or memory burden.
  • KLASS improves reasoning accuracy over standard greedy or Top-k decoding while halving sampling steps on GSM8K, MATH, HumanEval, and MBPP.The method reaches state-of-the-art results among diffusion samplers and also reduces perplexity while maintaining entropy in plain text generation.
  • KLASS achieves up to 2.78× faster sampling by more than halving diffusion steps through parallel unmasking of stable tokens.The sampler requires no additional training or external planner.

2 Related Works

Prior work established discrete diffusion foundations, masked diffusion language models, and iterative ancestral sampling. Recent acceleration methods reduce sampling cost through distillation or training-free certainty-based heuristics.

  • Foundations: Discrete diffusion research developed forward-backward processes for discrete state spaces, including continuous-time formulations and discrete score-matching objectives.These works frame discrete diffusion analogously to continuous diffusion while adapting training objectives to discrete spaces.
  • Foundations: Masked diffusion models substantially improved language-generation performance and narrowed the gap with autoregressive models, while LLaDA demonstrated scaling and reasoning abilities.The cited studies describe both performance gains from simplified masked diffusion and later scaling behavior in language domains.
  • Discrete diffusion samplers: Ancestral sampling generates text by repeatedly applying reverse denoising from a fully masked sequence across discrete timesteps.SUBS parametrization governs how predictions unmask tokens while preserving already revealed tokens.
  • Accelerated Sampling of Discrete diffusion models: Acceleration research reduces sequential sampling cost through distillation-based low-NFE training and other approaches adapted from continuous diffusion sampling.These methods train models to operate with fewer function evaluations rather than relying solely on the original sampling schedule.
  • Accelerated Sampling of Discrete diffusion models: Recent training-free samplers use model certainty, convergence, token position, entropy bounds, or Top-2 confidence gaps to guide multi-token decoding.The cited methods include confidence-aware decoding, staged certainty-based sampling, entropy-bound unmasking, and Top-2 confidence-gap selection.

3 Preliminaries

Masked diffusion models learn reverse denoising over sequences with absorbing masks, then discretize time and sample backward from a fully masked sequence. Simplified models estimate masked tokens while preserving already unmasked tokens.

  • The forward process uses an absorbing mask with a predefined monotonically decreasing schedule αt, and the reverse posterior is parameterized by a neural network.
  • Simplified masked diffusion models estimate only masked tokens while maintaining unmasked tokens throughout generation.
  • Training minimizes the negative ELBO, whose continuous formulation is invariant to the noise schedule under the stated framework.
  • Parallel decoding with KLASS un masks tokens only when both predictive confidence is high and the probability distribution is stable across steps.
  • Inference initializes a fully masked sequence, samples backward over discretized timesteps, and produces a complete sequence after T steps.

4 Method

KLASS selects tokens for unmasking using confidence and temporal stability measured by token-level KL divergence. It adaptively unmasks stable tokens and falls back to confidence-based selection when needed.

  • The confidence score is the largest predicted vocabulary probability, with higher values indicating greater certainty about a token.
  • The KL score measures divergence between a token’s predicted distributions at consecutive timesteps, with low values indicating consistent and more reliable estimates.
  • Correct reasoning samples consistently have lower KL scores than incorrect samples across models and datasets, motivating KL as a sampling signal.
  • At each diffusion step, KLASS un masks selected stable tokens and otherwise unmasks a fixed number u of positions with the highest confidence.
  • KLASS defines stable tokens as positions with confidence above τ and recent KL divergences below ϵKL.

5 Theoretical Rationale

The theoretical rationale argues that incorrect tokens cannot remain dynamically stable as their context is progressively resolved. KLASS therefore delays unmasking until tokens exhibit dynamic stability, supporting improved generation quality.

  • Theoretical rationale: Incorrect tokens cannot remain dynamically stable as context is progressively resolved under the paper’s well-trained-model perspective.The theoretical discussion frames this claim using conditional approximation and margin assumptions.
  • Theoretical rationale: A token that is wrong at the near-optimal context but correct at the target context must become dynamically unstable somewhere along the context path.Its average per-step KL is bounded away from zero along that path.
  • Theoretical rationale: KLASS delays unmasking until tokens exhibit dynamic stability, which the paper connects to improved generation quality.

6 Experiments

Experiments evaluate KLASS across reasoning, text, image, and molecular generation, showing improved quality, faster sampling, and low computational overhead across diverse settings.

  • Reasoning tasks: 40–70% fewer sampling steps yield up to 2.78× wall-clock speedups while KLASS improves accuracy over standard greedy decoding across reasoning tasks.KLASS also outperforms confidence-based Top-2 decoding, which reduces steps but degrades accuracy.
  • Text generation: KLASS improves generative text quality over existing discrete diffusion samplers, with higher MAUVE and lower perplexity while maintaining comparable entropy.The evaluation uses multiple oracle models on unconditional text generation.
  • Image generation: KLASS improves image quality over MMaDA’s confidence-based sampler, yielding lower FID and higher IS under both decoding schedules.The comparison uses the same schedules and fairness controls without modifying the backbone or adding auxiliary guidance.
  • Molecular generation: KLASS reduces total sampling steps while maintaining target reward for conditional molecular generation on QED and ring-count objectives.Experiments use QM9 molecules conditioned on drug-likeness and number of rings.
  • Ablation studies: Combining confidence and KL thresholds is essential for optimal MATH accuracy, while single-criterion thresholds reduce performance.Accuracy remains stable near each model’s distinct optimum despite differing optimal threshold values.
  • Ablation studies: Parallel KLASS unmasking improves MATH accuracy by up to 4.8 points while cutting sampling steps by nearly 50%.Similar accuracy-efficiency trends hold on GSM8K compared with sequential variants using the same stable token pool.
  • Computational overhead: KL computation adds no additional forward pass, with memory overhead below 1.57% and latency overhead below 0.21% per decoding step.The computation is a lightweight post-processing operation on existing logits with linear cost in the number of masked tokens and vocabulary size.

7 Conclusion

KLASS reduces sampling effort while maintaining or improving accuracy, achieving state-of-the-art results on math and code reasoning benchmarks without additional training. Its applicability extends across multiple modalities, while future work includes alternative noise schedules and larger models.

  • KLASS substantially reduces sampling steps while maintaining or improving accuracy, achieving state-of-the-art performance on math and code reasoning benchmarks.The method requires no additional training.
  • KLASS generalizes across multiple modalities as a practical sampler for faster and more reliable masked-diffusion generation.
  • Future work includes extending KLASS to discrete diffusion models with alternative noise schedules and evaluating it with larger models.

A Theoretical Proofs

The theoretical material relates distributional change across diffusion contexts to total-variation and KL-divergence bounds, while the surrounding material describes alternative discrete-diffusion samplers and experimental procedures. The reported wall-clock comparisons show that KLASS can reduce runtime substantially against 256-step Top-k decoding.

  • Theoretical bounds: Total variation across consecutive conditional distributions is bounded through triangle inequality and related KL-based inequalities.The proof defines T_t as TV(P_t+1, P_t) and invokes Pinsker’s inequality to connect KL divergence with T_t^2.
  • Related samplers: Alternative discrete-diffusion samplers trade accuracy, bias, or computational cost through caching, τ-leaping, higher-order integration, exact simulation, or sequential unmasking.The cited methods include O(τ) discretization bias for τ-leaping and O(L) sequential events for first-hitting sampling.
  • Experimental procedure: KLASS uses a lightweight validation search that first estimates KL thresholds, searches confidence values, and then refines the KL threshold.The procedure uses a small validation set and negligible computation relative to training.
  • Runtime comparison: 47.4% and 16.1% are the reported wall-clock reductions for KLASS versus 256-step Top-k decoding on LLaDA and Dream, respectively.

D.2.1 Experiment details

The experiments evaluate KLASS across text, image, and molecular generation settings with specified models, baselines, decoding budgets, and conditional-generation objectives. Results report maintained target rewards, reduced steps, and improved reasoning accuracy under selected configurations.

  • Reasoning evaluation: KLASS consistently improves reasoning accuracy while reducing sampling steps, with gains of 6.22–8.00 points and up to 79% fewer steps at temperature 0.At temperature 0.2, reported accuracy gains range from 0.69 to 5.10 points with 39%–71% step reductions.

E.1 Performance on reasoning tasks

Across reasoning benchmarks, KLASS outperforms confidence- and entropy-based alternatives at fewer sampling steps, while its gains vary with temperature and history length. Hyperparameter tables document selected configurations and baseline accuracies for LLaDA, Dream, and MDLM.

  • Baseline configurations: The reported baselines include 39.63% accuracy for HumanEval with LLaDA and 48.64% for MBPP with LLaDA, both using 256 steps.
  • Sampler comparison: KLASS consistently outperforms Top-k Margin and Entropy samplers with fewer steps, including top results on LLaDA GSM8K and HumanEval in under half the usual iterations.Entropy can achieve the highest accuracy on MATH and MBPP at 256 steps, but its performance drops sharply with fewer steps.
  • History-length ablation: History length 2 provides the best reported accuracy-efficiency balance for LLaDA and the highest accuracy for Dream under selected KL and confidence thresholds.

F Examples of generated samples

The examples compare KLASS with baseline samplers on mathematical reasoning and long-form text generation. KLASS produces a correct mathematical solution and more coherent text than the illustrated MDLM baseline.

  • Long-form text coherence: The MDLM sample begins on-topic but degenerates through repetition, nonsensical phrases, and an unrelated spam link.The cited example illustrates a complete breakdown in long-form coherence.
  • Long-form text coherence: The KLASS sample on “urban sprawl” remains relatively coherent in contrast to the degraded MDLM sample.

G.1 Limitations

The paper identifies limitations in evaluation scale, hyperparameter tuning, and potential misuse. Larger-model benchmarks remain unavailable, KLASS adds search cost, and improved samplers may facilitate harmful generation if misused.

  • Evaluation cannot extend to challenging agentic-system benchmarks because larger discrete diffusion models are unavailable.
  • KLASS introduces minimal additional hyperparameter search cost relative to Top-1 sampling, although performance is robust around reported optima.Further tuning could still identify better settings.
  • Improved diffusion samplers could be misused to create harmful or misleading text, images, or novel molecules.The paper recommends accompanying sampling improvements with AI safety research and ethical guidelines.
Loading 2511.05664v2…