Breaking AR’s Sampling Bottleneck: Provable Acceleration via Diffusion Language Models

We thank the reviewer for the insightful comments and thorough assessment of our work.

Weakness:

• Thank you for the suggestion. We add a synthetic numerical experiment in the revised manuscript to validate our theory.

  • Model. we use a 1-D $K$-state Potts chain of length $L$ with coupling $J$ to model the distribution of text $X=(X^{(1)},\dots,X^{(L)})$. Concretely, $X^{(1)}\sim\text{Unif}([K])$ and $\mathbb P\\{X^{(i)}=y|X^{(i-1)}=x\\}=\exp(J1\\{x=y\\})/(\exp(J)+K-1), \quad \forall x,y\in[K].$ This construction allows us to compute explicitly the mutual information, the optimal mask predictor $p^\star(\cdot | X_t)$, and the densities of $p_{X_0}$ and $p_{Y_0|M}$.

  • Mask predictor. We use the explicitly computable optimal mask predictor $p^\star(\cdot | X_t)$ to run the sampling process.

  • Sampling error. Given the explicit densities of $p_{X_0}$ and $p_{Y_0|M}$, the expectation in the KL divergence, taken over both the mask schedule $M$ and the data distribution $p_{X_0}$, is approximated using Monte Carlo simulations.

  • KL divergence vs. mutual information. For fixed $T=10$, $K=10$ and $L=100$, KL divergence vs. mutual information is summarized in the following table:

    J	Mutual Information	KL
    0.000	4.80e-13	0.000
    0.333	1.18e+00	0.051
    0.667	5.53e+00	0.309
    1.000	1.43e+01	0.748
    1.333	2.86e+01	1.302

    One can verify that the KL divergence error increases linearly with the mutual information.

  • KL divergence vs. number of iterations $T$. With a balanced mask schedule, $K=10$, $L=100$ and $J=2$, $\log KL$ vs. $\log T$ is summarized in the following table.

    $T$	KL	$\log$ KL	$\log T$
    2	18.70	2.928	0.693
    4	8.85	2.180	1.386
    5	7.24	1.980	1.609
    10	3.10	1.130	2.303
    20	1.49	0.395	2.996
    25	1.16	0.151	3.219
    50	0.44	–0.818	3.912

    One can verify that the slope is very close to -1, demonstrating that the KL divergence error scales proportionally to $1/T$.

  • Conclusions. Therefore, these two results validate our established sampling convergence theory.

• Thank you for the suggestion. According to Theorem 1, one can see that $C_1 = T s_{\max}/\sum s_t \asymp 1$. In addition, a direct inspection of the proof of Theorem 2 shows that we may take $C_2 = \frac{1}{2}$ and $C_3 = \frac{1}{16}$. We will state these explicit constants in the revised manuscript.

• Thank you for bringing these references to our attention and for the detailed description, which is very helpful for improving this paper. We’ll incorporate them in the revised manuscript to make the introduction more accurate.

Questions:

• This is certainly a valuable point! Classical results from information theory show that when the text distribution has a Markov structure, the mutual information between any token with the rest of the sequence is substantially smaller than that in the full dependent scenario. For instance, when $(X^{(1)}, X^{(2)}, \dots, X^{(L)})$ forms a first-order Markov chain, then $I(X^{(1)}; X^{(-1)})=I(X^{(1)}; X^{(2)})$, which is typically far less than that in the general case. Because our sampling bound depends on the same mutual-information terms, it automatically tightens for such structured data. Investigating diffusion language models under explicit structural assumptions (e.g., Markov, low-order dependencies) is therefore a promising avenue, and we will add a brief discussion of this point in Section 5.

• Thank you for the suggestion. Our ongoing studies suggests that it is possible to extend our results to the entropy-based sampling and show its performance improvement over the random sampling. Specifically, in each iteration, we can rank the still-masked positions by their conditional entropy (which can be estimated using the learned mask predictor). Unmasking the lowest-entropy positions first reduces the per-step conditional mutual information (because mutual information is always upper bounded by entropy), thereby leading to a smaller final error bound. Our conjecture is that when the entropy-based sampling procedure is used, the conditional mutual information terms in our current error bounds can be replaced by these smaller conditional entropy terms, which yields improved sampling accuracy. Providing a rigorous theoretical proof for this adaptive strategy is an interesting direction for future work.

