Confidence-Based Decoding is Provably Efficient for Diffusion Language Models

Changxiao Cai; Gen Li

Source-linked AI summary

Confidence-Based Decoding is Provably Efficient for Diffusion Language Models

Changxiao Cai, Gen Li

arXiv:2603.22248v1Mar 23, 2026cs.LGcs.AIcs.ITstat.ML

TL;DR

DLM decoding efficiency is theoretically underexplained because adaptive confidence-based unmasking depends on generated tokens and prediction confidences. The paper introduces an entropy sum-based analysis and proves efficient KL-accurate sampling that adapts to low-entropy data without prior distribution knowledge or tuning.

Problem
Confidence-based decoding has strong empirical performance, but its adaptive sampling efficiency and error guarantees in DLMs remain theoretically limited.
Method
The paper analyzes an entropy sum-based strategy that greedily unmasks tokens until cumulative predicted entropy exceeds a threshold.
Results
The strategy achieves ε-accurate sampling in KL divergence with expected iteration complexity scaling as Õ(H(X_0)/ε).
Takeaways & Limitations
When target entropy is low relative to sequence length, entropy-based decoding provides sublinear sampling complexity and automatically adapts without prior entropy knowledge or hyperparameter tuning.
Takeaways & Limitations
The analysis assumes an optimal mask predictor and leaves the effects of training error on confidence-based sampling performance for future study.

Abstract

from arXiv · show

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.

1 Introduction

DLMs enable parallel, non-causal generation, but their decoding strategy determines the trade-off between sampling speed and generation quality. This work develops theory for confidence-based decoding and analyzes an entropy sum-based strategy with provable efficiency and adaptivity.

Motivation: DLMs generate multiple tokens in parallel and relax autoregressive left-to-right generation, offering potential sampling acceleration and support for bidirectional reasoning.Their non-causal structure also supports global constraints.
Motivation: Simultaneous unmasking introduces factorization bias by ignoring conditional dependencies, whereas one-token-at-a-time decoding removes that bias but sacrifices parallelism.The decoding strategy therefore controls the central speed-quality trade-off.
Decoding strategies: Confidence-based decoding selects easy, high-confidence tokens first so newly revealed tokens provide context for harder predictions.Confidence may be quantified using prediction entropy or related signals.
Theory gap: Theoretical understanding of confidence-based sampling efficiency remains limited because adaptive unmasking dynamics depend on generated tokens and their predicted confidence levels.Open questions include iteration counts, sampling error, and adaptation to unknown data complexity.
Approach: The paper studies an entropy sum-based strategy that greedily unmasks tokens until cumulative prediction entropy exceeds a threshold.This controls cumulative uncertainty while revealing as many tokens as possible per iteration.
Results: The analysis establishes KL-divergence sampling guarantees and shows sublinear iteration complexity when target entropy is low relative to sequence length.The strategy is distribution-adaptive without requiring prior entropy knowledge or hyperparameter tuning.

2 Preliminaries

The preliminaries define masked diffusion language models through forward masking and iterative reverse unmasking. A mask predictor supplies conditional token distributions, while the decoding strategy selects positions and controls the random number of iterations.

Forward process: Masked diffusion models add an absorbing mask state and progressively replace data tokens with masks in the forward process.The masking process ends when the entire sequence is masked.
Mask predictor: The mask predictor learns conditional distributions of masked tokens given unmasked context, with time-independent posteriors determined by that context.Training targets conditional marginals for each token position.
Mask predictor: Training maximizes the log-likelihood of masked tokens conditioned on unmasked tokens, equivalently minimizing KL divergence to the product of learned conditional marginals.Finite samples and random masking patterns approximate the training expectation.
Reverse process: Reverse generation starts from a fully masked sequence and repeatedly selects masked positions, samples their tokens simultaneously, and inserts the results.The final sequence is produced after a potentially random number of iterations.
Decoding strategy: With a fixed mask predictor, the generated distribution is completely determined by the decoding strategy's choice of unmasking sets.This makes decoding strategy a defining component of the sampling process.
Entropy-based decoding: Entropy sum-based decoding initializes a random token order and ends each iteration when the accumulated entropy exceeds a threshold or all tokens are unmasked.The algorithm records the total number of iterations and outputs the fully unmasked sequence.

3 Main results

The paper develops entropy-based decoding strategies with theoretical guarantees for DLM sampling. The entropy sum-based strategy controls cumulative uncertainty and achieves adaptive, entropy-dependent iteration complexity, while the maximum entropy variant offers a weaker guarantee with distribution-dependent parameter requirements.

