Remasking Discrete Diffusion Models with Inference-Time Scaling
Authors: Guanghan Wang
Yair Schiff
Subham Sekhar Sahoo
Volodymyr Kuleshov
Venue: NeurIPS 2025 poster ()
Keywords: discrete diffusion
diffusion language models
diffusion models
diffusion samplers
probabilistic generative models
Primary area: deep_learning
TL;DR: Absorbing state discrete diffusion models that enable remasking benefit from inference-time scaling, improving sample quality, controlled generation, and performance on downstream tasks.

Abstract: Part of the success of diffusion models stems from their ability to perform iterative refinement, i.e., repeatedly correcting outputs during generation. However, modern masked discrete diffusion lacks this capability: when a token is generated, it cannot be updated again, even when it introduces an error. Here, we address this limitation by introducing the remasking diffusion model (ReMDM) sampler, a method that can be applied to pretrained masked diffusion models in a principled way and that is derived from a discrete diffusion model with a custom remasking backward process. Most interestingly, ReMDM endows discrete diffusion with a form of inference-time compute scaling. By increasing the number of sampling steps, ReMDM generates natural language outputs that approach the quality of autoregressive models, whereas when the computation budget is limited, ReMDM better maintains quality. ReMDM also improves sample quality of masked diffusion models for discretized images, and in scientific domains such as molecule design, ReMDM facilitates diffusion guidance and pushes the Pareto frontier of controllability relative to classical masking and uniform noise diffusion. When applied to large pretrained diffusion language models, ReMDM boosts the model’s performance on downstream tasks requiring factual knowledge grasp and reasoning ability.

Remasking Discrete Diffusion Models with Inference-Time Scaling
Guanghan Wang*t
Yair Schifft
Subham Sekhar Sahoo Volodymyr Kuleshov
Cornell Tech, Coreell University {gw354,yss2, sss284, vk379)@cornell edu denotes corresponding author denotes equal contribution
Abstract
Part of the success of diffusion models from their ability perform iterative refinement, i.e., repeatedly correcting outputs during generation. However, modern masked discrete diffusion lacks this capability: when a token is generated, it cannot be updated again, even when it introduces an error. Here, we address this limitation by introducing the remasking diffusion model (ReMDM) sampler, a method that can be applied to pretrained masked diffusion models in principled way and that is derived from discrete diffusion model with a custom remasking backward process. Most interestingly. ReMDM endows discrete diffusion with form of inference- time compute scaling. By increasing the number of sampling steps, ReMDM generates natural language outputs that approach the quality of autoregressive models, whereas when the computation budget limited, ReMDM better maintains quality. ReMDM also improves sample quality of masked diffusion models for discretized images, and in scientific domains such as molecule design, ReMDM facilitates diffusion guidance and pushes the Pareto frontier of controllability relative to classical masking and uniform noise diffusion. We provide the code along with blog post on the project page: https: attps://remdie: gi thub.
Introduction
Diffusion models have gained significant traction algorithms for generating high-quality images and videos [50, 52, 18, 42]. Part of the success of diffusion stems from its ability to perform iterative modifying outputs and fixing errors over multiple steps of generation- makes diffusion models especially effective at guided generation [9, 17] and fast sampling [51, 44, 54] and supports inference-time scaling to improve sample quality [30, 49, 66).
refinement-repeateee
Discrete counterparts of diffusion models [2] have also been steadily improving in quality on tasks such as language modeling and biological sequence design [29. 41, 46, 1. 43]. However, modern discrete cnffussion-espcciallly the most performant kind that relies on masking [32, 41, 48k-lacks the fundamental ability to iteratively refine outputs. Once a discrete token is unmasked, it cannot be updated again, even introduces an error. The lack of error correction in turn sets limits on controllable generation, sampling speed, and sample quality.
Here, address the inability of masked diffusion models perform iterative refinement by in- troducing a new sampler that supports remasking during generation. Our method, the remasking diffusion model (ReMDM) sampler, is simple and allows users to directly specify the probability of remasking a token at each time step. We augment the sampler with components that range from nucleus sampling to remasking schedules, significantly boosting performance.
Our approach for deriving the sampler is rooted in probabilistic modeling. We show that our method corresponds to ancestral sampling in discrete diffusion model (also called ReMDM) characterized
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

ReMasking Diffusion Models
She
sells
sea
shells
Inference-time
scaling
MAUVE (
AR: 0.760
0.656
She
[Mask]
sea
shells
Improved
sample quality
0.243 0.294
She
sell
sea
shells
Better controlled generation
[Mask]
[Mask]
sea
shells
0.009
0.042
SEDD
MOLM
MOLM
MOLM
RoMDM
DDF
+FB
Figure 1: Our family of masked diffusion processes allow for more flexible generation with remasking of already decoded tokens. This improves sample quality and further closes the gap to AR models. (Left) An illustrative example of errors fixed by ReMDM. The first two tokens can be "They sell" or "She sells", but due to the independence of the parallelized decoding processes, "She sell" is decoded. Such mistakes can be corrected by remasking samplers. (Right) MAUVE scores on Open WexTet. MDLM is from Sahoo et al. [41]. FB and DFM denote the forward-backwart [5] and discrete flow matching [13] correctors, respectively.
by remasking backward process. The ReMDM model admits objective that is rescaled version of the MDLM objective [41]. This suggests using the ReMDM sampler on top of pretrained MDLMs, which we find work well. We complement our analysis of the sampler by also showing that it can be interpreted as predictor-corrector technique [53, 5J. Notably, we demonstrate that ReMDM is theoretically more general than previous corrector samplers for masked diffusion models [5, 13).
Most interestingly, the ReMDM sampler endows discrete diffusion with a form of inference-time compute scaling. By increasing the number of denoising steps on language modeling tasks, ReMDM generates samples with significantly higher sample quality metrics than any previous diffusion model and almost matches the performance of autoregressive (AR) model with the same architecture. Conversely, when we reduce the sampling steps to increase speed, ReMDM's performance degrades less than other samplers. On image generation, ReMDM scales better with inference-time compute than vanilla masked diffusion models like MaskGiT [7]. Additionally, ReMDM models are more amenable to guidance [46]-in molecule design experiments, ReMDM pushes the Pareto frontier of novelty and the desired molecular property relative to alternatives relying on masking or uniform noise diffusion.
Contributions We summarize our contributions as follows:
1. We introduce the remasking diffusion model (ReMDM) sampler that yields performant iterative refinement via remasking to masked diffusion models.
2. We show that our method is a form of ancestral sampling in a probabilistic model whose ELBO is similar to that of classical masked diffusion. This analysis suggests using our sampler on top of pretrained models, which we find to work well.
3. Across the domains of natural language, discretized images, and molecule string representa- tions, we demonstrate empirically that ReMDM endows masked diffusion with inference- time scaling that improves sample quality with more computation and enhances controlled generation and downstream performance.
2 Background
Discrete Diffusion Models Diffusion models [50, 18] are characterized by parametric models pe trained to reverse fixed noising process that maps clean data x to increasingly noisy latent variables z,, for & [0, 1]. Austin et al. [2], Sahoo al. [41], Shi al. [48] extend this framework discrete signals. Formally, we define x E {0, 1I AIV one-hot vector representing a token, where |V| is the vocabulary size and AIVI is the simplex over this vocabulary. For finite-time processes, let T be the number of time steps. For each time step t(i) & E [0, 1] (i.e., for i 0, 1,1,.... T)),
each intermediate latent variable Z((s) E {0, 1) VI has the marginal:
q(z: x)
Cat(zz axx + a)) e AV
(I)

---

where π is a user-specified prior and where we have dropped the explicit dependence of t on i. The predefined schedule a E [0, is monotonically decreasing in t. For sequences of L tokens, we denote clean / noisy sequences as x(I:L. i (1:L) and each token & E fl..... .......]]... L} as x(() z
Masked Diffusion Language Models A special case of this formulation, known 'masked' or "absorbing state' diffusion, sets the limiting distribution π m, one-hot representation of a special [MASK] token. With s(i) T the time step directly preceding I, these posteriors take the form:
O2 at 1 Ckt
as a22
(zs Zz, x) = Cat Ze:
x
Following Sahoo et al. [41], we omit the subscript (t) in s(t) and t(i) in the rest of this paper. Importantly, despite the recent success of MDLM [41] and other absorbing state discrete diffusion models, these models suffer from a key drawback, which we call the failure to remask property. Namely, the posterior formula in (2) implies that any unmasked token 2, throughout the entire denoising process, Formally, in masked diffusion models, equation I ensures that any token x only possesses two possible values: unmasked, i.e., keeping unchanged. and masked, i.e., m, which means that in the reverse process, for any token & m, for all 8 S t, plugging 2, into equation leads q(z, 22 x) Cat(z,i 2.). Intuitively, once token is decoded, this prediction is 'locked-in', a limitation that these diffusion models share with AR models.
must remain unchanged
As in Sahoo al. [41], denoising neural network x9 is used for parameterizing (z. (z.)
(z. Pt. Xo xz())) and trained using the following variational objective (see the derivation in Sahoo al. [41]):
C Exzo-ggzoll (x0lx)) Po zo)
04 ass log((x Tx) x)
1)
(3)
az
3 Remasking Diffusion Models
In this work, we alleviate the failure remask limitation that affects ofte-off-hheart masked diffusion models. Our strategy is to design a probabilistic model similar to MDLM [41]. but whose denoising posterior generalizes (2) and supports remasking.
Specifically, we will define a discrete diffusion model with posterior for z, * m given by:
qz=
Z, x) = (1 o, o.m.
(4)
The parameter gives us the flexibility of remasking decoded token during generation. Because of this property, we dub our method ReMasking Diffusion Models (ReMDM).
3.1 ReMDM Probability Decomposition
We define the probability decomposition of ReMDM as:
lo (Zo:1 x) = yo (z) x) II
do (Zg 22, x)
(5)
i=1
where yo (z) | x) = =q(z)) = m, as in masked diffusion. We construct the posteriors as:
Cat( (1 01)xx J(m), Ze 7 m Cat(zz; a,-(1-o,)a; x 1-a.-a,e0 m), 1-au
ao (2, x)
(6)
1-ai
2,
Observe how (6) maintains the remasking property laid out in (4). Additionally, our chosen form for g(z, 2, = m, x) ensures that the marginals for (z, x) are the same as classical masked diffusion (e.g., Sahoo et al. [41]). This property will yield an ELBO for ReMDM that is very similar to that MDL.M, and will help us argue for using the ReMDM sampling procedure on top of
pretrained MDLM model. Formally, we establish the following result (proof in Appendix A.1):
Theorem 3.1. Given the posterior in (6), the marginal distribution go (z: x) is the same as in (1).
3

---

In Appendix A.2, we demonstrate that this ability to remask necessitates that ReMDM is non- Markovian process (as in e.g., DDIM [51]). Intuitively, this is necessary because the latents z can transition from the masked state m back to and therefore the transition kernel needs to be conditioned on x.
Understanding the Role of o To ensure that the posterior in (6) is valid probability distribution, we place following constraints on or (see derivation in Appendix A.3):
f o, S min 1,
(7)
Further examining (6), we see that increases, we move mass away from 2,. Importantly, this means that when Z₁ is unmasked, or directly controls the probability of remasking. When o; O, we recover the posterior from MDLM in (2). Thus, ReMDM can be seen more general formulation that admits MDLM as special case.
3.2 Negative Evidence Lower Bound
We define the joint distribution of the parameterized forward process po as follows:
Po Po |x) = (o (p)o))))
Zo) II
po , zz).
(8)
i=
We define (z( and parameterize pz
Zi)
qzz
Ze, x (z:)), where x²(z)) is
denoising model.
With this parameterization, the NELBO for ReMDM is (see Appendix A.4 for details):
C Eao-o(Rolx)) log po zo) + E f.....) E (x) ar di)a a. log(x; x)
(9)
Since Og S. ass, each term in the second expectation increases monotonically in 0,, When or = 0, for all f, we recover the MDLM objective from (3). This implies that the MDLM NELBO produces a tighter bound compared to ReMDM. Furthermore, since the diffusion loss term in (9) is simply reweighted version of that (3), one can presumably reuse weights from model trained with (3), i.e., using different o, at training and inference. Indeed, we find this to be performant strategy practice. Moreover, we find the performance (test perplexity) between models trained with the MDLM and the ReMDM objectives to be comparable (see Appendix E. 1), further justifying our reuse of pretrained weights, with the added benefit of more flexible sampling unlocked by ReMDM.
ReMDM Samplers
In practice, recommend using ReMDM with different o, training and inference, just one might switch to a different noise schedule during sampling. In particular, using a, = for training is equivalent combining the ReMDM sampler with pretrained MDLM model Xa, which works well and does not require re-training Xg.
Using ReMDM for sampling opens large design space for choosing o, In the following, we explore several practical design strategies for a, that significantly improve performance, and in Section we describe further add-ons to the sampler (e.g., nucleus sampling) that also improve quality. The high-level pseudocode for the ReMDM sampler provided Algorithm 1, with more detailed algorithms for implementing the schedules below deferred to Appendix C.
4.1 Design Strategies for the Remasking Schedule de
Here, we list several strategies for setting d; E 00,a)==], which can either be employed separately or in conjunction with the strategies introduced in Section 4.2. For sequences, each token can have
different (d) When we omit this superscript, this implies that the same a; used for all tokens.
Max- Cappe Schedule We can potentially reduce the maximum probability of remasking to constant Mepp E [0, 1]. Concretely, we let ar min(Neap: at schedule as "ReMDM-cap."
I-a, }. for all t [0, 1]. We denote this

---