We are thankful for the reviewer’s valuable feedback and close attention to the manuscript.

• For the question regarding generality, this indeed requires more clarification. First, our theory applies to any target text distributions---we impose no distributional assumptions whatsoever---so we refer to ``general data distributions.’’ Second, we analyze the standard sampling procedure of diffusion language models without introducing additional constraints or restrictions for analysis (e.g., we allow arbitrary mask sizes per iteration). Hence, we describe our results as covering “general sampling procedures” of diffusion language models.

• Thank you for catching the typos. We have fixed them.

We appreciate the reviewer’s time and effort in providing helpful feedback.

Weakness:

• This is certainly a valuable point that needs more clarification. We do not assume a near-optimal mask predictor. Our analysis only requires access to a pre-trained mask predictor (of arbitrary quality) and provides an off-the-shelf framework for assessing the text generation quality of diffusion language models. We decouple the mask predictor training phase from the sample generation stage, showing that the final error in KL divergence can be decomposed into (i) a training error term $\varepsilon_{train}$ that captures how far the learned predictor is from the optimal one, and (ii) a sampling error term incurred during the sample generation phase. The central contribution of this paper is to characterize the sampling error in terms of the iteration number and the statistical structure of the text distribution. Crucially, this sampling term is independent of the quality of the mask predictor. To make this point more precise and avoid misunderstanding, we will revise our statement about the pre-trained mask predictor from “We assume access to sufficiently accurate mask predictors” to “We assume access to some mask predictors”.

• We add a synthetic numerical experiment in the revised manuscript to validate our theory.

  • Model. we use a $K$-state Potts chain of length $L$ with coupling $J$ to model the distribution of text $X=(X^{(1)},\dots,X^{(L)})$. Concretely, $X^{(1)}\sim\text{Unif}([K])$ and $\mathbb P\\{X^{(i)}=y|X^{(i-1)}=x\\}=\exp(J1\\{x=y\\})/(\exp(J)+K-1), \quad \forall x,y\in[K].$ This construction allows us to compute explicitly the mutual information, the optimal mask predictor $p^\star(\cdot | X_t)$, and the densities of $p_{X_0}$ and $p_{Y_0|M}$.

  • Mask predictor. We use the explicitly computable optimal mask predictor $p^\star(\cdot | X_t)$ to run the sampling process.

  • Sampling error. Given the explicit densities of $p_{X_0}$ and $p_{Y_0|M}$, the expectation in the KL divergence, taken over both the mask schedule $M$ and the data distribution $p_{X_0}$, is approximated using Monte Carlo simulations.

  • KL divergence vs. mutual information. For fixed $T=10$, $K=10$ and $L=100$, KL divergence vs. mutual information is summarized in the following table:

    J	Mutual Information	KL
    0.000	4.80e-13	0.000
    0.333	1.18e+00	0.051
    0.667	5.53e+00	0.309
    1.000	1.43e+01	0.748
    1.333	2.86e+01	1.302

    One can verify that the KL divergence error increases linearly with the mutual information.

  • KL divergence vs. number of iterations $T$. With a balanced mask schedule, $K=10$, $L=100$ and $J=2$, $\log KL$ vs. $\log T$ is summarized in the following table.

    $T$	KL	$\log$ KL	$\log T$
    2	18.70	2.928	0.693
    4	8.85	2.180	1.386
    5	7.24	1.980	1.609
    10	3.10	1.130	2.303
    20	1.49	0.395	2.996
    25	1.16	0.151	3.219
    50	0.44	–0.818	3.912

    One can verify that the slope is very close to -1, demonstrating that the KL divergence error scales proportionally to $1/T$.

  • Conclusions. Therefore, these two results validate our established sampling convergence theory.

  In addition, because the true data distribution is unknown in the real-world datasets, it is infeasible to compute the token-level mutual information or the KL divergence between the output and the data distribution. Hence, we restrict our numerical experiments to the synthetic data.

