Remasking Discrete Diffusion Models with Inference-Time Scaling

We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns and questions below.

Concern 1: Clarify on ReMDM’s mathematical novelty

Although ReMDM and DDIM share the same high-level idea of non-Markovian diffusion processes, ReMDM enjoys the following differences from DDIM:

• ReMDM operates on masked diffusion, whereas DDIM was proposed for continuous diffusion, and in their appendix, they only showed applications to uniform noise discrete diffusion.
• The derivation of ReMDM allows us to easily reason about the bounds of the $\sigma_t$ value, making the strategy exploration more rigorous.
• ReMDM proposes several design strategies, further boosting the empirical performance.
• ReMDM connects itself with existing predictor-corrector samplers and showcases that existing discrete predictor-corrector samplers are ReMDM’s special cases.

Concern 2: Large masked diffusion models like LLaDA are slower than AR

While certain works, such as LLaDA, indeed lag behind comparably sized AR models, several commercial dLLMs, such as Inception Lab’s Mercury [1] and Google’s Gemini Diffusion, achieve a throughput of over 1000 tokens / sec, outperforming state-of-the-art AR LLMs with the same parameter size. These models demonstrate the great potential of masked diffusion, and ReMDM, as a general framework, can be applied to any type of masked diffusion model, alleviating a key shortcoming of these models.

Concern 3: The magnitude of the performance gains on LLaDA.

Below, we provide a new set of experiments on the popular math benchmark MATH500, where LLaDA with ReMDM achieves better pass@1 results (with high statistical significance) when compared with LLaDA-8B-Instruct as shown in the table below.

We report the 95% confidence intervals of 5 runs with different random seeds. The fact that LLaDA with ReMDM outperforms LLaDA beyond the margin of error on MATH500 further justifies ReMDM’s practical significance.

Concern 4: Typos and writing issues

We thank the reviewer for pointing these out. We will polish the paper’s writing in our updated manuscript.

• Clarifying failure to remask property: We will add explicit clarification that in masked diffusion models $\mathbf{z}_t$ can only take on two possible values: 1) the value of the original token $\mathbf{x}$ or 2) the masked token $\mathbf{m}$. Plugging in the values of $\mathbf{z}_t = \mathbf{x}$ into Equation 2 directly leads to the failure to remask property. Please refer to our response to Concern 1 for Reviewer ep2E for a detailed explanation.
• Adding context for Eq. 3: This equation presents the training objective for masked diffusion models that has been previously derived in works such as Sahoo et al. [2]. Given the space restrictions we were not able to reproduce this derivation in our main paper, but we will include a detailed re-derivation in an appendix section of our updated manuscript.
• Indexing of latent variables: In terms of the indexing of the latent variables, this is simply a matter of convention: one can either number the latent variables starting from $0$ or from $1$. The notation we use follows that of Sahoo et al. [2], where $\mathbf{z}\_0$ represents the first latent variable that comes directly after the clean data $\mathbf{x}$. $\mathbf{z}\_0$ does not exactly equal $\mathbf{x}$ because in masked diffusion models’ definition of the noise schedule, $\alpha_0$ approaches $0$ and only equals $0$ as $T$ goes to $+\infty$ (see the line below Eq.4 in Sahoo et al. [2]).
• Awkward terminology of "complete forward process": We thank the reviewer for pointing this out. Eq.5 represents the ReMDM’s decomposition of the joint distribution. We define the posterior in the main paper because it is used in the derivation of the NELBO. The forward distribution $q_\sigma(\mathbf{z}_t \mid \mathbf{z}_s, \mathbf{x})$, however, is not used in deriving the NELBO, so we put it into the appendix to help readers understand ReMDM better. To clean up the wording, we will change Line 89 into “3.1 ReMDM Probability Decomposition” and Line 90 to “We define the probability decomposition of ReMDM as:”.
• Typo in number of Eq 6 / 51: We thank the reviewer for pointing out the typo in Lines 92 and 97. We will change Eq (51) to Eq (6).
• Defining $i$ in Eq 8: As a common practice in masked diffusion literature, we used $s$ as an abbreviation for $s(i)$ that is defined in Line 73, and used $t$ as an abbreviation for $t(i)$ that is defined in Line 66. We made note of this notational convention in Line 68, but, per the reviewer’s suggestion, we will reiterate it when introducing Equation 8 in our revised manuscript.