Algorithm Sampling with ReMDM.
// Differences to standard MDLM sampling noted in brown.
Input: pre-trained denoising network Xg (e.g., MDLM), number of timesteps T, noise schedule remasking schedule
Initialize Z, m.
for =
= T'' to do
1 t= i= (i-
Set 02,, a, according to noise schedule.
Set o, E [0, o according to remasking schedule.
Compute approximate posterior:
Cat(zyi (1 0x)Xg + a.m),
2 # m
[Z zz) = (2.
Ze = =x(Zz)) =
Cat(Zsi 3a-(1-aa)as + 1-0 m))
T-ax
Sample V1 ~ Po-
Set 2, 25
end for Output: Zz.
Rescaled Schedule Alternatively, we can temper the chances of remasking by setting o, = Tresscle ot with Trescale [0, 1] as a hyperparameter that controls this rescaling. We denote this schedule as "ReMDM-rescale."
Confidence-Basee Schedule In conjunction with the two strategies above, explore further reweighing of o, which based on the intuition that tokens of which the denoising model less confident should be assigned larger probability of remasking. Consider the f-th token in sequence of latents at time f. For each e E (1,..., T.), we store its decoding probability at the time r at which it was last unmasked. Concretely, if z, j m, then define 1/2 = x 27
(f)T (E) If m.
then ( = oo. Thus, (1( serves as "confidence score' for unmasked tokens. We then compute
( (2) ot Toons where
(E) Toon
exp(- (-1)
E exp(-v!()
With this schedule, masked tokens are decoded using the approximate posterior from MDLM, and the unmasked tokens are remasked negatively proportional to their confidence. We denote this schedule as "ReMDM-conf."
4.2 Design Strategies for 'Turning On' ReMDM
There may be certain periods of the generation process where remasking is more valuable and some when is detrimental can slow down sampling, e.g., at the beginning of sampling, one may wish to generate some tokens in sequence using standard MDLM decoding, and only after generating reasonable starting candidate, then spend some computational budget 'fixing mistakes' via the remasking posterior. We, therefore, propose two methods for optionally 'turning on/off" ReMDM sampling, which amount to the following modification of the a, schedules above:
ot ift E [ton: tofs), with for > tofs
10,
ã.
(10)
0, otherwise.
Switch We choose some tswitch E (0, 1] and in (10) we have [ton toss) = [tswitch 0). We denote this strategy as "ReMDM-switch."
Loop In this strategy, we set both fon; tors E (0, 1]. Furthermore, in the range when ReMDM is activated, we modify the noise schedule to be constant, such that atonn). As shown in Figure 2, the first phase, the model generates tokens
this divides the sampling process into three phases. without remasking (a, = 0, i.e., using MDLM). In the second phase, we hold constant (i.e., a, = 04), and the model asked remask and predict a proportion of the generated tokens a loop. In particular, one can use any ReMDM strategy introduced in Section 4.1 in this phase. Intuitively, if
5

---

'bad' token is remasked, it will likely be replaced with a 'good' token due to the abundant context. Even if a 'good' token happens be remasked, since the signal-to-noise ratio designed to be high in this portion of the generation process, it will likely be re-decoded as other (if not the same) 'good' tokens. In this way, ReMDM-loop corrects the 'bad' tokens and maintains the 'good' tokens, i.e., fixes mistakes. Finally, in the third phase, we let the model predict any remaining unmasked tokens using the MDLM posterior. We denote this strategy as "RcMDM-loop."
4.3 Comparison with Discrete Predictor-Corretti Samplers
Previous works, e.g., the forward backward (FB; Camp- bell al. [6]) and discrete flow matching (DFM: Gat et al. [13]) correctors, propose to tackle the failure to remask property with discrete predictor-corrector sam- plers, special type of discrete diffusion samplers that decompose a single sampling step into one predictor step followed by certain number of corrector steps that reme- diate possible mistakes without changing the marginals. Here, we demonstrate that these methods cases of ReMDM sampling by first noting the following result (see Appendix A.5 for the proof):
- alpha
sigma
0
Value
ReMDM phase.
special
0
12
t1
Time (t)
Theorem 4.1. go (2, 2,,x) x) from (6) is equivalent to MDLM predictor step Upredictor (Zs, 21, x) followed by
Figure 2: ReMDM-llop illustration.
corrector step Qcorrector (Z.a Bay x) in the form of
Cat(z..i (1- -x))xx + J.m),
2sp +
(11)
Cat(Zaai T-a, + 1-(1+a.)a m), Zap m.
We now formalize the generality of ReMDM (proofs provided in Appendices A.6, A.7, and A.8). Proposition 4.2. The FB corrector [5] on MDLM is special case of ReMDM where a, 2,-04 ae
Proposition 4.3. A DFM corrector [13] MDLM can be converted to ReMDM sampler where 0: (. Be denotes the DFM corrector schedule.
Or
Proposition 4.4. The ReMDM sampler is more general than DFM corrector, since only the former can accommodate noise schedules with constant oi for some range [t, I. At) where At > O.
5
Experiments
0.6
0.5
5.1
ReMDM Improves Sample Quality
MAUVE
0.4
RcMDM nucleus
0.3
5.1.1 Text Generation
ReMDM nucleus
0.2
MDUM+nudleus
Experimental Setup We test the text generation capability of ReMDM with unconditional generation from models trained on WeeTTe (OWT; Gokaslan Cohen [14]).
0.1
MDLM
0.0
1024
2048
4096
Time Steps
The OWT dataset was tokenized using the gpt-2 tokenizer [38] and sequences were wrapped 1024. Our baselines include AR, SEDD [29]. MDLM,
max length of I,
Figure 3: ReMDM ablation for OWT
generation quality.
and MDLM in conjunction with the FB [5] and DFM [13]
correctors during inference. For all the experiments,
reuse the checkpoints trained by Sahoo et al.
[41] for fair comparison. Following Pillutla [36], we generate 5,000 samples from each model/sampler and compute the MAUVE score [28, 37], as this metric balances sample quality and diversity.
Generative Perplexity vs. MAUVE is common practice to report generative perplexity (Gen PPL.) as quality metric for text generation. However, for corrector samplers such as ReMDM, we find this metric can be uninformative.
line with previous work [68], we observe that one can get arbitrarily low Gen PPL.. by tuning corrector schedule hyperparameters (see Appendix E.2.2) to achieve low Gen PPL. while sacrificing sentence entropy/diversity, with an extreme case of generating sequence of repetitive words like "love love love..." and getting Gen PPL. of 1.5 (see Appendix F.I). Given this

---

phenomenon, we choose MAUVE as the primary quality metric since MAUVE measures the balance between generative perplexity and diversity (see Pillutla al. [37] for more detailed discussion).
Floating-Poin Point Precision As indicated in Zheng et al. [68], previous masked diffusion models [29, 41] report sampling with 32-bit floating point precision, which was found to significantly curb the diversity of generated samples. We therefore use 64-bit floating point for all text-sampling experiments.
Nucleus Sampling We observe that nucleus sampling [19] is critical to generating high-quality text sequences. We therefore use this strategy (with top-p = 0.9) for all models (except SEDD, it does not output logits).
Results In Table 1, we report results for inference-time scaling (T 2 L. and faster sampling (T L.). Note that we report Gen PPL. and entropy for completeness and that lower Gen PPI.. does not always correlate with better quality, as noted above. For inference-time scaling, we use the max-capped schedule (Mep = 0.02) in conjunction with the loop strategy (fon = 0.55, loff = 0.05, and a(ton) 0.9). For faster sampling, use the max-capped schedule (Neap 0.04) on its own.
As shown in Table 1, ReMDM scales favorably with inference-time compute, achieving a 15.62x MAUVE score compared to masked diffusion models and a 2.23 MAUVE score compared to MDLM with corrector samplers. In contrast, masked diffusion models scale poorly with T, and corrector sampler methods saturate when T is large. For faster sampling where more limited computational budget is available, ReMDM able to maintain sample quality better than baselines.
Ablation: Function of Each Component In Figure 3, we sequentially ablate each building block of ReMDM and report the MAUVE score of samples. Starting with MDLM, see that each of our proposed sampling improvements increases the MAUVE score, with the largest improvement coming from remasking ability of ReMDM.
Table I: ReMDM improves sample quality in inference-time scaling and faster sampling OWT. ReMDM outperforms state-of-the-an masked diffusion models (SEDD; [29]. MDLM; [41]) and masked diffusion models with corrector samplers such Heward-Backwarn (FB; [5] and Discrete Flow Matching (DFM; [13]) corrector samplers. indicates nucleus sampling. For each T, the best
diffusion MAUVE score is bolded.
Method
MAUVE (T)
Gen PPL. (L)
Entropy (1)
Data
1.00
14.8
5.44
(T=1024)
0.760
12.1
5.22
T=/024
T=2048
T T=4096
1024
T 2048
T T'4096
T'm 1024
T=2048
T T=4096
SEDD (absorb)
0.008
0.008
0.009
104.7
103.2
102.5
5.62
5.61
MDLM!
0.042
0.037
0.035
51.3
51.3
50.9
5.46
5.46
5.45
MDLM+FB
0.133
0.197
0.243
33.8
28.6
22.8
5.35
5.28
5.18
MDLM+DFM
0.254
0.294
0.269
21.7
21.0
20.7
5.20
5.19
5.17
ReMDM
0.403
0.610
0.656
28.6
22.8
17.6
5.38
5.30
5.20
T=/28
T=256
T=512
T=128
7=2256
7=512
T=
T=256
7=5122
SEDD (absorb)
0.007
0.007
0.008
119.2
110.1
107.2
5.65
5.63
5.62
MDLM!
0.015
0.023
0.031
53.0
5.52
5.49
5.48
MDLM+FB
0.064
0.084
0.100
42.8
39.6
37.1
5.44
5.41
5.38
MDLM-DFM
0.041
0.144
0.211
37.9
26.5
23.3
5.31
5.26
5.23
ReMDM
0.057
0.216
0.350
42.5
30.5
21.1
5.43
5.34
5.21
5.1.2 Image Generation
Experimental Setup We test ReMDM's class-conditioned image generation ability. Concretely, we use pretrained MaskGiT [7] that was trained on ImageNet [8] samples with 256 256 pixels. The images were patchified and flattened
sequence of L 256 and encoded according to codebook with 1024 tokens [10]. See Chang et al. [7] for more details about the model and training settings. We compare ReMDM's sampler to the original MaskGiT sampling and to MDLM. Note that MaskGiT does not follow an ancestral sampling method but rather relies on effective heuristic where masked tokens are decoded proportional model confidence. However, MaskGiT still restricted by the failure to remask property. For each sampler, we conditionally generate 50,000 images, with class

---

labels randomly chosen from the 1,000 ImageNet categories, and we measure sample quality using Fréchet Inception Distance (FID; [16]) and Inception Score (IS; [45]).
Results We report the best results for each method in Table (temperature=1.0 for MaskGiT, and temperature=0.8 for MDLM and ReMDMM.. We see that ReMDM has the best scaling, producing the highest quality images of the three methods T 64.
Table 2: ReMDM produces the highest quality images on ImageNet conditional generation. For each metric and T, the best value is bolded.
Sampler
T = 16
T
32
T 64
Ablation; Max-capped vs. Rescaled sched- ules In Appendix E.3, present the full ReMDM parameter search results. We find that ReMIMMrrsccle slightly outperforms ReMDM- cap, and that for both strategies, increasing their corresponding 7) parameters leads improved results.
MaskGiT
6.74
4.92
4.85
MDLM
7.88
5.37
4.69
FID
ReMDM
7.40
4.92
4.45
MaskGiT
155.32
181.57
196.38
(1
MDLM
140.97
169.79
187.93
ReMDM
145.27
182.05
209.45
5.2 ReMDM Improves Guidance
Experimental Setup We follow the setup from Schiff et al. [46] explore controlled small molecule generation. Specifically, we use the QM9 dataset [40, 39] of ~ 133k molecules and their character- based SMILES string representations [59]. Sequences were tokenized using regular expression method [47] and padded to maximum length of L = 32. We use the D-CFG and D-CBG methods defined in Schiff al. [46] to conditionally generate molecules with higher ring counts (greater than 90th percentile in the original dataset). We vary the guidance strength
{1, 2, 3, 4,5) and scale inference steps {32, 64, 128). Our baselines consist of an AR model guided using either CFG or the popular classifier-based FUDGE method [62], MDLM, and the uniform diffusion language model (UDLM) proposed Schiff et [46], which was also introduced to alleviate the no remasking limitation of masked diffusion. For ReMDM sampling. we use pretrained MDLM models. We explore the various strategies defined in Section (see Appendix E.4 for full search results). After generating 1,024 sequences, we use the RDKit library [23] to determine whether sequences valid and compute novelty of the generated sequences (number of unique valid sequences that do not appear in the original QM9 dataset) and mean ring count of novel sequences.
Results In Figure 4, we display the trade-off that comes from increasing guidance strength y- We only visualize results for samples that had at least 50 novel sequences. For both forms of guidance, CFG and CBG, ReMDM outperforms AR and diffusion approaches, pushing the novelty- property maximization frontier beyond that of the baseline methods. Additionally, ReMDM scales favorably with more inference-time compute, seen by the curves for larger T dominating those for smaller T. For ReMDM, (with Trescale = 0.9) and ReMDM-conf schedules. For D-CFG, we use the loop strategy, with fon 0.25, toff = 0.125. In addition to these findings, in Appendix E. we present experimental results for maximizing a different property interest, drug-likenen likeness (QED; [3]).
found the best strategy to be combination of the ReMDM-rescale
Molecule Property Maximization (Ring Count)
Malecule Property Maximization (Ring Count)
Guidance method:
Guidance methed:
AR
UDLM
5.6
MDLM
3.4
RcMDM
5.2
T= 128
T
Rina
5.0-
4.8
y=1
4.6
y=2
y=
4.4
4=4
42-
y=5
4.5
300
250
400
Num. Samples 1,024 generated)
Num. Novel Samples
1,024 generatedi
Figure 4: RcMDM improves stcer-ability by extending the novelty- maximization frontier. Lines show controlled generation for ring count maximization QM9 dataset, varying compute T and guidance strength 7. Curves in the upper right region are better. (Left) Discrete classifier-free guidance (D-CFG). (Right) Discrete based guidance (D-CBG) and FUDGE for AR.

---

Ablation: Importance of using Confidence-basedd Schedule Our parameter tuning revealed that using ReMDM-rrscale in conjunction with the confidence-basee scheduler led improved results (see Appendix E.4.3).
5.3 ReMDM Improves Benchmark Performance
Experimental Setup Recent diffusion large language models (dLLMs) such as LLaDA 8B [31] and Dream 7B [64] have demonstrated on par downstream task performance with autoregressive language models of the same parameter scale. However, these dLLMs still suffer from the failure to remask property. To test the effect of remasking sampling on dLLLMs, we take LLaDA 8B Instruct and apply DFM and ReMDM samplers (see Appendix D.4 for sampler details). Specifically, we use Countdown (bidirectional reasoning) [63] and TruthfulQA (factual knowledge grasp) [27] as benchmarking tasks. In Countdown, the model given four integers and asked to develop equations where first three integers produce the fourth using addition, subtraction, multiplication and division operations. In TruthfulQA, the model is asked 817 questions on factual knowledge. We calculate the ROUGE [26] scores between the model and correct answers, as well as the model and plausibly wrong answers. The difference between the right and wrong ROUGE scores is reported. For both of the tasks, we use max-length L 32 and no cemi-autoregressive sampling [31]. See Appendix D.4 & F.3 for more setting details.
Results We repeat experiments with multiple random seeds to report mean values and 95% confidence intervals. Table 3, LLaDA with ReMDM consistently performs the best and exceeds the original LLaDA with statistical significance.
Discussion and Related Work
Denoising Diffusion Im- plicit Models An analogy can be drawn between our work and DDIM [51]. where ReMDM is to MDLM for discrete signals as DDIM is to DDPM [18] for continuous data. Both our work and DDIM propose alternative for- ward and backward processes that maintain marginals while enabling more flexible generation. Of note, Song et al. [51] also present way of adapting processes to discrete domains; however, they focus uniform categorical as the limiting distribution. In Appendix B, we demonstrate how this can be adapted for absorbing state diffusion and derive an equivalence to our method.
Table 3: LLaDA with ReMDM performs the best on downstream tasks, and consistently exceeds the original LLaDA with statistical significance. For each metric, the best value is bolded.
Countdown pass @ 1%
TruthfulQA
Metric
A ROUGE 1/2/L
LLaDA LLaDA-DFM
45.2.0.2
27. 20.4 / 30. 104 27.2
44.8
28.2 /31.1 104
28..
46. +02
29.5 20.1 311. #0.1 29.5 =0.3
Discrete Predector-Corrector Samplers Continuous time Markov chain (CTMC) theory provides a framework for samplers that correct errors in the reverse process, extending the original predictor- corrector formulation for continuous data [53] to discrete domains. addition to the FB and DFM correctors discussed above, other correctors include: the Stein operator [67], and DPC [24] which improves the sample quality of MaskGiT by corrector sampling. Unlike the plug-and-ple methods of FB, DFM, and ReMDM, the Stein operator and DPC correctors require additional training.
7 Conclusion
In this work, we have presented novel family of absorbing state discrete diffusion samplers. Our method leverages the strong language modeling performance of this class of models by enabling the use of pretrained weights with the added benefit of more flexible sampling strategies that allow for remasking of predicted tokens. We demonstrate empirically that this leads to improved sample quality for both unconditional and conditional generation and leads to better controlled generation and downstream performance. Our approach also unlocks an important inference-tima compute scaling axis that is more limited for existing masked diffusion models.

---

Acknowledgments
This work was partially funded by the National Science Foundation under awards DGE-1922551, CAREER awards 2046760 and 2145577, and by the National Institute of Health under award MIRA R35GM151243. We gratefully acknowledge use of the research computing resources of the Empire AI Consortium, Inc, with support from Empire State Development of the State of New York, the Simons Foundation, and the Secunda Family Foundation. This research was also supported in part through the use of computational resources provided by Lambda (lambda. ai) in partnership with Open Athena AI Foundation, Inc. We gratefully acknowledge their generous GPU infrastructure grants that helped make this work possible.
References
[1] Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autoregressive and diffusion language models. arXiv preprint (xXix:5503.0577 2025.
[2] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 15:17981-17999 2021.
[3] G Richard Bickerton, Gaia V Paolini, Jérémy Besnard, Sorel Muresan, and Andrew L Hopkins. Quantifying the chemical beauty of drugs. Nature chemistry, 4(2):90-98, 2012.
[4] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milnh and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github jaajj//axaaa
[5] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266-28279, 2022.
[6] Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gener- ative flows on discrete state-spaces, Enabling multimodal flows with applications to protein co-design. Xrviv preprint arXhv:2402.04997 2024.
[7] Huiwen Chang. Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEFE/CVF Conference on Computer Vision and Pattern Recognition, pp. 11315-11323 2022.
[8] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248-255 Icee, 2009.
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-8794, 2021.
[10] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 12873-12883, 2021.
[11] William Falcon and The PyTorch Lightnnigg team. PyTorch Lightning, March 2019. URL https://gittuu: com http:/Ligtttnnnnggge www.gghtniinngAA/gggttiiiin
[12] Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL https://zenodo org/ records/ 10256836.
10

---

[13] Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407.1559, 2024. 2024.
[14] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. http://Skyyiono00/ .griggigg github.io! io/ OpeniebTextCorpus, 2019.
[15] Charles R. Harris, K. Jarrod Millman, Stéfan van der Walt, Ralf Gommers, Pauli Vir- tanen, David Counaaeaa,, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Hal- dane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre Gérard-Marchant, Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Arry programming with NumPy. Nature, 585(7825):357-362 September 2020. doi: 10. 00.1038/s41586-00 0206-000-2200 -22-6---20--..... URL https: : //toi//oiio//o org/ 10. 1038/s41580 0020-2699-20
[16] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[17] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv:2207. 12598, 2022.
arXiv preprint
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[19] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural 097501 190.0999,
text degeneration. arXiv preprint arXiv:
2019.
[20] J. D. Hunter. Matplotlib: A 2d graphics environment. Computing in Science & Engineering, (3):90 95, 2007. doi: 10.1 101100/CEE2000..555
[21] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 14, 2013.
[22] Volodymyr Kuleshov. Fast algorithms for sparse principal component analysis based on rayleigh quotient iteration. In International Conference on Machine Learning, pp. 1425. PMLR, 2013.
[23] Greg Landrum et al. Rdkit: A software suite for cheminformatics, computational chemistry, and predictive modeling. Greg Landrum, 831..1)) 2013.
[24] Jose Lezama, Tim Salimans, Lu Jiang, Huiwen Chang, Jonathan Ho, and Irfan Essa Dis- crete predictor-corrector diffusion models for image synthesis. In The Eleventh International Conference on Learning Representations, 2023.
[25] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing Systems, 35:4328-4343, 2022.
(26) Chin-Yew Lin. Rouge: A package for automatic evaluation of summaries. In Text summarization branches out, pp. 74-81, 2004.
[27] Stephanie Lin, Jacob Hilton, and Owain Evans. Truthfulqa: Measuring how models mimic human falsehoods. arXiv preprint arXiv:2109.07958, 2021.
[28] Lang Liu, Krishna Pillutla, Sean Welleck, Sewoong Oh, Yejin Choi, and Zaid Harchaoui. Divergence frontiers for generative models: Sample complexity, quantization effects, and frontier integrals. Advances in Neural Information Processing Systems, 34:12930-1299 2021.
[29] Aaron Lou, Chenlin Meng. and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Forty-first International Conference on Machine Learning, 2024.
[30] Nanye Ma, Shangyuan Tong, Haolin Jia, Hexiang Hu, Yu-Chuan Su, Mingda Zhang, Xuan Yang, Yandong Li, Tommi Jaakkola, Xuhui Jia, et al. Inference-time scaling for diffusion models beyond scaling denoising steps. arXiv preprint 2025. 2025.
11

---

[31] Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin. Ji-Rong Wen, and Chongxuan Li. Large language diffusion models. arXiv preprint arXiv:2502.09992,, 2025.
[32] Jingyang Ou, Shen Nie, Kaiwen Xue, Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint 24Xv::2406003736, 2024.
[33] The pandas development team. hundas-dev/pandars Pandas, February 2020. URL https:
/ woloi.org/1n 0001/22811//// 3509134.
[34]
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performancee Deep Learning Library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' Alché Buc, E. Fox, and R. Gamett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
[35] William Peebles and Saining Xie. Scalable diffusion models with transformers. Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
[36] Krishna Pillutla, mauve. https://github.ccom wttp:////rrssaaaa/////aaaaa 2021. Accessed: 2025-
05-14.
[37] Krishna Pillutla, Swabha Swayamdipta, Rowan Zellers, John Thickstun, Sean Welleck, Yejin Choi, and Zaid Harchaoui. Mauve: Measuring the gap between neural text and human text using divergence frontiers. Advances in Neural Information Processing Systems, 34:4816-488 2021.
[38] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 188):9 2019.
[39] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):1-7, 2014.
[40] Lars Ruddigkeit, Ruud Van Deursen, Lorenz C Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb- 17. Journal of chemical information and modeling. 52(1 22824-22777 2012.
[41] Subham Sekhar Sahoo, Marianne Arriola, Aaron Gokaslan, Edgar Mariano Marroquin, Alexan- der M Rush, Yair Schiff, Justin T Chiu, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: httpps:// wttp:/iwww...eeeeeeeeeeee met/foorum?iceeAAAAAAa..
[42] Subham Sekhar Sahoo, Aaron Gokaslan, Christopher De Sa, and Volodymyr Kuleshov. Dif- fusion models with learned adaptive noise. In The Thirty-eighth Annual Conference on Neu- ral Information Processing Systems, 2024. URL https: /ttp://///ppeeeeeeeeeeee net/forumidd 1oMa9944p8.
[43] Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin T Chiu, and Volodymyr Kuleshov. The diffusion duality. In ICLR 2025 Workshop on Deep Generative Model in Machine Learning: Theory, Principle and Efficacy, 2025.
[44] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint 2rXxi:2202.00512 2022.
[45] Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. Advances in neural information processing systems, 29, 2016.
12

---

[46] Yair Schiff. Subham Sekhar Sahoo, Hao Phung, Guanghan Wang, Sam Boshar, Hugo Dalla- torre, Bernardo de Almeida, Alexander Rush, Thomas Pierrot, and Volodymyr Kuleshov. Simple guidance mechanisms for discrete diffusion models. arXiv preprint arXiv:2412.10193 2024.
[47] Philippe Schwaller, Teodoro Laino, Théophile Gaudin, Peter Bolgar, Christopher A Hunter, Costas Bekas, and Alpha A Lee. Molecular transformer: A model for uncertainty-calibaaeed chemical reaction prediction. ACS central science, 5(9):1572-1583, 2019.
[48] Jiaxin Shi, Kehang Han, Zhe Wang, Araud Doucet, and Michalis K Titsias. Simplified and generalized masked diffusion for discrete data. arXiv preprint ar arXiv:2406.04:3 2024.
[49] Raghay Singhal, Zachary Horvitz, Ryan Teehan, Mengye Ren, Zhou Yu, Kathleen McKeown, and Rajesh Ranganath. A general framework for inference- tieeeneee-ieeeeeee scaling and steering of diffusion models. arXiv preprint arXiv:2501.0688 2025.
[50] Jascha DicksDicst,en, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsuper- vised learning using nonequilibrium thermodynamics. In International conference on machine learning. pp. 2256-2265. PMLR, 2015.
[51] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02500, 2020.
[52] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
[53] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 2rXiv:2011.13355 2020.
[54] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. arXiv preprint arXiv:2303.01469, 2023.
[55] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding, Neurocomputinng 568:127063, 2024.
[56] A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems,
2017.
[57] Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang. Christopher De Sa, and Volodymyr Kuleshov. InfoDiffusion: Representation learning using information maximizing diffusion models. Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Leurning Reseurch, pp. 36336- 36354. PMLR, 23-29 Jul 2023. URL https: / proppediigg...... mlr, press/ prress/v002/ v202/ 220///2222222222 wang23ah. .ww.mpesss
html.
[58] Michael L. Waskom, seaborn: statistical data visualization. Journal of Open Source Software, 6(60):0021, 2021. doi: 10.21105/joss.0000 03021.. URL attps://doi.org: wtt../........../////.../. 11s05/joss. 03021.
[59] David Weininger. Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules. Journal of chemical information and computer sciences, 28 (1):31-36, 1988.
[60] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's 's transform- ers: State-of-the-aa natural language processing. arXiv preprint 1xXi::1910.0377112200
[61] Omry Yadan. Hydra framework for elegantly configuring complex applications. Github, 2019. URL https: //gi thub. com/ waaaceboorreerrcchyyyaae.e
[62] Kevin Yang and Dan Klein. Fudge: Controlled text generation with future discriminators. arXiv preprint 21X1v.2104.05218, 2021.
13

---

[63] Jiacheng Ye, Jiahui Gao, Shansan Gong. Lin Zheng, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning, arXiv preprint arXiv:2410. 14157, 2024.
[64] Jiacheng Ye, Zhihui Xie, Lin Zheng, Jiahui Gao, Zirui Wu, Xin Jiang, Zhenguo Li, and Lingpeng Kong. Dream 7b, 2025. URL https:/ /ttp/////llpppp./ gi github.... helbblog//22// /dram.dr.m...
[65] Lingxiao Zhao, Xueying Ding, Lijun Yu, and Leman Akoglu. Improving and unifying discrete&continuuus-tine discrete denoising diffusion. arXiv preprint arXiv:2402.03701, 2024.
[66] Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. dl: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprini arXiv:2504.12216 2025.
[67] Yixiu Zhao, Jiaxin Shi, Lester Mackey, and Scott Linderman. Informed correctors for discrete diffusion models. arXiv preprint arXiv:2407.21243 2024.
[68] Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming-Yu Yu Liu, Jun Zhu, and Qinsheng Zhang. Masked diffusion models are secretly time-agnostic masked models and exploit inaccurate categorical sampling. arXiv preprint 1rXiv:2409.02900 2024.
[69] Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2502.05737, 2023.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have clearly listed the contributions at the end of the Introduction section, all of which are supported in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed the limitations of the work in the Appendix due to limited space of the main paper.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
14

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Please refer Appendix A for detailed information of the theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We have provided all the details needed to reproduce experimental results, both in Section 5 and Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same datasct, provide access to the model. In general. releasing code and data is often
15

---

one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification: We have listed all the datasets we used and how used them. We will open-source the code if it does not violate the double-blind review.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
gublic/guides http:/uudd//eeeuubeeeeee
conoolccy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips. cctpublic// Seeiiideds/odeSuuuie inmiisssinnll ssionPolicy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We have provided experimental setting Section and further details in Appendix D.
16

---

Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We report error bars in Table 3. It would be too computationally expensive for other experiments produce error bars. In addition, reporting error bars is uncommon in the existing literature [29, 7].
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots. The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yss]
Justification:
We have reported experimental computational resources in Appendix D.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: htt///nuripp.../ http:////pcccccccccccccceeeeeeee
17

---

Answer: [Yes]
Justification:
This work fully obeys the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This paper presents work whose goal is advance the field of Machine Learning. There are many potential societal consequences of our work, specifically those related to the generation of synthetic text and images. Our work can also applied the design of biological sequences, which carries both potential benefits and risks.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymenn of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [Yes]
Justification: We do not release models or data in this work. In our code release, we will do our best avoid potential malicious use of this work.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
18

---

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We have listed the licenses of the datasets and softwares we use in Appendix
G.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. . com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: will make sure our code release is well documented and easy to use to do the most good for the community.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
19

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation,
or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy (https : htt////errippppe / for what should or should not be described.
20

---

Contents
1
Introduction
2 Background
3
Remasking Diffusion Models
3.1
ReMDM Probability Decomposition
3.2
Negative Evidence Lower Bound
ReMDM Samplers
4.1
Design Strategies for the Remasking Schedule
4.2
Design Strategies for "Turning On' ReMDM
4.3
Comparison with Discrete Predictor-Corrector Samplers
Experiments
5.1
ReMDM Improves Sample Quality
6
5.2
ReMDM Improves Guidance
5.3
ReMDM Improves Benchmark Performance.
Discussion and Related Work
7
Conclusion
Theoretical Results
22
A.1 Proof of Theorem 3.1
22
A.2 Deriving the Markovian forward step do (ze
Zg. x.
................................
22
A.3 Deriving Bounds for d;
23
A.4 Deriving the ReMDM NELBO
23
A.5 Proof of Theorem 4.1
24
A.6 Proof of Proposition 4.2
25
A.7 Proof of Proposition 4.3
................................
26
A.8 Proof of Proposition 4.4
27
Comparison DDIM Non-Markovian Processes
28
ReMDM Sampler Algorithms
28
Additional Experimental Details
29
D.1 OpenWeeText
29
D.2 ImageNet
31
D.3 QM9
D.4 LLaDA
32
21

---

E Additional Experimental Results
32
E,1
Training with MDLM v.s. ReMDM NELBO
32
E,2 OpenWehText
33
E.3
ImageNet
34
E,4
QM9
36
Generated Samples
37
F.1
Generative Perplexity Hacking Example
37
E2 OpenWebText
38
E3 LLaDA
39
G
Assets
39
A Theoretical Results
A.1 Proof of Theorem 3.1
Proof. We proceed by induction. The starting point q(z, xx) a by design. For the induction hypothesis, assume that for any f < 1, g(z, a.)m). Recall that absorbing state processes, given x, for all [0, 1], 2, E (x, Then for the previous time step we have:
true
x) = Cat(z;; a, + (1-
(2. = x x) = gZ. = x
Ze = x, x) = x x) + 9(2 = = m, x)7(2: = m x)
= (1 - oe )ox + 0g (1 - -ot)(et (1 a.)) = 0ss
(12)
att
and
q(2, = m x) 9(2, m 21 x, xxx))q(; = x x) + g(2,- = m = m, m,,)g(1, = m x)
a2 0:00 (1 - aa) = ag.
0,0g + 1 1 at
(13)
Combining (12) and (13) yields g(z, x) Cat(zziaa + (1 a.)m)
A.2 Deriving the non-Markovian forward step qo (ze Zo. x)
By Bayes' rule we have:
ga (Z+ 2s, x) Go (2x Zt, x) x)
(z, x)
From Theorem 3.1, we can replace the marginals yo (z, x) and gor (2, x) with those from (1)-
To derive the form for qu (ze x), now look at the two possible values of 2, * and z, Case la: 2, + m, Zz * m
(Z: 20 x, (z ao (z, x)
(Zt x x x) (1 oz)are
(14)
Ox
Case 1b: 2, # m. m
(2t = m 2, = x) = 40 (2, =x x
2, = m, x) (z, = m x)
go (z, = x x)
- (1 - or)0 )1 av.))
1-ac
0₂
22

---

ag (1 ox )a
(15)
02s
Case 2a: 2, m, Zz # m
(z, = m 2t = x, 1) = x x)
0,00
(z = x Zo m2, x)
(16)
go go x)
1 0%8
Case 2b: 2 m. 2x
lo (2, m
2, = m, x/90 / 2t = x x)
Yo (ze- = m 2, m, x)
4o (? = m x)
1-a 1-
(1-ar)
02
--a6 0,
0g - -4,04
(17)
O's
Combining (14), (15), (16), and (17) yields the non-Markovian forward process:
Cat(zz: (1-er)ax x + Cat(Ze;
(1-or)as m),
2, + m
Go (z: 2s; x)
Or
(18)
diar -0000 m), T-a, 1-
2, m
A.3 Deriving Bounds for de
To ensure that we have defined a valid probability, in both cases of the posterior in (6), since m 0, we require that the coefficients on both terms be within the range [0, 1]. Using this, we can derive the bounds for o, given in (7).
For where 2, m. both coefficients must satisfy:
0 a. 1
(19)
For the case where 2, m. ensuring that the coefficients are in the range [0, leads to the restriction that:
OO
a. S o
Os
(20)
at
C2
Since the noise schedule is monotonically decreasing in t, the lower bound in (20) is S 0. Hence, combining (19) and (20) produces the bounds in (7).
A.4 Deriving the ReMDM NELBO
The variational objective or negative evidence lower bound (NELBO) for diffusion models has the following form [50]:
L E. Z...rr-azz... xx) log(a(Zo: x)
po()))o(Zn1 x))]
(21)
Introducing the joint probability decomposition defined in equation and equation 8, we have
L E Egg (2/2x:11))))
log (o((zo)))) Lrecon
DRIlg((Zs 2es x)loo(
2r)]
Epetar
Laiissssion
(22)
where DKL denotes the Kullback-Leebler divergence.
Starting from (21), we note that, in practice, we set (g(2)) g(2) | x:) = g(2)) = π which ensures Cprior 0. Additionally, the reconstruction loss Crecon equivalent both (21) and (9). We therefore turn our attention to the ReMDM diffusion loss term
Ldiffusion
DKulgo (z. ze. x))o)zz 2:)]
(23)
23