Entropy sum-based decoding strategy: The entropy sum-based strategy adaptively enlarges batches in low-uncertainty regions and shrinks them when prediction uncertainty is higher.It scans positions in random order and stops when the cumulative entropy exceeds threshold η or all tokens are unmasked.
Entropy sum-based decoding strategy: The decoding procedure determines unmasking sets from conditional entropies, then independently samples newly unmasked tokens from their corresponding conditional distributions.The resulting sampling process has a Markov-chain representation over partially unmasked sequences.
Theoretical guarantee: Under an optimal mask predictor, the entropy sum-based strategy achieves ε-accurate sampling in KL divergence with expected iterations governed, up to logarithmic factors, by H(X0)/ε.Here H(X0) is the entropy of the target data distribution.
Theoretical guarantee: When H(X0) ≪ L, the entropy sum-based method has sublinear iteration complexity and requires neither prior entropy knowledge nor hyperparameter tuning.Its adaptivity exploits low-complexity structure in the target distribution.
Effect of entropy criterion: The maximum entropy-based strategy stops when one token exceeds an entropy threshold and imposes a maximum batch size, rather than controlling cumulative batch entropy.It still achieves sublinear iteration complexity, but its guarantee is weaker and its parameters require knowledge of H(X0).
Effect of entropy criterion: Both entropy criteria achieve theoretical guarantees, but their comparison shows that the entropy criterion critically affects iteration efficiency and distribution adaptivity.The paper identifies algorithm-dependent lower bounds for different entropy criteria as a future direction.

4 Analysis

The analysis expresses the sampler’s KL error along its adaptive unmasking trajectory, then decomposes it into mutual-information terms controlled by the entropy threshold. This yields the stated KL guarantee after choosing the threshold as a function of ε and sequence length.

Trajectory representation: The realized iteration count and unmasking sets are deterministic functions of the sampled sequence and permutation.This characterization allows the sampling trajectory to be analyzed by conditioning on the permutation and previously revealed tokens.
Trajectory representation: At each iteration, the sampled distribution factorizes over newly unmasked tokens according to their true conditional marginals.The data distribution instead retains the true conditional joint distribution for the same unmasking set.
KL decomposition: The expected KL sampling error decomposes into per-iteration discrepancies between conditional joint distributions and products of conditional marginals.This isolates the factorization error introduced when multiple tokens are unmasked together.
KL decomposition: The joint-versus-product KL discrepancy is decomposed into mutual-information terms involving tokens unmasked in the same iteration.The resulting terms measure dependence after conditioning on already unmasked tokens and the permutation.
Bounding the error: The entropy threshold bounds the per-iteration contribution, leading to a final expected KL error bound.The proof uses the fact that the running entropy threshold permits at most one token with entropy exceeding η within an iteration.
Bounding the error: Setting η = ε/(4(log2 L + 1)) yields the desired ε-accurate sampling guarantee.The threshold choice converts the intermediate error bound into the theorem’s target accuracy.

5 Discussion

The discussion establishes that entropy-based unmasking can accelerate sampling and adapt to unknown data complexity, while identifying several directions beyond the current analysis.

Entropy-based unmasking can achieve iteration complexity sublinear in sequence length when the data distribution has low entropy.
The strategy adapts to the unknown intrinsic complexity of the data distribution without prior knowledge or hyperparameter tuning.
Future work includes analyzing deterministic ordering instead of the random order used in this work.
The analysis could be extended to confidence measures beyond entropy, including top-token mass and the mass gap between the two most likely tokens.
The current analysis assumes an optimal mask predictor, leaving the impact of training error on sampling performance unresolved.
Lower bounds on iteration complexity remain an open direction for characterizing fundamental limits of confidence-based sampling acceleration.

A.1 Proof of Lemma 1

The proof of Lemma 1 derives mutual-information bounds by reordering random variables, applying the chain rule, and using monotonicity of mutual information.

The proof introduces notation for suffixes and complements of the random-variable sequence to organize the chain-rule calculations.
Reordering the variables so that a selected variable appears last allows the same KL-divergence identity to be applied to that variable.
The argument uses mutual-information monotonicity after identifying each reordered variable and its conditioning context.
Summing the resulting bounds over all variables yields the claimed bound in equation (22), completing the proof.

A.2 Proof of Lemma 2

The proof of Lemma 2 compares the original sampling process with an auxiliary process that omits one token, tracking insertion positions, unmasking sets, and mutual-information contributions.

The auxiliary process unmasks every token except the selected token and follows the entropy sum-based procedure under the reduced permutation.
Before the selected token is unmasked, the original and auxiliary processes have identical unmasking sets and sampled tokens.
Conditioned on the other tokens and permutation entries, the selected token’s unmasking iteration is determined by its insertion position.
A high entropy for the selected token can cause the original iteration to terminate earlier, making its unmasking set a subset of the auxiliary set.
The proof groups insertion positions by iteration and relative position, then controls the resulting sums with dyadic size envelopes.
Mutual-information monotonicity and a chain-rule bound control the accumulated contribution, yielding the claim in equation (29).

A.3 Proof of Lemma 3

The proof of Lemma 3 expresses the data and sampled densities using the realized iterations and unmasking sets, then relates the resulting expressions to the entropy of the target distribution.