Question 1: Report ReMDM-conf results in text generation

For the inference-time scaling setting ($T\geq L$), we combine the confidence-based schedule with ReMDM-loop and have the following results:

For the fast sampling setting ($T<L$), we combine the confidence-based schedule with ReMDM-cap and have the following results:

The sample quality of confidence-based samplers generally lags behind that of the reported ReMDM results in the main paper. We will include these confidence results in our ablation table (Table 6) in our revised manuscript.

Question 2: Why does $\sigma>0$ yield better Gen PPL.

Despite having the same name, $\sigma$ in ReMDM and $\sigma$ in DDIM do not mean the same thing. In Eq.59, we show the mapping relation between $\sigma^{ReMDM}$ and $\sigma^{DDIM}$. When $\sigma^{DDIM}=0$, the corresponding $\sigma^{ReMDM}$ is $1-\alpha_s$ but not 0. In fact, the core result of DDIM is that non-Markovian samplers outperform the original diffusion sampler when $T$ is small, and this result is consistent with ReMDM (shown in the lower half of Table 1).

Question 3: Comparison between ReMDM and LLaDA samplers

For the masked diffusion model baselines and ReMDM in our experiments, except the LLaDA ones, we used ancestral sampling. Concretely, for each token $\mathbf{z}\_t$ at time step $t$, we randomly sample it from the posterior distribution $p\_\theta(\mathbf{z}\_s \mid \mathbf{z}\_t)=q(\mathbf{z}\_s \mid \mathbf{z}\_t, \mathbf{x}\_\theta(\mathbf{z}\_t))$ (where for ReMDM we use $q_\sigma(\mathbf{z}_s|\mathbf{z}_t, \mathbf{x})$ defined in Eq.6).

• We inherit the ancestral sampling from Sahoo et al. [2], which makes sure there are no mask tokens in the final generated sequence. After the decoding process, $\mathbf{z}\_0$ obeys the following marginal distribution: $cat(\mathbf{z}\_0; (1-\alpha_0)\mathbf{x}+\alpha_0 \mathbf{\pi}) \approx cat(\mathbf{z}\_0;\mathbf{x})$. To further handle extreme cases where there still exist masked tokens in $\mathbf{z}\_0$, as in similar works like Sahoo et al. [2], we use a final noise removal step, in which unmasked $\mathbf{z}\_0$ tokens remain the same and masked $\mathbf{z}\_0$ tokens are decoded based on the model generated distribution $\mathbf{x}_\theta(\mathbf{z}_0)$.
• For $T>L$, we simply define $t(i)$ as $\frac{i}{T}$ and plug it into ReMDM’s posterior. Due to the continuous property of $\alpha_t$, ancestral sampling allows $T>L,$ unlike LLaDA’s MaskGiT style sampler.

Question 4: Ablation studies on $\sigma$ for ReMDM-cap and ReMDM-rescale

• In Figure 6, we reported different $\sigma$ results for ReMDM-cap on text generation, and in Figure 8, we reported different $\eta$ results for ReMDM-rescale on text generation.
• In Table 7, we reported different $\sigma$ results for ReMDM-cap and ReMDM-rescale on image generation.
• In Figure 10-25, we reported different $\sigma$ results for ReMDM-rescale on the QM9 guidance experiments.

The takeaway message is that a larger $\sigma$ tends to improve generative perplexity but hurts diversity. A good $\sigma$ value should find a balance in the tradeoff.

---

References

[1] Inception, et al. "Mercury: Ultra-Fast Language Models Based on Diffusion." arXiv preprint arXiv:2506.17298 (2025).

[2] Sahoo, Subham, et al. "Simple and effective masked diffusion language models." Advances in Neural Information Processing Systems 37 (2024): 130136-130184.

