We thank the reviewer for their valuable feedback. Below, we address the reviewer's questions and concerns.

I'm curious about the distribution of the effective k chosen by EB sampler (across various datasets and sampling steps). If it is easy to produce such plots, I encourage the authors to include them in the appendix.

Answer:

We thank the reviewer for the suggestion, we will add histograms of effective-k values to the appendix (Unfortunately we cannot attach those here as per the rebuttal text-only responses policy).

Relatedly, it would be awesome to include some qualitative examples where such effective-k varies significantly for the EB sampler. Such plots and examples will benefit future researchers with a better intuition of the model's prediction distribution and further solidify the paper's claim on two types of errors (local model error and joint correspondence error).

Answer:

We refer the reviewer to the response to reviewer uZV5 where we included an example of the sampling process for an example from the GSM8K benchmark. Again, due to visualization constraints in the rebuttal, we will only add those in the final version of the paper.

Section 3.2: For joint correspondence error, the standard transformer architecture implies the predictions are not truly independent, so I feel unsure about the independence claim. It would be good to provide some concrete examples to show the failure modes.

Answer:

An MDM produces a marginal distribution per masked token when evaluated at the current state (as described in section XXX). When we unmask multiple tokens simultaneously, the newly unmasked tokens are sampled independently of each other, but importantly conditional on the current state. The joint dependence error comes from failing to account for dependency in this state-conditional sampling.

We will provide qualitative text examples of generation with a large $gamma$ that shows failure modes due to sampling independently conditional on the current state.

Figure 3: The legends can be improved. It is not clear to me where the current plot shows various k.

Answer:

We will add a comment clarifying that the plot is for 512 tokens generation length. Hence, the points at 512 NFE correspond to $k=1$, at 256 NFE to $k=2$ etc.

Equation 2: appeared twice in both the sum and max, which is a little confusing.

Answer:

We will clarify this equation that the sum and max terms are separate, and that the max is not within the sum.

Why does top-k perform quite well on maze and sodoku? What data property is making the gain less prominent than math and reasoning benchmarks?

Answer:

One explanation might be that unlike text, where there is more than a single answer to a specific prompt, for Sudoku and mazes, given the initial state, the solution is unique. Hence, the perfect model should be able to solve these problems in a single function evaluation as all tokens are deterministic. We believe that this property allows top-k method to also work relatively well.

We thank the reviewer for their valuable feedback. Below, we address the reviewer's questions and concerns.

Entropy Threshold Sensitivity and Tuning Guidelines: How sensitive is the method's performance to the entropy threshold hyperparameter across different tasks and datasets?

Answer:

In all the plots for language modeling the thresholds were: $[0, 0.001, 0.01, 0.1, 0.25, 0.5, 0.75, 1.0, 1.5]$, where performance is roughly maintained until a critical threshold value is reached and performance begins to decline.

For all benchmarks LLaDa has an optimal $\gamma$ between $[0.25,0.5]$ while Dream’s optimal $\gamma$ is in the range of $[0.1,0.25]$, indicating that the threshold is more model dependent rather than task dependent.

In the revised version of the paper we could add plots with $\gamma$ as the x-axis to better present its sensitivity and behaviour across datasets and models if the reviewer finds it to be helpful.

NFE Control and Predictability: Unlike top-k sampling where the relationship between k and number of function evaluations (NFE) is direct, EB-Sampler's NFE depends on the adaptive entropy-based decisions. How can practitioners predict or control the computational budget when using your method?

Answer:

From the figures, any small value of entropy threshold tends to work very well and performance across the number of steps tends to be flat until the threshold gets large enough. We further observe that the optimal threshold is roughly the same across benchmarks for each model, but the obtained NFE speed up varies. This implies that the required NFE budget may be task dependent, so our suggestion for practitioners is to gradually increase the threshold until performance starts to drop. Given the observed monotonic nature with respect to gamma, tuning this single hyperparameter is a simple process.

Similarly to adaptive solvers for other applications (such as ODEs), the adaptivity enables error control rather than NFE control.