For a fixed permutation, the realized iteration count and unmasking sets are determined by the sampled sequence and permutation.
Taking expectation over the target-distributed sequence yields the corresponding density identity for each fixed permutation.
The proof similarly expands the sampled distribution’s entropy using conditional token probabilities at each iteration.

A.4 Proof of Lemma 4

The proof of Lemma 4 establishes a correspondence between terms indexed by (k, q) on the left-hand side and terms indexed by (s, k, q) on the right-hand side. It verifies the required indicator condition and completes the matching argument.

Each left-hand-side term indexed by (k, q) is matched to a right-hand-side term indexed by (s, k, q), completing Lemma 4.The proof concludes after establishing this term-by-term correspondence.
The indicator 1{Drel_s(q) ≥ D_s} is shown to equal 1 because Drel_s(q) is at least q and q is at least D_s.This uses the fact that token i is unmasked at iteration s and that the preceding tokens must be unmasked before it.
The proof handles both cases, s = k and s < k, and obtains the needed inequality in either case.For s = k, the inequality is immediate; for s < k, it follows from the recursive threshold relations.
If s < k, the recursive definition forces D_{s+1} = 2D_s when D_{s+1} > q ≥ D_s.The argument also uses q ≤ D_{k}^{¬i} to conclude that D_s = D_{k}^{¬i}.

B.1 Proof sketch of Theorem 2

The proof sketch of Theorem 2 analyzes the maximum entropy-based strategy by decomposing KL error into mutual-information terms, bounding those terms with the entropy stopping rule, and controlling the iteration count.

Step 1: Decompose KL error into sum of mutual information: The proof first decomposes expected KL sampling error into mutual-information terms over tokens unmasked later in the same iteration.For each token, the relevant term measures its correlation with later-revealed tokens conditioned on earlier revealed tokens and the permutation.
Step 1: Decompose KL error into sum of mutual information: Only tokens unmasked after a given token within the same iteration contribute to its mutual-information error term.The last token unmasked in an iteration contributes no KL error.
Step 2: Bound KL error by entropy-based stopping rule: The KL error is then bounded using the maximum entropy-based stopping rule and the associated per-token lemma.Substituting the lemma’s bound and summing over token indices yields a bound on the expected KL sampling error.
Step 3: Control expected number of iterations: The theorem follows after setting the parameters and combining the KL-error and iteration-count bounds.The parameter choice is introduced immediately before the final conclusion of the proof sketch.
Step 3: Control expected number of iterations: The expected iteration count is controlled separately by Lemma 8.The proof states that the resulting iteration count has the same order as a procedure unmasking exactly Smax tokens per iteration.

B.2 Proof of Lemma 7

The proof of Lemma 7 controls a token-level mutual-information term by comparing the original process with an auxiliary process that removes the token under analysis. The main difficulty is the dependence between that token’s insertion position and the later-unmasked set.

The main technical challenge is statistical dependence between L_i and token i’s insertion position in the permutation.This dependence makes direct expectation over the permutation difficult.
The auxiliary process removes token i and runs the maximum entropy-based procedure on the remaining tokens.Conditioned on the other tokens and their relative permutation, its unmasking sets are determined independently of token i’s insertion position.
The later-unmasked set L_i in the original process is bounded by a corresponding auxiliary unmasking set.This containment allows the target mutual-information term to be bounded using monotonicity.
The proof introduces a random iteration index t_i for token i and bounds the probability that token i is unmasked at iteration k.Lemma 9 supplies the needed conditional probability bound after fixing the other-token configuration and permutation order.
Applying the resulting bounds and Lemma 5 completes the control of the target expectation and proves Lemma 7.The final combination bounds the right-hand side of the intermediate inequality and finishes the lemma.

B.3 Proof of Lemma 8

The proof of Lemma 8 bounds the expected number of iterations by classifying iterations according to the stopping rule and separately controlling capped and high-entropy iterations.

The stopping rule divides iterations into capped iterations, high-entropy iterations, and the final iteration.High-entropy iterations stop when the stopping token’s entropy exceeds η, while the final iteration contributes at most one iteration.
The total iteration count satisfies T ≤ Tcap + Thigh + 1.Here, Tcap and Thigh count capped and high-entropy iterations, respectively.
Capped iterations are controlled because each such iteration unmasks exactly smax tokens.This reduces their count to a quantity determined by the sequence length and the cap.
For a high-entropy iteration, the stopping token’s entropy and the entropy-sum bound constrain the iteration’s contribution.The proof invokes the entropy bound from Lemma 3 to control the relevant sum.
The auxiliary process and the uniform insertion position of token i are used to derive the remaining probability bound.Conditioning on the relative order of all other indices makes token i’s insertion position uniform, while the original and auxiliary processes coincide before token i is unmasked.
Combining these bounds yields the iteration-count conclusion, with each iteration unmasking at most smax tokens in the auxiliary construction.The final step explicitly uses the construction’s per-iteration cap.