• Our lower bound targets the balanced mask schedules used in practice, where the number of unmasked tokens per iteration is of the same order. Admittedly, the bound can fail for deliberately unbalanced mask schedules---for instance, unmasking half the sequence in the first iteration and then proceeding token-by-token in the rest of iterations. Because balanced mask schedules are much more commonly employed in practice, we believe this exclusion does not compromise the practical insights of our results.

Questions:

• “The sampling error has a linear dependence on the mutual information” is the main result of this paper (upper bounds in Theorem 1 and Corollary 1, and lower bounds in Theorem 2), rather than an assumption.
• Thank you for the suggestion. A practical implication of our theory is to unmask, at each iteration, the tokens whose conditional dependence on the rest of the sequence is weakest. Concretely, Eq. (24) shows that the per-step contribution to the total sampling error is the conditional mutual information between a newly revealed token and the remaining tokens. Hence, selecting positions with smaller conditional mutual information minimizes that contribution. Moreover, because $I(X;Y|Z)\leq H(X|Z)$ for any random variables $X,Y,Z$, we can implement this strategy by ranking tokens by conditional entropy: at step $t$, use the learned mask predictor $\hat{p}_i(\cdot | Y_t)$ to estimate the conditional entropy of the token at position $i$. Unmasking the positions with the lowest conditional entropy provides a simple heuristic that approximates the mutual-information criterion without additional training or external estimates. We will add a remark on this point in the revised manuscript.
• The original “diffusion model approaches” means using diffusion models for language models. We have changed it to “diffusion language models” to avoid this potential misunderstanding.
• Reversal reasoning is one key practical motivation of diffusion language models (see e.g., [1]). This task spans many real-world scenarios, including cloze/fill-in-the-blank problems, and span infilling. For example, literature examples often ask students to recover the first half of a poem from the second half; language exams routinely require predicting a verb given its object.
• Thank you for the suggestion. We have added the reference [2] to show that practical diffusion language models indeed mask more than one token in average at each iteration.

[1] Nie, S., Zhu, F., You, Z., Zhang, X., Ou, J., Hu, J., Zhou, J., Lin, Y., Wen, J.-R., and Li, C. (2025). Large language diffusion models. arXiv preprint arXiv:2502.09992.

[2] Yu, R., Li, Q., & Wang, X. (2025). Discrete Diffusion in Large Language and Multimodal Models: A Survey. arXiv preprint arXiv:2506.13759.

We appreciate the reviewer’s thoughtful comments and careful evaluation.

• Thank you for the suggestion. We agree that it would be beneficial to add an explicit remark on $\varepsilon_{train}$, which reads as follows. Our theory shows that the overall error in KL divergence exhibits a two-stage behavior: when the number of iterations $T<C_1\sum_i I(X^{(i)};X^{(-i)})/\varepsilon_{train}$ where $C_1$ is the constant from Corollary 1, the sampling error term dominates and the overall error decays proportionally to $1/T$; once $T$ exceeds this threshold, the error plateaus and becomes almost insensitive to further increases in $T$. Deriving a sharp scaling law for $\varepsilon_{\text{train}}$ itself would require new techniques to analyze the model capacity of transformers and the sample complexity of their nonconvex training, which is currently unavailable and beyond the scope of the present work.