We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns and questions below.

Concern 1: Why does Eq. 2 imply the “failure-to-remask” property?

In the forward Markov chain process of masked diffusion models (defined by Eq.1), each single token only has two possible states:

1. It can remain unmasked, i.e., it remains as the original token $\mathbf{x}$, or
2. It can become masked, i.e., it transitions to the [MASK] token $\mathbf{m}$ from which it never changes again.

Therefore, in Eq.2, $\mathbf{z}_t$ only has two possible values: $\mathbf{x}$ or $\mathbf{m}$. When $\mathbf{z}_t$ is unmasked, i.e., when $\mathbf{z}_t = \mathbf{x}$, we can plug it into Eq. 2, and have $q(\mathbf{z}_s \mid \mathbf{z}_t = \mathbf{x}, \mathbf{x}) = cat(\mathbf{z}_s; \frac{\alpha_s - \alpha_t}{1 - \alpha_t}\mathbf{x} + \frac{1-\alpha_s}{1-\alpha_t}\mathbf{x}) = cat(\mathbf{z}_s; \mathbf{x})=cat(\mathbf{z}_s;\mathbf{z}_t)$.

This means that once $\mathbf{z}_t$ is unmasked, the posterior will keep it unchanged, i.e., the “failure-to-unmask” property.

Concern 2: Clarify the meaning of $q(\mathbf{z}_s \mid \mathbf{z}_t=\mathbf{x}, \mathbf{x})$

To clarify, this is an abbreviation inherited from the MDLM [1] paper and $q(\mathbf{z}_s|\mathbf{z}_t=\mathbf{x}, \mathbf{x})$ means the $q(\mathbf{z}_s|\mathbf{z}_t, \mathbf{x})$ value when $\mathbf{z}_t=\mathbf{x}$. As mentioned in the previous comment, there are only two possible values for $\mathbf{z}_t$, and the equation $q(\mathbf{z}_s \mid \mathbf{z}_t=\mathbf{x}, \mathbf{x})$ denotes one of the two cases.

Concern 3: Why distinguish between $\mathbf{z}_t=\mathbf{m}$ and $\mathbf{z}_t \neq \mathbf{m}$ in the posterior?

As mentioned in our comment to the first concern, there are only two possible values for $\mathbf{z}_t$, $\mathbf{x}$ (denoted as $\mathbf{z}_t \neq \mathbf{m}$) and $\mathbf{m}$ (denoted as $\mathbf{z}_t = \mathbf{m}$). Therefore, when we derive the posterior of masked diffusion models, we separate it into two cases and plug the two possible values of $\mathbf{z}_t$ into the posterior, significantly simplifying the formula (see Appendix A.2.1 in MDLM [1] for more details).

Concern 4: Usage of time steps in the posterior

To clarify, we do use time steps to guide the remasking process by adjusting the value of $\sigma$ according to $t$ (denoted as $\sigma_t$). In Section 4.2, we discussed different strategies for changing $\sigma_t$, and we found that adjusting $\sigma$ values according to the time steps helped improve empirical performance.

Concern 5: Compare with an entirely random remasking model

Entirely random remasking models should be a naive special case of ReMDM, where the $\sigma_t$ values for all t and all tokens are handcrafted. We provided detailed discussions in Section 4 on various design strategies for ReMDM. As shown in our ablation studies (Tables 6 and 7, Lines 280-281), carefully designed strategies consistently outperform naive strategies across various tasks.

As a further justification, we provide a new set of LLaDA results where. At each time step, the sampler randomly remasks one token and decodes two tokens using the original LLaDA confidence sampler. LLaDA with ReMDM outperforms the random remasking model beyond the margin of error.

Concern 6: Move more ablations to the main paper

We thank the reviewer for this suggestion. We have moved some of our ablation results to the appendix due to the page limit, and we plan to include more ablation studies in the main paper when more space becomes available in future versions.

Concern 7: Typos and paper formatting issues

