Confidence-Based Decoding is Provably Efficient for Diffusion Language Models

TL;DR

This paper provides the first theoretical analysis of confidence-based decoding in diffusion language models, showing it achieves efficient sampling by adaptively controlling unmasking based on entropy, especially effective for low-entropy data.

Contribution

It introduces a theoretical framework for confidence-based decoding in DLMs and proves its efficiency in terms of expected iterations, adapting to data complexity without tuning.

Findings

Achieves $ ilde{O}(H(X_0)/ ext{epsilon})$ expected iterations for $ ext{epsilon}$-accurate sampling.

Automatically adapts to data entropy, providing acceleration for low-entropy distributions.

Provides a foundation for designing more efficient decoding strategies in diffusion language models.

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.