Confidence-Based Decoding is Provably Efficient for Diffusion Language Models

2026-03-23

AI Summary

This work addresses the critical role of decoding strategies in sampling efficiency for diffusion language models, noting that existing confidence-based approaches lack theoretical grounding. The paper introduces the first theoretical framework for such decoding methods and proposes an adaptive strategy based on entropy accumulation thresholds. By leveraging information entropy to quantify predictive uncertainty, the method dynamically determines both the number and order of tokens to unmask at each step, eliminating the need for prior knowledge or hyperparameter tuning. Theoretically, under low-entropy data distributions, the approach achieves ε-accurate sampling within an expected number of iterations bounded by $\widetilde{O}(H(X_0)/\varepsilon)$, substantially improving generation efficiency.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Sampling/Simulation-based Search
• Reasoning under Uncertainty: Sequential Decision Making

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) models for language modeling, allowing flexible generation order and parallel generation of multiple tokens. However, this flexibility introduces a challenge absent in AR models: the \emph{decoding strategy} -- which determines the order and number of tokens generated at each iteration -- critically affects sampling efficiency. Among decoding strategies explored in practice, confidence-based methods, which adaptively select which and how many tokens to unmask based on prediction confidence, have shown strong empirical performance. Despite this success, our theoretical understanding of confidence-based decoding remains limited. In this work, we develop the first theoretical analysis framework for confidence-based decoding in DLMs. We focus on an entropy sum-based strategy that continues unmasking tokens within each iteration until the cumulative entropy exceeds a threshold, and show that it achieves $\varepsilon$-accurate sampling in KL divergence with an expected number of iterations $\widetilde O(H(X_0)/\varepsilon)$, where $H(X_0)$ denotes the entropy of the target data distribution. Notably, this strategy yields substantial sampling acceleration when the data distribution has low entropy relative to the sequence length, while automatically adapting to the intrinsic complexity of data without requiring prior knowledge or hyperparameter tuning. Overall, our results provide a theoretical foundation for confidence-based decoding and may inform the design of more efficient decoding strategies for DLMs.

Problem

• diffusion language models
• confidence-based decoding
• decoding strategy
• sampling efficiency
• theoretical analysis

Innovation

• confidence-based decoding
• diffusion language models
• entropy-based sampling
• theoretical analysis
• efficient generation