Limited Theoretical Analysis: While the authors decompose the total error into model error and joint dependence error, the theoretical framework only addresses how to reduce the joint dependence error through entropy-bounded unmasking. It remains unclear how their approach mitigates the model error component, which represents the fundamental limitation of the underlying masked diffusion model's predictions, potentially leaving a significant source of error unaddressed.

Answer:

Our approach relies upon model error proxies from past research, see Figure 2 in our paper, where it is shown that different model error proxies result in different performance. If improved model error proxies are introduced in the future, they can be easily swapped into EB-Sampler.

In Figure 9-12, What is the difference between the top-k and margin methods represented by solid and dashed lines in the legend?

Answer: It is a typo in the legend, solid line is always EB-Sampler and dashed lines are Top-k. We will fix it in the revised version.

Lack of Sampling Examples: The paper focuses heavily on quantitative metrics but provides no actual text samples or sampling process generated by the EB-Sampler.

Answer:

We will add actual text samples in the revised manuscript. Due to text-only replies in the rebuttal policy we write out an example of the sampling process from the GSM8K dataset, sampling 40 tokens in 17 function evaluations:

############ PROMPT ############
Q: Jean has 30 lollipops. Jean eats 2 of the lollipops. With the remaining lollipops, Jean wants to package 2 lollipops in one bag. How many bags can Jean fill?

A:

############ GENERATION ############


--------- Step 0 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 1
[' ']
<|mdm_mask|><|mdm_mask|> <|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 1 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 2
['0']
<|mdm_mask|><|mdm_mask|> <|mdm_mask|>0<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 2 ----------
unmasked 2 tokens in this step
total number of unmasked tokens: 4
['3', ' ']
<|mdm_mask|><|mdm_mask|> 30<|mdm_mask|> <|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 3 ----------
unmasked 4 tokens in this step
total number of unmasked tokens: 8
['2', ' =', ' ', '2']
<|mdm_mask|><|mdm_mask|> 30<|mdm_mask|> 2 = 2<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 4 ----------
unmasked 3 tokens in this step
total number of unmasked tokens: 11
['8', ' l', 'ops']
<|mdm_mask|><|mdm_mask|> 30<|mdm_mask|> 2 = 28 l<|mdm_mask|>ops<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 5 ----------
unmasked 3 tokens in this step
total number of unmasked tokens: 14
[' has', ' -', 'ollip']
<|mdm_mask|> has 30 - 2 = 28 lollipops<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 6 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 15
[' Jean']
 Jean has 30 - 2 = 28 lollipops<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 7 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 16
['.']
 Jean has 30 - 2 = 28 lollipops<|mdm_mask|>.<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 8 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 17
[' left']
 Jean has 30 - 2 = 28 lollipops left.<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 9 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 18
[' ']
 Jean has 30 - 2 = 28 lollipops left.<|mdm_mask|><|mdm_mask|><|mdm_mask|> <|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 10 ----------
unmasked 2 tokens in this step
total number of unmasked tokens: 20
['2', ' ']
 Jean has 30 - 2 = 28 lollipops left.<|mdm_mask|><|mdm_mask|><|mdm_mask|> 2<|mdm_mask|><|mdm_mask|> <|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 11 ----------
unmasked 7 tokens in this step
total number of unmasked tokens: 27
['8', ' /', '2', ' =', ' ', '1', '4']
 Jean has 30 - 2 = 28 lollipops left.<|mdm_mask|><|mdm_mask|><|mdm_mask|> 28 / 2 = 14<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 12 ----------
unmasked 2 tokens in this step
total number of unmasked tokens: 29
[' can', ' bags']
 Jean has 30 - 2 = 28 lollipops left.<|mdm_mask|> can<|mdm_mask|> 28 / 2 = 14 bags<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 13 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 30
[' Jean']
 Jean has 30 - 2 = 28 lollipops left. Jean can<|mdm_mask|> 28 / 2 = 14 bags<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 14 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 31
[' answer']
 Jean has 30 - 2 = 28 lollipops left. Jean can<|mdm_mask|> 28 / 2 = 14 bags<|mdm_mask|><|mdm_mask|> answer<|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 15 ----------
