Stability-Weighted Decoding for Diffusion Language Models

Abstract: Diffusion LLMs (dLLMs) enable parallel text generation by iteratively denoising a fully masked sequence, unmasking a subset of masked tokens at each step. Existing decoding strategies rely on static confidence metrics computed at a single denoising step, ignoring temporal history and often leading to premature unmasking of unstable tokens. In this work, we theoretically establish that a token's temporal instability, quantified by the KL divergence between consecutive prediction distributions, provides a strict lower bound on its mutual information with the remaining masked context, indicating that temporally unstable tokens are inherently unsafe to unmask. Based on this insight, we propose Stability-Weighted Decoding (SWD), a training-free, plug-and-play strategy that incorporates temporal stability into token scoring and acts as a universal modulator for arbitrary score-based decoding policies. Experiments on code generation and mathematical reasoning benchmarks demonstrate that SWD consistently improves generation accuracy across representative scoring metrics and selection policies, and exhibits exceptional robustness, maintaining a significant performance lead over standard baselines across varying acceleration ratios.

• The paper introduces Stability-Weighted Decoding (SWD), a training-free method that mitigates premature token unmasking by penalizing unstable predictions.
• SWD leverages temporal KL divergence as an information-theoretic metric, effectively enhancing accuracy across various decoding tasks.
• Empirical results demonstrate significant gains, such as boosting MATH500 accuracy from 10.8% to 35.0%, while reducing computational overhead under parallel decoding.

Stability-Weighted Decoding for Diffusion LLMs: A Formal Analysis

Introduction

The study "Stability-Weighted Decoding for Diffusion LLMs" (2604.17068) addresses a critical problem in the deployment of diffusion-based LLMs (dLLMs): the risk of premature token unmasking during parallel decoding. While dLLMs have emerged as competitive alternatives to autoregressive models, their inherent parallelism introduces vulnerability to unmasking tokens whose predictions remain unstable—leading to error propagation and compromised sequence quality. This paper introduces Stability-Weighted Decoding (SWD), a training-free, modular method for decoding that penalizes temporary instability, strictly grounded in information-theoretic principles.

Background and Problem Setting

dLLMs generate text by iterative denoising from a fully masked input, where, at each step, a subset of masked positions is unmasked based on some scoring function applied to the model's predictions. Common score metrics include per-token confidence, entropy, and margin. Existing decoding policies typically make unmasking decisions based solely on these metrics at individual timesteps, neglecting the temporal dynamics of prediction stability—a notable oversight given the interdependence of token values throughout the iterative process.

The authors highlight that static, snapshot-based scoring is susceptible to premature unmasking when tokens have high instantaneous confidence but exhibit rapid fluctuations (temporal instability) over the generation trajectory. This dynamic can lock in errors, negatively influencing downstream predictions and overall generation quality.

Information-Theoretic Justification of Temporal Stability

The core theoretical contribution is a formal equivalence between temporal instability and mutual information with unrevealed context. Specifically, the temporal KL divergence between a token's predicted distributions at consecutive steps proves to be a strict lower bound for its dependency on the remaining masked content. The paper's theorem establishes:

$$I(x_0^i; \mathbf{U}_{t+1} | \mathbf{x}_{t+1}) \geq \mathbb{E}[\mathcal{D}_{\mathrm{temp}}^{(i)}]$$

where $I$ represents the mutual information, and $\mathcal{D}_{\mathrm{temp}}^{(i)}$ is the KL divergence between the predicted distributions over token $i$ at steps $t$ and $t+1$.

This implies that tokens whose predictions are unstable (high KL divergence) remain entangled with masked positions and that unmasking should be deferred until stabilization is observed.

Figure 1: Comparison of conventional decoding relying on static confidence with SWD, illustrating the correction of premature commitments via temporal stability modulation.

Stability-Weighted Decoding: Methodology

SWD is introduced as a training-free, plug-and-play “stability modulator” over arbitrary score-based decoding policies. For each masked token, a base score (e.g., confidence, margin, entropy) is multiplied by an exponential penalty based on temporal KL divergence:

$$\mathcal{S}_\mathrm{SWD}(i) = \mathcal{S}_\mathrm{base}(i) \cdot \exp(-\lambda \cdot \mathcal{D}_{\mathrm{KL}})$$

The penalty strength $\lambda$ is a hyperparameter that manages the trade-off between stability-constrained decoding and aggressive parallelism. This construction remains universally compatible with existing selection strategies, including Top-$k$, thresholding, and entropy budgeting schemes such as EB-Sampler.

Figure 2: SWD achieves favorable accuracy/speedup trade-offs compared to confidence-only baselines, especially sustaining accuracy at high parallelism.

Empirical Evaluation

Experiments are conducted on two leading dLLMs: LLaDA-8B-Instruct and Dream-v0-Instruct-7B, across benchmarks including HumanEval, MBPP (code generation), GSM8K, and MATH500 (mathematical reasoning). SWD is assessed under diverse scoring/selection paradigms and compared to both baseline (static) and recent history-aware decoders such as KLASS and CreditDecoding.

Key Observations:

• Universal Gain: SWD consistently improves accuracy across metrics, models, and tasks. For example, on Dream-7B (MATH500), confidence-based SWD raises accuracy from 10.8% to 35.0%.
• Efficiency: SWD achieves these gains while reducing the number of function evaluations (NFE), hence lowering both inference time and compute overhead despite minor additional cost for KL computation.
• Robustness: SWD maintains generation quality under high-acceleration (parallel decoding) regimes that would otherwise degrade baseline model performance.
• Versatility: Benefits hold under alternative selection policies (e.g., threshold, Top-1) and blockwise semi-autoregressive decoding.

  Figure 3: SWD outperforms baseline approaches consistently under multiple selection strategies (threshold and Top-1), demonstrating generality across decoding policies.

• Stability Tuning: Ablation over $\lambda$ shows a smooth trade-off, validating a “sweet spot” balancing aggressiveness and stability.

  Figure 4: SWD’s superiority is robust across $I$0, outperforming the confidence baseline for a broad parameter range.

Implications and Future Directions

SWD formalizes a principled approach to parallel decoding in dLLMs, grounding speculative heuristics regarding prediction stability in rigorous information-theoretic bounds. Practically, SWD enables more aggressive parallelism without sacrificing quality, making dLLM-based systems more viable for real-world, latency-sensitive applications. Theoretically, it establishes a general framework that can be adapted to multimodal generation, infill tasks, and advanced re-masking schemes.

Furthermore, the SWD paradigm is orthogonal to the underlying model and can be combined or extended with other history-augmented or credit-assignment decoding methods. The approach also opens the door to more nuanced, adaptive modulation strategies, possibly leveraging richer temporal statistics or context-dependent risk budgets.

Conclusion

Stability-Weighted Decoding offers a unified, theoretically justified mechanism for reliable parallel decoding in diffusion LLMs. By explicitly suppressing the scores of temporally unstable tokens, SWD significantly enhances accuracy and robustness for both code generation and reasoning workloads, regardless of the baseline scoring or selection strategy employed. This method provides a plug-and-play paradigm that bridges rigorous theory and practical deployment, setting a solid foundation for future research in high-quality, efficient decoding for discrete generative models.