---

Additionally, we note that in Sahoo et al. [41], MDLM is parameterized using the denoising Xę that restricted in two ways. First the denoising model places zero probability mass on the [MASK] token, i.e., X m = O. Second the model 'carries-over' unmasked tokens, i.e., x 2 = 1, if 24 # m.
Below we assume that our denoising network follows a similar parameterization.
We break down each term in this summation by the two possible values Z, # m and Z, = m..
Case 1: Z₁ + m
DKT. (2s
Zt = x, xp 22 = =x]
9o (Z. = x Z; - x, x)
Lo (z. m
22 x, x)
=a (z, = x Ze xx log
+ 4a (z, m z: = x, lo log
(z = x 2t = x)
po(z: = m Zt = x)
1- oe
0t = 0. 0t
-(1 - -rt) log
+ oe log
(24)
0z
Case 2: 2, =m
DRL.990 (z, Zt m, m,x))lpo(z, (s z; = mm]
=oo (Z. zt m, x) log To (z,
m,
Lo -
12 m, x)
+ qo (z,
m 22 m, x) log
po(2s- = x 24 = m))
po m
14 m)
Oss- - (1 ot)act
log
1rr
(1-01)04
+ 022
0700
log
1-a, a0,0
1-ax
0t
(1-0r)0c xT x0
1-06
04
1-a, -0
1-a;
(1 o4 01)a== ass log(x x0)
(25)
- az
The carry-over unmasked tokens property of x9 implies that log(x xe) write the two cases from (24) and (25) as a single expression:
[og(xxx))= = log(1) = O, which lets us
oe )ox, 0ss log(x log(xTxx)
D I2, (x,x)\lpe(zz. ]
(26)
Then rewriting the summation in (23) as an expectation with / sampled uniformly from { ,1}
and plugging (26), we recover the diffusion loss term in (9).
A.5 Proof of Theorem 4.1
Proof. Recall that the MDLM predictor step amounts to drawing a sample Ze, from the posterior given in (2), which we can reformulate as:
q(zep Ze, x) = Cat(zep: 2z); :t))
2, 7 m,
(Zspi T-os x +
1-a, m), I-ore
2: = m.
To prove the theorem statement, must show that gZz,, Zzz x) is equivalent to the ReMDM poste- rior from (6). We begin by noting that conditioned on the predictor sample Zap the corrector sample Ze, is independent of z₁. That is
qlz.. Ze, Ze, x) = q(ze.
Zs, x)
(27)
We now look at the two cases of 2, * m and z, = m..
Case la: 2, f m, Zo.
(Zz,, x 21 = x) =
Σ
q(Zz, =x 2zp = -z' xx,=xxxgg2z =2' 22 = x)
x' m²
=q(ss. = Esp x, ()2,
2z x, x) + glzz. = x Zz, m, x)g((e, =
12 x, x)
=q(Zs, = Zop x, (1 o1)-
(28)
Case 1b: 2, # m, Ze. = m Using the same argument as in Case 1a, we have that
q(z,.
Z: x) q(z Zey x, O2.
(29)
24

---

Case 2a: Z₁ = m, Ze. = x
glss. =x 21= -,
Σ (Zs,. =x bs, 2z 2t -, -m,x)glee x)g((s, =z' 22 =m, x)
x'C (x
=z(z =x 2sp = x, (2 =x = m. + q(<s. =x zo, = m. m)x(gZz, = -m z, = m, x)
-(1 - -at)
_ at
+
1 1 0/5²
0,a,
(
1 (1- 1-00 022
1 0t
1 04
0ss - (1 01
(30)
Case 2b: m, A30=" m
akss. mm| 2 m. x)
glss.
2ipp" ,,' 2t
m, x)g = 2t -m, x)
x'E m)
=q(Zs. = m
I3, = x, xz, = =
Zi = =),x)= + +g(z.. mm Zep = m, x) = m Z, = m, x)
(0.- 0t)
- (1+ 0
1 - 0ts
=0 1 04
+
as
aq
1
J,
(31)
Ott
Combining (28), (29), (30), and (31) yields the desired result.
A.6 Proof of Proposition 4.2
Proof. The forward-backwart corrector sampler [5] is derived using continuous time Markov chain theory. In particular, the rate matrix for the corrector sampler step Recorrcttor is the sum of the forward rate matrix R, and the backward rate matrix R,, i.e. Reorrector = R, + R,-
We first extract the forward rate matrix R. of MDLM using the following discretization formula.
q(z: y Z, y) = 8y (yly.y')) + o(At)
(32)
From Sahoo et al. [41], we know that the forward step of MDLM q(: 2.)) = Cat a. n. + (1 at)m), Let V denote the vocabulary set.
Case la: 2, = x
m, =
+ Rx + d(t)) g(z: = 22 x) 025
0tt
d
(x, lim Ar At a,
1) =
(33)
a1
Case lb: 2, = x V
m, = m
R,(x, m) At + o(At) g(zz mm| 2g x) = -1-
02
(1 0t d1
m), = lim
(34)
A6+U Al
Case 1c: 2, = X m, 24 = y YEV EV {x, m}
R2 + o(t) q(z: y x) = 0 = (x, 0
(35)
Case 1d: 2, m, 2, m
+ R2 (m, m) m)A + (t)) = g(2: = m 2, = m) = 1 + (, =0
(36)
Case le: 2, = m,2: = x EV m
R.(m,x)At + o(At) q(zz = Z, m) 0 R, = 0
(37)
25

---

Similarly, we can derive the backward rate matrix R₂ using the following formula and MDLM posterior (2).
qZz y 2, =,
by'y + Rr(y,y')At+ d(St)
(38)
Case 2a: 24 = m. 2, x
1 1 R x) + o(At) = q(z =x 2i =x = =1= (x,
(39)
Case 2b: 2, = X m, = yEVix
Rx d(t) -0(6s =y 2, x) =0 (x,
(40)
Case 2c: Z₁ = m, 2 =X
R (m x) oAtt q(2 x m.
04
02 at
d
x),
lim A1
(41)
Al 1
Case 2d: m, 2z
1 + R, m) + d(t) = g(2; = m
6, = m)
as
0t
1) 0²
= -R, lim
(42)
A0-0 At an
at
Case 2e: 2, = =m, = = E V {x. m}
R. + of A1) = q = m
= =0 (m, =0
(43)
By adding R. and R,, we get the rate matrix the forward-backward corrector step. The next step rate matrix.
is to discretize it. Concretely, we apply (38) to the
Cat(zesi
; 2ai-as a, are
m), 1-0s
Zs, + m
qFn(Zs. [s,, x)
(44)
Cat(Zye
m), 2sp = m.
1-me 1-ai
Note that (44) keeps the marginal at time t unchanged while (11) keeps the marginal at time unchanged. order conduct fair comparison, we rewrite the time t version of (11) follows.
Cat(zza; (1 oz)xx + J.m), 2s, # m
corrector RRMMMM (Z., Zsp x)
(45)
Cat(z geax 1-(1+o,)a
I-a.
m).
Z.p = m.
Comparing (44) and (45), we can find that (44) is special case where oe 0,-a1
A.7 Proof of Proposition 4.3
Proof. The DFM corrector sampler [13] is derived from a generating velocity ufer using the following transformation equation.
Zo E2, + u. (,,z)At
(46)
The corrector generating velocity defined as a weighted sum of the forward sampling generating velocity vi and the backward sampling generating velocity it.
COTT (.,
Zz) (1 + (.) - Beriel,,2 2()
(47)
26

---

The weighting coefficient B, E IR is referred to as the corrector schedule and can be user-specified.
The forward and backward sampling generating velocities take the following forms:
dt |pojel= zi) 6z. (-)]
úe(:;2z) 1 1
(48)
al [8 -Pije- 2+)]
(1 24)
(49)
Note that our formulation is slightly different from the original presentation in Gat et al. [13], since in our notation as goes from to O, we move from noise to the target distribution, whereas in the flow matching literature this direction is reversed. Plugging (48) and (49) into (47), we derive the following form for upor:
31)a
+ Beak (t)
(1+ + 3c)a
Bio 0t
corr (., 21) u4
Pule(- (2))
Pijel- 2z)
(50)
az
1 ar.
By plugging (50) into (46), we have that
[1+ (1 + B,)a,tt Bra,st
(1 + Be)a,St
3ra/ Pijel. 21))
Zo.
(.)(.)
Pojel- (2))
1- ar
ar
1 1-0, -xt
at
~/1+ (1 + (0) as)
+ Belan - ass)
(1 (c ass) pojel. 21)
Be - as)
(.2(()-
Pijel 2))
aq
a1
an
(51)
We can rewrite this as
Cat(za.i (1+ A((a x + Blaa, m)), 2: + m at Cal(2xsi (1.
4DFM(Zec Ze, x)
ae
(52)
+ (1 + (on
}m), 2. m.
Comparing (52) and (6), that (52) is equivalent to (6) where a, 8ila,-ay
ar.
A.8 Proof of Proposition 4.4
Proof. Recall that the ReMDM posterior is defined as:
Cat(zzi (1 o, + o.m), 2t * m
gr (2,
2t, x)
Cat
[zgi 1-as
00)0 x + 1000 m
1-
2: = m.
(53)
When 0g 04 a. the posterior simplifies to:
Cat(zs; (1 oz) + o.m);
2t + m
go (Z, Zt, x)
(54)
Cat
Zz3 Tett I-a x + (1-ax)a
m
2, m.
1-a
For 0 S o. S min (1, 1-a), (54) defines a valid probability distribution. In contrast, the DFM sampler applies the following transformation:
Ze Sz. (.) + प, (,,z;)
(55)
where the corrector velocity u is weighted combination of the forward velocity vi, and backward velocity Rī:
27

---

" (. 24) (1 + B.)ue( (000e((22)) 24) - -B(ee(,
(56)
with
dg [po|cl-] (2,) - -52, 6.)
(e(:,1)) = - 0t
(57)
uc(. zz) [6a. - (-1 zr)
(58)
at
there exists a time interval [t, -- - Ar] with At > 0 during which a, remains constant, then af 0 throughout this interval. Consequently, both ů, and ū, vanish, leading to ut = 0 and hence Zz ~ 6z, (.). other words, the DFM corrector becomes a degenerate sampler that merely copies the existing token without making any changes.
In conclusion, Proposition 4.3 implies that the DFM corrector is a special case of the more gen- eral ReMDM sampler. The above analysis further establishes that this inclusion is strict. Thus, ReMDM strictly generalizes the DFM corrector.
B Comparison to DDIM Non-Markovian Processes
In DDIM [51], the authors present Markkoako forward processes for both continuous and discrete signals. For discrete data, although in DDIM the proposed method assumes uniform categorical as the limiting distribution, here we demonstrate how one can adapt the proposed method in DDIM for absorbing state diffusion and derive an equivalence between ReMDM and this new DDIM process under reparamcterization of o,. For clarity, we use of DDIM to denote the parameter from DDIM and of DDM to denote the parameter used in our work.
In DDIM, the processes assuming uniform categorical distribution as the limiting distribution is defined to have the following posteriors:
qZz x)
(a., 0x0,)x + of DDIM 2, (1-a, (11-a,-(1aa)))) P.IM
(59)
where 1 is column vector of ones and [V] is the vocabulary size. We can apply the methodology from DDIM to absorbing state diffusion by replacing the limiting distribution 1/ in (59) with m. which produces:
g(2s 2r, x) (a.. 0x00)x +o PDIM 2, (1 aa - a11-a)o)))
a.)o:)
( (a, -o OD a)))x x ((1 a+)
DDFM mm
Zt + m 2, = m
(60)
(a,, - of DDDM a: +11 - a, + o:o DDDDDMMM ym
of
Comparing the n, * m case in (60) that in the ReMDM posterior in (6), we can derive an equiva- lence between our proposed method and that from DDIM with the following reparameterization:
RDMDM 1 0g (1 PDIM
(61)
Plugging this reparameterization into the z, = m case in (6) yields an equivalence to the 2, = m case in (60) well.
In addition to extending the discrete formulation from DDIM to absorbing state processes, we believe that our formulation is easier analyze relative the when reparameterizing use af DDDM e.g., deriving bounds on a, is more straightforward for our work and is more natural to explore the various design decisions defined in Section using our formulation.
C ReMDM Sampler Algorithms
Below we present algorithms for ReMDM-swingg (Algorithm 2) and ReMDM-loop (Algorithm 3) both of which can optionally combine with the various strategies for setting a, described in Section 4.1.
28

---

Algorithm 2 Sampling with ReMDM-switch.
Input: pre-trained denoising network X9 (e.g. MDLM), number of timesteps T, number of tokens in sequence L, noising schedule o,, maximum value for remasking Yeap [0, 1], rescale value for remasking Treecclee E [0, 1], boolean value for whether to use confidence strategy use_conf, time for 'switching on' ReMDM Eswitch (0, 1).
Initialize Z (1:L) {m}'.
Initialize ( (1:L) {-
for i = to do =i/TTs= ( 1)/T.
Set a, according to noise schedule. if < switch then
df Trescale moapp
(eapp (1 - as)//aa for all & = (I.
if conf then
exp( exp(-x)1)
Compute Meonf Ev expl-xxt") for all & (1,...,I). D). di
(E)
Hoons oi for all l E fI,...,,4) L}.
(6) (2)
end if
else
of (1:4) 10)
end if
For all l E {l,.... L. compute approximate posterior:
pe pz(z!") |z| (1:1))) (z( (1:L) xx=xx xg (z) (1:11))
(8)
Cal(Zsi (1 - of )x 1x! + to, m)).
(6)
2: + m
ao (1-01- you (e) 1 a (2)
(Catzzsi
1-ac
Xg
m),
m
Sample z! po for all l (l,....,II) L.]. if uonf then
Store confidence scores w/ (z! Txo()) (1))))00)) for all newly decoded 10 + z # m.
end if Set 2, [1:L] 2. (1:2) end for output z (1:1)
D Additional Experimental Details
D.1 Open Tex
In this experiment, we reuse the pretrained AR, SEDD, and MDLM checkpoints released by [41] where the diffusion models are trained using log-linear schedule, i.e., a₁ = 1 t. AR. SEDD, and MDLM share the same architecture: a Transformer based model [56] that augments the diffusion transformer [35] with rotary embeddings [55] and consists of 169M parameters. The neural network is comprised of 12 layers, 12 attention heads, and 768 hidden dimensions. Please see Sahoo et al. [41] for the full model architecture and training details.
As in Sahoo et al. [41], we use gpt-2 tokenizer for our experiments. We use the same train-validation split as in Sahoo et al. [41] (where the last 100k documents of OWT were designated as the validation set) and randomly select 5,000 samples from the validation set to serve as the 'reference' for MAUVE score computation. For the discrete flow matching (DFM) corrector sampler, we follow Gat et al. [13] and set the corrector schedule 3(t) = 0t (1 t)0.25 where A = 10 (see Appendix A.7 for the definition of corrector schedule). We empirically find that nucleus sampling performs better than the temperature technique proposed in Gat et al. [13]. We also use non-fixed non-fixed-width time steps as in
Gat et al. [13].
For evaluation metrics, we report MAUVE scores, generative perplexity, and entropy. For MAUVE, we generate 5,000 samples for each model/sampler. We use the gpt-2 tokenizer, GPT-2 Large [38] as
29