unmasked 8 tokens in this step
total number of unmasked tokens: 39
['.', ' The', ' is', ' ', '1', '4', '.', '<|endoftext|>']
 Jean has 30 - 2 = 28 lollipops left. Jean can<|mdm_mask|> 28 / 2 = 14 bags. The answer is 14.<|endoftext|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>
---------------------------

--------- Step 16 ----------
unmasked 1 tokens in this step
total number of unmasked tokens: 40
[' make']
 Jean has 30 - 2 = 28 lollipops left. Jean can make 28 / 2 = 14 bags. The answer is 14.<|endoftext|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|><|mdm_mask|>

We thank the reviewer for their valuable feedback. Below, we address the reviewer's questions and concerns.

The overall novelty seems to be a little bit limited as it is a simple modification of the existing top-$k$ method.

Answer:

While the EB-Sampler is a simple modification to existing top-$k$ methods, it is new in MDM sampling literature and it provides unprecedented strong efficiency improvements at zero cost. Simplicity that brings strong improvements should be considered as a strong suit rather than a disadvantage.

We would also like to highlight the contribution of the mathematical framework we introduce that covers many other adaptive multitoken sampling schemes for (time independent) masked discrete diffusion, extending upon the adaptive MDM inference from Kim et. al. (2025).

For the theoretical part, the reviewer is wondering whether it is possible to perform some analysis on the proposed EB sampler to see how the final discrepancy between the target distribution and the distribution of the generated sample scale with respect to the discretized timestep. The authors might consider referring to the stochastic integral framework proposed and used in [1,2] to see how the techniques developed in these two papers may be applied to analyze the proposed EB sampler to better understand the its effectiveness theoretically. Other references that the reviewer thinks might be helpful are [3,4,5]. Specifically, [4] provides the first theoretical analysis of masked discrete diffusion models to the best of the reviewer's knowledge. Would it be possible for the authors to at least cite these papers and briefly comment on how the EB sampler might be analyzed theoretically while revising the manuscript?

Answer:

We appreciate the suggestions from the reviewer and will add citations to these papers, where we note [2] is already cited in related work.

We quantify the discrepancy between the target distribution and the model distribution via the KL divergence in Section 5, proving the decomposition into model error and joint dependence error in Appendix A. Unlike more general continuous-time discrete diffusion, we have no explicit time discretization error from our sampling scheme. Our MDM sampling dispenses with time entirely as the optimal solution for MDM predictions are clean data conditionals, and time-independent, as described in Section 2.2 (see Ou et. al. (2025), Zheng et. al. (2024)). Naturally many MDM’s, including LLaDa and Dream, are parameterized without time given this optimal solution. That time is not required for MDM is embodied within our framework of adaptive multi-token sampling in Section 5 that considers state transitions via unmasking as opposed to following the state of a discrete diffusion process over time. This framework is expressive enough to contain existing (top-k) as well as novel adaptive multi-token MDM sampling schemes, including EB-Sampler, while avoiding heavyweight mathematical machinery needed for more general continuous-time discrete diffusion.

EB-Sampler only bounds the KL divergence in this framework if there is zero model error amongst the unmasked tokens. A theoretical understanding of how EB-Sampler would behave under inaccurate model predictions could likely benefit from the more advanced techniques introduced in these references.

For the empirical part, on the other hand, the reviewer thinks it might still be necessary to include extra baselines like Autoregressive Models (ARMs) in Figure 5,6 and 7 as they are more commonly used nowadays in the tasks included in this paper. Would it be possible for the authors to add ARM as a baseline method for comparison?

Answer:

Our paper focuses on accelerating MDM sampling for SOTA large language diffusion models like LLaDa and Dream that have already been shown to be competitive with similarly sized ARM models in terms of performance on these tasks.

Would adding a point to the figures showing performance of similar sized ARMs sampled with full-NFE next-token prediction be helpful? If so, we will add it in the revised version.

Finally, there are also a few miscellaneous problems regarding the presentation and typos of the manuscript. For instance, both Figure 10 and 11 contain two plots for the same method "Top- , margin" - may the reviewer ask whether this is a typo or not? Also, the notation "arg Top2" used in line 113 might need to have some text-based explanation as this is not a commonly used notation.

