Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
Authors: Chen-Hao Chao
Wei-Fang Sun
Hanwen Liang
Chun-Yi Lee
Rahul Krishnan
Venue: NeurIPS 2025 poster ()
Keywords: Discrete Diffusion Models
Masked Diffusion Models
Generative Models
Primary area: deep_learning

Abstract: Masked diffusion models (MDM) are powerful generative models for discrete data that generate samples by progressively unmasking tokens in a sequence. Each token can take one of two states: masked or unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading to redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens to take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive a variational training objective and introduce a simple architectural design to accommodate intermediate-state inputs. Our method demonstrates superior performance across a diverse set of generative modeling tasks. On text data, it achieves a perplexity of 15.36 on OpenWebText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, comparable to leading continuous generative models.

Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
Chen-Hao Chao', Wei-Fang Sun', Hanwen Liang', Chun-Yi Lee", Rahul G. Krishnan
University of Toronto, Vector Institute
NVIDIA AI Technology Center
National Taiwan University
{chchao, rahulgk) torontooooooooooooooooooooooooo. edu, cylee@csie atu..du..e.
toronto...
Abstract
Masked diffusion models (MDM) powerful generative models for discrete data that generate samples by progressively unmasking tokens a sequence. Each token take of two states: masked unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive variational training objective and inputs.
introduce simple architectural design to accommodate intermediate-ste Our method demonstrates superior performance across a diverse set of generative modeling tasks. text data, it achieves perplexity of 15.36 on OpenWebTee WexText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on letge comparable to leading continuous generative models.
Introduction
Discrete data generation has been central focus of probabilistic modeling since the early development of neural networks [1+ The goal is build model pable of generating sequences of symbolic units, known tokens, which can represent words in natural lan- guage data or pixels in images. The field has been largely shaped by autoregressive models (ARM) (e.g.,
1,200
MDM
MDM-Prime
Number of Idle Steps
1,000
MDM-Prime
600
400
(466), which capture the distribution of sequence data by factorizing according to prespecified left-to- right order. Recently, promising results from [7] g have demonstrated that arder-agnostic models, such as masked diffusion models (MDM) (c.g., [0).. can be effectively extended large-scale generation tasks, opening new venue for discrete generative modeling.
200
500
1,000
1,500
2,000
Number of Sampling Steps
Figure 1: Number of idle steps during the reverse diffusion processes of MDM and MDM-Prime. The results are averaged over ten runs. é is the sub-token sequence length.
Masked diffusion models latent-variable generative models that introduce noise by progressively masking tokens within sequence over time. During the reverse diffusion process, masked tokens
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Unmasked States
Intermediate States
Masked State
Token Semanties:
bird
dog
frog
bird bird dog
dog frog
v frog
Smaskk
Token:
0
3
Sub-roken:
(C=
Sub-teken Sequenee:
2
00
00
m
m
I
III MII
(a) Standard Masked Diffusion
Masked Diffasion with Prime (Ours)
Sampling Process
State Transition Tree
Sampling Process Sequessod
State Transition Tree
Sequence
atpp-0
"miakk>
stpp-t)
<mask>
mask ratio
bird
cut
dog
V.
C.V.
bird
dog
frog
200"-3
3
200
0.0
0
10
les.
5
bird
bird
dog
dog
frog
cat
frog
Figure 2: An illustrative example of (a) standard MDM and (b) MDM-Prime. Each token and its corresponding sub-token sequence (constructed via base-b encoding) take one of three states: unmasked, masked, or intermediate. The masked and intermediate states serve as the latent representations produced by the forward diffusion process. This example contains C = 4 possible token classes, labeled as 'bird', 'cat', 'dog'. and 'frog'. l = indicates that each token is represented using two sub-tokens, and = VC = 2 denotes the number of classes per sub-token. The symbol m represents a masked token masked sub-token. The bottom-right sections of (a) and (b) illustrate the state transition trees. MDM- Prime supports transitions through intermediate states while retaining the ability directly reach unmasked states. The bottom-left portions depict the sampling process for a token sequence of length 4. In (a), an idle step occurs between steps 3 and 4. In contrast, (b) demonstrates sampling process without idle steps, which leads improved model utilization.
incrementally unmasked according to a predefined ratio. Each token takes one of two states, masked or unmasked. This binary representation introduces an inefficiency; the entire sequence often remains unchanged across consecutive sampling steps, causing the model to repeatedly process the same input. This phenomenon is illustrated in the green curve in Fig. [ which quantifies the number of such idle steps by simulating the reverse diffusion process of an MDM (). The figure shows that 37% of the 1,024 steps produce no update to the sequence during the reverse diffusion process. This inefficiency motivates investigation of redefining the diffusion process transform these idle steps into informative updates for improving the utilization of the model during generation.
We propose simple yet effective solution that allows each token to take intermediate state, which represents the interpolation between the masked and unmasked states, in the diffusion process. We refer to this method the Partial masking scheme (Prime), and denote MDM augmented with Prime as MDM-PriPn Prime represents each token as a sequence of sub-tokens using a base-b encoding, with masking performed at the sub-token level. Since sub-tokens during the forward diffusion process, this method introduces intermediate states that partially reveal token information. An illustrative example is shown in the top of Fig. [ where unmasked states '0-3' are first encoded as '00-11', and the intermediate states are obtained by masking one of the sub-tokens in the sub-token sequence. The intermediate states enable MDM-Prime to perform fine-grained denoising process. For example, as illustrated in the 'State Transition Tree' of Fig. 2 (b), four-choice prediction can be decomposed into two binary decisions during the sampling process (e.g., mm 1mm-ml-11)) Compared to standard MDM transitions (i.e., Fig. 2 (a)). MDM-Prime is capable of making predictions based on partially observed token information while deferring the final token revelation until later sampling steps. With its ability to transition through intermediate states, MDM- Prime demonstrates improved model utilization during sampling, as reflected in the reduced number of idle steps in Fig. (i.c., the purple curves). This in turn leads to enhanced performance. as later presented in Section The contributions of this work are as follows:
be masked independently
We propose MDM-Prime, generalized MDM framework that enables intermediate token transi- tions. This framework can be optimized through a variational upper bound, which approximates the negative log- likeelio.ooo.
We present a simple implementation of MDM-Prime built upon the standard MDM architecture. Our design requires only minor modifications to the input embedding layer of the standard MDM.

---

We demonstrate that MDM-Prime achieves superior performance on both image and text generation tasks. On the OpenWebTex, dataset (TT). MDM- Prime attains an evaluation perplexity of 15.36. outperforming ARM (17.54), MDM 7 (17.58). To the best of our knowledge, this is the first MDMM approach to surpass ARM without relying on the autoregressive formulation. Furthermore, MDMM-rime achieves FID scores [14] of 3.26 on CIFAR-I0 [15] and 6.98 on ImageNet-32 [16), demonstrating competitive performance comparable to leading continuous generative modeling methods [17H19].
g 10 12 212), and their hybrid variants [ 13
2
Background
We first provide background on continuous-time masked diffusion processes. Section 2.1 describes the forward process, while Section 2.2tioon222eee elaborates on the reverse process.
2.1 Forward Diffusion Process
Let X' é {0, {o..... C 1} denote set of C' tokens, and let = [x' .z"1e E X be sequence of tokens with length L, where the superseript indicates the index of each token. Let xn denote the sample drawn from data distribution Pdaaa which is defined
probability mass function (pmf) over XL_ Given a continuous time variable t E [0, 1]. the latent variable introduced by
the forward diffusion process is denoted as I₁ E (X' U {m})² - X4, where m represents the masked token. In addition, let (x) be the Kronecker delta function, which equals 1 if x z' and 0 otherwise. The forward diffusion process is performed through an element- conditional sampler q( (2o) II
((z](=]) constructed by interpolating between d-(-) and dug (-), defined
as follows 9 10 20
(<)1=²/) (1 au (z)))) + arg2g (zz).
where a E [0. 1] a strictly decreasing scheduling function with boundary conditions og = 1 and on Rs O. Intuitively, each perturbed token æt retains the original value s with probability oz, and is replaced by m with probability 1 at- time 1, the latent variable æ1 = [m, ........ , consists entirely of masked tokens, exhibiting randomness and revealing no information about do- This process can also be interpreted as decreasing the mutual information between z and zó over time according to a 121 20 (i.e., [ xn)) a, (x(), where I1 (z/) denotes the entropy of x/)) as explained in Appendix A.1.2 Based on Eq. (.. the forward diffusion process can be accomplished by first drawing data sample æy Pa (-) and then applying masking obtain I; g(-(2z)).
2.2 Reverse Diffuusio Process
Let and be two time variables that satisfy 0 < S 1. The reverse diffusion process is performed by iterating through p(æ,, (xtl2)))....)))))<))) starting from æ1. The distribution p(æ, xx) can be derived using the conditional distributions ((x,)x,)) and (xx, (x,, 2o). In particular, the transition distribution ((x//22 = I (xx) [x)) is defined to be absorbing (20] on the masked state (i.c., a masked token remains masked from l), and is derived from Eq. as follows [10]:
a.-ax ah (x:)
(a(2)) 8 (x:)
if x E X if x = m.
a( (
8my
Based on Eqs. (1) and (2), the posterior distribution q( lat, ma) = I q(ar) |ni, .t) can be derived using Bayes' rule and the Markov property of the diffusion process, and is expressed as 9 10:
8 (xr) &m (xt
if x E X.
((
a.- I-ar
(x+)
if x m.
indicates that, given zi and observing a masked token oi = m, the reverse process transitions the masked token to its original value z/ with probability a.-a: or retains its value with probability T-a.- For the unmasked token of E X,, its value remains unchanged for the remaining steps.
Eq. (3)
lll///99/ Since (x,) = Eytoro me [/(x=///<__..)00 |x,, =a)] 20), the reverse diffusion process can be performed by first drawing æg ~ (mr.) and then sampling a. ~ q(-12z ar)) Many recent works [0 20 22
24 choose to model p( (x:) po xz)), where denotes the model parameters. To facilitate
3

---

computational efficiency, this distribution is typically factorized as pe (2 - -I== po po x+) g 10 20 24). The parameter can be optimized by estimating the negative log- ligliielihood log po (an) using variational upper bound, written as 99 I0I:
Lan 0) = S-¹ di 1 0t
Σ log pa (a6) dl,
Reteo (20)
where 0f = a.. This objective specifies a cross-entropy loss, weighted by the coefficient T-o and satisfies (Zo 0) 2 log po (#t)
p 10]. In addition, since unmasked elements in 2, retain their values over time according to Eq. 3 the carry-over parameterization [9] (also referred as the mean parameterization in [TO] can be applied by explicitly setting the corresponding entries in Zo to the unmasked values in Tt. Formally, this is expressed as pe (z) ببg (zt) for position i
where =t X. This technique ensures that the model avoids redundant predictions unmasked elements, and therefore can effectively reduce Linh [9]
3 Methodology
In this section, we introduce a new MDM framework that incorporates the Partial masking scheme (Prime), referred to as MDM-Prime. We begin by introducing an invertible function for constructing sub-token sequences and define a novel masked diffusion process over these sub-tokens, which enables intermediate state transitions and reduces the number of idle steps (Section We then propose a novel parameterization that captures sub-token dependencies via joint probability distributions, and describe a simple model architecture design that operates over sub-token inputs (Section 3.2
3.1 Discrete Diffusion via Partial Masking
The core idea of Prime is to represent each token si with a sub-token sequence y WM Uy 6. allowing the creation of intermediate states during the element-wise masking of the forward diffusion process (i.e.. Eq. (T). Given target length l >> 1, this can be achieved through the use of an invertible function f. which maps each token z E to its base-b encoding yó E y (i.e., sequence of sub-tokens), where & fo,...
1} denotes the of sub-tokens with [VC] N.
The invertibility of ʃ enables the likeli- hood to be defined the set of sub-tokens V. For convenience, we slightly abuse no- tation and extend f to accept vector inputs: f(xo) A (=])), (10 1. Its inverse is denoted formation, the pmf of the data, Pamta (2o). can be cquivalently written Pdata f according to the change-of- varraabll principle. Rather than di- rectly modeling Vo, model yo through MDM, and reconstruct Zp using f-((zo)) (3)0). resulting in a parameterized pmf po())) = Po a f(zo). An illustrative example is provided in Fig.
X
Param. pmf
(
(V(Uo)
f( = [y].-.. vol yo- f-1 Under this trans-
....,,.. (vo.........
1.
(%o I.1 y/ (..)
Data pmf
Pdata
Pdata o ff (y))
Figure 3: Illustration of the data and parame- terized pmf with an invertible f. 'Param. pmf" denotes the parameterized pmf captured using an MDM with parameter f. In this example, e 2.
To learn (yo) using MDM, each data point Do first transformed into yo (wx)) and its corresponding latent variable y: is sampled from q(3+)0)) = II II; q(vi"" 12 according to
Eq. T by substituting x's with y's. The reverse diffusion process is parameterized by (o(vl))))
I po (a(a y.). and the model is trained to minimize the loss Lvs( 0) as in Eq. 4
('¹ 0/ Jo 1-
E (2r0) log po ( (pblee) dt.
Lvb (yo 0)
(5)
Section 3.2 provides details on the parameterization of po ( 9.). while Appendix A.2.1 shows that Eq. 5 defines variational bound that approximates negative log- likeliield (NLL).

---

(a) Independent
Softmax Distribution
Zeros
Vi (m,
(0,m) m)
(m, 1, m,
(0, 0,m)
000
000
000
000
001
001
001
001
Py (3/) v
010
010
010
010
011
011
011
011
100
100
100
100
(b) Joint
101
101
101
101
110
110
110
Figure 5: An illustration of the proposed carry-over parameterization technique. In this example, C 3. and b = 2. The condi- tional distribution pe (vály)) is defined in Eq. 7 Softmax distributions are formed by normalizing the correspond- ing logit outputs highlighted in yellow.
2
(1=3)
(f= 4)
Figure 4: Distributions modeled by MDM-Prime using (a) independent and (b) joint parameterizations. Models are trained on a two-dimensional synthetic dataset with at E [0.... 511]²² representing the coordinate the figure (512 x 512). Brighter regions indicate higher probabilities. Experimental details are offered in Appendix A.4.1
This new diffusion process operates on an augmented set & Du [m). For each token, the number of intermediate states introduced by f and the diffusion process is given by 1 = (b+1)² (0+1)=-(66 (oj-((+1 Proposition A.3 ensures that this number is always positive. This property allows MDM--rime to learn smooth transitions via a rich set of intermediate states. Moreover, Proposition [A.] and Eq. (A5) formally establish that the number of idle steps decreases as l increases, which guarantees an improved model utilization of MDM-Prime during the reverse diffusion process.
3.2 Parameterization
In this section, we discuss our implementation for pe yó ly:), which comprises a decoder for modeling the distribution of yó and an encoder for processing the input We begin by outlining the decoder design, followed by description of the encoder.
Decoder Design via Joint Probability. A straightforward approach to implementing (o(vé y:)
to factorize it as I[ Po this factorization allows for easy application of the carry-over parameterization [ 10 (discussed in Section 2.2), it introduces two key challenges: (1) an independence assumption and (2) potential
Po (vo [v:), modeling each component with a softmax distribution. While
generation of invalid samples. First, the factorized form pe (vbl3r) ve) = II (v6 yy)) imposes
an additional independence assumption across sub-tokens, preventing the model from capturing dependencies between them. As shown in Fig, 4 (a), increasing the sub-token sequence length leads to a deterioration in the sampling distribution due to this limitation. Second, since f is defined as injective function but not necessarily bijective (i.e., (f() I [)* some samples yo produced by such factorized model may not correspond to any valid To under the inverse mapping f-1 For instance, when encoding the GPT-2 [5] vocabulary of C 50, 257 classes with E 4 and b = [VC]- 15, the model may generate invalid sub-token sequences such y = (14, 14, 14, 14) during inference time, even though there no corresponding =6 = 50, 624 for decoding.
To address these two issues, propose model the joint distributions pe (vú (y:) of sequence of sub-tokens while explicitly zeroing out the probability mass assigned to invalid samples. This is achieved by parameterizing only the logits of base-b encoding y E f(X) that correspond to valid aó E X, which results in exactly entries in the logit outputs for each position i {1, ... ,L).
Based on the above joint probability design, to further support carry-over parameterization for MDM- Prime, the element-wise distribution should be defined as po (vb (y.) A 5, (vo') for all position
i.j where v y (following the end of Section 2.2 Since parameterize the joint distribution as
5

---

vo(vólue) po
yt). 330 3o vo
this condition is imposed on the marginal distribution as follows:
1. (10) po (o6
Σ M
Pe 4.1
Pe
1y)) - 8 بد (86
(6)
v ey
To meet this condition, the probabilities of y with any element yo that is inconsistent with y should be explicitly set to zero. The parameterized probability can thus be defined as follows:
exp(Ee(1b) ()))
if vó E V(vi),
po(vúlv1))
Neviw exp(Polur(ee))
0,
yb # V(vi).
where V(y() _ (y' 1ytt...... y E f(X) stt. (y w v (2 m)} denotes set of
outputs that is consistent with yt n. and Es ye x JixL + R is a scalar logit. Proposition A.4
guarantees the correctness of this parameterization (i.e., Eq. ( satisfies Eq. (6) Consider simple
example where C 7, 8 and b 2, then V(yi) corresponds to the following sets:
y (m, m, m), then V(ys) = {(0, 0, 0), (0,0, 0, 1), (0, 1, 0), (0, 1, 1), (1,0, (1, (1,1,, If y (0, m, m), then V(y(() {(0, 0, 0), (0,0, (0, (0,1,))
If y (0, m). then V(w(()
(0,
Note that (1, 1, is invalid since C= in this case. Some illustrative examples of (v( (y:) are provided in Fig. [ As the reverse diffusion process progresses, the number of unmasked sub-tokens in y increases, leading to substantial reduction in [V(y{)). This results in a decreasing number of candidate classes of y over time, and thus explicitly reducing uncertainty in the prediction task. From an implementation perspective, Eq. by v During the forward pass of the model, the logits of yó & (vy) excluded according to these filters. Further implementation details are provided in Appendix A.2.3 and Fig. A2
be efficiently derived using precomputed filters indexed
Encoder Design for Processing Sub-tokens. In contrast to the decoder, where the distri- bution over sub-tokens Vá E f(X) is repre- sented jointly using logit outputs with C' en- tries (see Figs. and (). the encoder receives noised inputs y that lie in the augmented set 5tt Since the size of this set, [5ae may grow with f and typically exhibits | C (also see Appendix (.2.2), creating an embedding lookup table for yi is impractical due to the resulting growth in the number of parameters in it. To address this issue, we propose to model each sub-token embedding separately (i.e.. creating a lookup table for individual vidd )), followed by merging operation to produce a token embedding.
RexL
RCxL
Lagiss
Loghts
Neural Network (0)
Neural Network (0)
DXL
Concatenation
12 Dx
R 2 Xxxx
Embedding Lookup Table (8)
Embedding Lookup
V Jext
IE E X (b) Standard MDM Architecture
(a) MDM-Prime Architecture
Figure 6: Comparison between (a) MDM-Prime and (b) standard MDM architectures. The embedding lookup table in (a) has fewer learnable parameters since [D] < | and D/L D.
In our approach, simple merging operation based on concatenation is employed. Let D denote the dimensionality of the token embed- ding vector. Each sub-token first embedded into a vector of size D/E. and the resulting embeddings are concatenated to form D-dimensional token embedding vector. This token embedding can then be processed by an arbitrary downstream neural network, which allows to reuse the standard MDM architecture. A comparison between this design and a standard MDM architecture shown in Fig. 6 Alternative merging strategies, such the Perceiver [25] cross-attention mechanism, are discussed and evaluated in Appendices A.3 and A.5.2 where we show that simple concatenation yields the best performance.
In summary, adapting standard MDM to MDM-Prime requires only minimal architectural modifica- tions on the embedding layer. This simple strategy preserves the overall architectural design of the standard MDM, enabling fair comparison with our baseline in the following experiments.

---

Table I: Zero-shot validation perplexities evaluated on seven textual datasets. Lower values corre-
spond to better performance. Methods marked with * incorporate autoregressive formulation. MDLM-Primms exhibits improved results on LAMBADA, PTB, and ArXiv.
LAMBADA
WikiText
PTB
LMIB
AG News
PubMed
ArXiv
ARM 9
51.28
25.75
82.05
51.25
$2.09
49.01
41.73
BD3-LM 13
$50.03
$31.31
$96.81
$60.88
$61.67
<42.52
$39.20
EDLM-cOAR* 12
$50.04
$28.31
$89.73
$60.23
$57.94
546.31
$39.02
SEDD [7
249.86
5.3428
$100.09
$68.20
$62.09
$41.89
$,38.48
EDLM-NCE 12
<46.92
$30.77
$93.21
$63.19
<60.02
<41.0
<36.63
MDLM 99
<47.52
$32.83
$95.26
$67.01
<61.15
<41.9
$37.37
MiLM-- (f=33 2)
$30.91
$27.93
$74.81
$55.50
$63.21
$33.32
$25,44
MDLM- M-rrime (f=3)
$27.75
$27.42
$60.13
$42.69
<62.58
$41.14
$24.71
MDLM- MDLM-Prime (f=4) 4)
$24.44
$23.86
$53.98
$38.02
$39.44
$48.64
$25.83
MDLM-Prinne M-Prime (f-
$25.80
$26.87
$62.62
$45.36
$60.16
$59.09
$25.19
MDLM-Prrme (f=8)
$25.23
$25.77
$53.77
$38.00
$44.89
$54.50
$25.79
4
Experiments
This section presents empirical evaluations to examine the effectiveness of the proposed method. We first report results in the text generation domain Section 4.1 Then, we provide comparisons on the image generation benchmarks in Section
4.1 Text Generation
Configuration. In this set of experiments, models are trained on the OpenWebText (OWT) dataset [T] The data are tokenized using the GPT-2 tokenizer (5), which defines 1, 024 and The Masked Diffusion Language Model (MDLM) 9 is adopted as our baseline, and the experimental setup is consistent with [9]. Prime with different & applied to enhance its performance, and our method is denoted MDLM-Prime in this section. We also include comparisons with several recent works 7 TO 1 12 [33, diffusion transformer (DiT-B) (26 with rotary positional embeddings (27]. The model contains 130M non-embedding parameters. No temperature applied to the output probabilities during training and test time. Detailed hyperparameters are offered in Appendix A.4.3) Additional results, including ablation studies and efficiency evaluation, are provided in Appendices A.5.2 and A.5.5
C' = 50,251 257.
[ In MDLM (9. MDLM-Prime. and all other recent methods [ the core model architecture, i.e.. the 'Neural Network (0)' component in Fig. 6 is
Improvements to Likelihood Evaluation. We eval- uate the models' ability to capture the data distribution using the perplexity (PPL) metric (28). Table 2 reports PPL on OWT, along with the idle step ratio (ISR), which is defined the proportion of idle steps relative to the total sampling steps and computed using Eq. (A6). We observe that increases, MDLM-Prime achieves lower PPL., with performance converging when & 2 4. Since ISR also converges when é > this trend suggests that ISR can serve as an indicator of improved likeli- hood modeling ability. We provide a further analysis of the relationship between performance and ISR, with guideline for selecting f, Appendix /A.2.4] Moreover, MDLM-P M-immee with > 3 outperforms ARM, MDM- based approaches 7 9 I0 [2], and their hybrid vari- ants [I3 12] by noticeable margin in terms of PPL, indicating that incorporating intermediate state repre- sentations allows MDLM-Prime to model data likeli- hood more effectively. Instead following recent ap- proaches 212 L3 that leverage autoregressive formulation to enhance MDM performance, MDLM- Prime maintains an order-agnostic framework while achieving superior performance on textual data.
Table 2: PPL and ISR evaluation on OWT. Methods marked with incorporate autoregressive formulation. The symbol + represents that lower values correspond to better performance. MDLM-Primm with & {> outperforms prior methods.
PPL (
ISR
ARM
17.54
BD3-LMM* 13
$20.73
EDLM-cOAR- 12
$17.58
SEDD
<24.10
GenMD4 [0
$21.80
EDLM-NCE 12
321.52
MDLM 9
$22.98
36.77%
MDLM-Prime (E
2)
$17.90
13.52%
MDLM-Prime (E
$16.36
4.97%
MDLM-Prima (f -4)
$15.62
1.83%
MDLM-Prime (E -6)
115.36
0.25%
MDLM-Prime (f= $15.48
0.03%
Improvements Generalizability to Unseen Text Data. then examine their generalizability to unseen textual datasets. To assess the models' generalizability
With the models trained on OWT, we

---

Table 3: FID and IS evaluation on CIFAR-10. The Table 4: FID and IS evaluation ImageNet-32.
arrow symbols 1 4 represent that higher / lower The arrow symbols t / represent that higher results correspond to better performance.
lower results correspond better performance.
CIFAR-10 10
ImageNet-32
Discrete
Discrete
Method
FID (C)
IS (T)
Method
(4)
(t)
MDM (NFE=512)
4.66
9.09
MDM (NFE=1,024)
7.91
11.60
MDM-Mixture (NFE=512)
4.80
9.22
MDM-Mixtue (NFE=1.024)
8.08
11.56
MDM-Prime (NFE=512)
3.26
9.67
(NFE=1.024)
6.98
11.65
PixelCNN 35
65.93
4.60
Continuous
D3PM Absorb
20
(NFF==1000)
30.97
6.78
QC-NCSN+ ++ 40
19.62
9.94
D3PM Gauss.
20
(NFE=1,000)
7.34
8.56
NDM 41
17.02
CTDD-DG [36
(NFE=1.000)
7.86
8.91
DDPM 18
16.18
Tau-LDR 22
(NFE=1,000)
3.74
9.49
MSGAN 42
12.30
Discrete FM 23
(NFE=1,024)
3.63
i-DODE (SP)
43
10.31
Continuous
i-DODE (VP)
45
9.09
NCSN 33
25.32
8.87
Stochastic Interp.
44
8.49
Continuous FM
38
6.35
Soft Trunc DDPM
45
8.42
11.82
Bit Diffusion
39
3.48
ScoreFlow (subV VP)
19
8.87
StyleGAN+ADA
17
3.26
9.74
ScoreFlow (VP)
19
8.34
DDPM
18
3.17
9.46
Continuous FM
38
5.02
across diverse text domains, we report PPL. on a suite of commonly used zero-shot benchmarks, including LAMBADA (29]. Wiki'Text (30), Penn Treebank (PTB) [31]. 1 Billion Word Benchmark (LMIB) [32). AG News [33, and Scientific Papers (PubMed and ArXiv subsets [34). The results reported in Table MDLM-Prime exhibits superior results on LAMBADA, PTB, and ArXiv, and achieves comparable performance to ARM on WikiText. While it underperforms ARM on AG News, the overall results highlight its superior generalizability across multiple domains. Furthermore, our ablation study in Appendix reveals that the carry-over parameterization plays an important role in enhancing zero-shot pertormance, offering improvements on both LAMBADA and PubMed.
4.2 Image Generation
Configuration. In this set of experiments, models trained and evaluated the CIFAR-10 [5 and ImageNet-32 [16] datasets. For both datasets, the dimensionality is set to L. = 32 x 32 x 3, with C 256 corresponding pixel intensity values. The core model architecture is adapted from the ablated diffusion model (ADM) [46), which contains 114M parameters and is the same as that used in (23). Sample quality is evaluated using the widely adopted Fréchet Inception Distance (FID) [4] and Inception Score (IS) (47 metrics. Experimental details are provided in Appendix A.4.2
Additional results are presented in Appendices A.5.5 A.5.6
Improvements to Sample Quality. We first compare MDM-Prime with varying val- ues of l against the baseline configuration (i.e..
17.5
ISR (%)
NFE-512
15.0
é= 1). Due to the relatively small number of classes in the image experiments, we addition-
10
12.5
10.0
ally explore the case of ( 2/3, where pixel values are merged into super- super-pixels, resulting in b = 2/$/256 = 4, 096 discrete classes in one of the experimental settings. As shown in Fig. 7 the configuration with 1 = 2 achieves the best FID scores and exhibits relatively low ISR compared to that of & = 1 and 2/3. While the settings with form comparably to the baseline, we observe
7.5
5.0
4-
2.5
2
2/3
2/3
4
4.096
256
4,096
256
Figure 7: FID and ISR evaluated under different CIFAR-10. NFE denotes the number of function evaluations during sampling.
3 and = 4 per-
that models with lower ISR are less sensitive to the number of function evaluations (NFF) during sampling, as reflected by the smaller FID performance gaps between NFE=128 and 512.
8

---

The benchmark results are reported in Tables 3 and which include two baselines, MDM and MDM-Mixture as well as several existing generative modeling approaches (17) 22 The MDM baseline corresponds to the standard configuration with & = 1, while Mix-urxttue extends this baseline by incorporating a mixture distribution using an auxiliary variable, similar to [48] In this comparison, MDM-Prime adopts é = 2.
23 35 45
As shown in the tables, MDM and MDM-Mixture are inferior to MDM-Prime. On CIFAR-IO, MDM-Prime achieves better results than the other discrete generative models while requiring fewer NFE, and attains performance comparable StyleGAN+ADA [7]. On ImageNet-32, MDM-Prime demonstrates improved performance over existing continuous diffusion and score-based based models (i.e., achieving an FID improvement of 1.36 over ScoreFlow (VP) [9]
18 19
45)
Imputation under Partial Masking. To fur- ther evaluate MDM-Prime's capability for condi- tional image generation, we conduct experiments
Condition
Imputation
on image imputation. In contrast to text generation, where the positional index j of a sub-token ( lacks explicit semantic interpretation, sub-token positions in the image domain have direct influ- ence on the resulting pixel values. Specifically, the first sub-token yo has the greatest impact on deter- mining the final pixel intensity. In this experiment, 'Condition' in
1.1
Figure 8: Imputation results with conditional images obtained from CIFAR-IO. 10. MDM-Prime generates visually coherent image variants. More examples are provided in Appendix AA.5.5
each conditional image (denoted Fig. () is first encoded into its corresponding sub- token representation. The first sub-token retained, while the remaining sub-tokens are masked and subsequently predicted by our model. Some generated examples are shown in the 'Imputation' section of Fig- [ The results demonstrate that our method can generate visually coherent images conditioned on the preserved sub-tokens, highlighting its effectiveness in controlled image synthesis.
5 Related Works
The general framework of discrete diffusion models was first introduced by (21] where the authors explored modeling binary data using a diffusion process. This framework was extended to real- world applications, such as text and image generation, by the authors in (0). who proposed various perturbation strategies to implement diffusion processes with discrete noise. The authors in [22] proposed generalizing this framework as Continuous- Timme Markov Chains (CTMC). Inspired by the success of scaling MDM for text generation 77). some works 9 10 49) explored simplifications of the training objective of discrete diffusion models with latent variables represented using masked tokens. Other studies 23 50) investigated learning approaches based on a broader class of latent representations formulated through flow matching.
Building on the theoretical foundation of MDM, several enhancement techniques 12 D
23 48
53 have been proposed. In [3), the authors proposed an interpolation between ARM and MDM to capture the left-to-right structure in textual data. In [12), an ARM was employed as an energy-based function to guide the sampling process of an MDM, resulting in improved performance. Other works [23] 51 52) modified the sampling process of MDM to selectively remask certain unmasked
predictions to enhance sample quality. In addition, the authors in [48 [33 explored distillation techniques designed to reduce the number of sampling steps while maintaining sample quality.
Another line of research has explored the use of diffusion models with continuous noise distributions (i.e.. Gaussian) for modeling discrete data. Representative methods include (39 (585 Among these works, Bit Diffusion [39 shares similarities with our approach. In their method, discrete data are first encoded into bit representations, and a continuous diffusion model (with Gaussian kernels) is trained to generate these encodings. The generated outputs are then quantized back into discrete tokens. However, due to its reliance on quantization, the model's 's likelihood becomes intractable, which leads to its inability to directly capture the distribution of discrete data.

---

Conclusion
Scientific progress has continually reshaped our understanding of what constitutes the most basic units of matter. Physicists initially believed that atoms were elementary units of matter. This view changed with the discoveries of the electron, the atomic nucleus, and eventually the development of the standard model [59], which describes fundamental particles, their interactions, and how they combine to form atoms. In the context of generative models, we proposed Prime, a method to decompose the elementary unit of discrete data-tokent totenstoken-into fine-grained subcomponents. MDM- Prime establishes principled framework for perturbing and reconstructing discrete data using sub-token representations. Experimental results on both text and image generation tasks demonstrated that sub-token representations provide a more expressive modeling paradigm. We believe that this framework holds potential for addressing real-world problems that require fine-grained and precise modeling of discrete data.
Acknowledgements
RGK is supported by a Canada CIFAR AI Chair and a Canada Research Chair Tier II in Computational Medicine (CRC-2022-00049) This research was supported by an NFRF Special Call NFRFR2022- 00526. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. The authors gratefully acknowledge the support from the National Science and Technology Council (NSTC) in Taiwan under grant numbers NSTC 1 0-12221---02-069-1 212-MY3, and NSTC Corporation and the NVIDIA AI Technology Center (NVAITC) for the donation of GPUs and access to the T supercomputer. Furthermore, the authors extend their gratitude to the National Center for High-Performance Computing for providing the necessary computational and storage resources. Finally, we thank David Pellow, Vahid Belazadeh Meresht, and the anonymous reviewers for valuable feedback.
NSTC 113-2221-E-O We also express our sincere appreciation to NVIDIA
References
[1] Iyya Sutskever, James Martens, and Geoffrey Hinton. Generating Text with Recurrent Neural Networks. In Proc. Int. Conf. on Machine Learning (ICML), 2011.
[2] Tomas Mikolov, Ilya Sutskever, Anoop Deoras, Hai Son Le, Stefan Kombrink, and Jan Honza Cemnocky, Subword Language Modeling with Neural Networks. 2011.
[3] Alex Graves. Generating Sequences With Recurrent Neural Networks. arXiv: 1308 0850 2013.
[4] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2017.
[5] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
[6] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. LLaMA: Open and Efficient Foundation [cs. CL., 2023.
Language Models. arXiv: 2302. 13971
[7] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution. In Proc. Int. Conf. on Machine Learning (ICML), 2024.
[8] 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: 2502. 09992 [cs. CL], 2025.
10

---

[9] Subham Sekhar Sahoo, Marianne Arrilla, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and Effective Masked Diffusion Language Models. In Proc. of the Int. Conf. on Neural Information Processing Systems (NeurIPS), 2024.
[10] Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K. Titsias. Simplified and Generalized Masked Diffusion for Discrete Data. In Proc. of the Int. Conf. on Neural Information Processing Systems (NeurIPS), 2024.
[11] Aaron Gokaslan, Vanya Cohen, Ellie Pavlick, and Stefanie Tellex Openwebtext corpus. http:// http:://Skyioon007..o wttpriioooooniiithbbboboo ho//ppenNeeTexxtCorpu... 2019.
[12] Minkai Xu, Tomas Geffner, Karsten Kreis, Weili Nie, Yilun Xu, Jure Leskovec, Stefano Ermon, and Arash Vahdat. Energy-Based Diffusion Language Models for Text Generation. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[13] 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. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[14] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs Trained by a Two Time-Scale Update Rule Converge to Local Nash Equilibrium. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2017.
[15] Alex Krizhevsky and Geoffrey Hinton. Learning Multiple Layers of Features from Tiny Images. Technical report, University of Toronto, 2009.
[16] Patryk Chrabaszez, Ilya Loshchilov, and Frank Hutter. A Downsampled Variant of ImageNet as an Alternative to the CIFAR datasets. arXiv: 170.008819 [cs. CV]. 2017.
[17] Tero Karras, Miia Aittala, Janne Hellsten, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Training Generative Adversarial Networks with Limited Data. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2020.
[18] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models, In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2020.
[19] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum Likelihood Training of Score-Based Diffusion Models. In Proc. of Ini. Conf. on Neural Information Processing Systems (NeurIPS), 2021.
[20] Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured Denoising Diffusion Models in Discrete State-Spaces. In Proc. of the Int. Conf. on Neural Information Processing Systems (NeurIPS), 2021.
[21] Jascha Sock-Dics,,ein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep Unsupervised Learning using Nonequilibrium l'ermodynamics. In Proc. Ins. Conf. on Machine Learning (ICML), 2015.
[22] Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and Arnaud Doucet. A Continuous Time Framework for Discrete Denoising Models. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2022.
[23] Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky T. Q. Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete Flow Matching. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2024.
[24] Kaiwen Zheng, Yongxin Chen, Hanzi Mao, Ming- Yu Liu, Jun Zhu, and Qinsheng Zhang Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Sampling. In Proc. Int. Conf. on Learning Representations (ICLR), 2024.
[25] Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General Perception with Iterative Attention. In Proc. Int. Conf. on Machine Learning (ICML), 2021.
11

---

[26] William Peebles and Saining Xie. arXiv:2212. 09748 [vs. 2022.
Scalable Diffusion Models
with Transformers.
[27] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. RoFormer: Enhanced Transformer with Rotary Position Embedding. 2xXiv:2104.0985
(cs.CL), 2023.
[28] Daniel Jurafsky and James H. Martin. Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition. 3rd edition, 2025.
[29] Denis Paperno, Germán Kruszewski, Angeliki I Lazaridou, Ngoe Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Proc. Annual Meeting of the Association for Computational Linguisties (ACL), 2016.
[30] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture 2016.
Models. arXiv: 1609. 07843
[31] Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewiez. Building a large annotated corpus of english: The penn treebank. Computational Linguistics (CL), 19(2):313-330, 1993.
[32] Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Philipp Koehn, and Tony Robinson. One Billion Word Benchmark for Measuring Progress in Statistical Language Modeling. In Proc. Conf. of the European Chupler of the Association for Computational Linguistics (EACL), pages 263-272. Association for Computational Linguistics, 2013.
[33] Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level Convolutional Networks for Text Classification. In Proc. of Int. Conf. Neural Information Processing Systems (NeurIPS), 2015.
[34] Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. A Discourse-Aware Attention Model for Abstractive Summarization of Long Documents. In Proc. Conf. of the North American Chapter of the Association for Computational Linguistics (NAACL), 2018
[35] Aaron van den Oord, Nal Kalchbrenner, Oriol Vinyals, Lasse Espeholt, Alex Graves, and Koray Kavukcuoglu. Conditional Image Generation with PixeICNN Decoders. In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2016.
[36] Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten. Unlocking Guidance for Discrete State-Space Diffusion and Flow Models. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[37] Yang Song and Stefano Ermon. Generative Modeling by Estimating Gradients of the Data Distribution. In Proc. of Inf. Conf. on Neural Information Processing Systems (NeurIPS), 2019.
[38] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow Matching for Generative Modeling. In Proc. Int. Conf. on Learning Representations (ICLR), 2023.
[39] Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning. In Proc. Int. Conf. on Learning Representations (ICLR), 2023.
[40] Chen- Chao, Wei-Fang Sun, Bo-Wun Cheng, and Chun-Yi Yi Lee. On Investigating the Conservative Property of Score-Based Generative Models. In Proc. Int. Conf. on Machine Learning (ICML), 2023.
[41] Grigory Bartosh, Dmitry Vetrov, and Christian A. Naesseth. Neural Diffusion Models. In Proc. Int. Conf. on Machine Learning (ICML), 2024.
[42] Mgoc-Trung Tran, Viet-Hung Tran, Ngoc-Bao Nguyen, Linxiao Yang, and Ngai-Man Cheung. Self-supervised GAN: Analysis and Improvement with Multi-class Minimax Game, In Proc. of Int. Conf. on Neural Information Processing Systems (NeurIPS), 2019.
12

---

[43] Kaiwen Zheng, Cheng Lu, Jianfei Chen, and Jun Zhu. Improved Techniques for Maximum Likelihood Estimation for Diffusion ODEs. In Proc. Int. Conf. on Machine Learning (ICML), 2023.
[44] Michael S. Albergo and Eric Vanden-Eiinden Building Normalizing Flows with Stochastic Interpolants. In Proc. Int. Conf. on Learning Representations (ICLR), 2023.
[45] Dongjun Kim, Seungjae Shin, Kyungwoo Song, Wanmo Kang, and II-Chul Moon. Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation. In Proc. Int. Conf. on Machine Learning (ICML), 2022.
[46] Prafulla Dhariwal and Alex Nichol. Diffusion Models Beat GANs on Image Synthesis. In Proc. of Inf. Conf. on Neural Information Processing Systems (NeurIPS), 2021.
[47] Shane Barratt and Rishi Sharma. A Note the Inception Score. Workshop on Theoretical Foundations and Applications of Deep Generative Models at Int. Conf. on Machine Learning (ICML), 2018.
[48] Satoshi Hayakawa. Yuhta Takida, Masaaki Imaizumi, Hiromi Wakaki, and Yuki Mitsufuji. Distillation of Discrete Diffusion through Dimensional Correlations. Machine Learning and Compression Workshop at the Int. Conf. on Neural Information Processing Systems (NeurIPS). 2024.
[49] 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 2025.
Data. arXiv: 2406. .
[50] Neta Shaul, Itai Gat, Marton Havasi, Daniel Severo, Anuroop Sriram, Peter Holderrieth, Brian Karrer, Yaron Lipman, and Ricky T. Q. Chen. Flow Matching with General Discrete Paths: A Kinetic-Optimal Perspective. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[51] Yixiu Zhao, Jiaxin Shi, Feng Chen, Shaul Druckmann, Lester Mackey, and Scott Linderman. Informed Correctors for Discrete Diffusion Models. arXiv 2407 21243 [cs.LG], 2024.
[52] Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, and Volodymyr Kuleshov. Remasking 2025, Discrete Diffusion Models with Inference- Time Scaling. arXiv: 00307
[53] Justin Deschenaux and Caglar Gulcehre. Beyond Autoregression: Fast LLMs via Self- Distillation Through Time. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[54] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categor- ical data, arXiv: 2211. 15089 [cs. [L), 2022.
[55] Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B. Hashimoto. Diffusion-LM Improves Controllable Text Generation. In Proc. of Int. Conf. on Neural Informa- tion Processing Systems (NeurIPS), 2022.
[56] Ishaan Gulrajani and Tatsunori B. Hashimoto. Likelihood-Basee Diffusion Language Models. In Proc. of Int. Conf. Neural Information Processing Systems (NeurlPS), 2023.
[57] Shujian Zhang, Lemeng Wu, Chengyue Gong, and Xingchao Liu. Language Rectified Flow: Advancing Diffusson Language Generation with Probabilistic Flows. In Proc. of Annual Conf. of the North American Chapter of the Association for Computational Linguistics (NAACL), 2024.
[58] Vincent Hu, Di Wu, Yuki Asano, Pascal Mettes, Basura Fernando, Björn Ommer, and Cees Snoek. Flow Matching for Conditional Text Generation in Few Sampling Steps. In Proc. of Conference of the European Chapter of the Association for Computational Linguistics (EACL), 2024.
[59] W. N. Cottingham and D. A. Greenwood. An Introduction to the Standard Model of Particle Physics. 2007.
13

---

[60] Weiwei Zhang, Jian Sun, and Xiaoou Tang. Cat Head Detection How to Effectively Exploit Shape and Texture Features. In Proc. of European Conf. Computer Vision (ECCV). 2008.
[61] Prajit Ramachandran, Barret Zoph, and Quoc Le. Searching for Activation Functions. arXiv:1718 171.::11000099411 [cs. NE). 2017.
[62] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In Proc. Int. Conf. on Learning Representations (ICLR), 2015.
[63] Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In Proc. Int. Conf. on Learning Representations (ICLR), 2019.
[64] Anji Liu, Oliver Broadrick, Mathias Niepert, and Guy Van den Broeck. Discrete Copula Diffusion. In Proc. Int. Conf. on Learning Representations (ICLR), 2025.
[65] Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: Non-linear Independent Components Estimation. 2015.
[66] Laurent Dinh, Jaseha Sohl-Dickstein, and Samy Bengio. Density Estimation using Real NVP. In Proc. Int. Conf. on Learning Representations (ICLR), 2016.
14

---

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: The main claims presented in the abstract and the introduction section reflect the contributions of this paper. The key theoretical results discussed in Section 3 The experiments in Section provide empirical justification for these claims. Finally, Section summarizes both the theoretical and empirical contributions of this work.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made 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: The limitations of this work are summarized in the Appendix A.6 The discussion covers the assumptions made in this 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-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.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
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 acknowledged
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 to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
15

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The assumptions and proofs of our theoretical results are presented in detail in Appendices A. and A.2
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 in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are 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 be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: This paper discloses the information needed to reproduce the experimental results. The experimental configurations, detailed hyperparameter setups, and hardware requirements for the experiments are elaborated in Appendix A.4
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 perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS 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 the architecture clearly and fully.
model architecture, the paper should describe
(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- 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 partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
16

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or 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, as described in supplemental material?
Answer: [Yes]
Justification: All the data used in this paper are publicly available. The code and instal- lation instructions are available in anonymous repository, with the link provided in Appendix A.4
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
conooninoolllloooooooolmoooimill for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" 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 NeurIPS code and data submission guidelines (https
/
CooeeSumeeeeeeeeeeeeee ssi wwoooonoolllllloooooooooooiiiill for more details.
The authors should provide instructions on data access and preparation, including how 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. If 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 much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The experimental configurations, detailed hyperparameter setups, and hardware requirements are elaborated in Appendix A.4
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 to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or 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: The results in Fig. A1 | show the mean values with shaded regions representing the variances. The caption of Table A3 offers the confidence interval. Results in Tables follow prior work by reporting performance from single run, due to the high computational cost of training multiple variants.
Guidelines:
17

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror 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 capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
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).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
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 or figures symmetric error that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The hardware requirements (e.g., the computational hardware configurations and the execution time) for each experiment elaborated in Appendix A.4
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 compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: All authors have reviewed the NeurIPS Code of Ethics and confirmed that the research conducted in this paper complies with
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 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]
18

---

Justification:
This paper discusses its potential impacts in Appendix A.7
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 why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment 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 partieular 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 an 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 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 (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how 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 models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper poses no risk for misuse.
Guidelines:
The answer NA means that the paper poses no 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.
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 is challenging, and many papers do not require this, but we 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 (c.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: The creators of assets are properly credited through citations, and the license is included in the asset (see Appendix A.4).
Guidelines:
The answer NA means that the paper does not 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 used and, possible, include URL.
19

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssippsiiiie 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-pa se-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 ' 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: The assets (e.g., code, installation instructions, and rumning commands) summarized an anonymous repository, with the link provided in Appendix [A.4] The datasets are publicly available, and the experiments are performed on them with the default setup.
Guidelines:
The answer NA means that the paper does release assets.
Researchers should communicate the details the dataset/code/mooel 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 is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing 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 details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethies, 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 the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not involve research with human subjects.
Guidelines:
20

---

The answer NA means that the paper does not involve crowdsoureing 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
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if 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.
Question:
Answer[[NA]
Justification:
LLMs are only used for editing grammar.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
21

---

A Appendix
In this appendix, provide additional analyses and experiments. Section theoretical properties of masked diffusion models (MDM). Section A.2 offers an analysis of the proposed Partial masking scheme (Prime). D.ction/A.3 MDM augmented with Prime (i.e., MDM- MDM-Prime). Section Section reports additional experimental results. Section Section|A.6jsummen Aummmueesss the limitations this work. Finally, Section A.7 discusses the potential impacts of this work. The following table of contents summarizes the structure of the main manuscript and this appendix.
A. presents some
a number architectural designs of
the experimental configurations.
Table of Contents
Introduction
2
Background
2.1
Forward Diffusion Process
2.2
Reverse Diffusion Process
3
Methodology
3.1
Discrete Diffusion via Partial Masking
3.2
Parameterization
4
Experiments
Text Generation
Image Generation
5
Related Works
Conclusion
10
Appendix
22
A.1
Analyses Masked Diffusion Processes
23
A.2
Analyses of the Partial Mask Scheme
25
A.3
Model Architecture Designs
28
A.4
Experimental Setups
29
A.5
Supplementary Experiments
A.6
Limitations
35
A.7
Broader Impacts and Future Works
35
22

---

A.1 Analyses of Masked Diffusion Processes
In this section, we examine two properties of MDM. In Section A.I.I we derive an analytical expression for the expected number of idle steps in the reverse diffusion process. In Section [A.122] we show that the mutual information between the latent variables and data exhibits linearly decaying trend with respect to the scheduling function over time.
A.1.1 Expected Number of Idle Steps
In Section T we show that MDM may have significant number of idle steps during the sampling process. In this section, we derive a closed-form formula for calculating the expected number of idle steps and provide the reason why Prime must reduce this number.
Proposition A.1. Let L be the token sequence length, T be the total number of discretized timesteps for the sampling process, and a, [0. 1] be strictly decreasing scheduling function in t [0. 1]. Suppose the sampling timesteps are indexed by k E {0, .T...TTT---... 1}, the expected number of idle steps y, i.e., the entire token sequence remains unchanged between consecutive steps, is given by:
T-1
Σ k=0
(a1- - -
(A1)
Proof. At reverse step - k+ + 1, token remains unchanged if it is already unmasked at =11 or it is masked 1 & while remaining masked & 1_ي According the forward
diffusion process, the probability that a single token remains unchanged is given by:
a A+1
01
(e) (a)- 01 -1-+)=
(A2)
0 1-
(8)
(it)
(126)
where (i) is the probability of observing an unmasked token time t 1 -
k (i.e., Eq. 1.. (ia)
the probability of observing masked token time k+1
4. (i.e., Eq. ). and (iii) is the
probability that the masked token remains masked at time t = 1
(i.e., Eq. For all tokens
T
to remain unchanged, the joint probability is - expectation, the expected number of idle steps be calculated by accumulating Eq. (A2) for all k E (0,... 1}. and is written as:
-1- Due to the linearity of
Σ
01 011 -0 -]
(A3)
Eq. A1 enables the computation of idle steps given the sequence length L, the number of discretized steps T. and the scheduling function at, Consider the commonly used linear schedule a = 1 into the formula yields:
1 -- [9 [0]. Substituting the difference a t+1 01 $
T-1 Σ K=
L =T -
1)
L
(A4)
where (i) holds when T' is large. For example, when T = 1,024 and L = 1,024, we obtain Te L/T ~ 1024 e x 376, indicating that approximately 37% of the sampling steps are idle in the discrete diffusion process.
[0, 1], the quantity in Eq. (A1) decreases as the token sequence length L increases. In MDM-Prime, the sequence length is extended by a factor > 1 due to the introduction of sub-token sequences, and the number of idle steps 7/ is given by:
Moreover, since (qI- 1 -a
7 Prime
T-1
Σ k=0
[1
(a) k-1 - 01 -1-+))
LXE
(A5)
23

---

MDM 15, 1824)
MDM Prime & 1021. 2=
MDM-Prime (L 1921 (=
MDM-Ptime' Z 1024. (-
Number of ldlk Stops
3,000
Theoretical (yy
Taovretical
Theureticl syeas Empineal
Theoretical (
2,500
Empercal
Empirical
Empuncal
2,000
1,500
1,000
500
0
1,000 2,000 3,000 4,000
1,000 2,000 3,000 4,000
1,000 2,000 3,000 4,000
1,000 2,000 3,000 4,000
Number of Sampling Steps
Number of Sampling Steps
Number of Sampling Steps
Number of Sampling Steps
Figure A1: Comparison between idle steps obtained from simulation results and analytical computa- tion. Solid curves and their corresponding shaded areas represent the mean and variance from ten independent simulation runs. Dashed lines indicate the theoretical values.
always smaller than yj in Eq. (1 Therefore, employing Prime must result in fewer idle steps. To support this observation and verify the correctness of Eqs. (AT) and (A5), Fig. [A]] compares our simulation results with the analytical values computed using 7) and Prime The results demonstrate agreement between the theoretical and empirical estimates.
Given identical T and Or, P/rmie Eq. A5
Based on the above definition of expected idle steps, we define the idle step ratio (ISR) as the proportion of expected idle steps relative to the total number of sampling steps T:
+ T-
(0
-0.-)-a)-;))) LXE
(A6)
ISR quantifies the model utilization in diffusion processes. evaluate we can fix the scheduling function de and choose large T (i.c., to approximate the continuous-time limit T oo). In our experiments presented in Section we set T as 1, 024.
A.1.2 Mutual Information Scheduling
In this section, we show that the MDM framework implements mutual information scheduling (i.e., T(x;; =/)) = a, (x))). This relationship was initially established in prior works [20 21], where linear scheduling function a = 1 - was considered
special case. In our work, we extend this
result to arbitrary scheduling functions as E [0, 1], as formally stated in Proposition A.2
Proposition A.2. Let a4 E [0,1] 1] be the scheduling function and q (x)))) be the distribution defined in Eq. 7). The mutual information between on and x; satisfies:
I [x a2 (x.).
(A7)
Proof: The mutual information can be decomposed as follows:
I(x); H(x)) H(x)(a))
(A8)
The following proof expands both entropy terms to derive Eq. (A7).
(i) (x/la))) Given zi the distribution over of is:
((>jه=ع)) = a dgg (xi) + (1 ar) 8,(1)).
(A9)
Its entropy can be expressed as the entropy of Bernoulli:
(a = H (Beroulli((a:) = - a4 log ag - (1 ar) log ar).
(A10)
(ii) H(x)): The element-wise distribution (x;)) can be written as follows:
pe(xt)=
Σ po(a)) ((1)116))
afEx
(A11)
a4 Pa (zf),
if X,
Σ
po(=6) (ord + +1 -aa)6
11 at
if/==
21 x
24

---

The element-wist entropy can be expanded using Eq. AAll as follows:
H(x)) = Σ
(e(x) log
ZEA
Σ WE
(xpo(x)) log(aspu(z))) - az) log(1 at)
(A12)
po(x)) log log 04 log az (1 a4) log(1 an)
log
04 H (x)) + ((
According to Eqs.
[ (A10) and A12). the mutual information is expressed as follows:
T(r):/)) = H(x;) H(x (g))
= =o, (26) + ())) ((Bernoulli(az)
(A13)
= o,ll(()))
A.2 Analyses of the Partial Mask Scheme
In this section, we provide thorough analyses of MDM-Prime. Section A.2. derives a formula for computing the negative log- ligeliieldoodd (NLL).. Section A.2.2 shows that Prime yields a positive number of intermediate states. Section A.2.3 examines the correctness the carry-over parameterization for Prime. Finally, Section A.2.4 offers guideline for selecting the target length L.
A.2.1 Negative Log Likelihood Calculation
In Section 3 we introduce an invertible function f X - yEX transform between two vectors, 3o and 3/o- In this section, we detail the computation of the expected negative log-likelihood (NLL) of MDM after applying this transformation. We begin by revisiting the derivation of Eq. and then extend it to our proposed objective in Eq. (5).
The expected NLL of 26 expressed as Epent(nn) |- log (=))) which can be approximated using variational upper bound [9] [0 12) as expressed as follows:
Epena (20)) log (o
S E (so)
[log po By(we)loo)) [x:)]
dl
(A14)
ag
E Epon(wo) (20)
Byg (aa)
Σ log (o (xx)) di
at
where (i) derived from (a( ax))
T Po (1 (x)). Due the introduction f. the expected
NLL expressed by MDM-Prime is written as E, pan (mo)
log Pa o f(zo)), where o f(xo) is the This expectation can be estimated via Eq. (5). as
parameterized pmf as discussed Section 3.1 derived below:
E Epan (mo)) 1 log Po f(xo)] = Epun of
(wo) (wo)
log po (y60)
E of (((
[oo yy)) db
at
of Egivi (Bo)
log po (o (v)) Σ
Epanos of ((o)
dt
as
(A15)
where (i) is due to pa(ooly:) _ po (v) yy) following the derivation in Eq.
A14
To sample
data points from the distribution Paata f-1, one can first sample æ: Pdat and then transform principle for probability distributions. This
Yo = (xo) according to the change-of-variabhe formulation enables to estimate the expected NLL of LDM-Prime.
25

---

A.2.2 Number of Intermediate States
In Section 3.1 we claim that the number of intermediate states can be quantified as |39| 12
(b+1) (C+1) and that this number is always positive. To verify this, we provide Proposition
(C+)
A.3
Proposition A.3. Let y {O. ... , 1} U {m} and X {0. C - U {m). Ler l > 1 be a positive integer and let VVC The number of intermediate states is positive integer:
D | >
(A16)
Proof. The number of original tokens with the mask token is:
X
C
(A17)
The total number of possible sub-token sequences (including the mask) is:
1 (b 1b+
6 ( be
+0)
+1 +
+1,
(A18)
where (i) is derived by the binomial theorem. Since and thus:
[CC by construction, have be 2 C, 16+1)= +
(A19)
By rearranging this equation, we conclude that (b + 1 (C + 1) > 0. As result,
|Sa - (b+ (C+
(A20)
To illustrate how the number of intermediate states (i.e., 1 ( may increase with &, we define
a function (e, C) = (b + 1) (C' + and evaluate its value under different f. As concrete example, consider C 256 and & E {2, 4, 8). Substituting these values yields (2, 256) 32, (44, 256) = 368, and M88, 256) = 6304. This example highlights that Prime can produce substantial number of intermediate states by selecting é. This growth of M(,, C) motivates our design choice for the embedding lookup table discussed in Section 3.2 infeasibility of directly modeling the embeddings for y
where we mention the
A.2.3 Carry-over Parameterization
we introduced the carry-over parameterization for MDM-Prime. In this section, we justify this approach by presenting Proposition and discussing its practical implementation.
In Section 3.2
Proposition A.4. Let Vyyl) denote the set of yo such that each 1. is consistent with y E ), i.e.,
Viyl & (x' = -.."....... ........' 1 Given the parameterized distribution po(vé (y)) expressed as follows:
f(X) s.t. (y = ) (vi =m)
(A21)
exp(Eolgélve) polyblue) Evievcof Exp(Est v.Jj ifub V(vi).
(A22)
it $ V(yi).
the marginal distribution polyo 2.d (y)) of Eq. A22) satisfies the carry-over condition for all position i where 3 E y: po (ylo 3:0 (y:)
Po Po(v ... 1o 1y))) =5, (ub))
(A23)
No
Key
Proof. Given i, such that vi e V, according to Eq. A21 y0 non-zero probability:
v is necessary condition for
po(volze) >
Vo i.j
y
(A24)
26

---

m
One
Zero
Softmax: Distribution
Lookup
Marginal Distribation
Marginal Distribation
000
000
000
000
Apply Logits
000
000
000
001
001
001
001
and
001
001
001
Logical AND
010
Normalie
011
011
011
011
011
011
011
100
100
100
100
100
101
101
110
101
101
101
110
101
%
103
110
Uio
1.3
Filerr
Hilter
Filler
Combined Filter
P% (v0 M))
po(polM Ma)
Po (V,))
(a) Computation Procedure
(b) Carry-over Condition
Figure A2: An illustration (a) the computation procedure of the carry-over parameterization and (b) the corresponding carry-over condition (i.e., Eq. (A23)) the marginal distributions. In this example, 7, = 3, and b = 2 (i.e., binary encodings). In (a), the filters are precomputed and stored in a lookup table indexed by ut These filters exclude logits corresponding to yó with yo that are inconsistent with v E ) (i.e., yt" = 1 and yt = 0). For each position j E {1.... L), a filter is queried from the lookup table using 26" Then, a combined filter is then constructed by applying a logical AND operation across filters. Finally, the softmax distribution pe (y) (:) is established by normalizing the filtered logits. In (b), the resulting distribution po yó yy) satisfies the carry-over condition for vt 1- y' as defined in Eq. A233
In other word, if 2 * v for any j is observed, then po (vbl.))) = 0. This indicates that the marginal
= 2b ". - w0th ey Pe vó y1) = 0 for any yo + 1.j
Since Egeee pe(yo i y.) = E Ent po (olvblv ve) = by Eq. A22 Therefore, the marginal distribution corresponds to the Kronecker delta function as follows:
po (0
1- (y.) = 1 for fo Vi
1, if A. 4.j
Po (1 45 ye) 0,
d 3. = =,,,=000))
(A25)
if yo Vi
To understand this parameterization method, examine its implementation in practice. The goal ey
.6. 6--1 No 166t) ...., .....iiey Po(yá 3:)) is to ensure that the marginal distribution Po(v) (y.) =En'' ... ..--
zero whenever the candidate state yó contains any Mn i.j that inconsistent with ye (i.e., when 5. * vt" as discussed in the proof of Proposition A.3). Since the marginal probability is a sum over probabilities, it follows that each individual probability po (0( yt) should be zero for yó with yor that is inconsistent with i-j To enforce this, we use filters with C entries to exclude the output logits of invalid yd For example, suppose C = 7, 3, and b 2. If y! we create l filters (i.e., Filters 1-3) as follows:
(ol 1.1 1.2 v. (1,m, 0).
For yi" 1, Filter excludes yó where the first position is not 1.
For y! m, Filter excludes no state since no condition needed be satisfied this position.
For yt 0. Filter excludes yó where the third position is not O.
We then combine these element-wise filters using a logical AND operation, since any violation of an element-wise condition leads to an invalid vi that has inconsistency with y/ The resulting combined filter excludes y & Vvy)), retaining only y V(v(). Subsequently, we apply softmax over the logits corresponding to yh E Vvy() to define the probability distribution pe po(v) (o..)l.t) An
illustrative example of this procedure is depicted in Fig. A2
For efficient implementation, we precompute and cache lookup table with size [ [(x) (b+1) 1) C, which contains all possible filters for o During the forward pass, we query this lookup table using y² to retrieve the corresponding filters, which are then applied to the logits to zero out the probability of y & Voyt).
27

---

()OWT
(2) WilsTTRt
(3) LMIB
(4) PubMed
56
17.4
54
CIFAR-IO
16.2
64
15.0
36
56
8 2
(==
=
E m.
RD
4.8
(5) LAMBADA
(6) PTB
(7)AG News
(3) ArXiv
64.5
25.75
4.0
28.5
32
27.0
00
25.00
24.75
8
6
Figure A4: (1)-() PPL of MDLM-Prime with different l evallated on OWT and seven zeroshot datasets. (9) FID of MDM-Prime across varying & evaluated on CIFAR-10. Lower values indicate better performance. Red stars highlight the values corresponding to the elbow points of the ISR curves (see Fig. A3 These points yield near-optimal performance across most evaluations.
A.2.4 Analysis of the Target Length
Prime introduces a parameter that controls the target length of sub-token sequences. This section analyzes its impact on performance and provides a practical guideline for selecting an appropriate value
Section 4 demonstrates that, although increas- ing generally correlates with improved perfor- mance, the relationship is not strictly monotonic, ie., larger values of do not always lead to bet- ter results. As larger values of typically induce more intermediate states (see Section A.2.2), we hypothesize that the associated increase in learn- ing complexity may degrade performance under fixed model capacity (e.g., model size). This ob- servation motivates our development for identifying effective L.
OWT
CIFAR-10
17.5
12
15.0
12.5
10.0
ISR
7.5
5.0
2.5
strategy
6=2
8
2
Figure A3: ISR under different choices of f. The elbow points are highlighted using red stars.
Through empirical analysis across both image and text datasets, we find that selecting l near the el- bow point of the Idle Step Ratio (ISR) curve (i.c., Eq. (A6) yields strong performance. Fig. A3 presents ISR curves for various values of f, with elbow points indicated by red stars. The relationship between ISR elbow points and MDM-Prime's performance is depicted Fig. (4 where these points align with near-optimal results across most evaluations. Based on the above findings, we recommend 2 for image datasets and f for text datasets.
A.3 Model Architecture Designs
This section compares a number of architectural variants of MDM-Prime, with a focus on the design choices for the output logit layer and the input embedding layer. For reference, the standard MDM architecture is shown in Figs. A5 (ol) and (il), which illustrate its output logit and embedding layers, respectively. In standard MDM, the output layer produces logits with C' entries, while the embedding layer processes inputs x, E 8 produce L D-dimensional embeddings.
Output Logit Layer. In Section 3.2 we introduce two alternative designs for the output layer. The first design assumes independence among sub-tokens, leading factorized form of the probability distribution: pa Po(válen (y.) = 1 Po (yp ye ). This formulation can implemented using L logits, each with b entries, as illustrated Fig. A5 (o3). The second design (i.e., Fig. A5 (o2)) models a joint distribution sub-token sequences, producing L logits with C entries, consistent with the standard
28

---

Rooz
ReXL
Output Designs
C Cstries
C entries
R Rbxex
Neural Network (0)
Neural Nerwork (9)
Neural Network (0)
(of) Standard MDM
(o2) Joint Distribution
(o3) Independent Distribution
PXL
Dx
Input Designs
R Dxt
Concascaation
Cross Attention (8)
I 2 XtxL
K. E DxtxL
Qo E Bx
Emhedding Lookup (9)
Embedding Lookup
Embedding Lookup Table (9
J.rr. (13) Merging via Cross-Attention
Learnalle Queries
(12) Standard MDM
(i2) Merging via
Figure A5: Comparison of input and output layer designs in MDM and MDM-Prime. The top row illustrates three types of output layer designs: (ol) standard MDM, (o2) MDM-Prime with joint distribution, and (o3) MDM-Prime with independent distribution. The bottom row shows three embedding layer designs: (il) standard MDM embedding layer, (i2) MDM-Prime embedding layer with concatenation, and (i3) MDM- Prime embedding layer with Perceiver 25 cross-attention.
MDM architecture shown in Fig. A5 (ol). ). As discussed in Section 3.2 we adopt the joint distribution
design the factorized alternative (Fig. A5 (o3)) due to two key limitations of the latter: (1) the independence assumption, and (2) the potential to generate invalid samples. These issues hinder the model's ability capture complex data distributions (see Fig. 4 and limit its applicability to real-world generative tasks.
Embedding Layer. As discussed in Section 3.2 creating an embedding lookup table for y! is impractical due to the resulting growth in the number of parameters in it (also see Section A.2.2).
To address this, model sub-token embeddings by creating lookup table for y e y Since the input and output sequences differ in length, we introduce simple merging strategy based on concatenation, as illustrated in Fig. AS (i2). In this approach, each sub-token embedded into vector of size D and the sub-token embeddings are then concatenated to form D-dimensional token-level embeddings.
Fig. A5 (13) shows an alternative design based on Perceiver [25], which employs cross-attention with learnable latent queries to merge sub-token embeddings. In this design, each sub-token is embedded into a D-dimensional vector, producing are then used the key and value matrices in cross-attention layer, while a learnable query matrix Qo E D retrieves the token-level embeddings. The resulting outputs are D-dimensional token embeddings. Although this approach is well-established for processing high-dimensional inputs in transformer-baset models, introduces substantial computational overhead [25). Moreover, empirical results in Appendix A.5.2 show that it does not lead to performance gains. Based these observations, we adopt the simple concatenation- merging strategy in MDM-Primer s embedding layer.
sub-token embeddings. These sub-token embeddings
A.4 Experimental Setups
this section, we provide additional implementation details and hyperparameter settings used in our experiments. The configurations for the two-dimensional synthetic experiment, image generation experiment, and text generation experiment are presented in Sections A.4. A.4.2. and A.4.3
respectively.
29

---

A.4.1 Tvo-Dimensional Example
Dataset. The experiment in Fig. 4 is conducted synthetic dataset constructed by converting pixel values from an image in the Cat dataset [60] into a two-dimensional probability histogram. The image is first cropped to 512 x 512 pixels and converted to grayscale. A probability distribution is then constructed according to the normalized pixel intensity values. In this example, X = where L = 2 and C' 512. The sample Zo E X' represents the coordinate of the figure (512 512).
Training and Implementation Details. The network architecture is a four-layered multilayer percep- tron (MLP) with hidden dimension 512 and Swish [6]] as its activation function. The models are trained using the Adam optimizer [62] with a learning rate of 10-3 and a batch size of 4, 096. The training is performed on a single NVIDIA A40 GPU with 48 GB memory.
A.4.2 Image Generation
Datasets and Evaluation Methods. The experiments described in Section 4 are conducted on the CIFAR-10 10 (I5) and ImageNet-32 [16] datasets. For both datasets,
32 x 32 and C' == 256. The CIFAR-10 10 training set contains 50, 000 images, while the ImageNet-32 dataset comprises 1, 281, 149 training images and 49, 999 validation images. The sample quality is assessed using the Fréchet Inception Distance (FID) [4] and Inception Score (IS) [47). implemented via the torchmetrics image. f fid and torchmetrics .mmge.. inception libraries, respectively. The corrector steps [22] 23 are adopted during sampling. For CIFAR-10, FID is computed using the training set as the reference distribution, whereas for ImageNet-32, the validation set serves as the reference distribution. This evaluation protocol is consistent with prior works 17 22 23 35445)
Training and Implementation Details. Our model architecture follows 23 to adapt the ablated diffusion model (ADM) [46]. which has a U-Net structure with a symmetric design. It comprises a downsampling module, a bottleneck module, and an upsampling module. Each module contains multiple residual convolutional blocks, with attention layers incorporated at selected resolutions. The downsampling and upsampling modules use an embedding dimension of 96 and channel multipliers of [3, 4, 4), respectively. The depth of each block is set to 5. For attention layers, the number of head channels set to 64. The network optimized using the Adam optimizer [62 with B1 = 0.9, B2 = 0.999, and a learning rate of 1 x 10-4. The scheduling function is defined as third-order polynomial, given by a (1 - z)3. Following prior work 20 23], the coefficient of in the loss term is omitted during training to improve sample quality. As for sampling, temperature scheduling (23]. corrector steps [23] and timestep scheduling [I0] are incorporated into both MDM and MDM- --iiieeeette to enhance performance. The sampling parameters are selected via grid search (see Section 1. The model is trained with a batch size of 512 for 4, 250 epochs on CIFAR-10 100 and 1, 0000 epochs on ImageNet-32. The training is performed on eight NVIDIA L40 GPUs with 48 GB memory.
1-ax
A.4.3 Text Generation
Datasets and Evaluation Methods. The training is performed on the OpenWebText (OWT) dataset [I]. Text data is tokenized using the GPT-2 tokenizer [3], which defines L = 1, 024 and C = 50, 257. Since some sequences in OWT exceed the maximum length L, they are con- catenated with the <eos> token as separator and then wrapped into segments of length 1, 024. In addition, the first and last tokens of each sequence are set to <eos>. As OWT does not include an official validation split, we follow the procedure in [9] by reserving the last 100, 000 samples for validation. The above data preprocessing and evaluation setup is consistent with 19.
Training and Implementation Details. Following prior works 7 2
12 our model archi-
tecture is a diffusion transformer (DiT) (26) with rotary positional embeddings (27). The model consists of 12 DiT blocks with hidden dimension of 768 and 12 attention heads, consistent with the function is defined as O₁
used in 9.. A dropout rate of 0.1 is applied throughout training. The scheduling t. The network is optimized using the Adam optimizer (63 with 31 = 0.9, 32 = 0.999, and linear learning rate warm-up from 0 to 3 x 10-4, The model is trained with a batch size of 128 for million steps, corresponding to total of 262 billion tokens [9] To facilitate efficient training, the model is first trained without the carry-over parameterization for 900K iterations, and then enables this mechanism in the final 100K iterations (see Section A.5.3) The
training is performed on eight NVIDIA L40 GPUs with 48 GB memory.
30

---

Table A1: shot validation PPI. evaluated on seven textual datasets. Lower values correspond better performance. Methods marked with * 'ncorporaae an autoregressive formulation. MDLM- Prime with carry-over parameterization achieves improved results on LAMBADA and PubMed.
Carry-over LAMBADA
WikiText
PTB
LMIB
AG News
PubMed
ArXiv
ARM [9]
51.28
25.75
82.05
51.25
52,09
49.01
41.73
MDLM 9
<47.52
$32.83
$95.26
$67.01
$61.15
$41.89
$37.37
MDLLM-rrimx (f
234.61
$31.94
$91.85
$65.25
363.74
$63.51
$25.95
MDLM-Primm (E
$30.91
$27.93
<74.81
$55.50
$63.21
<33.32
$25.44
MDLM--rinns
$34.55
$31.00
$55.46
$41.15
$62.51
$171.74
$29.44
MDLM-Prims
Prime (E 3)
$27.75
<27.42
<60.13
<42.69
<62.58
<41.14
<24.71
MDLM-Prine Prime (f 4)
$25.79
$23.52
$49.90
$37.70
$59.65
2200.96
$23.19
MDLM-PPP M-Prime (E
$24.44
$23.86
$33.98
$38.02
$59.44
<48.64
$25.83
MDLM-Prinee (E 6)
$25.89
$25.86
$52.42
$37.65
$56.68
<497.39
$24.99
MDLM-PPIE M-Prime (f
<25.80
226.87
<62.62
<45.36
360.16
$59.09
$25.19
MDLM- (E
$32.20
$26.01
$53.27
$38.05
$65.63
$218.99
$28.63
MDL.M-rrime
$25.23
$25.77
$53.77
$38.00
$64.89
$54.50
225.79
A.5 Supplementary Experiments
In this section, we present additional experimental results. Section A.5.1 discusses a technique for tuning the hyperparameters based relative FID values. Section A.5.2 provides ablation studies on the embedding layer designs and the carry-over parameterization method. Section A.5.3 reports the time cost per training iteration under different values of f. Section[A.5.. compares the training curves of MDLM, MOLM-Prrie, and ARM. Section A.3.5 presents qualitative results for both the text and image generation tasks. Finally, Section A.5.6 provides visualization of the sampling process of MDM-Prime.
A.5.1 Hyper-Parameter Tuning based on Relative FID
Since computing the FID score is computationally inten- sive, tune hyperparameters (e.g.. batch size, dropout rate, scheduling coefficients at, and sampling parame- ters) using FID scores estimated from smaller number of samples. We refer to this approximation as relative FID. Empirically, correlated with the standard FID computed on the full dataset. Fig. A6 illustrrrs this trend, comparing relative FID computed with 2,048 samples to standard FID com- puted with 50,000 samples. In our experiments, use relative FID to guide hyperparameter search for both the baseline methods (i.e., MDM and proposed method, and report their best-performing results the benchmark comparison.
13
find that relative FID is positively
(DD150,00 000
CIFAR-10 10
Imag MeeegeNet-333 32
15
18
27
and our
relative (2,048 samples)
Figure A6: FID versus relative FID.
A.5.2 Ablation Study
Carry-over Parameterization. In this subsection, eval- uate the effectiveness of the carry-over parameterization. Ta- ble A2 compares the performance of MDLM-Prime trained with and without carry-over on the OWT dataset. It is observed that incorporating the carry-over parameterization consistently leads to slightly improved performance across different values of C. In contrast, when the model is evaluated using unseen data from zeroshot datasts, the difference becomes significant. comparison is offered Table AT This parameterization method consistently offers performance gains on LAMBADA. addition. Mrimee without the carry-over parameteriza- tion has noticeably inferior performance (i.e., difference >> 30)
Table A2: Perplexity (PPL.) evalua- tion OWT.. The symbol repre- sents that lower values correspond better performance.
PPL (4)
Carry-over t=
/
118.04
$17.90
f=
<16.43
$16.36
{= 4
<15.67
<15.62
{= <15.43
<15.3
8=8 <15.48
115.45
on specific text domains such as PubMed. We hypothesize
that the carry-over parameterization, by explicitly removing predictions of yo inconsistent with yt.
31

---

Table A3: The average runtime calculated based on ten training iterations. The results are reported in seconds per iteration. The 95% confidence intervals of the results are around 0.05.
Table A4: FID scores evaluated under different NFE on the CIFAR-I0 and ImageNet-32 bench- marks. Lower values correspond to better perfor- mance. MDM-Prime trained with
Runtime (scc./iter.)
CIFAR-10
NFE
128
256
512
Carry-over
MDM
7.55
5.00
4.66
= 1, = 50, 257
1.06 e-1
e-1
MDM-Prime
5.22
3.73
3.26
= 2, = 225
1.06
1.29 e-1
3,
b 37
1.11
c-34
ImageNet-32
{ = 4.
b = 15
e-10-11
e-43
NFE
128
256
512
= 6,
b 7
L.13
c.75
MDM
9.55
8.24
8.12
= 8, b == 4
1.07
2.12e-1
MDM Prime
7.85
7.61
7.31
Table A5: Ablation results on the validation perplexities of MDLM-Prime = 6). The model is trained and evaluated with or without the carry-over parameterization. Lower perplexity indicates better performance.
Carry-over
Dataset
Train Evaluation
OWT
LAMBADA
WikiText
PTB
LMIB
AG News
PubMed
ArXiv
< <55.43
$25.89
$25.86
$52.42
$37.65
$56.68
$497.39
$24.99
$15.41
$24.36
$24.25
$49.77
$37.35
$56.48
$63.58
$24.81
/
<15.36
<25.80
$26.87
$62.62
<45.36
160.16
$59.09
$25.19
effectively reduces the uncertainty in the model's outputs and therefore enhances generalization to unseen domains.
Embedding Layer Designs. In this subsection, we compare the performance MDM-Prime when trained using two embedding- strategies: the concatenation- based design (denoted as (o2)) and the cross-attention-base design (denoted as (o3)), both of which are detailed in Appendix [A.3 We first evaluate these approaches on the image generation task using the CIFAR-10 dataset with
2. The concatenation- design (o2) achieves a superior FID score of 3.26, compared to the cros-attentioon-basel design (o3), which yields an FID score of 3.98. The results indicate that the concatenation- based design offers better overall performance across both modalities.
A.5.3 Runtime Evaluation
Table A3 presents a comparison of the time required to optimize the model parameterized with and without the carry-over condition (i.e., Eq. (6)) under different choices of f. We observe that larger values of incur higher computational costs for the setup with carry-over condition. The increase in runtime arises from the filter lookup table calculation required by the carry-over parameterization (see Section A.2.3). Due to the 48 GB memory limitation of our available GPUs, parallelized querying of the filter lookup table (with batch size of 128) often results in out-of- memory errors. To mitigate this issue, we implement the lookup operations sequentially our text experiments, which introduces additional computational overhead as l increases. To maintain runtime comparable to the baseline, we first train the model without the carry-over parameterization for 900K iterations, and then enable this mechanism during the final 100K iterations. Alternatively, the carry-over parameterization can be applied only at inference time. As shown in Table A5 this setup-training without the carry-over parameterization but evaluating with it-achieves performance comparable to, or slightly better than, that of the model trained and evaluated with the carry-over parameterization, while avoiding the higher training cost. Nonetheless, to preserve the mathematical formality of our proposed framework, we adopt the carry-over parameterization during both training and evaluation.
As for inference efficiency, although idle sampling steps in masked diffusion models can be executed efficiently using cached model output probabilities (since the model is timestep-independent ()), determining the unmasked elements (i.e., resampling) during these steps still incurs computational overhead, even without extra forward passes. Consequently, the sampling cost of MDLM grows substantially when targeting a higher effective NFE (i.e., NFEs that cannot be cached). As shown
32

---

Table A6: Runtime and effective NFE evaluation of MDLM and MDLM-Prime (l 4).
MDLM
MDLM- Prime (L = 4)
Effective NFE
768
1,003
1,023
768
1,003
1,023
Discretized Timesteps
1,685
25,000
500,000
772
1,022
1,044
Runtime (sec.)
8.87
23.30
159.81
8.72
12.31
12.58
ARM
MDLM
MDLM-Prime (8 (4)
Validation Porpexityy (PPL)
4.5
Training Loss (NLL)
40
4.0
30
20
2.5
1.0
10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Steps (million)
Steps (million)
Figure A7: The training and evaluation curves of ARM, MDM, and MDM-Prime (E 4).
in Table A6 MDLM requires 500,000 discretized steps to achieve an effective 1,023 NFE, severely limiting sampling efficiency. In contrast, MDLMM-rrmme (E 4) attains the same effective NFE with only 1,044 discretized steps, highlighting its substantially improved efficiency.
A.5.4
Training Curves
Fig. A7
compares the training and evaluation curves of MDLM, MDLM-Prrmm (e= 4), and ARM. MD M MMMM-M--Me--- (E 4) achieves the lowest evaluation PPI. among the three models.
A.5.5 Sample Quality Evaluation
Image Generation. In this subsection, we provide additional quantitative and qualitative results assess sample quality in image generation tasks. Table A4 presents comparison of FID scores between MDM-Prime and its baseline (i.e., standard MDM). The results show that MDM-Prime outperforms MDM across different number of function evaluations (NFE). In addition, Fig. A19 provides a number of imputation examples, while Fig. A200 presents qualitative results showcasing multiple uncurated samples generated by MDM-Prime.
Text Generation. In this subsection, we present additional sample quality evaluations for text gen- eration tasks. Fig. A presents the generative per- plexity (Gen PPL) results for ARM, MDLM, and MDLM-Prime. Gen PPL calculated by evalu- ating the perplexity of generated samples using a pretrained large language model. To ensure a comprehensive assessment, we report Gen PPL obtained from both order-agnostic models (c.g., LLaDA-SB and order-specific models (e.g.. GPT-2 Large (5] ). Following (24], we adopt 64- bit floating-poin precision enhance sampling accuracy. We observe opposite trends depending on the pretrained model: MDLM and MDLM- Prime outperform ARM under LLaDA-based Gen PPL., whereas ARM achieves lower Gen PPI. than MDLM and MDLM- Primee when evaluated using GPT-2.
ARM
MDLM
MDLM-Prime
100
(LLaDA)
80
Gen PPL (GPT-2)
90-
70
50
40
40
500
1,000 1,500
2,000
1,000 1_500
Number Sampling Steps
Number of Sampling Steps
Figure A8: Gen PPL evaluated using LLaDA- 8B (left) and GPT-2 Large (right). Lower values correspond to better performance. MDLM is evaluated using their officially released code.
33

---

2)
4]
(2 = 8)
Figure A9: Visualization of the sampling processes of MDM-Prime with é (top), l (middle), and l (bottom). The models are trained on CIFAR-10.
In addition, to qualitatively evaluate MDLM-Prime,, we present both unconditional and conditional generation samples. Unconditional generation results are shown in Fig. A18 For conditional generation, we provide the model with prefix (i.e., Fig. A13) and suffix (i.c., Fig. |A texts sourced from an online article describing Rabindranath Tagore's 's poems (link). and assign the model to generate the middle content. Figs.
some samples generated using MDLM-Prime
A.5.6 Visualization of Sampling Processes
Unlike continuous diffusion models (e.g., 19 37)), visualizing the sampling processes of MDM and MDM-Prime is nontrivial due to the presence of masked tokens and sub-tokens. To qualitatively analyze the generation behaviors MDM and MDM-Prime we develop a visualization technique that illustrates the evolution of samples throughout the diffusion process. Instead of directly visualizing the latent variables (i.e., I, for MDM or yE for MDM-Prime), we visualize the model's predictions of the final unmasked sample. Specifically, show 2o where (- (y:) for MDM-Prime.
Po(: (æx)) for MDM, and æo ff (yo)
Fig. A9 shows the evolution of samples zo generated by MDM-Prime trained with different values of on CIFAR-IO. For each setup, we capture a snapshot every 25 timesteps, with the total number of function evaluations (NFE) fixed at 500. The snapshots displayed from left to right, depicting the progressive refinement from noisy initialization to the final output.
The text generation processes are presented in Figs. A10 A12 Since the coarse-to- torrse-tt-finee transitions of 2p are less visually disceriible text generation, we additionally illustrate the denoising progression using token-win wise masked ratios. In the figures, darker shades of blue indicate a higher degree of masking, while lighter colors suggest that a token is closer to final predicted state. The results demonstrate that MDLMMB M-Prime with larger values of enables a more fine-grained denoising process. In contrast, the original MDLM employs a binary masking scheme, where each token is either masked or unmasked. The results thus highlight the property of the proposed Prime method.
34

---

A.6 Limitations
core assumption in MDM is that tokens (i.e., z/.... .x)) are conditionally independent given the latent representation Ty i.e., s6 1 zy 2. for + j. Under this assumption, many recent factorize the conditional distribution as po( = II (o(a) xt). which leads the sum of log-pree probability terms in the evidence upper bound (i.e., Eq. (4). While this
works 8 10
20
22
24
factorized parameterization offers significant advantages in sampling and training efficiency, some
recent studies 12 (64) have shown that it may degrade performance due its inability to model inter-token dependencies.
In this work, adopt the conditional independence assumption prior studies 10 20
22
24 and extend it our sub-token formulation. Specifically, we assume yó I. yō y: for J. and accordingly factorize the conditional distribution as po (Bol))) I Po yo yv). As discussed in Section 3 further factorizing each sub-token sequence prediction as po 3r I[ Pa (v6' (3b) leads to clear drop in performance. Therefore, we retain the inter-token factorization (i.e., polyoloo )) strike balance between accuracy and computational efficiency.
po(yo yo (1)) 1o
A.7 Broader Impacts and Future Works
This paper investigated whether discrete data can be effectively represented and reconstructed using sub-token representations, We proposed MDM-Prime simple yet effective instantiation of this idea. Given its superior performance across both image and text generation tasks, MDM- holds potential for positive societal impact. To support broader scientific contributions, outline two potential directions for extending the proposed framework and guiding future research this area:
Learnable Transformations for Discrete Data. MDM-Prime adopts base-b encoding invert- ible mapping to extend discrete data into longer sequences of sub-tokens. Analogous to normalizing flows in the continuous domain (e.g., [65] (6]). which parameterize invertible transformations using carefully designed model architectures, we anticipate that this discrete transformation can likewise be parameterized and optimized during training encode the semantics of each token. As invertible modeling in the discrete domain remains underexplored, it presents opportunities to advance discrete generative modeling.
Capturing Inter-token Dependencies. Our current decoder implementation follows prior works (8] 10 20 224244, which assume conditional independence between tokens (see Section A.6). While
recent approaches [12 64) have proposed methods to relax this assumption, they require training an additional autoregressive model guide the sampling process of MDM, resulting in substantial computational overhead. Hence, developing more efficient method for MDM to model inter-token joint distributions represents promising direction for future work.
35

---

0.0
0.5
1.0
Masked Ratio
t 1.000
0.875
= 0.750
Bata
0.625
381
0.500
0.375
2311
teath
0.250
= 0.125
Ban GOLL'S
Always
vo14
MILL
sTaund
The
0.000
egiin
Figure A10: Visualization of the sampling processes of MDLM. The masked ratio is measured on a per-token basis, with higher values indicated by darker shades of blue. The samples generated with prefix and suffix presented in Figs. A13 and A14 respectively. Further experimental details are shown in Section A.5.5
36

---

0.0
0.5
1.0
Masked Ratio
t == 1.000
7M
t 0.875
Unb
Will
t 0.750
X3T
0.625
55
Mill
Hewr
t 0.500
Lactark
Febr
[ 0.375
5073
UN
t 0.250
1jxx
00
202
t == 0.125
180
refere
t 0.000
600g
ACTUTO
Figure A1l: Visualization of the sampling processes of MDLM- (l 2). The masked ratio is measured on per-token basis, with higher values indicated by darker shades of blue. The samples are generated with prefix and suffix presented in Figs. AI3 and A14 respectively. Further experimental details are shown in Section A.5.5
37

---

0.0
0.5
1.0
Masked Ratio
1.000
0.875
0.750
= 0.625
2210
0.500
onter
1al
= 0.375
The
0.250
Be's
The
tigh
ed
= 0.125
1
1004
The
majes
TM4
= 0.000
5
Figure A12: Visualization of the sampling processes of MDLM-Prime (6= 4). The masked ratio measured on per-token basis, with higher values indicated by darker shades of blue. The samples are generated with prefix and suffix presented Figs. AT3 and A14 respectively. Further experimental details are shown in Section A.5.5
38

---

timeless poems by Rabindranae anath Tagore
Born on May 7, 1851 the Bard Bengal, Rabindranath Tagore has Inspired generations people through his writings, poetry and thoughts. Tagore was much ahead his time and his works were loved only in India first published in 1910 and later translated and published into English in 1912, prestigious Nobel Prize in Literature in 1913 for "his profoundly sensitive. fresh and beautiful verse, which, with consummate skill, he has made his poetic thought. expressed in his own English words. part of the literature of the West." Infact. Rabondranath Tagore was the first non Eun-Eun to ever win a Nobel Prize!
world.
much-acclaimed
"Gitanjali' him
Remenbering Tagore his 168th birth anniversary today. here List down of his timeless poens that continue to resonate his creative charm and are still as relevant. These poems reflect upon different moods and are a nust read for all. Read on!
Figure A13: Prefix text used for conditional sample generation.
Waiting
The song came to sing
remains unsung this day.
have spent my days in stringing unstringing my instrument.
The time has not come true.
the words not been rightly set: only there the agony of wishing in my heart
have not seen his face. listened
only I have heard his gentle footsteps from the road before ny house....
But the lamp has not been lit and cannot him into ny house:
live in the hope of ineting with hin; but neeting is not yet
Friend
Art thou abroad on this stormy night thy journey love, my friend?
The sky groans like one in despair.
have sleep tonight. Ever and again I open my door and look out on the darkness, friend!
nothing before ne. wonder where lies thy path!
what dint shore the ink-black river,
by what far edge of the frowning forest.
through what mazy depth of gloon thou threading
thy course to cone to me. my friend?
Figure A14: Suffix text used for conditional sample generation.
39

---

Travis's you
Son
time you ask
When I plead my secret
you, you won't hear the light.
It 1s the quiet decker of ny 1ife of love.
When I went which said
returns, and went
the world will tell ne the same time.
Do you like to think this like
will have song in the Adverse.
will have story for you
Figure A15: Conditional sample generated using MDLM-Prime (6 2).
Dreaming Our thoughts are spark the depths of life, but They really seen to be no new.
Dreaming
stranger. That is sonething new to me. makes me nore than could.
That makes think more than could have.
always think with my light.
And so see the edge of heaven and even-sappiness and hope.-
magic of ny mind, can go along my lightness.
Figure A16: Conditional sample generated using MDIM-Prime (6 = 4).
The Promise
think about ny life when make me the vessel of life.
Garnings. deeply
hinge of my it's the most tranquil, exposed world.
It is the very noment of such thing walk up my every tine night Leave me.
lifetine.
come.
begin mind and mind.
The passion and death of the night do not think of the other away.
Figure A17: Conditional sample generated using MDLM-Prime (E 6).
40

---

<| |Jndoftext argunent that such modeling is ane of the world's thinkers for clinate change.
going exanine then," she said. "Each day, there companies that only 26 -- 48 people approve of fossil fuel enissions. The huge zonal wave that ice is releasing into the atmosphere. Everything here is.having is back and it is coming."
"eoppee read the website.
For the political pro-energy group, SoundEarth and the Climate Collective think that it will lead pollution, which could help nake the greenhouse issue a premier
to a stronger understanding educational problen. The studies expected increase the potential respond climate change. This why of applications that can damage infrastructure. up in trees, up in trenches.
further: "For example. plants world need to encourage the kinds
in towers."
"We are most aggressive skeptics." Bec Lindy safd. They claim that the pollution is used from storage the emissions, tines, swarp energy use amounts of rain in mountains, caused storms. The trees also host public debate over whether it's used when placing organic synchtas and plants to actually remove enissions from the minerals of which it 1s created, as well as changing water to be natural.
to take out high-cost water straight from fossil fuels. Carbon diexide companies. This begins
gravity. through large
While nany of those skeptics are even aware that such pollution is from scientific study, the Baker Effect, they may be receiving new level of threat fron state.
Beegy said this is something she's hopeful about and suspect that it will continue to grow, especially after solar wind projects shrink the risk of air pollution. The solar Effect has its course changed.
affect climate." Lindy said. "We have models called climtate feedback that are able predate the climate on. So if we need to, it forces actively reduce output to actually increase carbon pollution. They use some of the efficiency нe already study as sources of governnent resources.
"We see public concern knowing
of
Conners working on supplementary Energy Initiative Decamp to study new complete refrigerators the University California that will charged and responsive. House inatenaes their solar panels will likely serve over 33,088 panels. consisting of residential, "basement cast" scenarios that could be significantly set up if rug leaks in the open-up.
"The University is moving forward with the capital energy." Lindy said.
"For years, we haven't worked in the BLy's first exper agrees Professor Mark Clerock, directar of the National Institute of Relioess and Turiative Research Unit at Hampshire University. "So whether it takes the hours per second test is six-figure! Even 2 niles wouldn't t have been predicted."
Professor Jeanne Morul Institute of California California State adds, stops saying this will only accelerate anyway. Any skeptics say the BLy's changes sheet this ght give even more freedon to Colorado.
In approving the research, the Democle Climate Council published statement that the additional valume. Therefore nagnf tude study "taught mixing severe weather events. forests, comparison." Professor Thonas Nuirena. severe consequences of weather could weaken aiming to do the research this year. when the
Professor A at West Texas's Gulf Observatory, gets through extrene weather. he attributes one of the biggest implications of his discoveries. "It is the poisoning of our income I don't see then taking any action and still don't t have the opportunity to see pollution. I don't think thetr compantes have anything to start doing." he to deal
saíd. "We" '11 for them betore whether that is thought locally entirely on what cones with power would be an investigation."
again on whether its power - but
see what consequences
We are thinking NOE to the chemicals, nethane that comes from traditional coal emitting plants. the standard statement is funny.
it. can't afford Department of Energy. But, government of the coalition has tried to extend those limits to 95 percent. It shows credit to Bob Marshall measures used by his governments decades ago. It raises some questions they cannot use plants nade from traditional. coal but can do get a product at nigh cost. Should they get more noney from the subsidies or others from their locally grown stations or the food grains required to comply? Their limits would be least the biggest in the world.
someone wants try do
environnental damage
The Harshall says that they send farns in rads-</eedoftext|>
Figure A18: Unconditional samples generated by MDLM-Prime with 1,024 sampling steps.
41

---

Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Condition
Imputation
Figure A19: Uncurrated imputation results generated by MDM-Prime with CIFAR- 10 images as conditions.
42

---

(a) CIFAR-10
(b) ImageNet-32
Figure A20: (a) CIFAR-10 and (b) ImageNet-32 samples generated by MDM- -rime with NFE=512.
43


Reviews:
[{"review_id": "pJjWrsFAR7", "reviewer": "Reviewer_uEyH", "summary": "The authors propose a partial masking scheme (Prime) to address redundant and perhaps suboptimal computation in the masked diffusion model unmasking generative process, wherein some tokens are sampled early and then copied over and over while others are unmasked to be consistent with those already chosen. To challenge this setup, instead of the original masking/unmasking scheme, they introduce latent variables that are ``partially masked\" tokens which allows the model to gradually sample a consistent set of tokens rather than specialize/commit some dimensions too early and copying them over the rest of the diffusion process. They propose a variational training scheme for handling these latent partial masks (\"intermediate states\") and report perplexities and FIDs on text and image data.\n\nWhile the two goals (1. minimizing steps of doing nothing 2. introducing a richer latent denoising space) may seem orthogonal at first, I think the idea is roughly that time spent doing could be time spent throwing extra computation into a more expressive model, and a more fine-grained latent space is one way to do this.", "questions": "- [conceptual clarification] do I understand correctly that by explicitly defining the mapping between original tokens and subtokens, partially-masked subtoken sequences essentially represent a set of possible tokens that the current noisey state for this token can be mapped to? And this set of \"tokens to denoise to\" shrinks as we unmask more subtokens? Therefore, this is sort of a diffusion over the space of sets of tokens?\n\n- [detailed clarification of NFEs vs sampling steps vs number of modeling calls] If i understand correctly, it seems like one needs a joint model over subtokens given state $y_t$. Is it right that this joint model factors over tokens ($y_0^i$ independent of $y_0^j$ given $y_t$) but is correlated across subtokens $y_0^{i,i}$ dependent with $y_0^{i,j}$ given $y_t$? As a result, in the likelihood bound, there is a marginalization over subtokens consistent with a given observed partial masking, is that right? Is that what equations (6) and (7) get at? If so, any considerations computationally ? Does this marginalization govern that $\\ell$ should be kept somewhat low? \n\nSpecifically, If multiple model calls are made in the marginalization sum, then what exactly do the NFEs in the tables mean? Does 500 NFEs mean 500 denoising steps, where each denoising step might call the model many times to marginalize? In either case, **the relationship between NFEs, number of denoising steps, and number of model calls should be clarified. If multiple model calls are made per what the table refers to as \"1 NFE\", then this should be made more prominent in the text.**. Clarification will be appreciated.\n\n- [exploratory question] On an intuitive level, it seems that which original tokens to \"pair\" via shared subtoken subsequences is potentially an important detail that could affect the learning problem. e.g. if \"cat\" is 0001 and \"bird\" is 0000 and the partial mask is 000m then the model knows from [000m] that the word  is \"cat\" or \"bird\" and statistical strength can be shared in sentences like \"Lucy has a pet [000m]\". But less such strength could be shared if it were \"bird\" or \"chair\". Any particular thoughts on automated ways to assign the mapping? How was it done in your work?\n\n- [exploratory question] Any thoughts about data-driven ways to set $\\ell$ ?", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\n\n- Generalize masking and unmasking to partial masking , where a partial mask  is a base b ssequence of length $\\ell$ of \"sub-tokens\".\n- adapting MDMs to MDM-prime requires some embedding changes but not much architecture or diffusion formalism changes\n- Good PPL/FIDs per number of function evaluations\n\nWeaknesses\n\nSee questions.", "was_revised": true, "final_justification": "I maintain my positive score for this careful and insightful work.", "rebuttal": "We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.\n\n---\n \n**Questions**\n\n**Q1**. [conceptual clarification] **(a)** do I understand correctly that by explicitly defining the mapping between original tokens and subtokens, partially-masked subtoken sequences essentially represent a set of possible tokens that the current noisey state for this token can be mapped to? And this set of \"tokens to denoise to\" shrinks as we unmask more subtokens? **(b)** Therefore, this is sort of a diffusion over the space of sets of tokens?\n\n**Response**. \n**(a)** Yes, this is a valid and correct interpretation. When some sub-tokens are masked, the noisy state corresponds to a subset of possible original tokens consistent with the unmasked parts. As more sub-tokens are revealed, this subset shrinks, narrowing down the candidate tokens. \n**(b)** Yes, this process can be intuitively viewed as diffusion over sets of tokens. Each intermediate state (i.e., partial sub-token sequence) defines a set of possible tokens, and the diffusion moves from larger sets (more uncertainty) to singleton sets (fully denoised tokens).\n\n---\n\n**Q2**. [detailed clarification of NFEs vs sampling steps vs number of modeling calls] \n**(a)** If i understand correctly, it seems like one needs a joint model over subtokens given state $y_t$. Is it right that this joint model factors over tokens ($y_0^i$ independent of $y_0^j$ given $y_t$) but is correlated across subtokens dependent with given $y_t$? \n\n**(b)** As a result, in the likelihood bound, there is a marginalization over subtokens $y_0^{i,i}$ consistent with a $y_0^{i,j}$ given $y_t$ observed partial masking, is that right? Is that what equations (6) and (7) get at? If so, any considerations computationally? Does this marginalization govern that $\\ell$ should be kept somewhat low? \n\n**(c)** Specifically, If multiple model calls are made in the marginalization sum, then what exactly do the NFEs in the tables mean? Does 500 NFEs mean 500 denoising steps, where each denoising step might call the model many times to marginalize? In either case, the relationship between NFEs, number of denoising steps, and number of model calls should be clarified. If multiple model calls are made per what the table refers to as \"1 NFE\", then this should be made more prominent in the text. Clarification will be appreciated.\n\n**Response**.\n**(a)** Yes, this is correct. Eq. (7) defines a joint distribution $p_\\theta (\\boldsymbol{y}^i_0| \\boldsymbol{y}_t)$ over sub-tokens at position $i$ given $\\boldsymbol{y}_t$.\n\n**(b)** No, the likelihood bound does not require explicit marginalization over sub-token combinations. Instead, the distribution $p_\\theta (\\boldsymbol{y}_0^i | \\boldsymbol{y}_t)$ in Eq. (7) is defined as a softmax over the set $\\mathcal{V}(\\boldsymbol{y}_t^i)$, and this formulation is guaranteed to satisfy the marginalization condition in Eq. (6), as shown in Proposition A.4.\n\nTo construct the softmax over $\\mathcal{V}(\\boldsymbol{y}_t^i)$, we apply a filtering mechanism that masks out the logits entries corresponding to inconsistent $\\boldsymbol{y}_0^i$ based on the observed $\\boldsymbol{y}_t^i$. This procedure, as detailed in Appendix A.4 and illustrated in Fig. A2 (a), performs logical AND operations across $\\ell$ binary filters to the logits, which does not require additional model evaluations.\n\nRegarding the computational cost, larger $\\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\\ell = 1$) and MDLM-Prime with varying $\\ell$, along with the corresponding perplexity (PPL).  Based on these results, we recommend $\\ell=4$ as an ideal balance between efficiency and performance:\n\n|  | Runtime (sec.)  |  \n| - | - |\n| $\\ell=1$ (PPL: 22.98) | 1.14e-1  |\n| $\\ell=2$ (PPL: 17.90) | 1.29e-1  |\n| $\\ell=3$ (PPL: 16.36) | 1.34e-1  |\n| $\\ell=4$ (PPL: 15.62) | 1.43e-1  |\n| $\\ell=6$ (PPL: 15.36) | 1.75e-1  |\n\n**(c)** Each NFE corresponds to a single denoising step, which includes one forward pass of the model. The marginalization condition is handled via filtering and does not require additional model calls as described above.\n\n---\n\n**Q3**. [exploratory question] On an intuitive level, it seems that which original tokens to \"pair\" via shared subtoken subsequences is potentially an important detail that could affect the learning problem. e.g. if \"cat\" is 0001 and \"bird\" is 0000 and the partial mask is 000m then the model knows from [000m] that the word is \"cat\" or \"bird\" and statistical strength can be shared in sentences like \"Lucy has a pet [000m]\". But less such strength could be shared if it were \"bird\" or \"chair\". Any particular thoughts on automated ways to assign the mapping? How was it done in your work?\n\n**Response**.\nWe appreciate the reviewer’s insight regarding the impact of sub-token assignments. In our current implementation, we directly apply the base-b encoding to the original token indices (i.e., [0, ..., 255] for images and [0, ..., 50,256] for text). An interesting extension for future work is to reorder token indices based on semantic similarity, which could be automatically determined by a word embedding model, prior to base-b encoding. This would allow related tokens (like \"cat\" and \"bird\") to share more sub-token prefixes, enhancing the model’s ability to support fine-grained imputation in textual settings, analogous to the sub-token completion task demonstrated for images in Fig. 8.\n\n---\n\n**Q4**.  [exploratory question] Any thoughts about data-driven ways to set $\\ell$?\n\n**Response**.\nAn effective approach to setting $\\ell$ is to monitor the Idle Step Ratio (ISR), which can be calculated using Eq. (A5). As shown in our experiments in Appendix A.2.4, the ‘elbow point’ in the ISR curve often aligns with optimal performance. Therefore, $\\ell$ can be selected by gradually increasing its value and choosing the smallest point at which ISR plateaus. This selection method can be performed efficiently prior to training a model."}, {"review_id": "dwGTOlboFq", "reviewer": "Reviewer_ZpLq", "summary": "This paper introduces MDM-Prime, an enhancement to Masked Diffusion Models (MDMs) designed to address computational inefficiency arising from \"idle steps\" where the model's input remains unchanged during the sampling process. The core idea is to replace the binary masked/unmasked states with a more granular system that represents each discrete token as a sequence of sub-tokens using a base-b encoding. Masking is then performed at the sub-token level. The method achieves a lower perplexity on the OpenWebText dataset than previous MDM and autoregressive models.", "questions": "None", "limitations": "- For the text generation task, performance is evaluated solely based on perplexity. Augmenting the evaluation with other automated metrics or a human study would substantially strengthen the claims of superiority.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is well written. The illustrations are nicely plotted.\n- The motivation for the work is clear and compelling. The concept of \"idle steps\" provides a strong intuition for the inefficiency of standard MDMs. The proposed solution of partial masking through sub-token decomposition is an elegant and simple way to create a more fine-grained and efficient denoising process.\n- MDM-Prime requires only minimal architectural changes to a standard MDM. \n- The method achieves impressive results, notably outperforming autoregressive models in text perplexity (15.36 PPL vs. 17.54 PPL) on the OpenWebText dataset.\n\nWeaknesses:\n- The connection between improved model utilization and superior likelihood estimation is not thoroughly explored. While the reduction in \"idle steps\" is the primary motivation, the paper does not provide a deep analysis of why this leads to such a dramatic improvement in modeling capabilities, to the point of outperforming autoregressive models. \n- The paper does not report the impact of MDM-Prime on training and inference time. Handling sub-tokens and a more complex parameterization should introduce overhead. A direct comparison of throughput against the baseline MDM is necessary for a complete picture of the method's practical efficiency.", "was_revised": true, "final_justification": "The authors fully addressed my concerns in the rebuttal, therefore I would like to increase the rating from 4 to 5.", "rebuttal": "We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.\n\n---\n\n**Comments**\n\n**C1**. The connection between improved model utilization and superior likelihood estimation is not thoroughly explored. While the reduction in \"idle steps\" is the primary motivation, the paper does not provide a deep analysis of why this leads to such a dramatic improvement in modeling capabilities, to the point of outperforming autoregressive models.\n\n**Response**. We thank the reviewer for the comment. In Section 1, we provide the core intuition behind our method, and in Appendix A.2.4, we analyze how the target length $\\ell$ (which controls the number of idle steps) affects the model’s likelihood estimation as measured by perplexity.  While our current focus centers on providing empirical evidence and intuitive motivation, we acknowledge that a deeper analysis would constitute a valuable contribution and consider this an important direction for future work.\n\n---\n\n**C2**. The paper does not report the impact of MDM-Prime on training and inference time. Handling sub-tokens and a more complex parameterization should introduce overhead. A direct comparison of throughput against the baseline MDM is necessary for a complete picture of the method's practical efficiency.\n\n**Response**.\nA runtime comparison between the MDLM baseline and MDLM-Prime with varying $\\ell$ is presented in Appendix A.5.2 and Table A3  (the throughput can be derived by taking the inverse of the reported runtime values). Larger $\\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\\ell = 1$) and MDLM-Prime with varying $\\ell$, along with the corresponding perplexity (PPL).  Based on these results, we recommend $\\ell=4$ as an ideal balance between efficiency and performance:\n\n|  | Runtime (sec.)  |  \n| - | - |\n| $\\ell=1$ (PPL: 22.98) | 1.14e-1  |\n| $\\ell=2$ (PPL: 17.90) | 1.29e-1  |\n| $\\ell=3$ (PPL: 16.36) | 1.34e-1  |\n| $\\ell=4$ (PPL: 15.62) | 1.43e-1  |\n| $\\ell=6$ (PPL: 15.36) | 1.75e-1  |\n\n\\\nFor inference, we perform additional experiments to measure the generative perplexity (Gen PPL) under approximately aligned sampling time. The corresponding number of non-cacheable function evaluations (N-NFE) is also presented in the table. The results indicate that MDLM-Prime exhibits lower Gen PPL when N-NFE is higher.\n\n- MDLM\n| Runtime (sec.)  | 2.89 |  6.02  | 8.87  |\n| - | - | - | - |\n| Discretized Timesteps | 262 | 644 | 1,685 |\n| N-NFE  | 256  |  512  | 768  |\n| Gen PPL ($\\downarrow$) | 112.0 | 99.0 | 87.0 |\n\n\n- MDLM-Prime ($\\ell=4$)\n| Runtime (sec.) |  2.83 |  5.99  | 8.72 |\n| - | - | - | - |\n| Discretized Timesteps | 230 | 500 | 772 |\n| N-NFE  | 230 |  500  | 768  |\n| Gen PPL ($\\downarrow$) | 123.0 | 90.5  | 77.0 |\n\n---\n\n**C3**. For the text generation task, performance is evaluated solely based on perplexity. Augmenting the evaluation with other automated metrics or a human study would substantially strengthen the claims of superiority.\n\n**Response**.\nWe appreciate this valuable suggestion. We selected perplexity as the primary metric in the main manuscript due to its widespread adoption and its suitability for the evaluation of likelihood-based models such as MDM-Prime. In addition to standard perplexity, we also report generative perplexity (Gen PPL) in Appendix A.5.3, which evaluates the quality of model-generated samples through pretrained large language models. We agree with the reviewer and consider the reviewer's suggestion regarding additional automated metrics or human evaluation an interesting and promising future research direction."}, {"review_id": "4PInBUa3fG", "reviewer": "Reviewer_47UQ", "summary": "This paper introduces a partial masking technique which let each token in the masked diffusion model goes beyond just masked or unmasked state. In particular, each token is decomposed to subtokens, where the subtokens can be masked individually, making the original token partially masked. There is a bijective mechanism that allows the translation between token and subtokens. Empirical results show improvement over traditional MDM.", "questions": "N/A", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n\n1. The idea of partial masking via sub token is novel and elegant. It makes a lot of sense to have partial masking for discrete diffusion models, as the continuous diffusion model has many intermediate state, but it is non trivial to design a mechanism to do so. The design of sub tokens and bijection mapping is a clever way to tackle the problem.\n\n2. The empirical result is strong. It surpasses MDM by a large margin on perplexity and FID.\n\n3. The paper is well written and the idea is presented nicely. The figures are also very helpful.\n\nWeaknesses:\n\n1. Some discussion on the computational complexity is needed. Does it add any overhead to training and inference?", "was_revised": true, "final_justification": "The author addressed my concern and I will maintain my score for weak accept.", "rebuttal": "We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.\n\n---\n\n**Comments**\n\n**C1**. Some discussion on computational complexity is needed. Does it add any overhead to training and inference?\n\n**Response**.\nA runtime analysis of MDLM-Prime with varying $\\ell$ is presented in Appendix A.5.2 and Table A3. Larger $\\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\\ell = 1$) and MDLM-Prime with varying $\\ell$, along with the corresponding perplexity (PPL).  Based on these results, we recommend $\\ell=4$ as an ideal balance between efficiency and performance:\n\n|  | Runtime (sec.)  |  \n| - | - |\n| $\\ell=1$ (PPL: 22.98) | 1.14e-1  |\n| $\\ell=2$ (PPL: 17.90) | 1.29e-1  |\n| $\\ell=3$ (PPL: 16.36) | 1.34e-1  |\n| $\\ell=4$ (PPL: 15.62) | 1.43e-1  |\n| $\\ell=6$ (PPL: 15.36) | 1.75e-1  |\n\n\\\nFor inference, we perform additional experiments to measure the generative perplexity (Gen PPL) under approximately aligned sampling time. The corresponding number of non-cacheable function evaluations (N-NFE) is also presented in the table. The results indicate that MDLM-Prime exhibits lower Gen PPL when N-NFE is higher.\n\n- MDLM\n| Runtime (sec.)  | 2.89 |  6.02  | 8.87  |\n| - | - | - | - |\n| Discretized Timesteps | 262 | 644 | 1,685 |\n| N-NFE  | 256  |  512  | 768  |\n| Gen PPL ($\\downarrow$) | 112.0 | 99.0 | 87.0 |\n\n- MDLM-Prime ($\\ell=4$)\n| Runtime (sec.) |  2.83 |  5.99  | 8.72 |\n| - | - | - | - |\n| Discretized Timesteps | 230 | 500 | 772 |\n| N-NFE  | 230 |  500  | 768  |\n| Gen PPL ($\\downarrow$) | 123.0 | 90.5  | 77.0 |"}, {"review_id": "zTRTXkAZ8Q", "reviewer": "Reviewer_f3At", "summary": "This paper introduces MDM-Prime, a partial masking scheme designed to improve Masked Diffusion Models (MDMs). The core innovation is to allow tokens to exist in intermediate states between fully masked and unmasked, which the authors conceptualise as decomposing tokens into finer \"sub-token\" representations. This approach facilitates a more fine-grained denoising process by enabling the model to make predictions based on partially observed information. Empirically, Prime demonstrates superior performance across tasks, significantly outperforming previous MDM and autoregressive models on the OpenWebText benchmark and achieving competitive results on image generation datasets.", "questions": "- It would be helpful to include a sampling algorithm in the paper for clarity. Specifically, does the posterior $q(x_s^{i,l} \\mid x_t^{i,l}, x_0^{i,l})$ follow the same formulation as Eq3? Additionally, how is sampling from $p(x_s \\mid x_t)$ performed in practice? Given that $p(x_0^i \\mid x_t)$ is no longer independent across sub-tokens, we do not have a closed-form expression for $p(x_s^i \\mid x_t^i)$. In this case, is the sampling procedure carried out by first sampling $x_0^i \\sim p(x_0^i \\mid x_t)$, followed by sampling $x_s^{i,l} \\sim q(x_s^{i,l} \\mid x_t^{i,l}, x_0^{i,l})$?\n- It would be helpful to include the training curve in the paper. Does the training cost increase as $ l $ increases? Additionally, how does the training cost of MDM-Prime compare to that of MDM?\n- The paper does not report the model size. Could you clarify whether the model size is similar to that of the baselines?\n- While MDM includes more idle steps, these can be practically skipped as the model's output is not time-dependent, meaning logits do not need recomputation. In this regard, given the same sampling steps, MDM-Prime has a higher computational cost compared to MDM due to having fewer idle steps. This raises a critical question: given the same computational budget, like flops, at the test time, does MDM-Prime still maintain its performance advantage over MDM?\n- In Figures A3 and A4, there are cases where $ l < 1 $ for CIFAR-10. Could you clarify what this means? My understanding is that $ l = 1 $ corresponds to the original MLDM parameterisation without sub-tokens.\n- Everything has two sides. While MDM-Prime currently performs very well in terms of likelihood, what do you think are its downsides?", "limitations": "Overall, the paper is well-written and the results are impressive, particularly the achievement of outperforming autoregressive models in language modelling. However, some experimental and implementation details are missing, and the paper lacks theoretical analysis to support the proposed method. I would be inclined to raise my score if the authors can address these concerns.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\nOverall, the paper is well-written and presents very simple yet effective ideas. It is very promising that the proposed method outperforms auto-regressive models in the openwebtext dataset and achieves comparable performance to diffusion on image modelling.\n\n**Weaknesses**\n\nMy main concern lies in the absence of theoretical analysis explaining why partial-mask modelling yields better performance compared to MDM. While the use of sub-tokens enables smoother transitions, the parameterization of p(y_0|y_t) remains independent across different tokens. Does the proposed approach provide a tighter lower bound on the likelihood?", "was_revised": true, "final_justification": "The paper introduces a new mask diffusion model that significantly improves performance. However, it is still unclear why MDM-prime theoretically has a better likelihood. Therefore, I think 4 is an appropriate score.", "rebuttal": "We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.\n\n---\n\n**Questions**\n\n**Q1**. My main concern lies in the absence of theoretical analysis explaining why partial-mask modelling yields better performance compared to MDM. While the use of sub-tokens enables smoother transitions, the parameterization of $p(y_0|y_t)$ remains independent across different tokens. Does the proposed approach provide a tighter lower bound on the likelihood?\n\n**Response**.\nWe thank the reviewer for the comment, and are pleased to share some preliminary thoughts about the potential relationship between the likelihood bounds (Eqs. (4) and (5)).\n\nBoth objectives aim to approximate the same target, $-\\log p\\_{\\theta}(\\boldsymbol{x}\\_0)$, using a discrete diffusion process with the same scheduling function $\\alpha\\_t$. The key difference lies in the expectations:\n\n$$\\mathbb{E}\\_{q(\\boldsymbol{x}\\_t \\| \\boldsymbol{x}\\_0)} \\left[\\sum\\_{i=1}^L \\log p^{\\text{MDM}}\\_{\\theta}(x\\_0^i \\| \\boldsymbol{x}\\_t) \\right] \\\\, \\\\, \\\\, \\\\, \\text{  and  }\\\\, \\\\, \\\\, \\\\, \\mathbb{E}\\_{q(\\boldsymbol{y}\\_t \\| \\boldsymbol{y}\\_0)}\\left[\\sum\\_{i=1}^L \\log p^{\\text{Prime}}\\_{\\theta}(\\boldsymbol{y}\\_0^i \\| \\boldsymbol{y}\\_t)\\right].$$ \n\nA possible direction for theoretical exploration is to analyze the per-step difference between these expectations: \n\n$$\\mathbb{E}\\_{q(\\boldsymbol{y}\\_t \\| \\boldsymbol{y}\\_0)} \\left[ \\sum\\_{i=1}^L \\log p^{\\text{Prime}}\\_{\\theta}(\\boldsymbol{y}\\_0^i \\| \\boldsymbol{y}\\_t)\\right] - \\mathbb{E}\\_{q(\\boldsymbol{x}\\_t|\\boldsymbol{x}\\_0)} \\left[\\sum\\_{i=1}^L \\log p^{\\text{MDM}}\\_{\\theta}(x\\_0^i \\| \\boldsymbol{x}\\_t) \\right] \\triangleq d.$$\n\nIf this difference ($d$) is consistently positive, it could suggest that MDM-Prime provides a tighter bound. However, since $p^{\\text{MDM}}\\_{\\theta}$ and $p^{\\text{Prime}}\\_{\\theta}$ are parameterized by distinct neural networks operating on different input spaces, establishing this relationship analytically is nontrivial and requires specific assumptions about the model families.\n\nTo investigate this empirically, we estimated the expectation difference using trained models at a number of discretized timesteps. The results show positive values as follows: \n\n| $t$ | 0.25 | 0.5 | 0.75 | 1.0 |\n|-|-|-|-|-|\n| $d$ | 199.38  | 208.80  | 59.25  | 11.47  | \n\nWe leave a rigorous theoretical analysis that fully explains this behavior for future work.\n\n---\n\n**Q2**. **(a)** It would be helpful to include a sampling algorithm in the paper for clarity. \n**(b)** Does the posterior $q(x\\_s^{i,l} \\| x\\_t^{i,l}, x\\_0^{i,l})$ follow the same formulation as Eq. (3)? \n**(c)** Additionally, how is sampling from $p(x\\_s \\| x\\_t)$ performed in practice? Given that $p(x\\_0^i \\| x\\_t)$ is no longer independent across sub-tokens, we do not have a closed-form expression for $p(x\\_s^i | x\\_t^i)$. In this case, is the sampling procedure carried out by first sampling $x\\_0^i \\sim p(x\\_0^i \\| x\\_t)$, followed by sampling $x\\_s^{i,l} \\sim p(x\\_s^{i,l} \\| x\\_t^{i,l}, x\\_0^{i,l})$?\n\n**Response**.\n**(a)** We thank the reviewer for the suggestion and will enhance our description regarding the sampling process in Section 3.1 in an updated manuscript.\n\n**(b)** The posterior $q(y\\_s^{i,j} | y\\_t^{i,j}, y\\_0^{i,j})$ follows the same formulation as Eq. (3) and is defined as follows:\n\n$$q(y\\_s^{i,j} \\| y\\_t^{i,j}, y\\_0^{i,j})=\\begin{cases}\n\\delta\\_{y\\_t^{i,j}}(y\\_s^{i,j}) & \\text{if $y_t^{i,j}\\in \\mathcal{Y}$} \\\\\\\\\n\\frac{1-\\alpha\\_s}{1-\\alpha\\_t} \\delta\\_\\texttt{m} (y\\_s^{i,j}) + \\frac{\\alpha\\_s-\\alpha\\_t}{1-\\alpha\\_t} \\delta_{y\\_0^{i,j}} (y\\_s^{i,j})  &  \\text{if $y\\_t^{i,j}=\\texttt{m}$}\n\\end{cases}$$\n\n**(c)** Yes, the sampling procedure is carried out by first sampling $\\boldsymbol{y}\\_0^i \\sim p\\_{\\theta}(\\boldsymbol{y}\\_0^i \\| \\boldsymbol{y}\\_t)$, followed by sampling each sub-token $y\\_s^{i,j} \\sim q(y\\_s^{i,j} | y\\_t^{i,j}, y\\_0^{i,j})$. \n\n---\n\n**Q3**. **(a)** It would be helpful to include the training curve in the paper. **(b)** Does the training cost increase as $\\ell$ increases? Additionally, how does the training cost of MDM-Prime compare to that of MDM?\n\n**Response**.\n**(a)** Due to NeurIPS requirements, we are unable to add figures during the rebuttal phase, but we would be pleased to include the curves in the final manuscript version. We provide the following table that compares the evaluation perplexity (PPL) of MDLM-Prime and MDLM at various training steps:\n\n| Training Steps |  50K | 100K | 250K  | 500K  | 1,000K  |\n|-|-|-|-|-|-|\n| MDLM | 29.21 | 27.18 | 24.81 | 23.84 | 22.98 |\n| MDLM-Prime ($\\ell=4$) | 25.92 | 19.45 | 18.32 | 16.87 | 15.62 |\n| MDLM-Prime ($\\ell=6$) | 22.72 | 18.54 | 16.77 | 15.95 | 15.36 |\n\n**(b)** A runtime analysis of MDLM-Prime with varying $\\ell$ is presented in Appendix A.5.2 and Table A3. Larger $\\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\\ell = 1$) and MDLM-Prime with varying $\\ell$, along with the corresponding perplexity (PPL). Based on these results, we recommend $\\ell=4$ as an ideal balance between efficiency and performance:\n\n|  | Runtime (sec.)  |  \n|-|-|\n| $\\ell=1$ (PPL: 22.98) | 1.14e-1  |\n| $\\ell=2$ (PPL: 17.90) | 1.29e-1  |\n| $\\ell=3$ (PPL: 16.36) | 1.34e-1  |\n| $\\ell=4$ (PPL: 15.62) | 1.43e-1  |\n| $\\ell=6$ (PPL: 15.36) | 1.75e-1  |\n\n---\n\n**Q4**. The paper does not report the model size. Could you clarify whether the model size is similar to that of the baselines?\n\n**Response**.\nFor image experiment, MDM-Prime is with 114M parameters, which matches the MDM baseline. For text experiment, MDLM-Prime is with 130M parameters, which is the same as the MDLM baseline. We’ll include this information in the final version.\n\n---\n\n**Q5**. While MDM includes more idle steps, these can be practically skipped as the model's output is not time-dependent, meaning logits do not need recomputation. In this regard, given the same sampling steps, MDM-Prime has a higher computational cost compared to MDM due to having fewer idle steps. This raises a critical question: given the same computational budget, like flops, at the test time, does MDM-Prime still maintain its performance advantage over MDM?\n\n**Response**.\nFor the image experiment, we follow [1] to adopt the ablated diffusion model (ADM) architecture, which is conditioned on the time embeddings and the sampling updates are not cacheable (i.e., updates cannot be skipped). Therefore, MDM-Prime still maintains its performance advantage over MDM according to Tables 3 and 4 in the manuscript.\n\nFor the text experiment, we perform additional experiments to measure the generative perplexity (Gen PPL) under approximately aligned sampling time. The corresponding number of non-cacheable function evaluations (N-NFE) is also presented in the table. The results indicate that MDLM-Prime exhibits lower Gen PPL when N-NFE is higher.\n\n- MDLM\n| Runtime (sec.)  | 2.89 |  6.02  | 8.87  |\n|-|-|-|-|\n| Discretized Timesteps | 262 | 644 | 1,685 |\n| N-NFE  | 256  |  512  | 768  |\n| Gen PPL ($\\downarrow$) | 112.0 | 99.0 | 87.0 |\n\n- MDLM-Prime ($\\ell=4$)\n| Runtime (sec.) |  2.83 |  5.99  | 8.72 |\n|-|-|-|-|\n| Discretized Timesteps | 230 | 500 | 772 |\n| N-NFE  | 230 |  500  | 768  |\n| Gen PPL ($\\downarrow$) | 123.0 | 90.5  | 77.0 |\n\nWe would also like to highlight that MDLM requires large numbers of discretized timesteps to achieve high N-NFE due to the existence of idle steps (see Eq. (A4) in Appendix). Since determining the unmasked elements (i.e., resampling) in the idle steps causes computational overheads (even without forward passing the model), the sampling cost of MDLM becomes increasingly higher with higher N-NFE:\n\n- MDLM\n| N-NFE |  768  | 1,003 | 1,023 |\n| - | - | - | - |\n| Discretized Timesteps  |  1,685  | 25,000 | 500,000 |\n| Runtime (sec.) | 8.87 | 23.30 | 159.81 | \n\n- MDLM-Prime ($\\ell=4$)\n| N-NFE |  768  | 1,003 | 1,023 |\n| - | - | - | - |\n| Discretized Timesteps  | 772  | 1,022 | 1,044 |\n| Runtime (sec.) | 8.72 | 12.31 | 12.58 | \n\n---\n\n**Q6**. In Figures A3 and A4, there are cases where $\\ell<1$ for CIFAR-10. Could you clarify what this means? My understanding is that $\\ell=1$ corresponds to the original MLDM parameterization without sub-tokens.\n\n**Response**.\nYes, $\\ell = 1$ corresponds to the standard MDM configuration. For cases where $\\ell < 1$ appears, this refers to *super-tokenization*, where multiple pixel values are grouped into a single token before applying MDM. For example, $\\ell = \\frac{2}{3}$ indicates that three 8-bit RGB channels are merged into two 12-bit super-tokens, resulting in a larger token class size (e.g., from 256 to 4,096).\n\n---\n\n**Q7**. Everything has two sides. While MDM-Prime currently performs very well in terms of likelihood, what do you think are its downsides?\n\n**Response**.\nWe thank the reviewer for the question. The current MDM-Prime has the following limitations:\n\n(i) Non-monotonic Performance Scalability with $\\ell$: As demonstrated in our experiments (Appendix A.2.4), increasing $\\ell$ generally leads to better performance. However, the trend is not strictly monotonic, as larger $\\ell$ does not always yield improvement. Although we can properly select this value based on the elbow point method presented in Appendix A.2.4, we believe that the development of advanced training techniques to enable improved performance scalability with larger $\\ell$ constitutes an interesting future avenue.\n\n(ii) Fixed Design of the Sub-tokenization Function $f$: When extending MDM-Prime to applications such as text imputation, the choice of the sub-tokenization function $f$ and its inverse $f^{-1}$ becomes critical (see Q3 from Reviewer uEyH). The grouping of semantically related tokens may improve generalization in these tasks, and we consider this a promising direction for further investigation.\n\n---\n\n**References**\\\n[1] Gat et al. Discrete Flow Matching, NeurIPS 2024."}]

Author rebuttal:

Meta-review: This paper introduces partial masking for masked discrete diffusion, by decomposing each token into a small set of sub-tokens to define intermediate partially masked states. All reviewers agreed that the approach is simple, the paper is clear and the results are strong. There are some weaknesses, a particular one being that there is not a clear theoretical reason for why partial masking should improve results. Other comments/concerns raised by reviewers were generally resolved satisfactorily, with some additions to paper promised  (sampling algorithm, training curves, parameter counts, better comparisons). These will all improve the paper. I am recommending an accept, and strong encourage the authors to add all the promised items in the camera-ready.
Decision: Accept (poster)