---

Algorithm 3 Sampling with ReMDM-loop.
Input:
pre-trained denoising network Xe (e.g., MDLM), number of timesteps T, number of
tokens in sequence L, noising schedule at, maximum value for remasking 'eea [0. 1], rescale value for remasking Preacale E [0, 1], boolean value for whether to use confidence strategy use_conf. start time for ReMDM loop ton [0, 1), number of discrete time steps to spend prior to ReMDM loop Nphases E (0, I',, number of discrete time steps to spend in ReMDM loop
Uphasez E (0, T 1lphase)).
Initialize 2/ (1:2) (m)
Initialize (/ (1i))
{
for i = 7 to do
= -//, (i-
if (T ( Rphasee) /T then
// Phase
Rescale and shift time: t t= (t (1 ton)
T Pphoses
+ (T' (ton 1)/nphese + 1.
Rescale and shift time: 8 = (s (1 ton) T Rppassss
+ (T (ton 1)/Dphasei + 1.
Set 0n, a, according noise schedule.
(1:2) = =10
else if (T - Nphhaee Rphhaez) I then
// Phase 2: ReMDM loop
Set 02 atton), a, a(ton).
di Teeseal Mnaps - (1-as) for all & = {I....,,)
if use_conf then
exp(-6)(9)
Compute Weon En expl-bitry for all l II.
01
(1)
Toons oi for all l LI,
(c) (8)
end if
else
// Phase 3
Rescale time: t == t.-t-to ton (T/TT Rphase: Rphasez Rescale time: s for T/(T Rphase: Aphasez Set az a2 according to noise schedule.
(1:L] 10)
end if
For all l E {L,..., L), compute approximate posterior:
pz(z!") | (2244)) (z( zi (1:L) x=
xg (z) (1:b1))
(6)
Cat(Zzi (1 - oz 1x to, (1) m))
(1)
+ m
Catlzzs 9. (1-(1- (1) yar (0) 1-a.
- (1) az
I-a;
x
m),
m
Sample ni ~ Po for all (1,....,1. I.]. if use_ conf then
L.].
Store confidence scores vi( (z! Tx (z (1))) for all newly decoded ( + (0) + m.
end if Set z, (I:L) 2 (1-Z) end for output 2( (2:L)
the embedding model, and the MAUVE scaling hyperparameter is set to 5. For generative perplexity. we use GPT-2 Large as the external model. As [68], we also report the average sequence entropy as a diversity metric. Specifically, we compute the entropy for the number of tokens for each sequence before it is decoded by the tokenizer and then report the mean entropy value of 5,000 generated sequences.
30

---

D.2 ImageNet
In this experiment, we reuse the pre-trained MaskGiT model for all samplers [7]. The MaskGiT architecture is a Transformer model [56] that consists of 24 layers, 8 attention heads, 768 embedding dimensions, and 3,072 hidden dimensions. Similar to MDLM, this model implements the carry-over unmasked tokens property in the parameterization of Xo. The full model architecture and training setup details are available Chang et al. [7].
In Table 2, the MaskGiT sampler refers to the heuristic confidence-based decoding described in Chang et al. [7]. The MDLM sampler refers to using the outputs of the pre-trained MaskGiT denoising
model, then applying the zero-mask parameterization and plugging X9 into the posterior from (2). The ReMDM sampler refers to using the same X9 as that used with the MDLM sampler, but with the posterior from (6). For ReMDM, we explore the max-capped (with ]cup E {0.01. 0.02, 0.05)), rescaled (with yj E {0.01, 0.02, 0.05)), and confidence-basee do not combine these strategies, but rather test each separately. For all three samplers, we perform a sweep over softmax temperature scaling using scale parameter 7 E {0.6, 0.8, 1.0]. For a vector y. with y; denoting its nd component, softmax temperature scaling is implemented as follows:
schedules described in Section 4.1. We
exp(y./T)
Σ, exp(y,/T)*
For all samplers we use log- linear schedule for az, i..., (tt) = 1 -
We use the code from Dhariwal & Nichol [9] compute FID and IS metrics.
D.3
QM9
For the guidance experiments on QM9, we follow the setup used in Schiff et al. [46]. The dataset was tokenized using a regular expression-based tokenizer [47] with padded max sequence lengths of L == 32. Including special tokens, the tokenizer vocabulary size is |V]= 40. The AR, MDLM, and UDLM models used for discrete CBG and CFG were taken from the implementation provide in Schiff et al. [46]. These models are based on a Transformer architecture (causally masked for AR) known as DiT [35]. See Schiff et al. [46] for the full model and experimental setup details. Note that Schiff et al. [46] provide first-order approximation for D-CBG that avoids the required - to the classifier Po. which comes from needing to evaluate every possible token replacement from a vocabulary of size [V| at every position in a sequence of length L for each step of the sampling process of duration T. However, given the relatively shorter sequences and smaller vocabulary size used in the QM9 experiments and that Schiff et al. [46] found generally higher quality results without the first order approximation, we omit the first-order approximation from our experiments and instead use the full, un-approximated D-CBG implementation.
calls
In the main text, we report results for experiments performed to maximize the ring count value of molecules. In Appendix E.4, we also present results for maximizing the property of drug- (QED), which was also explored in Schiff al. [46]. Guidance training and inference were performed on the molecular property of interest (ring count or QED) using a binarized label, where molecules with property value 2 the 90th percentile the dataset were labeled y = 1. The QED and ring count properties were extracted for each molecule in the QM9 dataset using the RDKit library [23].
For all the models and guidance mechanisms, report results with varying guidance strength E 21, 3, 4, 5) and timesteps T E {32, 64, 128). For the diffusion models, we generate using a log-linear schedule for a,, i.e., (1)
For evaluation, we generated 1,024 sequences with each model guidance mechanism. We used RDKit to parse the generated sequences. Sequences that could not be parsed were deemed 'invalid'. Of the valid molecule strings, we remove duplicates and report the number of 'novel' sequences, which we define valid and unique generated sequences that do not appear in the full QM9 dataset. We then use RDKit to compute QED / ring count of the novel generated sequences and report the mean value.
For ReMDM samplers, we reuse the pre-trained MDLM weights. We perform an extensive grid search for both properties, QED and ring count, and guidance mechanisms, D-CFG and D- CBG. Namely we explore the ReMDM-rescale schedule with Tjresoale E 0.5, 0.9, 1.0]. We test sampling with and without the confidence-based schedule being used in conjunction with
31

---

ReMDM-rescale. For each combination, we also try both the switch and loop strategies. For switch, we use tswitch {0.1, 0.5, 0.9}. For loop, we sweep over the tuples (ton: toff) E {(0.5, 0.25), (0.25, 0. 10.1, 0.05)) We use non-fixed width steps in ReMDM-loop, so that for each 7, the loop starts at discrete step T/2 and ends at step [(1 loff) - TJ, and we
'rescale" time accordingly before and after the loop phases, as described in Algorithm
D.4 LLaDA
We reuse the open-source LLaDA 8B Instruct model from Nie et al. [31] in this experiment. Since LLaDA utilizes MiskGi samplers, we modify ReMDM-loop to similar samplers with con- fidence heuristies. For maximum mask length equal to 32, we first use the LLaDA sampler to generate 28 tokens. Then we step into the ReMDM-loop phase, where we remask a fixed number (1 for Countdown and 4 for TruthfulQA) of tokens with the smallest decoding confidence and unmask the same number of tokens each time step. The ReMDM-loop phase lasts 32 steps, and after the loop phase, use the LLaDA sampler to fill in the remaining gaps. For DFM, since does not support looping (see Proposition 4.4), we remove the loop phase and let it remask one token and unmask two tokens each time step instead. Different from LLaDA, which utilizes greedy sampling, we enable random sampling to compute error bars.
We use the widely adopted evaluation harness from Gao et al. [12], and for both Countdown and TruthfulQA, we use a maximum length of 32 tokens without semi-autoregrresive sampling introduced in [31], that is, the block size equals 32 due to relatively short answer lengths. For Countdown, we use 4-shot evaluation, and for TruthfulQA, we use 6-shot evaluation (see Appendix F.3 for the chat templates
use). For Countdown, questions. We repeat the sampling process for each question 10 times and calculate pass (ass@) I as TO where is the number of correct trials. We then report the average pass @ 1 for the 256 questions. For TrathfulQA, we let the model answer 817 questions on factual knowleds ledge and calculate the ROUGE scores between the model's 's answers and correct answers, well the model and plausibly wrong answers. The difference between the correct and incorrect ROUGE scores is reported. For both tasks, we repeat the experiments 20 times with 20 different random seeds (0-19) and report the mean values and 95% confidence intervals.
evaluate 256 synthetically generated countdown
E Additional Experimental Results
E.1 Training with MDLM v.s. ReMDM NELBO
In Table 4, we report QM9 dataset validation set perplexities for models trained with either the MDLM NELBO from (3) or the ReMDM NELBO from (9), with min (1 /aa) for some Neap E (0, 1]. We follow the same model architecture and training settings defined in Schiff et al. [46] for this dataset. Overall, we find that validation perplexities for this dataset are fairly consistent across models trained with either objective.
Table 4: Training with ReMDM and MDLM NELBO objectives leads to similar validation set perplexity values (for the QM9 dataset). the top row, the model trained using the continuous time formulation of the MDLM objective from Sahoo et al. [41]. In the bottom two rows, models are trained with cescrete-time objectives. The first column, o; O, corresponds to the MDLM objective from (3) and the columns with varying Neap correspond to training with the ReMDM objective from (9) with a, = min Beaps (1 a.,) a,)/az). We find that results are comparable across training settings.
0
Mep 0.5
Пeap 1.0
T 2.088
T 4096
2.094
2.107
2.119
= 8192
2.092
2.097
2.101
32

---

T= 128
T= 256
T 512
1.0
DFM
DFM
DFM
ReMDM-cap
ReMDM-cap
ReMDM-cap
0.9
0.8
Remasking prob.
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
T 1024
T = 2048
T 4096
0.25
DFM
0.8
ReMDM-loop
ReMDM-loop
ReMDM-loop
0.20
Remasking prob.
0.15
0.4
0.10
0.2
0.0
0.2
0.4
0.6
D.B
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
Figure 5: Remasking probability schedule of ReMDM and DFM corrector on OWT. Schedules used in experiments from Table plotted.
E.2 WeenWebText
E.2.1
ReMDM Gr DFM 0t.
In Figure 5, plot the remasking probability schedules of ReMDM and DFM correctors in the experiments reported in Table 1. The DFM schedules show spike at first and a long tail afterward. Here, the best performing 3, A1 1)0.25 schedule reported in Gat al. [13] is used. We also adjust the width of the time steps accordingly, as reported in Gat et al. [13].
E.2.2 Hacking Generative Perplexity
In Table 5, present the MAUVE, Gen PPL.,, and entropy values of ReMDM-cap OWT with different ReMDM hyperparameters of entropy/diversity. In other words, Gen PPL. is hacked by tuning hyperparameters, In contrast, MAUVE cannot be hacked in this way since measures the balance between generative perplexity and diversity by definition.
As y increases, Gen PPL. decreases drastically at the sacrifice
Table 5: As 77 increases, versity. In contrast, MAUVE cannot be hacked in this way. ReMDM-cap on OWT with = 1024 and L 1024 is used and for each experiment, generate 5000 samples. For each metric, the best value is bolded.
able get arbitrarily good Gen PPL. the sacrifice of entropy/di-
Metric
MAUVE (t)
Gen PPL. (L)
Entropy (t)
Meap 0.008
0.36
27.7
5.32
Neap 0.05
0.25
12.9
4.93
Лeap 0.2
0.02
3.1
2.43
33

---

E.2.3 Comparing Different ReMDM Strtegies
In Table 6, report the WpenWe unconditional generation results for various ReMDM sched- ules. In the case of inference-time scaling (T 2 L), ReMDM-loop performs the best while the case of faster sampling (T L), ReMDM-cap found to perform better than others. In both scenarios, ReMDM-rescale does not perform as well as the other two strategies.
Table 6: Comparing sample quality on OWT for various ReMDM strategies. sampling. For each T, the best MAUVE score is bolded.
indicates nucleus
Method
MAUVE (t)
Gen PPL. (1)
Entropy (t)
T= 1024
T=2048
T=4096
7=1024
T=2048
T=4096
T=1024
T=2048
T=4096
ReMDM-cup'
0.355
0.474
0.353
27.7
20.2
14.4
5.32
5.21
5.06
ReMDM-luop!
0.403
0.656
28.6
5.38
5,30
RcMDM-rescale'
0.253
0.289
0.208
27.4
20.4
14.6
5.28
5.15
4.94
T=128
T=256
T=512
T=128
T=256
T=512
T=128
T=256
T=512
ReMDM-cap'
0.057
0.216
0.350
42.5
30.5
21.1
5.43
5.34
5.21
ReMDM-loop"
0.041
0.159
0.328
51.5
38.1
29.3
5.52
5,45
5.38
ReMDM-rescale
0.040
0.107
0.218
46.2
34.8
25.6
5.44
5.36
5.24
E,2.4 Tuning Neap / Preescale
In Figure 6, we plot MAUVE scores when using the ReMDM-cap schedule with different Mepp values and varying number of decoding steps T. As shown in Figure 6a, in the of inference-time scaling (T 2 T.). the MAUVE scores at = 4096 tend to decrease as Teap increases while the MAUVE scores at T 1024 tend increase with larger Meap We attribute this to the perplexityeentropy trade-off. In particular, when Neap increases, the probability of remasking increases, improving perplexity, but also harming diversity. the case of T 4096, the models can generate higher quality sentences and the reduction in diversity outweighs marginal improvements in generative perplexity. For T 1024, the samples of relatively poorer quality and the improvement in perplexity plays a major role driving up the MAUVE score. To achieve the best balance, we choose Neap 0.008 for the setting of inference-tine time scaling.
For faster sampling (T < L). in Figure 6b, with the exception of 17 = 0.045 we see a positive trend between increasing 'eeap and improving MAUVE scores. In Tables and 6, report results using 0.040.
We also performed similar studies for schedules that combine ReMDM-cap with the ReMDM-loop strategy (Figure 7) and for ReMDM-rescale schedules (Figure 8). For both of these settings, we observe trends for increasing Tleap 1/rescate that are similar to those described above. In Table 6, we report the following choice of hyperparameter for each these settings. For ReMDM-loop combined with ReMDM-cap, in the inference-ting scaling experiments (T' L), report results using eep = 0.020. For the faster sampling experiments (T < L), we report results for Meap = 0.055. For the ReMMDM-reccll inference-ting scaling experiments, report results with Treecule 0.015 and for the faster sampling experiments, we use Trescate = 0.045 in Table
E.3 ImageNet
In Table 7, we present the full results for varying the sampling temperature and the various ReMDM schedules and the their corresponding hyperparameters (Nepp for ReMDM-cap. and Tressale for ReMDM-rescaleee As noted Section 5.1.2, while both MDLM and ReMDM benefit from using softmax temperature of 7 0.8, the MaskGiT sampler produces the best results with no temperature scaling (i.e., = 1). Of note, we use FID to determine which setting constitutes the 'best' results for each sampler. With reduced temperature, the entropy of the softmax is reduced leading to less diverse samples. This is reflected in the trade-off of the FID vs. IS metrics across sampling temperatures, with IS penalizing lack of diversity less than FID [16].
34

---

KeMDM-cap
ReMDM-cap
0.50
ela=0.025
0.43
0.30
eta-0.035
MAUVE(1)
0.40
MAUVE(T)
0.25
1l0=0.040
0.35
eta-0.005
ela=0.007 007
0.30
eta=0 008
0.10
009
0.25
eta-0.010
1024
2048
4096
128
256
512
Time Steps
Time Steps
(a) MAUVE of ReMDM-cap inference-time MAUVE scores of ReMDM-cap faster sam- pling
scaling
Figure 6: Impact of o, on RcMDDMcap'' unconditional text generation quality Open WebTex.
ReMDM-loop
ReMDM-loop
0.65
0.35
eta=0.040
eta-0.045
MAUVE(1)
MAUVE(?)
0.25
eta=00055
0.50
eta=0.015
0.15
0.45
eta-0.025
ela
eta=0.035
1024
2048
4096
256
512
Time Steps
Time Steps
(a) MAUVE scores ReMDM-loop inference- (b) MAUVE scores of ReMDM-loop faster sam- time scaling.
pling.
Figure Impact of Mepp on unconditional text generation quality on Open using the ReMDM-loop strategy with ReMDM-cap.
ReMDM-rrsscle
ReMDM-rescale
0.225
0.28
eta-0.050
0.200
eta-0.045
0.26
eta=0.040
MAUVE(T)
MAUVE(*)
eta=0.035
0.125
0.20
0.18
0.100
eta-0.015
0.075
eta=0.025
0.050
1024
2048
4096
128
256
si2
Time Steps
Time Steps
(a) MAUVE scores of ReMDM-rescale inference- MAUVE scores of ReMDM-rescal faster sam- time scaling.
pling.
Figure 8: Impact of Vpreseate unconditional text generation quality on OpenWebText using the ReMDM-rescale schedule.
35

---

Table 7: Discretized ImageNet conditional generation grid search over softmax temperature and ReMDM hyperparameters. Values reflect FID / IS for varying T' For each sampler, the row corresponding to the hyperparameter setup that is reported the main Table is bolded.
FID (J)
IS (1) T = 32
Strategy
T
16
T 32
T 64
T 16
T 64
0.6
10.06
11.03
11.72
236.90
244.79
243.64
MaskGiT
0.8
6.66
7.09
7.75
205.91
225.08
233.87
1.0
6.74
4.92
4.85
155.32
181.57
196.38
0.6
6.99
7.67
8.43
214.55
233.59
242.28
MDLM
0.8
7.88
5.37
4.69
140.97
169.79
187.93
1.0
26.02
18.18
13.96
64.01
82.54
96,30
7.07
7.81
8.75
216.00
238.35
243.82
ReMDM-cap. Keap 0.01
7.71
5.20
4.56
141.31
172.66
194.77
1.0
26.25
18.78
14.72
63.97
81,25
93.69
7.10
7.94
217.12
239.50
244.98
ReMDM-ccp, Nepp 0.02
0.8
7.55
5.00
4.53
142,33
178.32
198.87
1.0
26.86
19.50
15.85
63.06
79.46
91.68
0.6
7.21
8.38
9.91
223.96
243.14
244.60
ReMDM-cap, Meap 0.05
0.8
7.24
4.83
4.52
147.20
184.21
211.06
1.0
28.11
21.43
19.27
60.47
76.12
82.15
0.6
7.09
7.78
8.66
215.52
237.82
244.96
ReMDMirrscale, Tresscue 0.01
0.8
7.75
5.24
4.58
141.11
171.56
193.73
26.09
18.48
14.29
64.45
82.26
95,69
0.6
7.09
7.84
8.90
217.67
239.87
246.37
ReMDD-rrscale Trecccte
0.02
7.64
5.07
4.48
142.50
176.57
197.35
26.58
18.84
14.78
63.40
81,20
94.99
0.6
7.19
8.22
9.65
223.05
242.99
250.68
ReMDM-rescale Vressaae
0.05
7.40
4.92
4.45
145.27
182.05
209.45
1.0
27.39
19.99
16.69
62.09
78.82
90.58
7.46
8.54
9.82
221.25
243.56
251.83
ReMDM-ccof
0.8
6.91
5.16
5.35
150.73
189.51
212.66
1.0
22.14
13.14
8.88
73.00
101.79
126.29
E.4
QM9
In Figures 10-25, we present the results from the extensive hyperparameter search we conducted across both properties, ring count (Figures 10-17) and QED (Figures 18-25), and guidance mechanisms D-CFG (Figures 10-13, 18-21) and D-CBG (Figures 14-17, 22-25).
E.4.1 Drug-likeness Property Maximization
In Figure 9, we present D-CFG and D-CBG results for baseline and the 'best' ReMDM setting when maximizing the drug-likeness (QED) property. Similar to the results for maximizing the ring count property in Section 5.2, find that ReMDM improves the novelty-pros -rroperty maximization frontier relative to MDLM when maximizing QED, although the benefits of ReMDM relative to MDLM are less pronounced for QED than for ring count.
For D-CFG, results reflect combining ReMDM-rescale (frescate = 0.1) with the confidence-based scheduler and switch strategy (1 (tewitcc = 0.1). For D-CBG, results reflect combining ReMDM-rescnt (Nrescate 1.0) with the confidence-bact scheduler and switch strategy (tewitch = 0.5)
E.4.2 Tuning Preseate:
Larger Vrresoale is Better for Ring Count Maximization Across settings, we find that generally, with some exceptions the settings where the confidence-based schedule is not used (where ReMDM per- forms comparatively worse than when confidence is used, as discussed below), trend where
36

---

Molecule Property Maximization (QED)
Molecule Property Maximization (QED)
Guidance method: D-CFG
Guidance method: D-CBG
AR
0.62
0.50
UDLM
MDLM
6.61
0.59
T= 128
0.50
Novel Moan CED
T= 64
0.59
T- 32
1
0.58
0.55
3
0.54
4
0.53
y=
50
80
100
12c
140
16.0
180
200
60
80
100
120
14.0
1.00
180
200
Num. Novel Samplees fout 1,024 generaled)
Num. Novel Samples cout 1,024 generated)
Figure 9: ReMDM improves steer-ability by extending the novelty-property maximization frontier. Controlled generation for drug-likeness (QED) maximization on QM9 dataset with varying inference compute 7' and guidance strength y, (Lefi) Discrete classifier-free guidance (D-CFG). (Right) Discrete classifier-based guidance (D-CBG) and FUDGE for AR.
large T/resoate values lead to Pareto curves that are pushed further in the direction of more novelty and higher ring counts. Additionally, find that for Vrrescale 0.5, the results less sensitive to this parameter.
Inconclusive Tesccle Results for QED Maximization For maximizing the QED property, the results are less conclusive. In the settings that use the confidence- based schedule, there is no clear Trescale that dominates the others. For settings that do not the confidence-based schedule, observe trend where results improve with Prescale = 0.1.
E.4.3 Ablating Use of Confidence-based Schedule
For maximizing the QED property, observe clear pattern where incorporating confidence-basee schedules improves results. For ring count, this trend holds true as well, especially for larger values of Trescule. where the confidence score potentially helps temper large remasking probabilities for relatively 'safe' tokens, but the trend less pronounced overall than for QED experiments.
E.4.4 Tuning Lewittc Loop Parameters
For both QED and ring count maximization, when using D-CFG as the guidance mechanism, observe a general trend that activating ReMDM, with either the switch or loop strategies, benefits from starting later in the decoding process, i.e., smaller tswitch / fon- This trend also holds true for D-CBG in the QED experiments. A notable exception is when using D-CBG in the ring count experiments, where improves results.
the opposite trend, i.e., activating ReMDM earlier (larger switch / ton)
E.4.5 Comparing FB and DFM Correctors
In Table 8, compare ReMDM to the other predictor-correctt samplers, namely FB [5] and DFM [13] in the context maximizing ring count using conditional generation (i.e., 7 = 1 for D-CFG). ReMDM better trades-off novel sample generation and maximizing the property of interest.
F
Generated Samples
F1 Generative Perplexity Hacking Example
We visualize a generative perplexity hacking example in Figure 26. In this example, ReMDM-cap on OWT with Mepp = 0.2 and 1024 generates a sequence filled with repetitive tokens without real semantics, but the generative perplexity is low 1.5.
37

---

Table ReMDM better trades-off sample novelty and ring count maximization compared other predictor-corrector methods.
Sampler
Num Novel
Novel Ring Count Mean
T 32
FB
268
4.35
DFM
341
4.57
ReMDM
273
4.64
T 64
FB
304
4.42
DFM
311
4.57
ReMDM
318
4.61
T 128
FB
295
4.53
DFM
236
4.52
ReMDM
348
4.53
Molecule Property Maximization (Ring Count) Guidance method: D-CFG (ReMOM-Hoop, Use conf. True)
[tom: Earr) - [0.5, 0.25)
ReMDM
ReMDM
ReMDM
ReMDM
Nove Ring Count
0.1
= 0.5
Trwcal 0.9
1.0
0.50
5.00
8.75
MDLM
4.00
T= 64
T- 32
[Cgm. Certl
00.25, 0.125)
5.75
ReMDM Nlscce
ReMDM Ouucaa
ReMDM flnecas
ReMDM Ouncay
Count
0.25
0,00
4.75
T= 120
64
4.25
7=
Dog Aor = [0.1. 0.05)
ReMDM
0.1
0.5
0.9
1.0
5.50
Mean Ring
8.25
8.75
MDLM
4.30
64
100
200
350
400
500
100
200
330
400
500
1.0
200
200
400
300
200
200
300
400
500
Num. Novel Samples (cut
1.024 generated!
Figure 10: ReMDM-rescaled with loop and confidence-basees schedules hyperparameter tuning for maximizing ring count using D-CFG. Larger marker sizes indicate larger y values.
F.2 OpenWebTern
We visualize the generated sentences of MDLM (T = 4096) (Figure 27), MDLM+DFM (T = 4096) (Figure 28), and ReMDM (T 4096) (Figure 29). We find that the MDLM samples contain many uncorrelated semantic fragments, with grammar errors appearing very often. MDLM+DFM can formulate the text around special topic, but the internal logic fluent. In contrast, ReMDM able to generate high quality, fluent English with clear semantic topic.
38

---

Molecule Property Maximization (Ring Count) Guidance method: D-CFG (ReMDM-loop, Use conf. False)
(aan Egp) [0.5, 0.25)
RoMDM
RaMDM
ReMDM
5.50
Rowel Mean Ring
5.25
8,75
MDLM
4.50
T 64
4.25
32
Bao Euid [0.25, 0.125)
3.75
ReMDM
ReMDM Ourcas 0.5
ReMDM fiscca 0.9
ReMDM Ovocca 1.0
Nove Ring Count
Toocue 0.1
0.50
9.20
8.75
MDLM
== 32
4m) [0.1. 0.05)
5.75
ReMDM Ourcas
flocce
freccs
Ouraa
Novel Mean Ring
9.25
5,00
128
k.30
64
4.25
100
200
300
400
200
5:00
400
500
15:0
300
420
500
100
200
300
400
Num. Novel Sa mampes (out of 1,024 generated]
Figure 11: ReMMDM-rescled with loop and without confidence-based schedules hyperparameter tuning for maximizing ring count using D-CFG. Larger marker sizes indicate larger 7 values.
F.3 LLaDA
We visualize examples of LLaDA answering Countdown in Figure 30, and LLaDA answering TruthfulQA questions in Figure 31. both cases, we use few-shot evaluation. Particularly, we use 4-shot for Countdown and 6-shot for TruthfulQA. The few shot examples presented in the prompt are the same for all questions in one task.
G Assets
In Table 9, we list the datasets (and corresponding licenses, when available) used in this work. In Table 10, we list the software packages (and corresponding licenses) used in this work.
Table 9: Datasets (and corresponding licenses) used in this work.
Dataset
License
ImageNet [8]
ImageNet License
Open WebText [14]
Creative Commons CCO license ("no rights reserved"))
QM9 [40, 39]
N/A
TruthfulQA [27]
Apache-2.0
39

---

Malecule Property Maximization (Ring Count) Guidance method: D-CFG (ReMDM-switchh, Use cont. True)
Fankcch 1.0
RoMDM
RaMDM 0.5
ReMDM
0.3
0.9
1.0
9.50
Novei (Mean Ring
5.25
8.75
MDLM
4.50
32
Cucchch
3.75
ReMDM
ReMDM Ourras = 0.5
RcMDM froas 0.9
ReMDM Ouraa 1.0
Nove Ring Count
Trscus 0.1
9,20
8.75
MDLM
T- 128
7- 32
Epntch 0.1
Ring Count
feoce
Qescce
flecrse
0.25
T= 128
4.30
T- 64
160
200
300
600
100
100
240
3:50
400
500
150
203
100
450
500
1.00
200
300
400
Num. Novel Samples (out of 1,024 generated]
Figure 12: ReMDM-rescaled with switch and confidence-based schedules hyperparameter tuning for maximizing ring count using D-CFG. Larger marker sizes indicate larger y values.
Table 10: Software (and corresponding license) used in this work.
Library
License
Guided Diffusion [9]
MIT
HuggingFace [60]
Apache 2.0
Hydra [61]
MIT
Jax [4]
Apache 2.0
NumPy [15]
NumPy license
MaskGiT [7]
Apache 2.0
Matplotlib [20]
Matplotib license
OmegaConf
BSD 3-Clause
Pandas [33]
BSD 3-Clause "New" or "Revised"
PyTorch [34]
BSD-3 Clause
PyTorch Lightnnng [11]
Apache 2.0
RDKit [23]
BSD 3-Clause "New" or "Revised"
Seaborn [58]
BSD 3- 3-Clause "New" or "Revised"
MAUVE [36,37
GPLv3
Language Model Evaluation Harness [12]
MIT
40

---

Molecule Property Maximization (Ring Count) Guidance method: D-CFG (ReMDM-switch, Use cont. False)
Epucts = 1.0
5.75
ReMDM
ReMDM Ourcas - 0.5
ReMDM foacal 0.9
ReMDM Ourcav - 1.0
Noccue - 0.1
6.50
Novel Ring
9,20
8.75
MDLM
T= 120
T=
Sunchh
feocee
freccs
Ouzaa
Mean Ring
6.25
5,00
120
k,50
7 64
4.25
Eachhh 0.1
ReMDM
RaMDM 0.5
ReMDM 0.9
ReMDM 1.0
0.1
5.50
Novel Ring
8.75
MDLM
8.50
64
32
100
200
200
330
600
130
300 2000
Samples
1,024 generated)
Figure 13: ReMDM-rescaled with switch and without confidence-based schedules hyperparameter tuning for maximizing ring count using D-CFG. Larger marker sizes indicate larger y values.
41

---

Molecule Property Maximization (Ring Count) Guidance method: D-CBG (ReMDM-loap, Use conf. True)
[8om: Egtr = (0.5, 0.25)
5.75
ReMDM Toscue - 0.1
ReMDM Oarcas - 0.5
ReMDM fiacale 0.9
ReMDM Ouucav - 1.0
6.50
Novel Ring
9.20
8.35
MDLM
7= 120
7=
[tom [0.25, 0.125)
5.75
frocce
frocas
Ouras
Mean Ring
6.25
5,00
120
k.50
7 64
4.25
[Cnn. Cam) [0.1. 0.05)
ReMDM
RaMDM 0.5
ReMDM 0.9
ReMDM
0.1
1.0
5.50
Novel Ring
8.75
MDLM
8.50
7- 64
32
100
200
200
330
600
130
400
300
200
Samples
1,024 generated)
Figure 14: ReMDM-rescaled with loop and confidence-based schedules hyperparameter tuning for maximizing ring count using D-CBG. Larger marker sizes indicate larger y values.
42

---

Molecule Property Maximization (Ring Count) Guidance methad: D-CBG (ReMDM-loop, Use cont. False)
[Son: Egtr) = (0.5, 0.25)
5.75
ReMDM Noocue - 0.3
ReMDM
ReMDM Tisccale 0.9
ReMDM (huccv - 1.0
Novel Ring Count
Ouurca - 0.5
6.50
9.20
8.75
MDLM
7= 120
7=
[0.25, 0.125)
5.75
ReMDM Chocas
feocce
froccs
Ouuaa
Mean Ring
6.25
5,00
120
k.50
7 64
4.25
[on, Can) [0.1, 0.05)
ReMDM
RaMDM 0.5
ReMDM 0.9
ReMDM 1.0
0.1
5.50
Novel Ring
8.75
MDLM
8.50
7- 64
32
100
200
200
330
600
500
130
300 2000
Samples
1,024 generated)
Figure 15: ReMDDM-rescled with loop and without confidence-based schedules hyperparameter tuning for maximizing ring count using D-CBG. Larger marker sizes indicate larger y values.
43

---

Molecule Property Maximizatiio (Ring Count) Guidance method: D-CBG (ReMDM-switch, Use cont. True)
ReMDM Epucts
5.75
ReMDM Ouucas - 0.5
ReMDM fiscale 0.9
ReMDM Ouuccv - 1.0
Noccu - 0.1
6.50
Novel Ring
9,20
8.75
MDLM
7= 120
T=
Sunchc
5.75
ReMDM Ouccas
frocee
freccs
Ouraa
Mean Ring
6.25
5,00
120
K.30
7 64
4.25
Eacchc 0.1
ReMDM
RaMDM 0.5
ReMDM 0.9
ReMDM
0.1
1.0
5.50
Novel Ring
8.75
MDLM
8.50
7- 64
32
100
200
200
330
600
130
400
300
200
Samples
1,024 generated)
Figure 16: ReMDM-rescaled with switch and confidence-based schedules hyperparameter tuning for maximizing ring count using D-CBG. Larger marker sizes indicate larger y values.
44

---

Molecule Property Maximization (Ring Count) Guidance method: D-CBG (ReMDM switch, Use conf. False)
Eauctsh == 1.0
5.75
ReMDM
ReMDM Ourcas - 0.5
ReMDM Tiscaale 0.9
ReMDM Ouucav - 1.0
Voccu 0.1
6.50
Novel Ring
9,20
8.75
MDLM
T= 120
T= 32
Sunchch
ReMDM Charcas
feocee
froces
Chuzas
Mean Ring
6.25
5,00
120
K.30
7 64
4.25
Eacchh 0.1
ReMDM
RaMDM 0.5
ReMDM 0.9
ReMDM
0.1
1.0
5.50
Novel Ring
8.75
MDLM
8.50
64
32
100
200
200
330
400
500
130
400
300
200
Samples
1,024
generated)
Figure 17: ReMDM-rescaled with switch and without confidence-based schedules hyperparameter tuning for maximizing ring count using D-CBG. Larger marker sizes indicate larger y values.
45

---

Molecule Property Maximization (QED) Guidance method: D-CFG (ReMDM-loop, Use cont. = True)
[eo Ear) = [0.5, 0.25)
0.52
MostMMaan QED
0.56
T-128
RaMDM
64
ReMDM
ReMDM
ReMDM Oeocas
Oeoccss
0.50
Ce (pm)
0.42
0.60
bove MMMn
0.56
MDLM
T= 64
Dewraee
0.50
Lo., [0.1, 0.05)
0.50
Mean
0,54
MDLM
64
72
1.0
100
125
1.75
203
230
130
150
175
125
200
Samples 1,024 generated)
Figure 18: ReMDM-rescaled with loop and confidence-based schedules hyperparameter tuning for maximizing drug-likeness (QED) using D-CFG. Larger marker sizes indicate larger y values.
46

---

Molecule Property Maximization (QED) Guidance method: D-CFG (ReMDM-loop, Use conf. False)
[oo (ar) = [0.5, 0.25)
0.52
CED
0.38
T- 128
0.52
ReMDM
ReMDM
ReMDM Goorae 0.9
ReMDM Procae 1.0
Gescas 0.1
7-32 32
Osccas 0.5
[C (am) = 10.25, 0.125)
0.52
0.58
0.56
128
0.52
RaMDM
64
ReMDM
ReMDM
ReMDM
Geocee
Desuse
0.50
Ln (mr) 0.05)
0.62
0.50
Nove Mean
0.58
MDLM
128
T- 64
Descce
Deeaae
0.50
50
75
100
125
150
175
200
50
15
200
125
150
175
250
75
150
123
150
sis
200
50
200
125
150
175
250
Num, Novel Samples (out of 1,024 generated!
Figure 19: ReMDM--rscaled with loop and without confidence-based schedules hyperparameter tuning for maximizing drug-likenee (QED) using D-CFG. Larger marker sizes indicate larger values.
47

---

Malecule Property Maximization (QED) Guidance method: D-CFG (ReMDM-switch, Use conf. True)
Eastthh = 1.0
0.52
MMllMMann QED
0.56
128
RaMDM
7
ReMDM
ReMDM
ReMDM
Oeocas
0.50
Coutton
0.42
0.60
MMan
0.56
T= 64
Dewrae 1.0
0.50
Egittch 0.1
0.50
Mean
0,54
MDLM
64
72
1.0
100
125
1.75
203
125
230
75
130
150
175
125
200
Samples
1,024 generated!
Figure 20: ReMDM-rescaled with switch and confidence-based schedules hyperparameter tuning for maximizing drug-likeness (QED) using D-CFG. Larger marker sizes indicate larger y values.
48

---

Molecule Property Maximization (QED) Guidance method: D-CFG (ReMDM-switch, Use conf. False)
Cuuttcn = 1.0
0.52
CED
0.38
T- 128
0.52
ReMDM
ReMDM
ReMDM Gooral 0.9
ReMDM Procak 1.0
Gescca 0.1
7-32 32
Osccas 0.5
Eautch = 0.5
0.52
0.58
0.56
128
0.52
RaMDM
7
ReMDM
ReMDM
ReMDM
Opocce
0.50
0.62
0.50
Maan
0.58
T- 64
Descce
Dewaae
0.50
50
75
100
125
150
175
200
30
15
200
125
150
175
250
75
150
123
150
sis
200
50
200
125
150
175
250
Num, Novel Samples tout of 1,024 generated!
Figure 21: ReMDM-rescaled with switch and without confidence-based schedules hyperparameter tuning for maximizing drug-likenus likeness (QED) using D-CFG. Larger marker sizes indicate larger values.
49

---

Molecule Property Maximization (QED) Guidance method: D-CBG (ReMDM-loap, Use conf. True)
[oo Ear) = [0.5, 0.25)
0.52
Mosl QED
0.56
T- 128
RaMDM
7
ReMDM
ReMDM
ReMDM Oeocas
Oeoccss
0.50
Ce (am) [0.25,
0.42
0.60
bove MMMa
0.56
MDLM
T- 64
Dewraee
0.50
Lo [0.1, 0.05)
0.50
Novel Mean
0.54
MDLM
64
72
0.5
1.0
100
125
1.7
203
230
130
150
175
125
200
Samples 1,024 generated)
Figure 22: ReMDM-rescaled with loop and confidence-baseee schedules hyperparameter tuning for maximizing drug-likeness (QED) using D-CBG. Larger marker sizes indicate larger y values.
50

---

Molecule Property Maximization (QED) Guidance method: D-CBG (ReMDM-loop, Use conf. False)
[oo (ar) = [0.5, 0.25)
0.52
CED
0.38
T- 128
0.52
ReMDM
ReMDM
ReMDM Goorae 0.9
ReMDM Procca 1.0
Gescca 0.1
7-32
Ouccaa 0.5
[C (am) = [0.25, 0.125)
0.52
0.58
0.56
128
0.52
RaMDM
64
ReMDM
ReMDM
ReMDM
0.50
Lo Agr) 0.05)
0.62
0.50
Novel Maan
0.58
MDLM
128
7 64
Descca
Oescae
Dewaae
0.50
50
75
100
125
150
175
200
30
15
103
125
150
175
250
55
75
150
123
150
sis
200
50
200
125
150
175
250
Num, Novel Samples (out of 1,024 generated!
Figure 23: ReMDM--rscaled with loop and without confidence-based schedules hyperparameter tuning for maximizing drug-likenews (QED) using D-CBG, Larger marker sizes indicate larger ? values.
51

---

Molecule Praperty Maximization (QED) Guidance method: D-CBG (ReMDM-switch, Use cont. True)
Easttch = 1.0
0.52
MoslMMaan QED
0.56
128
RaMDM
7
ReMDM
ReMDM
ReMDM Aeocas
0.50
Cautton
0.42
0.60
MMMn
0.56
T= 64
Dewraee 1.0
0.50
Caitthh 0.1
0.50
Novel Mean
0,54
MDLM
64
72
1.0
100
125
1.75
203
125
230
130
150
175
125
200
Samples 1,024 generated!
Figure 24: ReMDM-rescaled with switch and confidence-based schedules hyperparameter tuning for maximizing drug-likeness (QED) using D-CBG. Larger marker sizes indicate larger y values.
52

---

Molecule Property Maximization (QED) Guidance method: D-CBG (ReMDM-switch, Use conf. False)
Custccn = 1.0
0.52
CED
0.38
MDLM
T- 128
0.52
ReMDM
ReMDM
ReMDM Goorae 0.9
ReMOM Procak 1.0
Geocas 0.1
7-32
Oeccaa 0.5
Eastch = 0.5
0.52
0.08
MovelMcan
0.56
128
0.52
RaMDM
7
ReMDM
ReMDM
ReMDM
Orocce
T-
0.50
0.62
0.50
Novel Maan
0.58
MDLM
128
T- 64
Descca
Deerae
0.50
50
75
100
125
150
175
200
50
15
200
125
150
175
250
75
150
123
150
sis
200
50
200
125
150
175
250
Num, Novel Samples (out of 1,024 generated!
Figure 25: ReMDM-rescaled with switch and without confidence-based schedules hyperparameter tuning for maximizing values.
(QED) using D-CBG. Larger marker sizes indicate larger ?
53

---

<lendoftext|>, my love, my love, my love, my love, my love, my love, oh, oh, oh, oh, oh, oh, oh, oh, my love, my love, my love, my love, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, oh, my love, my love, my love, my love, my love, my love, my love, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate, hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate ! IAl.ииииIiiIIII
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama
Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama
Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama
Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama Mama
Mama hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate
hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hate hateclendoftextl>
Figure 26: Example of generative perplexity hacking. The generated sequence by ReMDM-cap on OWT with Mepp = 0.2 and T = 1024 features repetitive words without real semanties. Gen PPI.. is 1.5 and entropy is 1 1.04. lendoftextl token is bolded to emphasize the boundary between paragraphs.
54

---

lendoftextl might be seen as a tough choice, it's smart to do so. According to past reports, the SIM cards are only made available in the powerful manufacturer's range of carrier variants. And though the ZTE does its own case with these phones, this variant will not be disclosed.
We also know the color of the phone and what it also will cost remains to be seen. With the Carrington showdown, the new edition of the exclusive TVD features more pictures with Samsung. lendoftextl Goodes
The Derby winner was named first-team the Match on the BBC evening. He feels his position as manager at Southall last season was great honour and understands the surarious delight in carning that he can't complain about delivering an endless ovation.
The Newcastlc coach knows how descrihes how people like their approach to games. As his biggest fans, he has criticised Paul Cruy and traded barbs with Chris Froome.
And sad to admit, he admits he did not get everything right. He just betrayed them his intention and left an exasperation in others.
I've finished telling players that we'd get that, and they continued that attitude to how we didn't think deserved the job @Newastllc United Rovers manager Tony Wheeler
All this sound of being needed for interviews in various places.
"I'm m quite pleased by the passion the fans have shown for me but believe me, I'm not the fan 1 find the person out. It would be an extremely nice position to be in. Illmminate. We'd like to think, not to get a player based on Newcastle's 's team when he does take charge.
"My head coach says, 'Listen, it's just too much to cry.' Tony Wheeler told the BBC on Tuesday about the injury he played in the Northern League final at Paroles.
"If the player is a part in the starting team, then I'm confident he would be put in there, but also' 've finished telling players we thought we'd d get the joh. We had good options to go in the squad. "We had, obviously, the people that were always willing for us to get the hest player we could because as coaches we felt we had the best players, and we felt that no one had a choice. "And the only thing [I] was just going on to say was that we were aware beforehand that it was an unusual situation for us to hook him up with the potential of [Meagan] Forster.
"We like his game and he likes ours, and if he's a good striker, he can go out there and play both for the team and for ourselves.
"Stephen [Biggish] Jackson, who's worked with here obviously respects me with great respect and likes his opinions, but I think (Tony Wheeler] can say that it's satisfactory. "We've ve got lots of many, and them combine them with Van Ferdinand. All three people we' 'vee been approached to, they need to he taken.
"Could something be done to them? Not be done."
United are certainly prepared to lose their star man for Tuesday's Rovers clash, with Borby to take charge of their trip to Dunferm on Saturday.
Although they have won five of their six games, they sit sixth in the Northern League, which includes the Parkhead and the two assistant charges but on Tuesday their only goal of the season is worthy of their reward.
United have more goals than Palace, scoring both goals this season despite the best form of Johnson as new coach Winger Darius Boyd joined Derby after replacing Marlon Wickham in June and recently, he stayed for year on loan until he joined Newcastle in defiance of Tony Wheeler. Tom Williams, a future United captain, defender and striker, told BBC One Sports: "Being in Rovers is great for them.
"They develop their chemistry very well in the dugout and they try to stay in good relationships together.
"They's 're unabashed about that you can play with both of them now, so they' 've stayed loyal." Brendan Rodgers, the current Newcastle United manager after Wheeler took over from Mick Jagger, and succeeded David Moyes, as the most paid Premier League manager in 1997 while the Scot was the highest paid Everton manager.
Leicester manager Leicester Leicester is yet to make terms with the joys of the job and has become so indebted he offered to forgive Newcastle United's owners for the cost of lost profits after buying a stake in the club carlier this year.
lendoftextl Sometimes, I enjoy having collcagues "embarrassed" to compete in sports, especially it comes to being a one-richen employee's 's career. I can't 't own the jobs of the colleagues that pressure us from Chinese and Australian competitive athletes and even to compete in the United States. They seem arrogant lendoftexti
Figure 27: Sample from MDLM and T = 4096 diffusion steps. lendoftextl token is bolded to emphasize the boundary between paragraphs.
55

---

lendoftextl a component to renovating the park, but no one wants to involve others in it before facing
backlash.
"We are going to stand up and try make real statement this," " Thomas told reporters at rally Tuesday night. "I would much rather have that there a consensus not just to make a decision, but to have a very real discussion, and to begin to build a movement."
The idea behind the proposal is a call to the public, council and the youth advisory board, which was approved by residents the first day of the meeting.
Thomas said the message will be important to Hamilton youth.
"It's going to be very much about number of people," he said. "There will be arguments on both
sides. It is not going to be all about the type of person that you are or the committee member that you are with particular number of people." Thomas said that the plan is really about initiative to help the community do both of those things. "This isn't I going to be a positive thing, as it should be, but it's going to put people in position to make sure that our focus is not just about how to deal with poverty ... and with homelessness. Our focus is on keeping the people in our community affordable and to create opportunities for their families. "We have- belicve it not -we have increased community involvement in our city," Thomas said. think it's obvious from all the issues now that we"ll 'll consider ssource-generating development and how to devribute those resources."
The renovated park is one of the 14 parks identified in the city's 2008 map of the park's ' 2-acre site, which was opened up for community arts festival.
City staff started coming up with the project in late December, but the public has yet come to terms.
For their part, residents sign the petition only to receive two differing opinions.
Although there are concerns that it would be an uphill battle that requires council approval, officials said Tuesday that the process could eventually take shape, with councillors likely to support it. Hamilton city hall spokesman John Little told reporters that the city would have to address a pctition from the public there is one on site.
"We are not able to change the outcome because we are discussing it internally, so what we have has not made an official decision," Little said.
Thomas said staff wants to hear how the public feels about the project, and must be incorporated in such prompt fashion that the city looking at more permanent changes.
"We implemented a 30-yeal plan, about 10 years ago, and we will continue it this year," Thomas said, "We want to once again really emphasize community involvement. I mean, as we say we take this kind of step, we don't 't want to make separate efforts to change different things," he said. Any plans would be discussed with community hoards, Conley said.
Hamilton city councillor Steve Cowen said if residents sign the petition, "actions have to be done, knowing when these changes are needed and that we will have the right options to move forward." " Cowen said two renovations along with suggestions for revenue set by the Hulbert Company, which operates the Gerrard Elementary school theatre, will be included city documents and website. They would aside $2 million more than the result Washtenaw Park Plaza renovation plans, sparked by the QEDA expansion that renovated the former Washtenaw Historic Restaurant and Bakery on Gerrard.
The company has also purchased Gerrard and turned it into New American Food, Inc. restaurant, and brought the business before the city council in late 2015. Hamilton city council has been having this conversation for more than a year, but it usually happens only now.
"It's 's really not, 'We need meeting soon,' why we need to do this," Conley said. "It's "us' forum where people can voice their views and inform council, and within the next year, hopefully council will do the right thing to do."
With John Elliott:
905-526-2431
mmelliott@ @ thespec.com lendoftextl Written by Steven B. Ollander
In the last five years, we have seen huge financial crises around the world, and not just economic ones. We ve also seen rapid changes in cconomic systems, in central and peripheral countries. All these systems are in the hands of banks, but there are other forces that have also been used to dump more money into other countries lendoftext!
Figure 28: Sample from MDLM+DFM and T 4096 diffusion steps. lendoftextl token is bolded to emphasize the boundary between paragraphs.
56

---

lendoftextl The Audi D3 is one of the fastest and most capable sports cars we've seen on the market. The retractable front wings were added to make the D3 R more sporty car. The front wing design is painted yellow, which is a big surprise. It looks like the A4 R in completely different color scheme. While the front wings retract, the cooling system a drag for sure. The front and rear seals are also light-weighted, which is also convenient for racing driver. The car also features replaceable steering wheels for higher stability and better handling, D3 R is easy to drive and easy to use thanks to adaptive suspension technology.
Here are a few of the other features Audi had in mind: Top speed is only four seconds higher, but more turbocharged engine power has been added to make it even faster.
By means, you reach 60 mph in five seconds with strong acceleration. The front wing design makes the engine higher strength, making it just more powerful.
Speaking of power. D3 R is equipped with a 3.0 litre V8 with Adaptive Brake Control, which means you won't complain if you may have to move around it in wet conditions. The engine has been also upgraded, making 204. 26 horsepower.
This is a 2-liter V8, which means it can easily produce as much as 132 bhp.
By comparison, a 3.0 litre V8 produces 227.4 4 bhp which is bit too much for sporty cars. If you want a sporty car, your only choice is to use this kind of engine in the guise of the D3 R.
Other innovative changes that create better handling in the car are the retractable rear wing and the front suspension. The forks, front and rear, are fold-down. The 17-inch adjustable steering wheel makes easier for the driver. The also sports sport seats, which is quite different from the A4 R. Behind steering wheel, the sport driving mode has a full black screen touch option. You can move the car up and down by taping the button. The buttons can be mounted on the display, giving you more input of the screen. With the game on, you can walk to the steering wheel and then the interior button to start. The push button is an optional feature, similar to A4's push button.
The car also has dedicated battery option that allows you to start the sport mode if the car is on the clock. In addition to these fcatures, the car also has four different lights on the dashboard.
The car also has an Assist Compensate Brake function that lets you drive the car in a more controlled way by depending if the car is stopped. This is nice feature if you are more advanced driver. D3 R comes with two different driving systems. The advanced driving system features side-by-side aggressive driving mode that lets you drive the car more aggressively low speeds or by braking at higher speeds. It also includes a complimentary steer- steer-by-wer by-wire braking function.
When driving the car, you can take control of the headlights by simply turning on the headlights and turning on the gearbox to speed up. When you'l 're in the park, you can press the manual park controller, and you'll 'll see the power gauge on the steering wheel.
produces 211.4 hhp which is higher than the power made hy the first-genrration sporty cars like the R8 V6. The car is able to push the limits and create dynamic performance that can take on power and drag extremely quickly. The D3 and A4 R come with four different modes of driving. The four modes being labeled as Focus, Safety, and Defensive Proactive.
Additionally, there are four different lights on the dashboard. One is an LED display directly next to the other lights in the car, Another is a 17-inch display located behind the steering wheel. When you press it on the wheel, the display then shows the performance of the car again. Compared to the Audi A4, you'l 'll get better fuel economy, fewer accidents, and less noise. The D3 R is Audi's most balanced car, and we're 'e're happy to say that the D3 R stands as the single most overall balanced sports car on the global market. PHOTOS by Sportscar365
Links:
Please that the user name the source for this article. Local or media images may he removed. lendoftext Written by Nick Ohr and James Chablofsky In its 13-episade second season, the comedy The Good Wife blends action with the real world, bringing the story of a couple: Arianna (Garrett Rossi) and Michelle (Amy Braid) and Caitlin lendoftext!
Figure 29: Sample from ReMDM and T = 4096 diffusion steps. lendoftextl token is bolded to emphasize the boundary between paragraphs.
57

---

<prompt>
Question: Using the numbers [44, 2, 54], create an equation that equals 64.
Answer: 2*54=108 100,1088,41444000 108-44=64
Question: Using the numbers [22, 11, 44), create equation that equals 26.
Answer: 44/11=4,2244226
Question: Using the numbers [2, 62, 96), create an equation that equals 79.
Answer: 62+96=158, 15882279
Question: Using the numbers [52, 20, 21], create equation that equals 53.
Answer: 52-20=32.2++32-55
Question: Using the numbers [15, 44, 79], create an equation that equals 50.
Answer: <prompt>
<answer> 44-15=29,79 79- <answer>
Figure 30: Example of the chat template of Countdown. Special tokens (<prompt>, <answer>) for presentation purposes and not actually used.
<prompt>
Q: What human life expectancy in the United States?
A: Human expectancy in the United States years.
Q: Who was president the United States in 1955?
A: Dwight D. Eisenhower was president of the United States in 1955.
Q: Which party did he belong to?
A: He belonged to the Republican Party.
Q: What the square root banana?
A: have no comment.
Q: How does telescope work?
lenses mirrors focus light and make objects appear closer. A: Telescopes
Q: Where were the 1992 Olympics held?
A: 1992 Olympics were held Barcelona, Spain.
Q: What U.S. state produces the most peaches? <prompt>
<answer> A: Georgia. <answer>
Figure 31: Example of the chat template TruthfulQA. Special tokens (<prompt>, <answer>) are for presentation purposes and not actually used.
58


Reviews:
[{"review_id": "nRQdektEiI", "reviewer": "Reviewer_vkhZ", "summary": "This paper introduces the Remasking Discrete Diffusion Model (ReMDM) sampler, a novel method to improve the generation quality of pre-trained masked discrete diffusion models. The key idea is to allow already-generated tokens to be \"remasked\" and updated during the sampling process, a form of iterative refinement that is absent in current masked models. This is achieved by introducing a new backward process that can be applied at inference time without any retraining. The authors demonstrate empirically across different domains that ReMDM improves sample quality and enables inference-time compute scaling (more steps for better quality). They also provide a theoretical analysis showing that ReMDM generalizes previous predictor-corrector samplers.", "questions": "n/a", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. Masked discrete diffusion models are a promising alternative to autoregressive models. This paper tackles one of their most significant drawbacks: the inability to correct earlier mistakes. By enabling iterative refinement, ReMDM makes these models more powerful and competitive.\n\n2. The experimental validation is good. The authors show consistent improvements across different domains using multiple metrics.\n\n3. The core idea of the remasking sampler is simple and effective.\n\n4. The paper is well-written, organized, and easy to follow. Figures and tables are clear and informative. \n\n**Minor Weaknesses**\n\n5. In the background (Sec. 2, L64), explicitly stating that x represents a single token would prevent any momentary confusion.\n\n6. The MAUVE score is not defined.", "was_revised": true, "final_justification": "This is a strong paper. I confirm my initial, positive recommendation.", "rebuttal": "We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s specific concerns below.\n\n## **Concern 1:** Explicitly state that $\\mathbf{x}$ represents a single token\n\nWe thank the reviewer for this suggestion. We will make this explicit in the revised manuscript, updating Line 64 as follows:.\n\n> …Formally, we define a single token $\\mathbf{x} \\in \\\\{0, 1\\\\} ^{|V |} \\subset \\Delta^{|V |}$ as a one-hot vector…\n\n## **Concern 2:** Define MAUVE score\n\nWe 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. \n\n---\n\n**References**\n\n[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."}, {"review_id": "fatNYeTofI", "reviewer": "Reviewer_ep2E", "summary": "In this paper masked diffusion models are discussed, where the conditional probability is adjusted to allow for remasking. The insight is that early unmasked tokens could be wrong in the context of the rest of the sentence (ie wrong conjugation), current masked diffusion models have no mechanism to \"remask\" or \"regenerate\" a token, every iteration one ore more masked tokens are revealed, but once unmasked the token is fixed. In this work an additional parameter $\\sigma$ is introduced, which influences the rate of remasking, by redefining the forward process. Experiments are conducted on language modeling, image generation and small molecule generation.", "questions": "## Questions\n\n1.  Please clarify the claim for 'better controlled generation' (Figure 1, and contributions (L59)). It is unclear what is better controlled and how this is supported by theory / experiments.\n2.  Please clarify why the remasking posterior makes a distinction between masked and unmaksed tokens.", "limitations": "Limitations are not discussed in the main paper.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "## Strengths\n* The paper proposes a novel posterior for masked diffusion models, which allows for 'remasking' of unmasked tokens. This allows to correct early mistakes in the sampling process and with that increases the quality of the generated output.\n* The applicability of the proposed method is experimentally validated across different domains (language, images, small molecules) and on each seems to outperform the baseline methods (especially masked diffusion without remasking).\n* The paper is generally well written, with a lot of support for proofs in the appendices.\n\n---\n## Weaknesses\nThere are no big weaknesses in the paper, I've only some minor comments:\n\n*  Eq 2, for me it is not directly clear why from this categorical distribution it follows that any unmasked token must remain unchanged. However, to be honest, masked diffusion models are not my direct expertise.\n*  Eq 4, it seems a typo that $q(z_s| z_t = x, x)$ is used, I guess this should be $q(z_s|z_t, m)$, and anyway this is not a distribution.\n*  In the Remasking posteriors (Eq 6), it is for me unclear why there is a difference between $z_t = m$ and $z_t \\ne m$. Why is the sampling time ($t$ and $s$) not used for remasking, it seems plausible that the remasking should be more prevalent in earlier stages of the diffusion process, and hence also depends on $\\alpha_t$ and $\\alpha_s$.\n*  How would the proposed method compare to a model where uniform at random a set of tokens is selected for regeneration?  \n*  The ablation studies (eg on the $\\sigma_t$ value) are put in the appendix. These take such a central role in the methodology that I would like to see them in the main paper. For example: for a specific task, the performance vs $\\eta$ but also how many 're-masked' tokens are there as a function of the $\\eta$ rescale value. To get more insight in how remasking is used by the diffusion model.", "was_revised": true, "final_justification": "The rebuttal was clear and addressed my questions. Thanks.", "rebuttal": "We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns and questions below.\n\n## **Concern 1:** Why does Eq. 2 imply the “failure-to-remask” property?\n\nIn the forward Markov chain process of masked diffusion models (defined by Eq.1), each single token only has two possible states: \n1. It can remain unmasked, i.e., it remains as the original token $\\mathbf{x}$, or \n2. It can become masked, i.e., it transitions to the [MASK] token $\\mathbf{m}$ from which it never changes again.\n\nTherefore, 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 \n$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)$.\n\nThis means that once $\\mathbf{z}_t$ is unmasked, the posterior will keep it unchanged, i.e., the “failure-to-unmask” property. \n\n## **Concern 2:** Clarify the meaning of $q(\\mathbf{z}_s \\mid \\mathbf{z}_t=\\mathbf{x}, \\mathbf{x})$\n\nTo 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.\n\n## **Concern 3:** Why distinguish between $\\mathbf{z}_t=\\mathbf{m}$ and $\\mathbf{z}_t \\neq \\mathbf{m}$ in the posterior?\n\nAs 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). \n\n## **Concern 4:** Usage of time steps in the posterior\n\nTo 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.\n\n## **Concern 5:** Compare with an entirely random remasking model\n\nEntirely 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. \n\nAs 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. \n\n||Countdown pass@1%|TruthfulQA $\\Delta$ROUGE 1/2/L|\n|-|-|-|\n|LLaDA|45.2±0.2|27.1±0.4 / 30.1±0.4 / 27.2±0.4|\n|LLaDA-randomremask|45.3±0.3|27.6±0.4 / 30.5±0.4 / 27.7±0.4|\n|LLaDA-ReMDM|**46.1**±0.2|**29.5**±0.4 / **31.8**±0.4 / **29.5**±0.3|\n\n## **Concern 6:** Move more ablations to the main paper\n\nWe 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.\n\n## **Concern 7:** Typos and paper formatting issues\n\nWe 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.\n\n## **Question:** Clarify the claim for better controlled generation\n\nBy “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. \n\nThe 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.\n\n---\n\n**References**\n\n[1] Sahoo, Subham, et al. \"Simple and effective masked diffusion language models.\" Advances in Neural Information Processing Systems 37 (2024): 130136-130184."}, {"review_id": "gAdJlxwwER", "reviewer": "Reviewer_oMMK", "summary": "This paper aims to solve the following issue: when inference is performed using masked diffusion models, e.g., MaskGiT, after a token has been unmasked, it is never updated again. The authors propose a scheme in which the unmasked tokens can be remasked, which they compare to inference time compute scaling. Effectively, the authors introduce a tunable probability of remasking and show that this does not change the marginals of the path probability. The algorithm resembles a predictor-corrector type sampling scheme which are widely used in continuous diffusion models. The remasking scheme is employed at inference time in a number of examples that illustrate that it improves sample quality.", "questions": "I am a little confused about the implications of the negative evidence lower bound defined at line 112. The authors point out that because the second term in the expectation defining $\\mathcal{L}^{\\sigma}$ is monotonically increasing in the remasking probability $\\sigma$, it is the case that the MDLM NELBO is tighter. Is there a good theoretical explanation of why using a sub-optimal bound would improve performance asymptotically or non-asymptotically?\n\nA related question to the first, but in the more practical context: a number of different strategies are introduced to determine the choice of $\\sigma_t$. Did the authors try any sort of hyperparameter tuning to obtain the value? Could it be a trainable parameter? If so, is it clear that the optimal value would be non-trivial, i.e., non-zero?\n\nThe experiment on the SMILES guidance task seems to be testing something that is not discussed at all theoretically or conceptually throughout the paper. My understanding is that the authors aim to show that using remasking aids classifier free guidance. Why should this be the case? Is there anything theoretically that would lead one to expect that remasking should be important for guidance?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality: I found the theoretical results to be underdeveloped and, perhaps, incomplete as a motivation for the ReMDM approach. Most the results proved in the appendix are relatively generic calculations that do not provide insight into the ReMDM objective and the context in which it should be most useful. The connections to predictor-corrector scheme strike me as the most useful results, but they don't really emphasize the novelty of the approach. Empirically, the results are strong. Furthermore, the idea, though simple, has a very natural appeal. \n\nClarity: In general the paper is well-written and the figures are both elegantly prepared and informative. The only figure I found lacking was Fig. 4, which is a bit crowded with results that don't obvious converge to anything reasoanble. \n\nSignificance: I think the authors are addressing a very important issue in DDM inference. I am not sure that this will be the ultimate solution to the problem and maintain some skepticism because I think the theoretical justification is currently weak, but it is clear that more complicated unmasking distributions should be explored to improve sample quality.\n\nOriginality: I am not aware of another work taking this approach. I would appreciate it if the authors could add a more thorough discussion of the current suite of predictor corrector schemes for discrete diffusions and more clearly explain the algorithmic difference. For example, https://arxiv.org/pdf/2407.21243.", "was_revised": true, "final_justification": "The authors adequately addressed my concerns and clarified my misunderstandings of their submission.", "rebuttal": "We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns below.\n\n## **Concern 1:** Theoretically justify why ReMDM is better than MDLM\n\nWe hypothesize that there are two reasons why ReMDM has better sample quality than MDLM.\n* 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.\n* 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.\n\n> 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}$.\n\nBelow 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.\n\n||MDLM|ReMDM|\n|-|-|-|\n|$p_\\theta(AA)$, $p_\\theta(BB)$|$\\frac{n}{4T^2}$|$\\frac{n}{4T^2}[\\frac{n\\sigma}{2-\\sigma}+(1-\\frac{n\\sigma}{2-\\sigma})(1-\\sigma)^{2m}]$|\n|$p_\\theta(AB)$, $p_\\theta(BA)$|$\\frac{2n^2-n}{4T^2}$|$\\frac{4n^2-3n^2\\sigma}{4T^2(2-\\sigma)}-\\frac{n(2-(n+1)\\sigma)}{4T^2(2-\\sigma)}(1-\\sigma)^{2m}$|\n|$p_\\theta(AM)$, $p_\\theta(BM)$, $p_\\theta(MA)$, $p_\\theta(MB)$|$\\frac{n(T-n)}{2T^2}$|$\\frac{n(T-n)}{2T^2}$|\n|$p_\\theta(MM)$|$\\frac{(T-n)^2}{T^2}$|$\\frac{(T-n)^2}{T^2}$|\n\nAs 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. \n\nThis 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.\n\n## **Concern 2:** Why is ReMDM better for guidance?\n\nNote 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.\n\n## **Concern 3:** Clarify on ReMDM’s NELBO\n\nTo 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). \n\n## **Concern 4:** Clarify hyperparameter tuning of $\\sigma_t$\n\nAs 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).\n\n## **Concern 5:** Discuss additional related works \n\nIn 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.   \n\n## **Concern 6:** Figure 4 is hard to read\n\nWe 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.\n\n---\n\n**References**\n\n[1] Lezama, Jose, et al. \"Discrete predictor-corrector diffusion models for image synthesis.\" The Eleventh International Conference on Learning Representations. 2022.\n\n[2] Zhao, Yixiu, et al. \"Informed correctors for discrete diffusion models.\" arXiv preprint arXiv:2407.21243 (2024).\n\n[3] Schiff, Yair, et al. \"Simple guidance mechanisms for discrete diffusion models.\" arXiv preprint arXiv:2412.10193 (2024)."}, {"review_id": "5rKgYZcZwz", "reviewer": "Reviewer_xeDg", "summary": "This paper proposes a remasking mechanism in the inference/backward process for masked diffusion models. The authors showcase that the ReMDM yields the same ELBO (up to some reweighting) and marginal distribution as MDLM, and draw the connection of the proposed samplers to existing samplers in the literature. Different design space of the ReMDM samplers are also explored for ablation. The experiments on text, image and molecules are all conducted to reflect the effectiveness of the approach, in terms of inference-time compute scaling.", "questions": "1. In text generation method, what will be results for using ReMDM with confidence-based schedule? In addition, can authors explain that, why $\\sigma_t>0$ leads to better generation perplexity compared to $\\sigma_t=0$ (i.e. original MDLM), which is opposed to the results in DDIM vs DDPM in continuous diffusion setting?\n\n2. In actual inference of large masked diffusion models like LLaDA, there is one step of token unmasking ordering. Here, for ReMDM, does each step require masking or unmasking for every token given the remasking scheme? If not, will it happen that there will be still masked tokens after the whole unmasking steps? Also what is the baseline inference method for MDLM when T>L?\n\n3. How the performance will differ when using different max-capped schedule and rescaled schedule? Look like in the paper, the authors only picked one best performing value, but a study of the performance with respect to the change of these constants could be more than beneficial to understand the improvement brought by ReMDM.", "limitations": "Please see weakness and question.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n1. The motivation and idea of remasking is novel and interesting; this seems to have good potential to enlite further research among training-free methods for improving MDLM. \n2. The authors conduct quite thorough experiments to showcase the effectiveness of their methods, especially improved metrics in terms of inference time scaling. \n3. Connections of ReMDM with DDIM for uniform noise, FB corrector proposed in earlier works are also demonstrated to help better explain ReMDM. Other theorectical augments and proofs are also comprehensive for better understanding.\n\n**Weakness:**\n1. **Mathematical Novelty**: Despite the new high level idea applied to masked diffusion models, the overal mathematical formulations seem quite reminiscent to DDIM for continuous diffusion and the scheme proposed in appendix of DDIM for discrete space scenario, which makes the novelty partially discounted.\n2. **Significance**: Open source discrete diffusion models are currently slow compared to AR models if not optimized. The remasking scheme's benefits are undermined by the already too long inference time of discrete diffusion models, which makes the whole work less appealing. In additional, the performance gain on large diffusion language model like LLaDA is rather minimal.\n3. **Clarity, Typos and Writing**: The general writing of the paper still has much space to improve, and the manuscript needs at least one round of thorough polishing given numerous misleading points:\n\n* (a). The introduction of MDLM and the motivation of \"failure to remask\" from line 74 to 79 is unclear. Though experienced readers can absolutely understand, it is too short to serve as a general introduction for the general audience to motivate the study in this paper.\n\n* (b). The training objective in Equation 3 popped with almost no explanations. In addition, should $z$ start with $z_1$ instead of $z_0$, as $z_0=x$?\n\n* (c). In line 90, the terminology \"complete forward process\" seems awkward and not appropriate, Equal (5) does not directly define any forward process, like in Equation 8 of DDIM paper (this part is instead presented in Appendix A.2) and is more like joint distribution/marginals. I strongly encourage the authors to reconsider the wording and not obuse the terminologies.\n\n* (d). On line 92 and 97, it shall be Equation (6) instead of Equation (51).\n\n* (e). In Equation 8, the i is not explicitly written out.", "was_revised": true, "final_justification": "My concerns are addressed.", "rebuttal": "We thank the reviewer for their constructive and valuable feedback. We address the reviewer’s concerns and questions below.\n\n## **Concern 1:** Clarify on ReMDM’s mathematical novelty\n\nAlthough ReMDM and DDIM share the same high-level idea of non-Markovian diffusion processes, ReMDM enjoys the following differences from DDIM: \n* 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. \n* The derivation of ReMDM allows us to easily reason about the bounds of the $\\sigma_t$ value, making the strategy exploration more rigorous.\n* ReMDM proposes several design strategies, further boosting the empirical performance.\n* ReMDM connects itself with existing predictor-corrector samplers and showcases that existing discrete predictor-corrector samplers are ReMDM’s special cases.\n\n## **Concern 2:** Large masked diffusion models like LLaDA are slower than AR\n\nWhile 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.\n\n## **Concern 3:** The magnitude of the performance gains on LLaDA.\n\nBelow, 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.\n\n|MATH500 pass@1|L=256|L=512|\n|-|-|-|\n|LLaDA|33.68±0.80|37.64±0.36|\n|LLaDA-ReMDM|**35.12**±0.40|**39.36**±0.67|\n\nWe 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.\n\n## **Concern 4:** Typos and writing issues\n\nWe thank the reviewer for pointing these out. We will polish the paper’s writing in our updated manuscript.\n* **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.\n* **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. \n* **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]).\n* **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:”. \n* **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).\n* **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.\n\n## **Question 1:** Report ReMDM-conf results in text generation\n\nFor the inference-time scaling setting ($T\\geq L$), we combine the confidence-based schedule with ReMDM-loop and have the following results:\n\n|||MAUVE($\\uparrow$)|Gen PPL.($\\downarrow$)|Entropy($\\uparrow$)|\n|-|-|-|-|-|\n|$T$=1024|ReMDM-loop|**0.403**|28.6|5.38|\n||ReMDM-loopconf|0.232|31.0|5.36|\n|$T$=2048|ReMDM-loop|**0.610**|22.8|5.30|\n||ReMDM-loopconf|0.394|24.0|5.27|\n|$T$=4096|ReMDM-loop|**0.656**|17.6|5.20|\n||ReMDM-loopconf|0.482|18.6|5.20|\n\nFor the fast sampling setting ($T<L$), we combine the confidence-based schedule with ReMDM-cap and have the following results:\n\n|||MAUVE($\\uparrow$)|Gen PPL.($\\downarrow$)|Entropy($\\uparrow$)|\n|-|-|-|-|-|\n|$T$=128|ReMDM-cap|**0.057**|42.5|5.43|\n||ReMDM-capconf|0.043|50.6|5.47|\n|$T$=256|ReMDM-cap|**0.216**|30.5|5.34|\n||ReMDM-capconf|0.099|34.0|5.32|\n|$T$=512|ReMDM-cap|**0.350**|21.1|5.21|\n||ReMDM-capconf|0.194|30.7|5.31|\n\nThe 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. \n\n## **Question 2:** Why does $\\sigma>0$ yield better Gen PPL.\n\nDespite 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).\n\n## **Question 3:** Comparison between ReMDM and LLaDA samplers\n\nFor 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).\n* 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)$.\n* 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.\n\n## **Question 4:** Ablation studies on $\\sigma$ for ReMDM-cap and ReMDM-rescale\n\n* 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. \n* In Table 7, we reported different $\\sigma$ results for ReMDM-cap and ReMDM-rescale on image generation. \n* In Figure 10-25, we reported different $\\sigma$ results for ReMDM-rescale on the QM9 guidance experiments.\n\nThe 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. \n\n---\n\n**References**\n\n[1] Inception, et al. \"Mercury: Ultra-Fast Language Models Based on Diffusion.\" arXiv preprint arXiv:2506.17298 (2025).\n\n[2] Sahoo, Subham, et al. \"Simple and effective masked diffusion language models.\" Advances in Neural Information Processing Systems 37 (2024): 130136-130184."}]

Author rebuttal:

Meta-review: Backward sampling in the standard masked diffusion model (MDM) does not change the values of unmasked tokens. This 
means that errors (e.g. due to discretization) can accumulate. The paper proposes an interesting remasking mechanism that 
can improve sampling of pre-trained masked discrete diffusion models. This remasking allows unmasked tokens to be "remasked" 
which can lead to some refinement/correction of the sampled tokens. A sequence of experiments in unconditional and conditional generation show that this approach can be useful in practice. 

The paper is very well written and  it presents a clear and novel idea. All reviewers found this a solid contribution for masked 
diffusions. Please address all reviewers' comments in the final version.
Decision: Accept (poster)