Answer:

Regarding Figures 10,11,12 - We thank the reviewer for pointing this out. It is indeed a typo in the legend, we will fix it in the revised manuscript. To clarify, solid lines are always EB-Sampler and dashed lines are Top-k.

We will also add an explanation for the $\arg \mathrm{Top}_k$ notation.

We thank the reviewer for their valuable feedback. Below, we address the reviewer's questions and concerns.

But I think the paper would benefit by including some baselines, based on using thresholds on metrics such as confidence, entropy and margin. The baseline could be the following:

1. Unmask all tokens having confidence greater than a threshold
2. Unmask all tokens having entropy less than a threshold
3. Unmask all tokens having margin greater than a threshold

I think the contribution of the paper is important, but these baselines would be useful for understanding why the method works.

Answer:

We appreciate the suggestion to explore additional members of the $\phi$ family (eq. (5) in paper). This threshold strategy, specifically for confidence, was introduced in concurrent research Wu et. al. (2025). As described in our paper, these alternative $\phi$'s can also leverage pre-trained masked diffusion models based on our multi-token adaptive sampling framework. We found that the suggested alternative $\phi$ performed similarly to EB-Sampler on MBPP but worse on GSM8K. In particular, for similar performance, EB-Sampler enables sampling more tokens on GSM8K.

For the different models we report NFE / Accuracy% in the table below for the optimal $\gamma$. For LLaDa we use confidence and for Dream we use entropy (see Tables 1,4,5,6 in paper).

In the revised manuscript we will incorporate these results as plots for all choices of error proxies for both models.

An interpretation of these results is that, even if model error is non-existent, an absolute threshold on token entropies/confidence could allow for an arbitrary amount of error to be introduced in sampling. Thresholding the entropy sum instead controls the overall error. This difference is more relevant when there is more variation in the token entropies. If all entropies were roughly the same, then an absolute threshold would behave similarly to an entropy sum threshold, which we observe does not occur on GSM8K.

How does the optimal threshold gamma for the method in the paper scale with the vocabulary size or the generation length?

Answer:

Generation length: We experimented with semi-autoregressive generation with varying block lengths for MBPP and GSM8K for both LLaDA and Dream. Results are in the tables below.

We observe that LLaDa maintains a stable optimal $\gamma$ values for different block sizes, while Dream seems to be more sensitive, requiring lower $\gamma$ values for smaller blocks on MBPP but higher $\gamma$ values on GSM8K.

We believe that the major factors affecting the optimal threshold values are the model itself and the task when blocks are very short. For large blocks such as 256 or 512, the optimal threshold seems stable across datasets for each of the models.

MBPP, LLaDa, semi-autoregressive generation

MBPP, Dream, semi-autoregressive generation

GSM8K, LLaDa, semi-autoregressive generation

GSM8K, Dream, semi-autoregressive generation

Vocabulary size: For the evaluated language MDMs in the paper, we do observe a difference in the optimal $\gamma$ value (LLaDa - $\gamma=0.25$, Dream - $\gamma=0.1$). Although LLaDa and Dream have different vocabulary sizes (126k and 151k, respectively), this alone is insufficient to determine the effect of vocabulary size. These models were trained differently, using distinct data and context lengths. A computationally intensive experiment involving training multiple models in a controlled environment would be necessary to isolate the influence of vocabulary size on the optimal $\gamma$.

Also, should it depend on some average entropy level of the predictions?

Answer:

We do observe that for the two language MDMs used in the paper the empirically observed optimal $\gamma$ is different: $\gamma=0.25$ for LLaDa and $\gamma=0.1$ for Dream, but they are consistent across benchmarks. The reviewer is correct though that potentially one could improve our results further by learning state-dependent entropy thresholds for EB-Sampler or more generally state-dependent $\phi$ functions as we discuss in the conclusions section. Tuning more than one parameter would require more substantive effort though, and we leave learning possibly even stronger adaptive unmasking schemes as extensions for future research.