We thank the reviewer for pointing out these mistakes. We will add an equation number to the NELBO formula and change Eq. (51) in Lines 92 and 97 to Eq. (6) in future versions. As for the experimental setup, we mentioned in the paper that we reused the pretrained checkpoint in Lines 202, 235, and 266.

Question: Clarify the claim for better controlled generation

By “better controlled generation”, we mean a better Pareto frontier of sample quality and the wanted feature. In the context of Figure 4, this refers to pushing out the curves to the upper right region. In Figure 4, ReMDM dominates all the other baselines, demonstrating that ReMDM has strictly better controlled generation than other guidance baselines.

The motivation for using ReMDM here is that standard masked diffusion models cannot change tokens once they are generated, and therefore are hypothesized to be less steerable. In contrast, ReMDM can avoid this issue and generate high-quality samples that better align with the desired conditioning property. Please refer to our comment on Reviewer oMMK’s Concern 2 for a more detailed discussion.

---

References

[1] Sahoo, Subham, et al. "Simple and effective masked diffusion language models." Advances in Neural Information Processing Systems 37 (2024): 130136-130184.

We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns below.

Concern 1: Theoretically justify why ReMDM is better than MDLM

We hypothesize that there are two reasons why ReMDM has better sample quality than MDLM.

• In practice, language models are not perfectly trained and can make mistakes. The failure-to-remask property of masked diffusion models makes it impossible to fix the errors generated by the model, thus making the sample distribution shift away from the target distribution as the generation process goes on. ReMDM, however, can remediate this problem by enabling remasking, since the number of token changes that ReMDM can make (denoted as $N$) can be larger than the sequence length $L$, while for MDLM, $N$ is always smaller than or equal to $L$ since predicted tokens cannot change, i.e., the failure-to-remask property.
• Even if the denoiser is perfectly trained, the parallel decoding nature of MDLM can lead to wrong token conjunctions (Figure 1, Leftmost), and ReMDM can remediate this problem. To formalize this, we consider the following example.

Example: Suppose we want to decode a two token sequence starting from an input “MM” (M denotes [MASK] token), and the vocabulary size is 2: A and B. “AB” and “BA” are correct sequences and both occur with equal probability in the true data distribution, while “AA” and “BB” are incorrect sequences, they have zero probability under the data distribution. We have a perfect denoiser, i.e. $x_\theta(A)=x_\theta(B)=0.5$, $x_\theta(AB \mid AM)=x_\theta(BA \mid BM)=x_\theta(AB \mid MB)=x_\theta(BA \mid MA)=1$. We use $T$ time steps and $\alpha_{t(i)}=\frac{i}{T}$.

Below we provide a table comparing the joint probabilities that one gets after $n$ steps of sampling from the MDLM posterior vs. performing $n$ sampling steps from the MDLM posterior followed by $m$ sampling steps from the ReMDM-loop posterior with a remasking probability of $\sigma$. These quantities can be computed by induction.

As long as $0<\sigma<\frac{2}{n+1}$, which is easy to achieve because for any $n<T-1$, $\frac{2}{n+1}$ is always smaller than $\sigma^{max}$ (defined in Eq.7), the incorrect sequence probability, i.e., $p_\theta(AA), p_\theta(BB)$ will monotonically decrease as $m$ increases, and the correct sequence probability, i.e., $p_\theta(AB), p_\theta(BA)$ will monotonically increase as $m$ goes up.

This observation justifies the two major claims of ReMDM: 1. ReMDM can improve sample quality 2. ReMDM benefits from inference-time scaling, i.e., the more ReMDM steps we take, the better masked diffusion model's sample quality will be.

Concern 2: Why is ReMDM better for guidance?

Note that guided sampling is equivalent to sampling from a distribution with unnormalized probability $p_\theta(x)p(y|x)^\gamma$ ($x$ for sampled data and $y$ for the conditioning class). Supported by our empirical results (Table 1, 2) and our comments on Concern 1, ReMDM provides a better model of $p_\theta(x)$. This, together with the fact that $p(y|x)$, which is either a separately trained classifier (classifier-based) or the denoising model itself (classifier-free), is shared across MDLM and ReMDM, explains why ReMDM provides a better model of $p_\theta(x)p(y|x)^\gamma$, i.e., better guidance.