• Thank you for the suggestion. We add a synthetic numerical experiment in the revised manuscript to validate our theory.

  • Model. we use a $K$-state Potts chain of length $L$ with coupling $J$ to model the distribution of text $X=(X^{(1)},\dots,X^{(L)})$. Concretely, $X^{(1)}\sim\text{Unif}([K])$ and $\mathbb P\\{X^{(i)}=y|X^{(i-1)}=x\\}=\exp(J1\\{x=y\\})/(\exp(J)+K-1), \quad \forall x,y\in[K].$ This construction allows us to compute explicitly the mutual information, the optimal mask predictor $p^\star(\cdot | X_t)$, and the densities of $p_{X_0}$ and $p_{Y_0|M}$.

  • Mask predictor. We use the explicitly computable optimal mask predictor $p^\star(\cdot | X_t)$ to run the sampling process.

  • Sampling error. Given the explicit densities of $p_{X_0}$ and $p_{Y_0|M}$, the expectation in the KL divergence, taken over both the mask schedule $M$ and the data distribution $p_{X_0}$, is approximated using Monte Carlo simulations.

  • KL divergence vs. mutual information. For fixed $T=10$, $K=10$ and $L=100$, KL divergence vs. mutual information is summarized in the following table:

    J	Mutual Information	KL
    0.000	4.80e-13	0.000
    0.333	1.18e+00	0.051
    0.667	5.53e+00	0.309
    1.000	1.43e+01	0.748
    1.333	2.86e+01	1.302

    One can verify that the KL divergence error increases linearly with the mutual information.

  • KL divergence vs. number of iterations $T$. With a balanced mask schedule, $K=10$, $L=100$ and $J=2$, $\log KL$ vs. $\log T$ is summarized in the following table.

    $T$	KL	$\log$ KL	$\log T$
    2	18.70	2.928	0.693
    4	8.85	2.180	1.386
    5	7.24	1.980	1.609
    10	3.10	1.130	2.303
    20	1.49	0.395	2.996
    25	1.16	0.151	3.219
    50	0.44	–0.818	3.912

    One can verify that the slope is very close to -1, demonstrating that the KL divergence error scales proportionally to $1/T$.

  • Conclusions. Therefore, these two results validate our established sampling convergence theory.

  In addition, because the true data distribution is unknown in the real-world datasets, it is infeasible to compute the token-level mutual information or the KL divergence between the output and the data distribution. Hence, we restrict our numerical experiments to the synthetic data.

• We appreciate this careful point. Lemma 1 is not ``missing” $\varepsilon_{train}$: the term $\varepsilon_{train}$ is fundamentally absent in Lemma 1 because this lemma analyzes the sampling error when the optimal mark predictor $p^\star$ is used. Error due to imperfect training is aggregated into $\varepsilon_{train}$ and handled in (17), which relates the distribution induced by the learned predictor $\hat{p}$ to that induced by $p^\star$. This separation decouples training from sampling, which allows Lemma 1 to focus on $KL(p_{X_0}\|p_{Y_0^\star|M})$. We will add a remark after Lemma 1 to make this point clearer.

• In the training stage, the mask predictor is parametrized such that tokens in the unmask set are predicted independently given the current context (see (4) and (5)). Therefore, the proof of Lemma 1 reflects this design: the gap between the factorized predictor and the true joint conditional over a block of tokens is controlled by conditional mutual information terms that precisely measure the statistical dependence among those tokens after conditioning on the revealed ones. This is why conditional mutual information naturally appears in the per-step error decomposition.

We thank the reviewer for the careful reading and the constructive feedback.

Weakness:

• Thank you for the suggestions. We will add plain-language explanations in the proofs. For instance, we will add the suggested summary for the KL divergence in (22); we will discuss the intuition for the KL divergence decomposition in (23); we will include an explanation for the conditional mutual information in (24).
  To aid navigation, we will also (i) restate the sampling update (6) and the factorized predictor (4), and (ii) remind readers of the masking sets $W_t$ and $D_t$ throughout the proof.

Questions:

• We appreciate this thoughtful point. A practical implication of our theory is to unmask, at each iteration, the tokens whose conditional dependence on the rest of the sequence is weakest. Concretely, (24) shows that the per-step contribution to the total sampling error is the conditional mutual information between a newly revealed token and the remaining tokens. Hence, selecting positions with smaller conditional mutual information minimizes that contribution. Moreover, because $I(X;Y|Z)\leq H(X|Z)$ for any random variables $X,Y,Z$, we can implement this strategy by ranking tokens by conditional entropy: at step $t$, use the learned mask predictor $\hat{p}_i(\cdot | Y_t)$ to estimate the conditional entropy of the token at position $i$. Unmasking the positions with the lowest conditional entropy provides a simple heuristic that approximates the mutual-information criterion without additional training or external estimates. We will add a remark on this point in the revised manuscript.