Concern 3: Clarify on ReMDM’s NELBO

To clarify, the NELBO introduced in Line 112 is not actually used, but rather serves to justify our use of MDLM’s pretrained checkpoint. This was explicitly stated in Lines 116-118. Besides, in Table 4 in the Appendix, we also demonstrated that training a new model using the ReMDM NELBO as opposed to that of MDLM yields similar test perplexities, further justifying our use of the MDLM checkpoints (stated in Lines 118-119).

Concern 4: Clarify hyperparameter tuning of $\sigma_t$

As mentioned in the comment on Question 1, we do not use the ReMDM NELBO to train the model but directly reuse the MDLM checkpoints. This means that we do not conduct any extra training, and ReMDM is a purely inference-time algorithm. Therefore, in our current setting, $\sigma_t$ is a hyperparameter and is not trainable because there’s no gradient flow with respect to it, so we cannot conduct back propagation. We did conduct grid search on the $\sigma_t$ values, and the ablation results were shown in Figures 6-8, Table 7, and Figures 10-25. Regarding its optimal value, we are certain that it should be non-zero, as ReMDM with $\sigma_t=0$ degenerates to MDLM (stated in Lines 113-114).

Concern 5: Discuss additional related works

In Section 6, we had a paragraph (Lines 314-319) discussing existing discrete predictor-corrector samplers and their connection with ReMDM. As written in this section, predictor-corrector samplers that do not require extra training, like FB and DFM, were shown in Section 4.3 to be special cases of ReMDM. Other works, such as DPC [1] and the Stein Operator [2] (i.e., the paper linked by the reviewer), require training a separate model. Specifically, the Stein Operator work requires training a new masked diffusion model using a Hollow Transformer parameterization. In contrast, we demonstrated that ReMDM can be applied to pre-trained models and involves only inference-time modifications. This is a key benefit of our work.

Concern 6: Figure 4 is hard to read

We thank the reviewer for pointing this out. To make this figure more accessible, we will add an arrow pointing to the upright in both subfigures to show that curves in the upper right region demonstrate a better tradeoff between generation quality and the wanted feature in our revised manuscript. As a further justification, in Figure 4, we displayed the Pareto frontier of sample quality and the desired feature, and we see that the red curves (ReMDM) dominate the curves from other models. Therefore, we claim that the red curves (ReMDM) are strictly better than all the other guidance baselines. The experimental setup and presentation follow similar works, specifically Schiff et al. [3], that have investigated how diffusion models navigate this Pareto frontier.

---

References

[1] Lezama, Jose, et al. "Discrete predictor-corrector diffusion models for image synthesis." The Eleventh International Conference on Learning Representations. 2022.

[2] Zhao, Yixiu, et al. "Informed correctors for discrete diffusion models." arXiv preprint arXiv:2407.21243 (2024).

[3] Schiff, Yair, et al. "Simple guidance mechanisms for discrete diffusion models." arXiv preprint arXiv:2412.10193 (2024).

We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s specific concerns below.

Concern 1: Explicitly state that $\mathbf{x}$ represents a single token

We thank the reviewer for this suggestion. We will make this explicit in the revised manuscript, updating Line 64 as follows:.

…Formally, we define a single token $\mathbf{x} \in \\{0, 1\\} ^{|V |} \subset \Delta^{|V |}$ as a one-hot vector…

Concern 2: Define MAUVE score

We thank the reviewer for pointing this out. Per the reviewer’s suggestion, we will add a section in the Appendix that more formally defines the MAUVE score as in [1] in the revised manuscript. In our initial manuscript, we found that the definition of the MAUVE score was difficult to describe within the limited space of the main paper. Therefore, we referred readers to the original citation and provided an intuitive explanation for our choice of using MAUVE as a metric that balances sample quality and diversity.

---

References

[1] Pillutla, Krishna, et al. "Mauve: Measuring the gap between neural text and human text using divergence frontiers." Advances in Neural Information Processing Systems 34 (2021): 4816-4828.