Simplified and Generalized Masked Diffusion for Discrete Data
Authors: Jiaxin Shi
Kehang Han
Zhe Wang
Arnaud Doucet
Michalis Titsias
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion
discrete
masked diffusion
absorbing diffusion
diffusion model
Primary area: diffusion_based_models
TL;DR: A simplified and generalized framework for training masked discrete diffusion models

Abstract: Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes.

Simplified and Generalized Masked Diffusion for Discrete Data
Jiaxin Shi', Kehang Han", Zhe Wang, Arnaud Doucet, Michalis K. Titsias Google DeepMind
Abstract
Masked (or absorbing) diffusion is actively explored as an alternative to autore- gressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parame- terization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained Open WebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion mod- els on pixel-level image modeling, achieving 2.75 (CIFAR- 10) and 3.40 (ImageNet 64 x 64) bits per dimension that are better than autoregressive models of similar sizes. Our code available at https:/ hagp::///ihhunn thub. com/ google-deepming
Introduction
Since their inception [1, 2, 3], diffusion models have emerged as the workhorse for generative media, achieving ofate-of-thh-art
tasks such as image synthesis [4, 5, 6], audio [7, 8) and video generation [9. 10, 11, 12, 13J. The majority of existing successes are for continuous state space diffusions. While diffusion models have been extended to discrete state spaces [1, 14, 15] and have been successfully applied to applications ranging from graph generation [16], text-to-sound generation [17] or protein design [18), they remain not as widely used as their continuous counterparts as they not competitive with autoregressive models in important domains such as text modeling. This has motivated the development of continuous space diffusion models where the discrete data are embedded in the Fuclidean space [19, 20, 21, 22, 23] the simplex [24, 25, 26, 27, 28). We believe that one of the reasons for the limited success of discrete diffusions is that they have been hindered by fairly complex formulations and training objectives. This paper is step towards closing this gap.
In this work, we focus on "masked" (or "absorbing") diffusions, a discrete diffusion formulation first presented by Austin et al. [14], and later explored by the literature from various perspectives [29, 30, 31, 32). We follow here a continuous-time framework which has proven very useful to improve the training and understanding of continuous state space diffusions [see e.g., 3, 33, 34]. We make several technical contributions which simplify the training of these models and improve significantly their performance. Our contributions
follows:
Using elementary arguments, we establish several properties for the forward process induced by this model and its corresponding time reversal, improving our understanding of this model class.
Equal contribution. Correspondence to: jiaxinstgoogle.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

We provide remarkably simple expression of the Evidence Lower Bound (ELBO) for masked diffusion models, showing that it corresponds to a weighted integral over time of cross-entropy losses. Similarly to continuous space diffusions [33], this objective can be rewritten in terms of signal-to-noise ratio and exhibits invariance properties.
We develop unifying understanding of previously proposed continuous-time discrete diffusion models [29, 32, 35], revealing the changes they made to our ELBO objective and/or model parame- terization. We show that these changes either lead expensive model evaluations, large variance in training, or breaking the consistency between forward and reverse processes.
On GPT-2 scale text modeling and pixel-level image modeling tasks, masked diffusions trained using our simple ELBO objective outperform previous proposals, leading to the best likelihood and zero-shot transfer performance among discrete diffusion models.
Finally, based on our simplified masked diffusion formulation, we propose generalized masked diffusion model that allows state-dependent masking schedules. This generalized masked diffusion model further improves predictive performance measured by test likelihoods.
Concurrent work by Ou et al. [36] and Sahoo et al. [37] derives similar simplified expression of the ELBO. Ou et al. [36]'s 's derivation relies on an observation similar to the one we made in Proposition 1.
Masked Diffusion
Consider sentence where we progressively replace each word with special mask token, transform- ing the sentence into a sequence of masks. Our goal to train a generative model that reverses this process, effectively turning sentence of masks back into meaningful text. More formally, assume our data consists of tokens from a finite discrete state space with m possible states, represented by and their corresponding one-hot vectors eo, ci...., cm-1- To accommodate the masking process, we augment this space with an additional mask state, denoted by the index m. The masking process transitions each token to the mask state at a random time. This process, known as the forward process, is applied independently to each token (e.g., each word), progressively converting the data into sequence of mask tokens. By learning reverse this masking process, create a generative model capable of producing coherent discrete data.
integers 0, 1, .....
Discrete-ting forward process. We start with the case of a single token and later expand to multiple dimensions. We define the forward process as a Markovian sequence of discrete random variables I; indexed by time f, where i runs from 0 to 1. Throughout the work, we abuse the notation such that Fg can be either an integer or its corresponding one-hot vector, whenever is clear from the context. We divide [0, 1] into T intervals, and let s(i) (i 1)/T, t(i) = i/T. Following Austin et al. [14], the state transition between [s(i), (()] is determined by a transition matrix of size an all-one vector of size m + 1, em
(m + 1) (m + 1): Q: = (1 B.)I + 8,1e where represents a one-hot vector where element at index m is 1. Each entry [Qi]je denotes the probability of transition from the state j to the state k:
[ = (()(()) k/ = j) = - 8.)5zk + 3:0kmm
This means that, with probability 1 - Bi, Ta(a) Is(i), otherwise it jumps to the mask state. Given the above transition matrix, the marginal distribution at time t(i) given xo is
(( Cat(z Q. mо) = z Q
Here, use Cat(x; p) denote a Categorical distribution where p is the vector of probabilities
of being in each category, and Q a I Q2 a,I + (1 (1-a;) az) 1e for a I (1 - By)) We
expect or to become very small or zero for a sufficiently large T such that g(z) (z0) for any zo will become a delta mass at the mask state.
Continuous-time limit. We can define continuous-time forward process by taking limit of the above discrete-time process. We first specify continuous function 3(t) such that B = (TTT)) We then let T oo in the discrete-timg time process and compute the limit of Q, (proved in Austin et al. 14, Appendix A.6, see also App. A) as
(t) lim Q. a + (1 1e where a e exp (- B(s)ds). T-400
(I)
2

---

at
0,0,/(1 a.)
1.0
0
linear
-5
geometric
0.5
cosine
poly2
15
poly0.5
0.0
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 1: Masking schedules in the literature: (Left) az (Right) weight of the cross-entropy loss w.r.t. t; Equations for these schedules are given in Tab. 4 in Appendix.
so that (( (zo) Cat(x;; Q(t) xo). For two arbitrary times, 0 < 1, the transition distribution that is compatible with the above marginal (i.e., ((22) Ex, (r,)) (zr))
g(ze(x))) = Cat(ar; (ss, t) x.). where (s.,t) ((s)) = 0: a,
a,
Note that Austin et al. [14] did not derive this explicit form of transition matrix between two arbitrary time and L, which appeared later in Zhao et al. [38] concurrently with our work.
Masking schedules. From the definition a,, we have that Ogg And similar to the discrete- time formulation, we would like be zero or very close to zero. We provide a summary of masking schedules from literature that satisfy these properties in Fig. The linear schedule was proposed in Sohl-Dickstein et al. [1] for binary variables and then re-derived by Austin et al. [14] from mutual information for discrete-time models. The geometric schedule at is plotted for Bmin = 10-5 and 8max 20. It was first used for continuous diffusions [3] and then for discrete by Lou al. [32]. The cosine schedule was originally proposed in MaskGIT [39], an iterative unmasking generative model inspired by diffusion. This schedule has the property of slowing down the unmasking process the beginning of the reverse generation. Aligning with their observation, we find that this results in lower chance of conflicting tokens being unmasked simultancously at the start of generation, thereby enhancing the overall generation quality.
Time reversal of the forward process given The analytic property of our forward process allows to compute many quantities of interest closed form. One such quantity frequently used in diffusion models is the time reversal of the forward process given zo: (x vo) for t. We derive it in App. as
gx [ut, mo) Cat(#s; Re (t, s) Tx), where ( (1, =
0g (2o em)
Oz
From the transition matrix (t, s) E ()(++))((++)) we can see the reverse process conditioned on zo has a very simple logic-if x, is a mask, with probability T-or it will jump to the state at time s, otherwise it will stay masked. Once is unmasked, it remains in the same state until the end.
Model and Objective
For discrete-time masked diffusion process, define generative model by approximately reversing the forward transitions using reverse model (o(x, One way to define this model is
po( A q(x,|=1, (00 (z,, +)),
(2)
where (2,, t) e m++1 a probability vector parametrized by a neural network fo with a sofumax applied to the output logits (note the m-th output is forced to 0 since the clean data cannot be masks):
softmax((f(x;;))) xr
It = m,
(2 t)
Te + m.
This is known mean-paraameterizaaioon since it leverages prediction model for the mean of æ0- A matrix-form depiction of (o(xal1))) is shown in Fig. (right). fact, we can select a time-invariant parametrization (æ (o(x))) as [36] showed that (2)) given &, is identical for any t.

---

Besides (=( we also need to specify p(ro| and the prior distribution p(2)(()) = p(z1). Following the practice in continuous diffusion models [33], we choose p( ((())))(1)) x (( And since q (20)= = 82,, for any Zo as 01 = O, we set p(r1) 2 dx,,,, see App. E.
We then write out the discrete-time diffusion model objective [1, 2], which is lower bound of the log marginal likelihood of data zo under the model p (known as the Evidence Lower Bound, or ELBO): log p(zo)= 2 Be (mn)) [lg (())) - (T( (ro))l(((11)) - Lr,
where Lr Ei= (orces) [KL(g(Zs()) (it) mu) (1) (m)))) )) For the above choices of the prior distribution, the term (T ( [()())))))))) becomes zero. Under the reverse model (2), the KI. divergence terms in Lr becomes (proof in App. D)
(L(g(x ( mo) lipo (xal<1))
0g2 0z 6... To log Ma (Zt, !),
which is a simple cross-entropy loss between the predicted logits and the clean data. In App. D, we show that LT Riemann sum and is lower bounded by the corresponding continuous integral:
a4 1- att
C00 P lim CT CT=
(mo) [0 To log μe (2 t) dt,
(4)
Je(1)
where af denotes the derivative of ar, with respect to f. Therefore, we can obtain ELBO that is tighter than that of any finite T by pushing T - cc. This ELBO can be further simplified by letting ((1) - 0. As result, Be(=aa)) /z0) [log (())) goes to O and the ELBO becomes -Loo.
For continuous state-space diffusions, the ELBO depends on the signal-to-noise ratio (SNR) at its endpoints but is otherwise invariant the noise schedule [33]. We establish here similar result for discrete diffusions. Consider choosing a = (Ae), where represents the sigmoid function (x) = T+e- this context, the log-SNR defined by A log T-o, log-SNR(t). By making change of variables in (4) make everything function of the log-SNR, we obtain
_a
S Jaco:
A1
T (A)lE (mo) [6=m,m z log (0 A)] dA.
where ūa(z; A) == (a(a t) and (=)) =- (( for t = log-SNR ( This shows that the only effect Og has on the loss through the values the SNR at the endpoints. Still, because we draw uniform samples of t to estimate the integral, the choice of masking schedule affects the variance.
Multidimensional data. In the previous sections, 2, assumed to be single discrete token. To extend the method multidimensional data, let ae be now a sequence (a ..x(() x (N)). where each element z/(h) represents a discrete token. We select forward process which factorizes across all N tokens: q(x,(2x) = -I {( (=("">"/)))) As a result, the forward marginals q(2, (zn) and reversal q(2,) (x,, also factorize. In this case, we define the reverse model as (=(x, & (m)
I ((=!") (m) (=("") M4 (x,,, t)), where (x(x,, t) is neural network that takes the full N tokens as input and outputs N probability v The n-th output (4
(1)) (xre, 1) is a prediction model for [rz], the mean value of the n-th token. Repeating above derivations gives
E[r6"" (
04
(NN
J0 1--a Eg(r(110)
[ _m (4) (n)) log n0 (ze. t) dt.
(5)
We term our simple masked diffusion model trained with loss (5) MD4 (Masked Discrete Diffusion for Discrete Data). A single step of MD4 training algorithm is described in Alg. 1 in Appendix.
4 Sampling
We ancestral sampling from our discrete-time reverse process for generation. We have found this yields slightly higher sample quality compared to other methods such as Euler discretization [29, 32]. For conditional generation tasks such as infilling, we find that the simple approach works best we keep the conditioning tokens unmasked throughout the generation process. A complete description of the sampling algorithm can be found in Alg. 2 in Appendix.
We intentionally choose the reverse model to factorize across dimensions because the true reverse transition (xx))x)) factorizes in the continuous time limit (as approaches t).

---

Linear dt
80
Linear an cosine grid
Tokens Revealed per Step
100
Cosine a;
60
Cosine (class cond)
40
60
40
20
Linear de
Cosine d;
50
100
150
200
250
Total Steps T
Steps
Figure Left: FID evaluation for 50k samples randomly generated from MD4 on pixel-level modeling of ImageNet 64 (numbers in Tab. 6). Right: Number of tokens revealed per generation step (T 256). Each image consists of 64 64 x 3 12288 tokens.
Impact of schedules and discretization. For comparing different sampling configurations, we primarily the FID score [40] image datasets our evaluation metric. We favor it text generative perplexity used in prior work [32], as the latter can be misleadingly reduced by lowering sample diversity [41]. We initially trained our model using the linear schedule, which achieves the best final ELBO overall; however, we found that sampling did not perform well with standard uniform discretization grid t(i) 4. We hypothesize that time discretization generating multiple tokens in single step. We then switched to the cosine schedule (Tab, 4) that slows down unmasking at the beginning of reverse process. This drastically improves the FID on ImageNet 64 64 from 70 to 17 for T 256 steps (Fig. 2, left). Building on this observation, we suggest using "cosine" discretization grid for sampling in models trained with linear schedule:
lead conflicts by
Th ((i) cos ( (1-
This induces the same discretization in a₁ as the cosine schedule with a uniform grid, leading to comparable sample quality, shown in Fig. 2 (left). In Fig. 2 (right), we plot the number of tokens unmasked per step for linear and cosine schedules with a uniform grid. We believe the cosine schedule performs better because it leverages information redundancy: with more tokens revealed, step.
remaining tokens become more predictable, reducing conflicts when unmasking them single
Although these findings were originally developed on images, we find them translate well text (see Fig. 10). we expect other techniques such as top-p sampling [41], classifier guidance [42, 43], and predictor-correst correctors [29, 44] to further improve sample quality of our models. While we reserve these for future work, we note that the JAX [45] implementation of categorical sampling implicitly truncates small probabilities, creating similar effect to top-p sampling. See App. G for details.
5 Relation to Existing Work
We discuss how to unify several existing masked diffusion models using our framework.
Continuous-Time Markov Chains (CTMC). To show the connection with the CTMC view presented Austin et al. [14], Campbell et al. [29], we can write out the forward and reverse masked diffusion using CTMC machinery. To see this, for short time At, given To, the Taylor expansions of our forward and reverse transition matrices at f. are
i At) = 1 + Q(t) ((At) for Q(t) a ((t((1))) -),
Rt - At) = T+ R (t) o for (t) 2 of
Cm Cm( 6m)
(8)
1-
where Q(t) and 1ea0tt are known as the transition rate matrices. Austin et [14] derived the same Q(t) in App. A.6 of their paper. However, they did not explore the reverse process
PPerplexity of generated samples scored by large language model such GPT-2.
5

---

continuous tietetuoou---ueeeeee objective. Campbell et al. [29] derived an alternative ELBO expression using rate matrices, which Kitouni al. [46] further simplified for absorbing diffusion. In App. H.I, we show how to recover their expression by separating out a constant from our ELBO expression (4) and applying a discrete Eineggaaionn-y--arte evaluations of the prediction model 1 t) to compute an inner summation. To circumvent this computational burden, they used a doubly stochastic estimate, However, this leads to significantly higher variance compared to the analytic cross-entropy (4) which only requires one pass of ua( t).
A key limitation of their expression is that it needs N
Please refer to App. H.2 for more details.
Score parameterization. While far used prediction model 14 (xee 1) for the mean of clean data given Tt (i.e., mean parameterization), one can choose other ways of parameterizing the reverse model. Lou et al. [32], Benton et al. [35] proposed parameterize the discrete "score" s(1,, 1)3 (e(e)]
g=Gi) and introduced score-based loss for discrete diffusions. In App. H.3, we provide
an alternative derivation of their loss which is simpler. We show the link between score and mean parameterizations through the following proposition.
Proposition 1 (Score Parameterization vs. Mean Parameterization). Let qt be the marginal distribu- tion of the masked diffusion defined in Sec. al time I. The discrete score s( ()g mask state It = m. and j + m can be expressed as
2:33) for (xx)
s(m,t) t)2 = at E| = mm]ea ej. which satisfies
Σ
s(m
Ott t)
1-a ax
di
jmm
Proposition (proved in App. H.3) implies that a reasonable score model for a mask state is
0t so t)j = 1 a po t)3-
(10)
Indeed, substituting (10) into the score-based loss of Lou et al. [32], Benton al. [35] recovers our objective (4). In Lou et al. [32], the score is parameterized as a neural network without enforcing the constraint in (9). This means the learned reverse model can be incompatible with the forward process. We find that our parameterization, which enforces the constraint, leads to more stable training and better results.
Any-order autoregressive models. The continuous-ting reverse process of our masked diffusion model can be viewed as an any-order autoregressive model (AO-ARM) [47]. To see this, we reorder the tokens according to the timing of their unmasking events in the reverse process. For all tokens, the cumulative distribution functions (CDFs) of unmasking times (T,, J are identical and satisfy
P(T, < 1) = P(=."") m) 1 dg As result, the ordering is uniformly random across all possible arrangements, and the token prediction during each unmasking event represents a prediction step in AO-ARMs. This connection was initially pointed out in Hoogeboom et al. [48, App. CJ. The relation between our simplified ELBO (5) and the AO-ARM objective is independently clarified by Ou et al. [36]. Despite this equivalence, our work demonstrates that the masking schedule ar introduces a new degree of freedom in the design of such models. Variations in a₁ can lead to different distributions of unmasking times, significantly impacting performance in diffusion-style parallel sampling under time discretization, as shown in Fig. 2.
Other related work. Due to space constraint, we defer the discussion on other related work, including MaskGIT [39], discrete flow matching [49], SDDM [30], Blackout diffusion [50] and SUNDAE [51], App. H.4.
Generalization to State-dependent Masking Schedules
Consider a scenario where some tokens hold more significance than others and we would like to unmask them earlier in the process. To achieve this, we introduce state-dependent masking schedules, where the probability of unmasking token depends not only time, but also on the token's value.
We first define the forward process for single token dt. Let os be a m + dimensional vector function, i.c., there is different function 0g for each possible value i of the token z;. Also, by

---

Mayor said
steps
that new plan
Mayor Bowser said meetings Commissioner Thursday that new plan will board
700 steps
Mayor Muriel Bowser said after meetings Commissioner on 850 steps
Thursday that new plan will be board in December
Mayor Muriel Bowser said after meetings with Commissioner Busby Thursday that new plan will on board December.
1000 steps
Figure 3: Iterative unmasking process for unconditionally generated sample by MD4. This visualization only includes a subsequence from a generated sequence of 1024 tokens. "?" represents masks. Masked tokens revealed sequentially: green (steps 500-700), yellow (700-850), and red (850-1000). Additional unconditional generation from MD4 can be found in App. K.5.
vector a we denote the element-win wise division of the two vectors. We define the forward transition a, as (xx)) = Cat(r;; (s, t) xx) where
((s, () diag
+
diag
Cre O22
1
and diag( (te) is diagonal matrix with the vector St in its diagonal. The probability of moving from current state J, to future state dt (either the as
ws mask) is determined by state-dependent
rate as Zx. while the marginal at time s given Zo
(<<,)<n)) Cat [(s)T2z) for Q(s) diag(os) (T diag(ax))
Further, for any time 0 I s < 1 holds that q( (z0) = Σx. q( (xx, (mv) so the above is valid continuous-tin Markov chain.
Given the forward conditionals and marginals, we can now compute the time reversal conditioned on is derived in App. 1.1. For m, we have
To. The full form of
g(x, lx, m, x0) [ - = m, To; x020
-=a
Te u. + (4-a))
To x... (11)
This suggests that the reverse model given 2; = m can be chosen as po m) 2 g( ( = m. oe (zz, t), (xx, where (z, t) a neural network that approximates E[xo|2,] while diag(po (22,, ()) approximates TF( [ara] diag(E[ar (xx))) We show in App. I.1 that the negative continuous-tingg time ELBO for the state-dependent rate case is
6x fx=f(( o4
[8m.,
(ro po t) + Fanrd log (o f))] dt.
(12)
Here, at is the elementwise derivative of az This generalizes the MD4 loss (4), which is recovered when a is scalar schedule times a vector of ones. For N tokens, the model further generalize similarly to Sec. 3 and the loss is given in (32). We call this generalized model GenMD4
To learn the token dependent masking schedule using ELBO optimization, we parametrize the m + 1 dimensional function a using the polynomial schedule (see Fig. 1) as a4, = 1 t and optimize each parameter w; 0.* The value of wy through the masking probability how fast the token with value i jumps to the mask state. Since in the loss (12) the distribution (xx))20) depends aę and thus the vector w, optimizing w poses discrete gradient estimation problem [see, e.g., 52]. Naive autodiff leads to biased gradients and pushes w towards zero because the gradients cannot propagate through the (discrete) samples drawn from q (z.). To fix this, used the REINFORCE leave-one-our estimator [53, 54] to compute low- wwriaariece unbiased gradients for optimizing W. Details are given in App. I.2.
Ctt determines
"We only need m learnable parameters for i = O,..., 1, since can never be the mask token. For the tinal mask dimension choose arbitrary fixed value such Wm 0.
7

---

Table Zero-shor unconditional perplexity on five benchmark datasets from Radford al. [57]. The numbers for other methods are from Lou et al. [32] except our reimplementation of SEDD Absorb. Our MD4 model achieves the best result on benchmarks except LAMBADA where is the second best. "The GPT-2 numbers are reported for the GPT-2 checkpoint pretrained on WebText instead of OWT thus is not direct comparison.
Size
Method
LAMBADA
WikiTex12
PTB
WikiText103
IBW
Small
GPT-2 (WebText)
45.04
42.43
138.43
41.60
75.20
D3PM
< 93.47
S 77.28
I 200.82
< 75.16
S 138.92
Plaid
S 57.28
X 51.80
X 142.60
S 50.86
S 91.12
SEDD Absorb
50.92
41.84
114.24
40.62
< 79.29
SEDD Absorb (rcimpl.)
S 49.73
S 38.94
3 107.54
S 39.15
S 72.96
MD4 (Ours)
< 48.43
34.94
102.26
S 35.90
S 68.10
Medium
GPT-2 (WebText)"
35.66
31.80
123.14
31.39
55.72
SEDD Absorb
42.77
31.04
87.12
29.98
61.19
MD4 (Ours)
I 44.12
S 25.84
66.07
S 25.84
S 51.45
7 Experiments
7.1 Text
Text is natural discrete data with rich structures. For comparison with prior work, evaluate likelihood on two datasets: text8 [55], a character-level text modeling benchmark, and OpenWebTees [56], an open clone of the unreleased WebText dataset used to train GPT-2 [57]. We also assess our model's performance on downstream tasks by training on Fine [58], high-quality dataset of fine educational text commonly used by the open-source community for comparing LLMs. Unless otherwise specified, linear schedule and cosine sampling grid are employed.
Open Text We train MD4 GPT-2 small (S) and GPT-2 medium (M) sizes Open Web- Text and evaluate zero-shot perplexity on five benchmark datasets used in Radford al. [57].
OpenWabt Text Evall Perplexity
Gaussian Diffusian-S
6x10
SEDD-5
We keep our evaluation setup the same as SEDD [32]. To ensure fair comparison, mented SEDD in our codebase. Our implemen- tation led slightly better results than those reported in their paper.
GenMD4-S
reimple-
4 10'00
MD4-M
3 10
www Man
2x101-
As in Tab. 1, our small model outperforms previous best discrete diffusion models on five tasks. We are also better than GPT-2 all tasks except LAMBADA where we are the sec- ond best method. When scaling up medium size, MD4 similarly beats SEDD and GPT-2 on 4 out of 5 tasks.
400
60c
800
1000
Training steps (1 unit 1000 steps)
Figure 4: Perplexity on OpenWebText (OWT) val- idation set during training. The final numbers are reported in Tab. in Appendix.
To confirm that the strong zero-shot performance stems from improved training, we plot perplexity on 2% Open WebText validation set in Fig. 4. Our models converge faster and have better final likelihoods than prior methods. We also observed that SEDD [32] has training instabilities, likely due to score parameterization breaking consistency between forward and reverse processes (Sec. 5). Although GenMD4 achieves lower perplexity than MD4, we observed that the learned ws can overfit to dataset statistics, making less effecttvv zero-shot transfer tasks.
We also assess our models' generation quality. Fig. shows a randomly selected, notably coherent sample from MD4-medium and its denoising process. Fig. 10 demonstrates MD4's 's text infilling ability and highlights a substantial quality gain when transitioning from uniform to cosine discretization (see Sec. 4). Despite MD4's strong performance on quantitative metrics like generative perplexity, we have placed these results Appendix Fig. due the metric's inherent unreliability, as noted in Sec. 4. We emphasize the more reliable FID-basck assessments found in our image experiments.

---

Table 2: Bits Per Character (BPC) on Text8 test set. All models use standard 12-layer transformers similar to GPT-2 small [57] except Discrete Flow which uses 8 layers.
Table Bits Per Dimension (BPD) on CIFAR-10
test set and Downsampled ImageNet 64x64 [63] validation set. All models in the table are trained without data augmentation.
Method
#arams BPD (L)
Method
BPC (1)
Autoregressive
PixeIRNN [63]
3.00
Contimuous Diffusion
Gated PixeICNN [64]
3.03
Plaid [22] (Our impl.)
S 1.48 1.41
PixcICNN++ [65]
53M
2.92
BFN [26]
PixeISNAIL [66]
46M
2.85
Any-order Autoregressive
CIRAR
Image Transformer [67]
2.90
ARDM [48]
< 1.43
Sparse Transformer [68]
59M
2.80
MAC [61]
< 1.40
Discrete Diffission
D3PM Absorb [14]
37M
Autoregressive
D3PM Gauss [14]
36M
$ 3.44
IAF/SCF [62]
1.88
Camphell [29]
36M
3.5.9
AR Argmax Flow [15]
1.39
Campbell [29]]2990]] Absorb
28M
3.52
Discrete Flow [59]
1.23
MD4 (Ours)
28M
2.75
Transformer AR [14]
1.23
Autorggressive
Discrete Diffusion
ImageNet 64x64
PixeIRNN [63]
3.63
Mult. Diffusion [15]
1.72
Gated PixelCNN [64]
3.57
D3PM Uniform [14]
< 1.61
Sparse Transformer [68]
152M
3,44
D3PM Absorb [14]
1.45
Routing Transformer [69]
3.43
SEDD Absorb [32]
3 1.39
Perceiver AR [68]
770M
3.40
MD44 (Ours)
< 1.37
Discrete Diffusion
GenMD4 (Ours)
S 1.34
MD4 (Ours)
198M
3.40
Text8. Following prior work [14, 32], we trained masked diffusion models text8 and evaluate the bits-per-character on the test set (details in App. J.1). As seen in Tab. 2, our models outperform previous discrete and continuous diffusion models, well state-of-the-art AO-ARMs which are closely related to discrete diffusion [48]. Our model is only beaten by an autoregressive (AR) transformer and the AR- RR-backbone Discrete Flow [59]. We believe this is because AR models only require learning fixed generation order thus better utilize model capacity. Text8's small vocabulary (26 letters and space) led us expect limited flexibility from
depende formulation. However, using the generalized objective in (12), GenMD4 achieved significantly better BPC than MD4, demonstrating the potential state-dependent diffusion for discrete data.
We train MD4 Fine Edu and evaluate its zero-shot accuracy on the Hellaswag dataset [60], popular common sense inference benchmark for LLMs. We di- rectly compared MD4 its AR counterparts transformers with identical configurations (ex- cept for causal masking) trained on the same data. Results are summarized in Fig. 5.
45
MD4-S
accuracy (%)
40
MD4-L
AR-S
35
AR-M
hellacwag
AR-L
MD4 demonstrates steady performance growth with increasing scale. While outperformed by AR models of the same size, the performance gap does not widen as model size increases. For example, AR-small reaches 30% accuracy in 50k steps, while MD4-small takes 200k steps (4x data efficiency difference). At the medium scale. AR achieves 37% in 270k steps, compared MD4's million steps.
30
25
10a
102
10a
1000
Figure 5: Hellaswag accuracy vs. training steps for MD4 and AR models at GPT-2 small, medium, and large scales.
7.2 Pixel-level image modeling
Unlike continuous diffusion which struggles with discrete data, we show that MD4, a discrete diffusion model, performs well on inherently continuous data, suggesting its potential for unifying

---

Figure 6: Non cherry-picked unconditional samples from MD4 trained on ImageNet 64x64, treating pixels as discrete tokens. More samples can be found in Fig. 9 in Appendix. The model is optimized for likelihood instead of visual quality-sec c.g., Kingma et al. [33] for samples from continuous diffusion model optimized similarly for likelihood.
modalities. We follow Austin et al. [14] and train MD4 on order-agnostic image data from CIFAR-I0 and downsampled ImageNet 64 x4x66 64 [63]. Each image is treated as set 256-valued discrete tokens, making the model agnostic to pixel proximity. We compare to other discrete diffusion and AR models with reported likelihood results on these datasets, although to our knowledge there are no published result on discrete diffusion for ImageNet 64 x
64 that directly model raw pixel space.
Tab. 3 summarizes our results. We establish new state-of-the-art for discrete diffusion models, outperforming previous work [14, 29] by significant margin. Our CIFAR- 10 result surpasses the best reported AR result. On ImageNet 64 64, our results are competitive with Transformer AR models that are 4x larger, as well as strong continuous diffusion model VDM [33]. Notably, despite lacking knowledge of the ordinal structure of pixel values, MD4 outperforms models trained with this inductive bias, including D3PM Gauss and Campbell et al. [29] where the noising distribution is discrete Gaussian that assigns larger probabilities near pixel values. To isolate the differences caused by training objectives, we also implemented the Campbell et al. [29] objective with the absorbing process, showing its high variance hinders learning even with our architecture.
We provide a random sample from our ImageNet 64 64 model in Fig. 6. More results can be found in App. K. In Fig. 2, plot the FID values of samples generated under different choices of schedules and discretization grids. We can see that the model with the linear schedule plus a cosine grid achieves an FID close the model with cosine schedule, both significantly outperform the linear schedule with uniform grid. We further trained a class-conditional model on ImageNet 64 x 64 that boosts the FID to around 7. Although these are not state-of- the-art art FIDs on ImageNet 64 x 64, we emphasize our models are optimized for likelihood instead of sample quality.
8 Conclusion
In this work, we revisit masked diffusion models, focusing on a flexible continuous-time formulation. Existing works in this area are not easily accessible to non-specialists and present ELBOs that are difficult to optimize, often resulting in performance that not competitive with continuous diffusions and AR models. The framework we propose provides a very simple expression of the ELBO as weighted integral of cross-entropy losses. Additionally, we propose a generalized masked diffusion formulation (GenMD4), where the masking schedule depends on the current state of the process, and derive its corresponding ELBO. On text data, our MD4 models outperform existing discrete and continuous diffusion models. For pixel-level image modeling, we significantly improve discrete diffusion results, outperforming similar-sic sized AR models and achieving comparable likelihoods to continuous diffusion models such as VDM. GenMD4 provides further improvements in terms of likelihoods over the state-independent case.
Although we have improved masked diffusion models, they still suffer from limitations. First, in some tasks such as text8, masked diffusions are not yet competitive with AR models. We conjecture that this is because AR models can better leverage model capacity since they only require learning one order. would be interesting to develop better architectures for discrete diffusions. Moreover, GenMD4 is promising, but it can easily overfit to the dataset, making it less effective for zero-shot transfer compared to simpler versions. Additionally, inference with a state-dependent schedule more challenging.
10

---

References
[1] Jascha using nonequilibrium thermodynamics. International Conference on Machine Learning, 2015.
Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning
[2] Jonathan Ho, Ajay Jain, and Picter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
[3] Yang Song, Jascha Sohl- Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poolc. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2020.
[4] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffuusinn models. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition. pages
2022.
[5] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint (xXiv:2204.06155, 1(2)3, 2022.
[6] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily I. Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan. Tim Salimans, et al. Photorealistic text-to- image diffusion models with deep language understanding. In Advances in Neural Information Processing Systems, 2022.
[7] Nanxin Chen, Yu Zhang, Heiga Zen, Ron J Weiss, Mohammad Norouzi, and William Chan. Wavegrad: Estimating gradients for waveform generation. In International Conference on Learning Representations, 2021.
[8] Zhifeng Kong. Wei Ping. Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021.
[9] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. In Advances in Neural Information Processing Systems, 2022.
[10] Ruben Villegas, Mohammad Babacizadeh, Picter-Jan Kindermans, Hernan Moraldo, Han Zhang, Moham- mad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan. Phenaki: Variable length vidco generation from open domain textual descriptions. In International Conference on Learning Representa- tions, 2023.
[11] Omer Bar Tal, Hila Chefer, Omer Tov, Charles Hermaan, Roni Paiss, Shiran Zada, Ariel Ephrat, Junhwa Hur, Yuanzhen Li, Tomer Michaeli, et al. Lumiere: A space-time diffusion model for video generation. arXiv preprint arXiv:2401. /2945, 2024.
[12] OpenAl. Sora. https: http://pppaaaaiccoeo conm/ndex/soor///r 2024.
[13] Fan Bao, Chendong Xiang, Gang Yue, Guande He, Hongzhou Zhu, Kaiwen Zheng, Min Zhao, Shilong Liu, Yaole Wang. and Jun Zhu. Vidu: a highly consistent, dynamic and skilled text-to-video generator with diffusion models. arXiv preprint arXiv:2405.0423 2024.
[14] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models discrete state-spaces. In Advances Neural Information Processing Systems, 2021.
[15] Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling. Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances in Neural Information Processing Systems, 2021.
[16] Clément Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang. Volkan Cevher, and Pascal Frossard. DiGress: Discrete denoising diffusion for graph generation. In International Conference on Learning Representations, 2023.
[17] Dongchao Yang, Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for texi-to-sound generation. JEEE/ACM Transactions on Audio, Speech, and Language Processing, 2023.
[18] Nate Gruver, Samuel Stanton, Nathan Frey, Tim GJ Rudner, Isidro Hotzel, Julien Lafrance- Vanasse, Arvind Raipal, Kyunghyun Cho, and Andrew G Wilson. Protein design with guided discrete diffusion. In Advances in Neural Information Processing Systems, 2023.

---

[19] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, al. Continuous diffusion for categorical data, arXiv preprini urXiv:2211. 2022.
[20] Ting Chen, Ruixiang ZHANG, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. International Conference on Learning Representations, 2022.
[21] Xiang Li, John Thickstun, Ishaan Gulrajani, Perey S Liang, and Tatsunori Hashimoto. Diffusion-LM improves controllable text generation. In Advances in Neural Information Processing Systems, 2022.
[22] Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood-based diffusion language models. In Advances in Neural Information Processing Systems, 2023.
[23] Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Q Weinberger. Latent diffusion for language generation. In Advances in Neural Information Processing Systems, 2024.
[24] Pierre H Richemond, Sander Dieleman, and Arnaud Doucet. Categorical SDEs with simplex diffusion. arXiv preprint arXiv:2210. 14784, 2022.
[25] Pavel Avdeyev, Chenlai Shi, Yuhao Tan, Kseniia Dudnyk, and Jian Zhou. Dirichlet diffusion model for biological sequence generation. In International Conference on Machine Learning, 2023.
[26] Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. arXiv preprint arXiv:2308.0700 2023.
[27] Kaiwen Xue, Yuhao Zhou, Shen Nic, Xu Min, Xiaolu Zhang. Jun Zhou, and Chongxuan Li. Unifying Bayesian flow networks and diffusion models through stochastic differential equations. arXiv preprint 1577, 2024.
[28] Guan-Horng Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror diffusion models for constrained and watermarked generation. In Advances in Neural Information Processing Systems, 2024.
[29] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. In Advances in Neural Information Processing Systems, 2022.
[30] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. In International Conference on Learning Representations, 2022.
[31] Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. arXiv preprint 2302.2300. 2023.
[32] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffiusion language modeling by estimating the ratios of the data distribution. In International Conference Machine Learning, 2024.
[33] Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. In Advances in Neural Information Processing Systems, 2021.
[34] Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems, 2022.
[35] Joe Benton, Yuyang Shi, Valentin De Bortoli, George Deligiannidis, and Amaad Doucet. From denoising diffusions to denoising Markov models. arXiv preprint 2rXiv:2211.0359) 2022.
[36] Jingyang Ou Shen Nie, Kaiwen Xue. Fengqi Zhu, Jiacheng Sun, Zhenguo Li, and Chongxuan Li. Your absorbing discrete diffusion secretly models the conditional distributions of clean data. arXiv preprint arXiv:2406.0377. 03736. 2024,
[37] Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint 24X.::2406.07522 2024.
[38] Lingxiao Zhao, Xueying Ding, Lijun Yu, and Leman Akoglu. Improving and unifying discrete and continuous-time discrete denoising diffusion. arXiv preprint arXiv:2402.03701, 2024.
[39] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
12

---

[40] Martin Heusel, Hubert Ramsauer. Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by two time-scale update rule converge to local Nash equilibrium. Advances in Neural Information Procexsing Systems, 30, 2017.
[41] Ari Holtzman, Jan Buys, Li Du Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2019.
[42] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:220../2598, 2022.
[43] Hunter Nisonoff, Junhao Xiong, Stephan Allenspach, and Jennifer Listgarten, Unlocking guidance for discrete state-space diffusion and flow models. arXiv preprint 1cXiv:2406.01577 2024.
[44] Yixiu Zhao, Jiaxin Shi, Lester Mackey, and Scott Linderman. Informed correctors for discrete diffusion models, arXiv preprint arXiv:2407.21243. 2024.
[45] James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: compos- able transformations Python+NumPy programs, 2018. URI. http: /ttp::////tttt thubbb....tt.....g thub. ctchub.com//.. ax- m1 jax.
[46] Ouail Kitouni, Niklas Nolte, James Hensman, and Bhaskar Mitra. Disk: A diffusion model for structured knowledge. arXiv preprint arXiv:23/2.05253, 2023.
[47] Benigno Uria, lain Murray, and Hugo Larochelle. A deep and tractable density estimator. In International Conference on Machine Learning, pages 467-475, PMLR, 2014.
[48] Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poolc, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Representations, 2021.
[49] Andrew Camphell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. In International Conference on Machine Learning, 2024.
[50] Javier E Santos, Zachary R Fox, Nicholas Lubbers, and Yen Ting Lin. Blackout diffusion: generative diffusion models in discrete-state spaces. In International Conference on Machine Learning, pages 9034 9059. PMLR, 2023.
[51] Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step- unrolled denoising autoencoders for text generation. In International Conference on Learning Representa- tions, 2022.
[52] Jiaxin Shi, Yuhao Zhou, Jessica Hwang, Michalis Titsias, and Lester Mackey. Gradient estimation with discrete Stein operators. In Advances in Neural Information Processing Systems, 2022.
[53] Tim Salimans and David A Knowles. On using control variates with stochastic approximation for variational bayes and its connection to stochastic lincar regression. arXiv preprint arXiv: 2014,
[54] W. Kool, H. V. Hoof, and M. Welling. Buy 4 REINFORCE samples, get a baseline for free! In Deepkl StructPred ICTR, 2019.
[55] Matt Mahoney. Text8.
net//c/textdata.. htt.://eeettaaa.aa........ Accessed: 2024-05-14
[56] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. OpenWebtTextCorpus. 2019.
http://Skyllono07, github. 10/
[57] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodci, and Ilya Sutskever. Language models are unsupervised multitask learners. OpenAl blog, 18))9, 2019.
[58) Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweh datasets: Decanting the weh for the finest text data at scale. arXiv preprint arXiv:2406. 17557, 2024.
[59] Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete flows: Invertible generative models of discrete data. In Advances in Neural Information Processing Systems, 2019.
[60] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi, Hellaswag: Can a machine really finish your sentence? arXiv preprint 1rXiv:1905.07836 2019.
13

---

[61] Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference any-order autoregressive models the right way. In Advances in Neural Information Processing Systems, 2022.
[62] Zachary Ziegler and Alexander Rush. Latent normalizing flows for discrete sequences. In International Conference on Machine Learning, 2019.
[63] Aäron Van Den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International Conference on Machine Learning, 2016.
[64] Aaron Van den Oord, Nal Kalchhrenner, Lasse Espeholt, Oriol Vinyals, and Alex Graves. Conditional image generation with pixelenn decoders. In Advances in Neural Information Processing systems, 2016.
[65] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelcnn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. In International Conference on Learning Representations, 2016.
[66] Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. Pixelsnail: An improved autoregressive generative model. International Conference on Machine Learning, 2018.
[67] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, 2018
[68] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
[69] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 53-68, 2021.
[70] Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in Neural Information Processing Systems, 30, 2017.
[71] Kehang Han, Kathleen Kenealy, Aditya Barua, Noah Fiedel, and Noah Constant. Transfer learning for text diffusion models. arXiv preprint 2rXiv:2401.17181, 2024.
[72] Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Newral Information Processing Systems, 2015.
[73] Peter W. Glynn, Likelihood ratio gradient estimation for stochastic systems. Communications of the ACM, 33(10):75-84, 1990.
[74] Ronald J Williams. Simple statistical Pradient-following algorithms for connectionist reinforcement lcarning. Machine Learning, 8(3-4):22-2555 1992.
[75] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEFF/CVF International Conference on Computer Vision, pages 4195-4205, 2023.
14

---

Table Masking schedule formulas.
Masking schedules
Qx
Cross-entropy loss weight of 1-a:
Linear
Polynomial
t
Geometric
exp (-Buim (Fmaa)
exp
8 mm
P
31-6 Pimxx
log Sam mas
exp
Cosine
cos( ( t])
tan 2 (1 11--)) t))
A Discrete-time derivation
We divide time from to into T intervals, and let s()) = (i 1)/T, t(1) = i/T. The forward transition matrix Q, E ((m++))x(m+))) (m is vocabulary size) at time t(i) is
j=k=
1-3
j=k m
[Qilı*
k= m.j # m
otherwise
or more compactly written as
== +
where 1 denotes an all-one vector of size m + 1, and is an one-hot vector of size + with the m-th element (recall that counting starts from 0) being one. We use an one-hot vector dt of length + 1 to denote the discrete state. The forward conditionals are defined as
( (1))) Cat(z Q7 x=(1)) 23(1) Qi2(()))
(13)
where Q. xa(() is the probabilities for each of the categories that F(()) can take. The marginal forward distribution at time t(ii given zy is
( = Cat(amn) Q mо) = 2 Q(
where Q = IT Q2 = IT 1
B₁) - - II Bÿ)) le To see what this leads to in
continuous time, we let B3 (({(s)) and T - 00:
II - B₂) exp Σ
log(1 - -B1)
j=1
j=l
B(((i)) +o(1/T))
cxp
T
j=1
((A)
T To exp
(-/ 3(s)ds).
We let Q(t) denote the limit of Q, in this case:
((t) exp ( (- f 3(s)ds) + (11-+((1))))) exp (-f B(s)ds)) Ie
I + (1 - aa)
Here define a & exp( So B(s)ds). And the marginal forward transition is g(z,(20) = Cat Q(1) 20) Txo) No Q())=; a x, (1 or TT.
(14)
15

---

B Continuous-time derivation
We consider a continuous-time Markov chain with transition rates
(t (Q: ()/L/T) - (1e) -)).
(15)
For simplicity, we let Q = le, T - T. The marginal forward distribution at time given To is ((za/zo) = Cat( Q(t) To), where
((t) = exp (% Q(s)ds) = exp (o / B(s)ds) = exp(3(t)Q).
Here we define B(1) _ Jo B(s)ds. The matrix exponential can be computed via eigendecomposition:
B(t)Q UAU
where
1
U 1 emen
1e 1mm
vn +
0-1 1 vn Iemem -lem
((t)(em I));
and thus exp(A) = 0z2 7 + (1 02y mm
((t)= U exp((A) a. + (1 a,)
A simpler derivation uses the following property:
==
Therefore,
@(t) cxp(3(t)Q)
I 3() + 1
+000++00+0000
5 20 + + +...
= =1+ =1+Q- - (1 ((t) + 2 3(1² 3(1 +...)Q
= = Q - -
= o, T + (1 a.
This marginal forward transition matrix at time coincides with the result (1) we get by taking the limit of discrete-time derivation.
Arbitrary discretization of the continuous-tin forward process. For the discrete-time process we have defined the per-step transition in (13). For the continuous titin process. we can derive the transition matrix (s,, t) & q( = =j(r,, = i) between two arbitrary time and t as the solution to the following differential equation (known as Kolmogorov forward equation)
d
dt ((s t) = ((s, tt) where Q(t) = 3(t)
with initial condition ((,, 8) = T. The solution given by
((s, t) = exp ((3(1) - ((=) = (())
Routine work (using the Woodbury matrix inversion lemma) shows that
Q(t) a; (1-
Plugging the result back, we get the forward transition distribution from
((x)(=r) = Cat ()x TT,))=))))))))))) x Q(1 1) where ((s t) A
(16)
0t
1 + (1 04 le
as
a
16

---

MASK
D0
1
1
1
1
1
0
a, a4
1-ar
1-0t
1-ag
a,- T-a
lP (Itt, t)
1-a,
l-ac
Q(xx= = =. 20)
q -|xx (0 t))
Figure 7:
The reverse transition probability and generative model. Left: g(2, -12; ;0)
in matrix form where first index is 2; and second index is 2g. Right: po(x, .2 = 2 q((,. I = Mo t)) also in matrix form.
C Time reversal of the forward process given Zo
The analytic property of our forward process allows to compute many quantities of interest in closed form. One such quantity frequently used in diffusion models is the time reversal of the forward process given rot ( We can compute it using (14) and (16) as
g( 20) = g( (x, )4(22(00)
( (rO)
a.r0 x 30
I. + m, Jt.
1-a1
1-ac
1-ax
Fs = m, 3t = m
(17)
2g 24
3t +
Visually, eqn (17) is a R(m+I)x(m++) matrix (Fig. 7, left) whose first index is and the second is Ts. The matrix is almost an identity matrix except the last row corresponding v, is the mask token. The last row means with probability of du the mask token gets unmasked to become Tp. and with probability of 1-a= remains masked.
T-ay
Alternatively, we can rewrite the above using reverse transition matrix R (t, E R(m++)x((++) as
q(Za 20) (tt s) T2x), where R = =I
04
Cm (2o em)
a-a
We also interested in what would happen in the infinitesimal time limit, i.e., when L Al and At + 0. Note that
Ot-A: at = - - +((I))
Plugging it into the original formula, get
0/
(L, - Al) 1 - 02
(To m Al + o(A.).
Comparing the above with the transition rate matrix R (t) definition
Re (t, At)
R At) + o(At),
have determined the transition rate matrix for the reverse process conditioned on Zn!:
af
Re (t) 1 ai
em (No em )
(18)
17

---

D Details of the ELBO
Using (17) and (3), we compute the KI. divergences between forward and reverse transitions
(L(g(z, (11 20,
(x))
(L (1) μi (2r,
(19)
Σ - ( 2o) log ((=_l=
21 m
0
3t + m
To Teex ex
0ss Ote
8x =m
xo c log
at
po t)
Tx
ck
kmm
Ott To log (2
6z
at
Note that 0 log 0 0. Alternatively, this result be easily obtained from the visual depictions of ( (24, and po (x, (x4) shown in Fig. 7. In this case, the reconstruction term becomes
(A(1)
E By(eeay (2o) plgg (000
Σ (ocy)((k()))))))))) log Σ jmm 40(1) (o (k-l)
(1) log (4111)
+ - an log m
(4(1)
(1 a)( log m.
The prior KL term can be computed as
[L(((())==)lll<xz) KLL6xx..m (m..m)) 0.
As usual, we take the continuous-time limit by letting T ->
Lo a lim Cr
T+000
a Oc(()) s(i) ((())
lim
30 Engie 2o) [osecoyi log (1( ((i))
T+000
s(1) t(i) 1 1-0((() (i(())
of
20 Eg(sslz0) 18 log po (x t)] dt.
((1) - at
E
Avoiding undefined KL divergence
When defining the forward process, often do not want on to be exactly 0, equivalently, X to be oo for numerical stability reasons. Instead, we set A₁ to be a finite value, and thereby 0i has small positive value. This has problem that the support of g( (zr) is no longer {m) and instead becomes {m, }.). As a result, the KL divergence between g(z1 (()) and p(21) is undefined because q(x1 (xo) is not absolutely continuous with respect to p(x1) 82,,m- To resolve the issue, modify the prior distribution p(x1) such that it has support over all m + 1 values. One such choice is letting
p(x1)
01
Σ jyhm
8ms + (1 a1 a1)mmmm
Then, the prior KL divergence term becomes
((21|_o) p(x1)
(L(g(z) (20) (20)))))))))
Σ 21
q(अ (20) log
a1d2,,00 + (1 a1)82,=m p(x1)
(an (11 + (1 a1)8m,mm log
on log
a1
(1 01) log
1 1- a1
01 a1/m
01
aj log m.
18

---

Details of Training and Sampling with MD4
E1 Training
Algorithm 1 A single step of training with MD4.
Input: data minibatch {a]] network ua(; t), masking schedule at for i 1,.... B do (in parallel):
t < mod(u u + t/B, 1), a ~ UO0, 1]
for n E [AV], mask out each token zo =;(()) independently with probability 1 at, to obtain aj.
T-ar, (ao for n E [NV if 2 m, compute weighted cross entropy loss و (ag ) log Ne (x) (x), 44)
Sum over all weighted cross entropy losses for mask positions and optimize via autodiff
F.2 Sampling
Algorithm 2 Unconditional and conditional generation (e.g., infilling) with MD4.
Input: Context sequence I" of length N, with masks indicating the target areas for generation Init: (1(1)) -- discretize((0, 1)). F((T) x
for i = T,T 1....,1 do t 1-)( 10,8888( 1(i-1) (n) draw x. Call T-or 1o (II, +
for n E [N], if dt
(m)
1-ax em) else m! +- 2t
(m)
I-a,
return Zn-
G JAX Categorical Sampling and Implicit Top-p
We noticed that the following equivalent implementation of Alg. 2 leads to significantly worse sample quality in JAX:
Algorithm 3 Variant of Alg. 2 that yields lower sample quality when implemented in JAX.
Input: Token sequence xf of length N, with masks indicating the target areas for generation discretize([O, 1J). Fe(r) x
Init: {1(i))
1, do for i = T,T 1, 1
t t t(1), < f(i-1) 1)
for n E [N] do (in parallel) draw U0,
if (m) = and < c.-ai then
I-ai
draw (n) Cal(ug (x,,))
(x)
clse
(n)
(x)
2 Iz
return Zo-
However, mathetically is equivalent to Alg. and should produce identical results. Our investigation revealed that the issue arises because Alg. 2 scales the output probabilities of μ4 by small factor 0.-a; as close f, causing some categories to have very low probabilities. JAX, however, I-au
implements categorical sampling using Gumbel argmax, which is less numerically stable than methods like binary search. As a result, categories with low probabilities are rarely sampled, even when their cumulative probability significant. In our experiment, we found that categories with probabilities below le-8 are rarely sampled out of a total of 50K categories. Thus, Alg. 2 implicitly performs top-p sampling (with dynamic p) under JAX's categorical sampling, yielding better sample quality than Alg. where μe is not scaled by small factor and has fewer categories truncated.
19

---

H Unifying Existing Masked Diffusion Models
H.1 The CTMC point of view
first prove a lemma that connects the forward and reverse transition rate matrices. This follows from the results in [29] but give proof for completeness.
Lemma 2. The forward transition rate matrix Qtt) and the reverse transition rate matrix (given mn) R (t) satisfy:
neth (1)s = x() (rjo(i]zo) (oo) for j + k.
(20)
2 (k) mo)
Proof
Consider the reverse transition from time + T to t. For j # k, Bayes' rule yields
g(3; = =| = k/ = j)
g( 5|æ147 k. k0)
q(Ee+t k/zo)
q(3: = 165r + (())kKT + +(())
(1 k((oo)
q(:: = =](z0)
8aj
Q(1)sKT + o(T).
g = k|zoo
Then, it follows from the definition of the transition rate matrix that (L) Q()) (io qajo (F|zz))
Proposition 3. We use the shorthand Ra(()k to denote the approximate reverse transition rate from the state is to j obtained by substituting our prediction model pa for ve in R (1)kj- Then, the continuous time objective (4) can be equivalently expressed as
C.0 =-f' Ba.o((k/oo) (2(1)
(t)() + Q0()ss log Ro))s + C,
(21)
jfk
where C' a constant independent of
Proof To rewrite objective Lo with the transition rate matrices, first go back (19).
There, instead plugging in the explicit form of R 8), we substitute with (8) which leverages the transition (t). (t). To simplify the notation, we assume I, k and use the shorthand Ro (t)xs & Rua(k) (R (t)kg. We then have
KL(g(Er-sslar) 10)p: KL(Cat(Zz; R (t, At) e*)) Cat(#,; Rua (t, f At) ex))
m
c + R (t) + of (t)))
ol 12 100 o(At))e; log
cg (I + Ro At + At)))
1=0
+ R (1) o(A))
(1 1 (1)xa At) log
1 Ro +
(= Al - d(A))
(Re0 (t) log
(o At ((A)) + d(tt)
J*
=R (t)ck Ro(t))x) Σ (t) log (1))g)A++
At +o(At) +((At).
20

---

For the last identity, we have used the fact that log + +) = I + o(x). To obtain Loo, we take the limit of LT as + oc, which is equivalent letting Åt 1/T We obtain
x lim
E (F (o) http:////LLaaaag///gg///////////
i=2
lim
E B (mo) (o)
( R (R Ro((()))
T mc
1:00 (1(1)) At )(1)) Ro (t()) At o(At) At d(tt
Σ R (t(3)) log j#
R (t)as dt. Ro(t)ks
(
(f)( (E( +
Σ (t) log
((1)
jtk
Note that Re0 (1) a constant matrix independent of 0. Absorbing all constant terms into C, we have
J. (kg0(k [oo
Loo
E
(1)a()) log Ro (U)ssss) dt C.
jykk
Next, we subtitute (t) with the forward transition rate using Lemma 2:
[(
5Q(1)) ( log Ro (U)ss) di di+ +
Lrx =
B (e())
(o(() mR= +
2:j0 (k(
jzk
fr
m
m
Ju())
(oo(((/rr)) (1))
((t)s) log Ro(() dt +C
k=
k=o jtt
(ko(k(a Ro (1)RK +
Σ> ((0)))4p0(())))) log R (1) dl de++0 CC 1=
((1)
k=0
where the last identity used the discrete analog to integration-by- by-paa (or summation-by-per Σx=o Ejtt k(i = Ek=o EzSx [k, Rearranging the terms then gives (21).
H.2 Differences from Campbell et al. [29]
Campbell et al. [29] used the first term of (21) as the training loss. A key limitation of this loss function is from the inner summation term
(()ss log Ro (() j
For single dimension case, the sum is analytically computable due the sparse structure of Rg (t)-if Dt = & is mask, the second term disappears; otherwise the only possible neighbor is mask. However, for multidimensional data, will represent all N neighbors in the forward process, i.e, the states get from mask out single unmasked dimension of z, k. Recall that Ro(t) is computed as substituting our neural network prediction model و4 (z for an R (t) Therefore, the summation together with Ro (4)uε requires N evaluations of sta(-). This is prohibitive since the neural network model is usually expensive. To resolve this issue, Campbell et al. [29] proposed to rewrite the sum
E)~(((1)) (1) [Z log Ro (t)xa where ((ik) = Q(1)a Z₁ 2g A Σ Q(())s
and estimate it through Monte Carlo. Taking into account the outer expectation under (nj (o) the computation of the loss then becomes a doubly stochastic estimate (using & ~ q:jo mn) and j ~ q(jlk) which suffers from large variance. In contrast, the form of our loss (4) only requires evaluating иe once for single stochastic estimation of the expectation w.r.t. q(xz(20).
21

---

H.3 Score parameterization
We provide simpler derivation of the score- booe-based loss [32, 35] below. We start from the form of the ELBO in (21) and rewrite it as
Loc
Ject)
Egg (k) sea)
Σ
(1tH ((RM00
R (t)as + 1 (t)xs log R (t)as ]dt.
R (t)ks
)
(22)
jzk
For the last identity we used the zero-row-sum property of transition rate matrix:
(² = ----rr *
If we plug (20) into (22) and reparameterize with score model
(23)
g( Iu (Iz))
we recover the score entropy loss function from Lou et al. [32], Benton et al. [35]:
Coo -Ta (k
Σ (())))((0((()) (so());
(]10(] (k(
log so (k), + sh (
)) dt, (24)
qe (
j7k
where (v) A ylogy- -g. Note that our derivation above is different and simpler than that of Campbell et al. [29] (which Lou et al. [32] is based on) since we leverage the conditional reverse transition rate given xp instead of the transition rate matrix of the reverse process. We can further simplify the loss with the following relationship between the conditional score and zo:
2n Q())ej
0tt xo ej for h: m.. +
(25)
(k(
Zo Q(1)ex 1 1-a --
Note that only the result under the case k = m is needed. This is because when 2t is un- masked, at any time between 0 and the state must stay unchanged and remain 20. As result, (I(((1)-AA/ (t, Pe (xc-)c = 0 for Ia * m. From (15), we know Q(1)² =
((t) 85k). Combining (25) and (24), we get
[( B(1) (Egnp(((e)))) Jmme Jan
Σ sa(k); 1 az To log so(k)) ( ) )dle. (26)
Lxx
j
Further, we can show the connection between (26) and (4) by reverting the score parameterization to mean parameterization using (23), equivalently so(2;); I (x) ej. By doing so, obtain
frr ((t) (Fago(tttoo)) [bmnk Σ so(k); 1 at
0f To log 1o (k:)
at dt.
L.o
aq
Observing that
Σ se(m); =
O2
(27)
1
jpm
conclude that this recovers the objective in (4). Interestingly, in Lou al. [32] the score parameterization is not constrained to satisfy (27). That means the learned reverse model might be incompatible with the forward process.
Below, we prove Proposition using the result from Eq, (25).
Proof of Proposition 1
(r( (m)
Ez, ()1(((/o))(Z0))
Exo x0 4:jo (200)
40o ((() 2rjo (m (r))
(m)
4: ( ())
Eng
a1
xi0 ej
T
aa
E E[rol mj] ee.
1 04
1 a2
22

---

H.4 Other related work.
MaskGIT [39]. MaskGIT is diffusion-inspired iterative denoising model for discrete image tokens obtained through models such as VQ-VAE [70]. Training of MaskGIT follows the steps: (a) Sample E [0, I]. (b) Given mask scheduling function (t), sample y(t) tokens to place masks. (c) For data zo of size (m + 1) x N and the partially masked state 2t, minimize the negative log-likelihood
! Ew: E [Enzg") =m (T log 1b (xz, t) dt.
Ma =
(28)
Our forward process satisfies (1)((()))))) = 11 Therefore, when we set the mask scheduling function as (() = 1 Or we obtain a loss similar to (5) without the - weighting. Note that there remains a difference in the sampling distribution of oy: in the masked diffusion forward process, tokens are sampled independently and do not necessarily yield exactly (1 ox) mask tokens time t, though the expected number is (1 az ). One might be interested in whether the uniform weighting can be recovered by selecting an appropriate schedule ar However, solving a such that af = 02 1 yields a = ce + 1 and there is no c that satisfies both og = 1 and a₁ = 0. This shows that training with the MaskGIT loss (28) may not be faithfully optimizing the model likelihood.
Discrete flow matching [49]. For the linear schedule 04. matrix (8) conditioned on To is:
our reverse transition rate
- 1 R (t) 1- em (m - -mm) =- - -mm
0f
This is the same as the conditional reverse transition rate used in Campbell et al. [49, Eq. (22) that their time / is reversed, and the rate matrix was therefore in the form 1800 (1) T- em (To cm)
SDDM [30]. Sun al. [30] proposed a becudo-likelhhood-iie objective for training discrete diffusion models that can also be applied to masked diffusion. However, their objective encounters the same challenge as Campbell et al. [29] - [ requiring N passes of the mask prediction model. To mitigate this, they introduced a new transformer architecture, which unfortunately leads to some performance degradation.
Blackout diffusion [50]. Santos al. [50] proposed "blackout" diffusion process that gradually diffuses images to black state. While this approach is similar to masked diffusion on binary data, key differences emerge when dealing with larger state spaces. their method, image pixel intensities gradually fade out, whereas ours directly transition to mask state. Our method offers more flexibility, being applicable to general discrete state spaces without requiring predefined structural relationships. It also demonstrates competitive performance in image generation, achieving this without knowing pixel value proximity.
SUNDAE [51, 71]. Unlike masked diffusion, SUNDAE uniformly corrupts data with random tokens in the vocab (known uniform discrete diffusion [14]). Additionally, it uses second loss term from cross entropy between clean data and 1-step unrolled model prediction. Similar ideas have been proposed in [72].
I Details for state-dependent rates
I.1 Derivations and time continuous limit
All derivations in this section assume that is a single token, while for N tokens the masked diffusion with state-dependent rates factorises across the N tokens. Learning from data of N tokens using variational inference is discussed in App. I.2.
23

---

Given the forward transition ((x))2,)) and marginal (20) derived in main text (Sec. 6) The reversal given zo is (2,, vo) = Cat R (t,, 8) xx) for
1-a4
Zo
j = m, k # m j = m, k = m
(L j)
1=
1-ai
20
8jk
j x m.
or alternatively can be written as
(x1a/Te. = g( x ()x
(Tr 20)
22 x. + (1
2 a. ee [a] xo00 + (1 (1-a, a, xo)em
(29)
ar zoto St + (1 - aj 2o) to)e,
To simplify this expression we consider the two cases: either xt = m (i.e. 3, mask) or ii # m where in the second case Tig For the case It = m, the denominator in (29) simplifies as
g(z:
m| m| = 1 a,
due to 2I 24 0 since xo F m, i.c. the observed token 20 cannot be a mask. Then given that Iy m the probability that T, = m is
1 a 3o
(1 a.)) xQ
1 a2
T
20
(30)
1 a (1 042) T x0
1 O2
while the remaining probability for Vs
Fg F m is
(a,, aa))TTa (a, oz) Tua
0i
Zu-
(31)
04 xp
(1 at)
T x
1 a₁
Then, combining (30) and (31) to write g( |a m. Fu) in unified way yields the expression (11) in the main Sec. In the second case, when I, = To F m, q( I * m, mo) from (29) simplifies dramatically and it becomes q<ra/0 # m, xo) = ze ra which a point mass that sets x, = Zt.
Derivation of the continuous-tima limit of the loss in (12). To simplify the notation, we let S.tt A Co- We first compute the KL divergence terms in the discrete-time ELBO as
(L(q(e, Fr (x)loo
Ex (11//11, mo) log gla, (m) po(2.1z))
It m + m 20 ex Eurooo log diag (Ze, t))ek
0
(1 - Eo,e) [ xo
= -aamm
Σ kmm
(1 - Eo,t) xo log (1 Ex,t) mo (2te, 1).
(1 - Eo,1) xo
= -ma.m
TT too log uo (2r t) + (1 (e,e) 20 log
(1 Eo,t) pa Malaze,
Let A, T t(1) s(i) for all i. Plugging 04- =====00 = a4 -0--00 of + o(At) into the above formula
and letting 7tt T-a." we get
ni
KL 1r, xo) (1pe
1 1 2o At + o(A))
Vt T Nono log mo (Z,, t) + (1 + 7/ no At log
1 1 71 Hol At + o(At)
Sa..m
hyT xooe log (r t) + 7o no At) (nr mnat 7t uo(ee t) o(Att) + +((At)
by a log mo( () + .o. xi Al 7t po( + o(Ai)]
+
Sarm No (roz) log (o t) + xo (1 +)) + o(At).
24

---

Therefore,
lim
7-000
T Σ
lim
By(Recc (m)) Boess.. Me (zuz)) log po (z)()) t(i)) + zu (())))+ + +(tt)
T-000
1=
=s" %/ 1o zzmm (2oz) log 1o (x2, + zo uo d))]d.
J(())
Letting t(1) + proves the result.
1.2 Training and gradient estimation
The model is applied to data consisted of N tokens where so (wg ...,...,, x0 ( and where each state in the masked diffusion is (x)....,,z) (N)). The reverse generated model has a factorizing N
transition conditional of the form II Pe (2.s
(1) (aze) where po (xx.")/=rr)) q(a:!") (n) .u (s ())
has form that depends on whether x;
m or
(n) # m. For the first case:
(1) pa (a(a p (z m, {xt }aZN) (
1 0g
(n) (Ie, (xe, TT ( 0t ar,
diag(p)) (x) (xr, (m)
or
14g
where Pe (n) (Ix, 1) = softmax (fo(x:))) a m + dimensional probability vector modelled by a NN (where the final value is constrained to be zero since 1. (I,, t) is reconstruction of z!") which cannot be mask, so in practice the NN classifier needs to have a softmax output only over the m actual token classes). Crucially, note that the NN classifier receives as input the full state Ia of all tokens, while additional time features to encode are also included. When 2((") * m the reverse (n) transition model is set to be pa(z,, (x (n) f m, {=/ (k) Jkpn) (m(")) Tx which matches precisely
)
((22 (m) | m, =r00" (m) = (az (n) (n) from the forward process.
The full negative lower bound for state-dependen rates and assuming N tokens is given by
(n) ((N) = . (1==/² Bg(er)zzo) [mreg" (2) (m) - Me (xe, t) + 20 (2x) log Mg 11 (Ie, t)) dt. (32)
=
Given that each ar, = 1 10--thh th;, the reverse model becomes
po(x!") (a())) * m, { (k)
(n) 1k = (c" #)T Mg (Iy, x) (0)T eoo
diag(up (1)
(IL, t))=
where is the m dimensional vector of all w,s. Note that the probability of (n) staying in the mask state, i.e., (n) = m depends on the full o, and it is given by log) pa (xx, t)
(n)
EnD e log H4
(n) (n) (xe, while the probability for (()) to take a certain non-mask token value
of ws the above (1 e log eo;log #) pe (xz, t);- The gradient wrt f is af,5 -w, and I-os;
loss is written as
C - f(
//² --
(g(xr/(ro) Ennz(">mm (1
ue (m) (I,, t) + (z( log po (n) (æx, (xx E)) dt,
where w is the vector of all w. 's. An unbiased gradient over the NN parameters 0 is straightforward to obtain since we just need to sample one time point and an ia ~ q(<z(2o) to approximate the integral and expectation and then use the gradient:
V
Σ RED
u
(n)
11, (xr. 1) + (m) (2(rr) loguu (x).
The gradient wrt the w parameters is more complex since these parameters appear also in the discrete distribution g(z)(20) which is not reparametrizable. To deal with this need REINFORCE
25

---

unbiased gradients [73, 74], and in our implementation we consider REINFORCE leave-one-out (RLOO) [53, 54] with two samples. Firstly, the exact gradient wrt w the exact loss is written as
-fo²
R (2)) lø( mo)] dt S/" (o (f V log ((Fe)zoo)
-
(33)
where
g =
Σ 1E: =m
(To (n)
(n) 1e (It, 1) + T (ab"))
(n) (=(")) log log
(") (xx, t)), f(1,, vo) = w g00, zo).
Note that g( rD) is a vector while f(n, va) scalar. The left term (33) is easy since it just requires sampling i and ~ g( (ro), while the right term is the REINFORCE term which could have high variance. For this second term we use RLOO with two samples z =a and construct the unbiased estimate
1
( w log ((x) (m)) - V ~ log q( vo)) [f(x), (af;
21
Thus, the overall unbiased gradient for we use is
21 [g(a), vo) + g(0 2o) + (Va log g()/ - V log (eg(((1=0)) [fer] mo) - [o)]]
Experimental Details
In all experiments, the model is trained with continuous time loss while samples are drawn from the discrete-time reverse model of 1000 timesteps unless otherwise noted. We used an exponential moving average factor 0.9999 for all evaluation including sample generation.
J.1 text8
We followed the standard dataset split as in Austin et al. [14), Lou et al. [32] and trained our models on text chunks length 256 for million steps with batch size 512. All models in the table used a standard 12-layer transformer architecture unless otherwise noted. Our transformer has also the same number of heads (12) and hidden dimension (784) as in Austin et al. [14], Lou et al. [32].
We used the continuous-time ELBO and drew one sample of t for each data to estimate the integral. To reduce the variance of training. we used the same antithetic sampling trick described in Kingma et al. [33] for continuous diffusion models. We used the linear masking schedule a 1 f and added a small shift e = 10 when / is close to and to ensure numerical stability. The shifted schedule is at (1 21 - f) + The shift leads to support mismatch between ((z) (æo)) and the prior p(æ1), leading to an undefined KL divergence term. We explain in app. how to modify the prior distribution to allow small uniform probabilities in non- non-mask states to mitigate this problem. The shift leads to a non-zero reconstruction term and KL divergence term for the prior distribution but both are of negligible scale
safely ignore them when reporting the ELBO.
We used cosine learning rate schedule with linear warm up of 2000 steps. We applied channel-wise dropout of rate 0.05 and used Adam W optimizer with learning rate 0.0003 and a weight decay factor of 0.03. Our model is trained on TPU-v5 lite for less than a day.
J.2
Deb
We kept 2% of the original training set for validation. Our small and medium transformer model have the same number of layers, heads, and hidden dimensions as in Lou al. [32] and our tokenizer was also kept the same with vocabulary size of around 50k. The training objective, masking schedule and other architectural choices were kept the same with the text8 experiment. We kept the training hyperparameters the same as text8 experiment except that we reduced the dropout rate to 0.02.
J.3 Fine Web-Edu
We kept the same training setup as the OpenWeeTeex experiments. Our transformer models have the same number of layers, heads, and hidden dimensions as those of GPT-2 models. We use the same GPT-2 tokenizer.
26

---

For Hellaswag evaluation, we concatenate question with each answer option, tokenize the concate- nated string, pad the length of 1024. The padded token sequence gets fed to our MD4 model's loss function for likelihood evaluation. We average 32 Monte Carlo samples to reduce variance. The answer with the highest likelihood estimate is the model's prediction.
J.4 Images
We used the same linear masking schedule as previous experiments all likelihood results. We used the same U-Net plus self-attention architectures from the continuous diffusion model described in Kingma al. [33] for CIFAR-10, except that we did not use Fourier feature inputs and added an additional input embedding layer with embedding size the same as the hidden dimension of the model. For ImageNet 64 64, we reduced the number of residual blocks (in one side of the U-Net structure) from 64 to 48 and added 12-layer diffusion transformer [75] with 768 hidden dimension and 12 heads in the middle.
For both datasets we used Adam optimizer and trained for 2M iterations. We used learning rate 0.0004, batch size 256, weight decay factor 0.01 for CIFAR-10 and learning rate 0.0002, batch size 512, weight decay factor 0.03 for ImageNet 64 x 64. The learning rate follows a cosine annealing after 100 warm up steps. Our CIFAR-10 model trained 32 TPU-v5 lite for 24 hours. Our ImageNet-64 x 64 model is trained on 256 TPU-v5 lite for 3.5 days.
As explained in Sec. 4, we have observed that the cosine schedule leads better sample quality so we used it to train a cheaper model for sample visualization. This model differs from the one that achieves best likelihood in that used 8 residual blocks (in side of the UNet structure) and 20-layer diffusion transformer in the middle. All other configurations are kept the same.
K
Additional Results
K.1 Sample quality evaluation by GPT-2
We use the GPT-2 large model to evaluate the perplexity samples generated by our model, following Lou et al. [32]. Results are shown in Fig. 8.
240
Generative Perplexity by GPT-2 Large
120
60
SEDD Small
MD4 Smell
MD4 Medium
GPT-2 Medium
256
1000
Decoding Steps
Figure 8: Generative perplexity evaluated by GPT-2 Large following Lou et al. [32]. We compare MD4 against the GPT-2 checkpoint (autoregressive baseline) and SEDD (the previous best discrete diffusion model this task) in generating 1024-token text sequences. We investigate the effects of two orthogonal factors on sample quality: model size and decoding steps. The numbers for GPT-2 and SEDD are from Lou et al. [32].
27

---

K.2
Perplexity on WpenWeb Text validation set
Tab. reports the final perplexity number achieved on OpenWebText validation set, corresponding to Fig. 4.
Table 5: Perplexity on OpenWebText validation set.
Size
Method
Perplexity (1)
Small
Gaussian Diffusion
27.28
SEDD Absorb (reimpl.)
24.10
MD4 (Ours)
22.13
GenMD4 (Ours)
21.80
Medium
MD4 (Ours)
16.64
K.3 FID evaluation of MD4 trained on ImageNet x4
We provide the FII numbers corresponding to Fig. 2 in Tab. 6.
Table 6: FID 50k samples generated by MD4 trained on ImageNet 64 64, corresponding to Fig. Top three rows show results from an unconditional model, while the bottom row is from model conditioned on class labels. Uniform discretization grid is used in Alg. 2 unless otherwise noted.
Timesteps T
Method
128
256
Linear du
193.81
128.18
72.94
50.21
Linear du, cosine grid
42.07
25.16
18.31
18.22
Cosine ax
47.46
23.84
17.8
18.74
Cosine ar, class conditional
30.75
11.39
7.13
7.8
K.4
Additional unconditional generation from MD4 trained on ImageNet 64x
We provide more unconditional generation results from our pixel-level modeling experiments on ImageNet 64 64 in Fig. 9.
K.S
Additional unconditional generation from MD4 trained on TipenWebTeer Text
Below we include two unconditioned text samples generated by our MD4 Medium model trained on OpenWeTText.
K.5.1 MD4-M unconditional sample 1024 tokens
like, don't have to be alive? Sometimes there are things that are too real and you're really supposed to experience them. that's good feeling. That is the scary thing. Not actually, being able to experience things, being able do these things, when you're doing them, which, for most people having to wake in a dream is something that seems the most significant, and then you think about it the next day. It's like the hope of the future, and you wake up right now thinking about it. What happens is,, then you have stop and think about it and then all of sudden, somebody always says, "You' re dreaming."
And sometimes I wonder if this is good time to teach your gut instincts to your actors when you're doing a show like this. Because even on this particular show, it feels like everyone's been through this all the time before, if even few years ago. I mean, you're doing show together, at least not on continuous development, you you're vet. mean, you should really be along.
28

---

Figure 9: More unconditional samples from MD4 trained ImageNet 64x64.
If you're 're not sure, well
vS: I'm working on that one.
Did any you guys feel that an instinct could work? thought, "Well, because you didn't do 'Deadwood' you should stop doing this." But when I read the story for the first time, I thought, think this is going to work." " What I can't picture is way to hold this apart.
29

---

we. When we wrote the first episode, we wrote script that we felt like me and myself would want to see. I knew that I wanted to be able something that was real, that you could see and just really step out of yourself. And then saw it. Then, you get rehearsing it and doing it. And then actually started shooting. I think I knew I didn't think it was going to be good. But, know it was good. And now people are talked about because it's not good enough.
VS: That's me. It's what we have to
something -- and wanted to able to take refuge
Growing up, you say that you just completely hated the show, "Lost." Isn't that what you wish for the end of the day?
VS: I don't like the concept.
And there's lot that you don't know about that, I think for to have had these ideas, if you didn't understand even that it was coming out of this world that doesn't exist, we might never get together.
It's so weird. This happened to happen at the same time?
vS: Yes. happened happen basically the same time.
Nobody's even show or had movie/come out of the movie,
VS: If I'm going to pretend I'm definitely not you and stuff, I don't think I'm going to swallow that. didn't expect it to do quite that long.
live through that
There are always things now that happen with 'Dedwwood' where you don't 't know where it's going to end up next time, but I think there occasions now where have to keep the fight, even if 'Lost' was pretty consistent in the mindset and the form.
VS: I'm glad that we did fight the odds, because we should have understood that almost sense of not that had showed up on the same day, we know we work in the same pieces, but a lot of stuff we don't know about. Some of it, language, and there are a lot of things where we take from them and we do this what they did because we want
there was direct link. But there
need to deal with. We also just have accept the
K.5.2 MD4-M unconditional sample 1024 tokens
the groups let recreational vehicles the meantime of fighting off the permit. "The purpose of the permit is to make sure that work with the NPS and made roadways and rest areas. We're not just scaring guys kind of messing around." Community plans to build an urban bike facility marched forward the ongoing staff meeting of the King County Commission.
the three roads that will stay open in
Trail will be finished just south of the Greenviex 5.
Instead of continuing with a pedestrian and bike trail to the MBTA's campus, these two trails could bridle the areas from Market to 14 and carry communitio time closer.
"This project will provide a car-free path to King County," " said Andrew Weed. It's been put the brakes on in the past several months, but there are those residents still skeptical.
"I've addressed some expressed some of their concerns. I don't think it's terribly reasonable from
community concerns that 've been raised. They've
30

---

transportation standpoint."
The trail had been set up to meet for more than year when the council approved funding a different proposal.
Mayor Muriel Bowser said after meetings with Commissioner Bushell on Thursday that the plan will be board in December.
"There's enough of finish for this project to roll out time, and we're going to get it done," Bowser said.
For the public, the campaign appears over.
"There was one meeting that I feel like lost at last night's 's meeting," said Shelley Potts, local resident.
Local resident Joel Grimy, who lives on Uman Road, met residents there as well.
And in other groups that rode through Mayor assistant Stacey Land and even son held fliers saying to look for light sign, and also met with Bowser's son, Deion Bowser, about a future plan also have dog park on the transit corridor.
Advocates at Brickley's event, many one waited least minutes during the start of the public meeting, said they expect at least another month from the Board of Commissioners, even after public hearing on Nov. 13.
"We've been trying to be talkative board where we are meeting in advance, being respectful of folks," Bowser said.
He considered that the proposal for the section of trail between the Greenview 5 and 3 "has to move on schedule. We have other historic preservation projects that would take over that."
But Chad Routledge, local advocate of the project, spoke against the mayor's plan.
"The mayor has sent new meeting to the public using the same route that resulted from the loud criticism and onslaught of complaints from the community committee back during the public hearing," " Routledge said.
The BDC doesn't have a particular plan-turns around for the end of the planned path, and says "nothing practical can happen right now." " But, she said the agency still "looking to make investments in facilities along the route."
And still there is another part of the trail that might be just much wish for the dogs, as cars: the district wants to go west foot couple blocks south, to make the trail safer for dogs.
"I feel that the accessibility of the trail is pretty important. think the education of the trail, and the uses along different routes are very important pieces of balanced outcome," " said Bushell.
Trams coming off Route 1
K.6 Conditional generation from MD4 trained on OpenWebText
We share conditionally generated text samples by MD4 Medium in Fig. 10 and observe that slow unmasking near 1, enabled by the cosine schedule, tends to help produce more consist and meaningful samples than uniform unmasking counterpart.
31

---

MD4-M schedule
skydlving is fun sport, pretty risky, You're getting season if something goes wrong and can happen you know, we know about season, especially Skydlving. anybody wins this year
Saturday you pretty sey: "This is what am going to be doing right now. It's just the simplest thing-thhet why always shampoo and shower three times a day.
get
the
Innor schedule
skydlivig a fun sport. is pure endurance excitement many people in many times could have something goes wrong and we can't make it, we have help another friend as we have 200
just keep your healthy as possible.yu never want to be ollythat is why lawways shampoo twice day and shower day.
cosine grld
lost. if
MD4-M cosine schodule
skyydiving fun sport, it's extremelly risky. You can have so many injuries one time and then one next time. There are many ways you can neuroconcusioons, especially from Skydiving, continuing rise every year
Though antibacterial products poison, skin needs chemical solution that protects from bacteria and spots that form within -thatt why always shampoo twice day shower three times day..
Figure 10: Conditionally generated text samples from MD4-M. Top: MD4-M trained with the linear schedule, sampled with uniform grid; Middle: MD4-M trained with the linear schedule, sampled with the cosine grid; Bottom: MD4-M trained with the cosine schedule, sampled with uniform grid. Context text shown blue, model-generale text in black.
K.7 Effect of discretization zero-shot perplexity
We carried out ablation study on the effect of discretization on zero-shot perplexity. Results are included in Tab. 7. Note that this is inference ablation with the same trained model (MD4-S trained with the continuou-time objective).
Table 7: Effect of discretization zero-shot perplexity.
Size
Timesteps
LAMBADA
WikiText2
PTB
WikiText103
IBW
Small
T 100
105.2
S 36.1
70.3
T 1000
48.5
35.0
< 102.5
S 35.0
S 68.4
T = 10000
S 48.4
34.9
< 102.4
< 34.9
< 68.2
T = do (continuous)
48.4
34.9
102.3
3 35.9
3 68.1
32

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main theoretical and experimental contributions claimed in the abstract and demonstrated in the paper. They reflect the paper's 8 contributions and scope.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes].
Justification: The limitations of our work are detailed in the very last paragraph of the paper (see Section 7).
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes].
Justification: All the theoretical results are proven in the supplementary material.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
We included all experimental details in App.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility 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, or other means that are appropriate to the research performed.
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) the contribution is primarily a new algorithm, the paper should make 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 a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
33

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to 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, described in supplemental material?
Answer: [Yes]
The datasets we used are all public datasets. Our code released at https: //g //github. com/ google-deepmine haappmiidd/d/e
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
oniooicc)))))) 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 NeurlPS code and data submission guidelines (https: / //nips. cctppblic..///cccccccccccccc...e gewuddee/oodeiieeeeeeeeeeeeeeeee
ss wswsnsoloo.. onoolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
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: Included in App. J.
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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can be 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: [No]
Justification: We follow the practice in prior work [33] to not include error bars, partly because the models are expensive to train.
Guidelines:
34

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-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 figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Included in App. J.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well 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 failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]-
Justification: After careful review of the NeurIPS Code of Ethics, is clear that the research presented in this paper conforms with the Code of Ethics in every respect.
Guidelines:
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
This paper is mostly theoretical and methodological. We do not see immediate societal impact this work. However, we acknowledge that large scale implementation of our algorithm might suffer from the same societal biases as any other generative models.
Justification:
11. Safeguards
35

---

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)?
Question:
Answer: [NA].
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines 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 original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We cited the dataset sources.
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 dataset.
The authors should state which version the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [NA].
Justification: This paper does not release new assets.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]-
Justification: The paper does not involve crowdsourcing nor research with human subjects. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
36

---

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 an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
Justification: The paper does not involve crowdsourcing nor research with human subjects.
37


Reviews:
[{"review_id": "MJWNOiHah6", "reviewer": "Reviewer_bny3", "summary": "Summary: This paper introduces a framework for masked diffusions that consolidates previous research on the topic and organizes it into a cohesive structure. The authors also present a generalized model within this framework, which enables the use of state-dependent masking schedules and optimization of scheduler parameters.", "questions": "1. GenMD4 has not been tested on image datasets. Could you please share the results of GenMD4 when applied to image datasets?\n2. Since introducing GenMD4 results in additional variance, what if all tokens share the same w (referred to as \"simplified-GenMD4\")? This would result in less variance. Given that GenMD4's performance is close to MD4, can simplified-GenMD4 achieve the same BPC?", "limitations": "The method is only applied to masked diffusions.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The GenMD4 framework offers a valuable approach to optimize the forward process. In earlier studies, forward processes were typically manually designed and set within the model. However, GenMD4 adjusts the forward distribution to align with the estimated distribution, thereby improving the forward process. This innovation may serve as a source of inspiration for developing more effective forward processes.\n\n2. This paper summarizes previous formulations of masked diffusion models and establishes the connections between them.", "weaknesses": "1. In line 90. The handling of $p(x _0|x _{t(1)})$ could be enhanced. Assuming $p(x _0|x _{t(1)}) \\propto q(x _{t(1)} | x _0)$ is equivalent to assuming that $q(x _0)$ is uniformly distributed. In reality, it should be treated the same as other $p(x _s|x _t)$.\n2. In line 114. When discussing multidimensional data, it is not straightforward to assume that the backward process factorizes across tokens. This is because the distribution $p(x _0)$ does not factorize across tokens. Achieving factorization necessitates a small time step dt, which may not be easily observable. Additionally, in the previous single-token scenario, dt dose not need to be small, indicating that one step is sufficient to model the distribution $p(x _0 | x _1)$. This aspect is crucial for multidimensional data and should be emphasized in a fundamental paper like this.\n3. In append F. The presence of a non-zero $\\alpha _1$ may result in the \"medium brightness problem\" [1]. However, there is no singularity when $\\alpha _1$ is zero if log-SNR is not introduced, and the time interval can be extended to [0, 1].\n4. In append G2. When applied to masked diffusion, $R_{kj}$ is zero when $ j \\ne k$ and $j \\ne m$. Given that $R_{kk} + R_{km} = 0$, $\\tilde{q}$ can only take on one value (m), resulting in no additional variance.\n5. In image experiments, MD4 employs masked noise, while $\\tau$LDR uses Gaussian noise. We recommend conducting experiments with the same noise scheduler to demonstrate conclusively that MD4 is superior. If the goal of this paper is solely to establish that masked noise outperforms Gaussian noise, we recommend explicitly stating this claim. Additionally, we advise detailing the sampling method, as variations in methodology can influence the quality of generated samples.\n\n[1] Common Diffusion Noise Schedules and Sample Steps are Flawed, Lin et al., 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time you’ve taken to review our work and for the positive and constructive feedback! We are glad that you found our paper \"offers a valuable approach to optimize the forward process”, “organizes prior work into a cohesive structure”, and \"may serve as a source of inspiration\" for future work. We respond to each of your comments below.\n\n\n## The handling of 𝑝(𝑥0|𝑥𝑡(1)) could be enhanced\nThank you for the suggestion. In our implementation t(1) is tiny ($\\\\alpha_{t(1)} = 10^{-4}$) and we expect p(x0|xt(1)) to be very close to q(xt(1)|x0). We have also tried treating p(x0|xt(1)) the same as other p(xs|xt) and did not observe significant differences.\n\n\n## Multidimensional data and backward factorization\nThank you for the suggestion. We agree with the reviewer that the true backward distribution becomes factorized only as $s \\\\to t$, which is why we adopt the continuous-time formulation throughout the work. We will follow the reviewer’s suggestion to make this more clear in the final version.\n\n\n## Non-zero $\\\\alpha_1$\nWe originally introduced the non-zero $\\\\alpha_1$ to resolve numerical issues in GenMD4, where the REINFORCE LOO gradient requires calculating $\\\\log \\\\alpha_t$. We have since then used the same schedule for MD4 as well. We agree with the reviewer that the non-zero $\\\\alpha_1$ is unnecessary for MD4. Although we have not found any “medium brightness” problems, we will still rerun the MD4 experiment with $\\\\alpha_1 = 0$ and update the results.\n\n\n## Appendix G2: Variance of Campbell et al. (2022)\nWe thank the reviewer for noting this. In the single dimension case, we agree that there is no additional variance via further simplification of the Campbell et al. bound. The additional variance comes from multidimensional data where the sum over $j$ becomes the sum over all neighbors in the forward process, i.e., the states that mask out a single dimension of $x_t$ that has not been masked yet. In this case, the Campbell et al. (2022) bound simplifies to\n$$\n\\\\beta(t) \\\\sum_{i: x_{t, i} \\\\neq m} \\\\log (x_{t, i}^\\\\top \\\\mu(x_t(i\\\\to m))_d\n$$\nwhere $x_t(i\\\\to m)$ denotes the state we get by masking out the i-th dimension of $x_t$. This equation requires $d$ function evaluations of the prediction model $\\\\mu$, thereby incurring additional variance if we estimate the sum with Monte Carlo. We will fix the discussion of Campbell et al. (2022) in the final version. \n\n\n## 𝜏LDR with masked noise\nWe followed the reviewer’s suggestion to perform an additional experiment using 𝜏LDR with masked noise. We made sure the neural networks and training hyperparameters used by 𝜏LDR (masked) is the same with our MD4 experiment. The results are as follows (for comparison we also included the MD4 results)\n\n\n| Method | BPD (CIFAR-10) | \n| :------ | ------: | \n| 𝜏LDR (masked) | <= 3.52 |\n| MD4 | <=2.78 |\n\n\nThe results show that 𝜏LDR with masked noise suffers from the high-variance objective even with the same networks and training hyperparameters from MD4.\n\n\n## Details of the sampling method\nPlease refer to our response to questions shared by reviewers.\n\n\n## Results of GenMD4 on image datasets\nWe ran GenMD4 on both image datasets with the same architecture and hyperparameters as in MD4. On CIFAR-10, we get 2.7749 (GenMD4) v.s. 2.7847 (MD4). On ImageNet 64x64, we get 3.4233 (GenMD4) v.s. 3.4273 (MD4). We see consistent improvements in test BPD over MD4 throughout training although the improvement is smaller than that of text experiments (1.34 v.s. 1.37 BPC on text8 & 21.80 v.s. 22.13 PPL on OpenWebText validation set).\n\n\n## Simplified GenMD4\nWe followed the reviewer’s suggestion to perform an additional experiment on text8 that only learns a scalar w for GenMD4 that is shared across all schedules. The result is 1.37 BPD which is almost the same as the result of MD4 and worse than GenMD4 with a vector w (1.34 BPD). We believe this is because the simplified GenMD4 reduces to MD4 with a learnable schedule (we can show this by choosing $\\\\alpha_t$ in (19) as a scalar schedule times an all-one vector) and MD4 is invariant to masking schedules.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4Lcb9YecYh", "reviewer": "Reviewer_Y9JA", "summary": "The paper proposes a streamlined and generalized framework for masked diffusion models, addressing the complexities and inefficiencies of existing models, including those based on Score Entropy Discrete Diffusion (SEDD). It introduces a continuous-time variational objective for masked diffusion models, simplifying the evidence lower bound (ELBO) to a weighted integral of cross-entropy losses. Additionally, the paper presents state-dependent masking schedules, enhancing the flexibility and performance of these models. The proposed methods demonstrate state-of-the-art results in text and image tasks, significantly improving likelihood and zero-shot transfer performance.", "questions": "- Could the authors provide more insights into the practical challenges faced during the implementation of the state-dependent masking schedules?\n- How does the proposed model ensure consistency between the forward and reverse processes, and how does this impact training stability compared to SEDD? \n- Could the authors provide a detailed and separate description of the training and sampling algorithms, similar to what is provided in the Appendix of the SEDD paper, to better and more easily understand the proposed method?\n- How sensitive is the proposed method to hyperparameter choices? Do multiple runs with the same hyperparameters yield consistent performance?", "limitations": "None", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper offers a novel theoretical formulation of the continuous-time variational objective for masked diffusion models, simplifying the training process and ensuring consistency between forward and reverse processes.\n- The introduction of state-dependent masking schedules provides a more adaptable approach, catering to the specific characteristics of the data and improving model performance.\n- The proposed methods achieve state-of-the-art performance in both text and image generative tasks, significantly enhancing likelihood and zero-shot transfer capabilities.\n- By reducing the ELBO to a weighted integral of cross-entropy losses, the paper makes the training and understanding of masked diffusion models more accessible and potentially more stable.\n- The paper includes comprehensive experimental validation on various datasets, demonstrating the robustness and superiority of the proposed methods.", "weaknesses": "-  Despite the theoretical simplifications, the practical implementation of state-dependent masking schedules can still be complex and computationally demanding. Specifically, obtaining the starting x_T is challenging, and since the sampling process lacks stochasticity, sampling cannot be done from the completely masked state.\n- The state-dependent models have a tendency to overfit to dataset statistics, which can limit their effectiveness in zero-shot transfer tasks.\n- While the paper demonstrates superior performance, a more detailed comparative analysis with other state-of-the-art methods, particularly regarding computational efficiency and training times, would provide a clearer picture of the advantages.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time you’ve taken to review our work and for the positive and constructive feedback! We are glad that you find our work \"offers a novel theoretical formulation\" and \"achieve state-of-the-art performance\" with \" comprehensive experimental validation\". We address each individual comment below.\n\n## Practical challenges in implementation of state-dependent masking schedules \n* Regarding computational cost, the state-dependent masking schedules and corresponding GenMD4 objective require only twice the computation of MD4 with the same network size. While slightly more expensive, this increase remains within affordable limits.\n* Regarding sampling, we clarify that x_T is a full mask state even for the state-dependent schedule case. This is because the learned schedule $\\\\alpha_t = 1 - t^w$ still satisfies $\\\\alpha_1 = 0$ regardless of the value of $w>0$. We also clarify that the sampling process is stochastic in each step, please refer to the sampling algorithm described in our response to questions shared by reviewers.  \n\n\n## Effectiveness of state-dependent models in zero-shot transfer tasks\nWe have included the results of GenMD4 on zero-shot transfer tasks in our response to shared questions. The results are mixed: on some datasets (Lambada, PTB) GenMD4 results in significantly better zero-shot PPL numbers than MD4 while on other datasets (WikiText) GenMD4 is slightly worse. We believe this indicates some degree of overfitting to the dataset statistics of OpenWebText. \n\n\n## Comparative analysis regarding computational efficiency and training times\nWe follow the reviewer’s suggestion to compare the training times of MD4/GenMD4 on OpenWebText using 128 TPUv3. We can see that MD4 is as fast as our reimplementation of SEDD while achieving better results. GenMD4 is slightly slower than MD4 due to the 2x computational cost.\n* Gaussian Diffusion: 3.5 steps / s\n* SEDD: 4.2 steps / s\n* MD4: 4.2 steps / s\n* GenMD4: 3.5 steps / s\n\n\n## How the model ensures forward/reverse consistency and impact on training stability v.s. SEDD\nFor simplicity let’s consider the single-dimension case. As we have shown in Proposition 1, the true score function for a mask state $x_t=m$ has the following relationship with the conditional mean of $x_0$:\n$$\ns(m, t)\\_j = \\\\frac{\\\\alpha_t}{1 - \\\\alpha_t} E[x_{0,j}|x_t=m]\n$$\nNote that here $x_0$ is a one-hot representation of the data therefore $\\\\sum_j x_{0,j} = 1$ and thus also $\\\\sum_j E[x_{0,j}|x_t=m] = 1$. This implies that\n$$\n\\\\sum_j s(m, t)_j  = \\\\frac{\\\\alpha_t}{1 - \\\\alpha_t}.\n$$\nSince the transition rate matrix of the true reverse process depends on the score, the above equation implies a constraint that the reverse process has to satisfy in order to be consistent with the $\\\\alpha_t$-determined forward process.\n\n\n\n\nAs we have explained in Sec. 4 and App. F3, we can interpret both our method and SEDD as learning a model $s_{\\\\theta}$ for the true score function $s$. The differences are\n* In MD4, $s_{\\\\theta}$ is parameterized as \n$$\ns_{\\\\theta}(m, t)\\_j = \\\\frac{\\\\alpha_t}{1 - \\\\alpha_t} \\\\mu_{\\\\theta}(m)\\_j\n$$\nwhere $\\\\mu_{\\\\theta}$ has a softmax output that produces a probability vector. Therefore, the score model also satisfies $\\\\sum_j s_{\\\\theta}(m, t)\\_j = \\\\frac{\\\\alpha_t}{1 - \\\\alpha_t}\\\\sum_j \\\\mu_{\\\\theta}(m)\\_j  = \\\\frac{\\\\alpha_t}{1 - \\\\alpha_t}$. \n* In SEDD, $s_{\\\\theta}$ is a free-form neural network model that outputs a real-valued vector, which has no guarantee it satisfies the constraint. This inconsistency between forward and reverse process leads to large variational gaps in the ELBO which we have empirically shown to lead to unstable training (see Figure 3).\n\n\n## Detailed description of the training and sampling algorithms\nPlease refer to our response above to questions shared by reviewers.\n\n\n## Sensitivity to hyperparameter choices and multiple runs\nWe test the sensitivity to multiple runs by re-training MD4-S with the same hyperparameter as the original one reported in the paper, where the original PPL on OpenWebText eval split is 22.126 and the re-trained MD4-S gets 22.134. In the pdf we uploaded for rebuttal, we include an analysis of the impact of masking schedule choice on sample quality, showing the benefits of a cosine schedule.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KTr6AwozjU", "reviewer": "Reviewer_Pt68", "summary": "The paper simplifies the mathematical formula for the absorbing state diffusion process. By doing so, the authors derive a continuous-time ELBO for masked diffusion models. Their method, MD4, achieves better perplexity scores than SEDD on text8 and zero-shot perplexity on numerous datasets.", "questions": "1. Clarification on D3PM experiments in Table 1 as mentioned in the \"weaknesses\" section in the reviews.\n2. Why did the authors decrease the dropout to 0.02 for OWT experiments and not set it to 0? Diffusion models are heavily regularized due to the randomness in the input to the model and oftentimes don't require additional regularization such as dropout. Hence, an intuitive or an empirical explanation would be helpful.", "limitations": "Yes", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "Simplifies the complex mathematical formulations for the absorbing state diffusion for D3PM.", "weaknesses": "Weaknesses:\n\n1. Weak empirical results\n    1. The zero shot numbers for D3PM in Table 1 look fishy.  There are only 2 differences between Md4 and Absorbing State D3PM:\n        1. Mathematical simplification. In the discrete case (Eqn. 6), even though MD4 features a Simplified functional form for the ELBO, it shouldn't give it any performance benefits in terms of perplexity since it is mathematically equivalent to D3PM.\n        2. The improvement in ELBO could be because of the continuous time formulation. However, VDM [1]  has shown that for gaussian diffusion, improvement from discrete (T=1000) to continuous time (T = $\\infty$) barely improves the likelihood by less than 1%. For this reason, I request the authors to perform eval on an already trained model and report the  perplexity numbers on text8 or OWT using Eqn (6) with T=100, 1000, 10000. If the numbers reported for D3PM in Table 1 are indeed correct, and if the entire improvement is coming from the continuous time formulation, then the discrete time MD4 should get a number that's comparable to D3PM's zero shot ppl numbers. \n        \n        Questions: How did they retrain D3PM? Did they use the same transformer backbone as MD4? Did they use the same model size and data pre-processing scheme? Did they use  uniform state or absorbing state diffusion process? The authors need to clarify this.\n        \n    2. CIFAR10 Experiments. The AR baselines use old transformer models hence the comparison isn't quite fair. Current SOTA diffusion models on Imagenet 32 achieve a NLL of 2.55 [2] which is far better than the absorbing state diffusion models. So, I'm unsure about the takeaway from Table 3. In the conclusion section, the authors claim that \"… on text and image data, the resulting masked diffusions outperform existing discrete and continuous diffusion models …\" which is factually incorrect given that their method largely underperforms against gaussian diffusion [1, 2].\n2. Limited evaluation of GenMD4. The authors mention that GenMD4 performs poorly on zero-shot tasks. I request the authors to quantify this poor performance by providing \n    1. Validation ppl numbers on OWT\n    2. zero-shot ppl numbers.\n\n[1] Kingma, D., Salimans, T., Poole, B. and Ho, J., 2021. Variational diffusion models. *Advances in neural information processing systems*, *34*, pp.21696-21707.\n\n[2] Sahoo, S., Gokaslan, A., Sa, C., Kuleshov, V., 2024.  Diffusion Models With Learned Adaptive Noise. arXiv:2406.07524", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the time and the feedback! We clarify a few points.\n\n##  D3PM results & differences between MD4 and D3PM\n* The reviewer is concerned if the D3PM results for zero-shot transfer tasks are comparable to MD4. We clarify that these D3PM results are from the SEDD paper (Lou et al., ICML 2024) which used the same setup as SEDD, including an absorbing state diffusion. This paper states:\n\"... baselines are ... D3PM (with Absorbing Transition) … We retrain both models … reuse our model architecture and match hyperparameters (i.e. model size, training specifications).\"\nTo ensure direct comparability with MD4, we reimplemented SEDD and reproducing their results (see \"SEDD Absorb (reimpl.)\" in Table 1). This ensures a fair comparison between the models.\n* The reviewer states that MD4's simplified ELBO form shouldn't provide benefits compared to D3PM as both are ELBOs. However, we emphasize that this simplification is crucial for a well-engineered, numerically stable ELBO implementation, leading to significantly better performance. To illustrate, we analyze D3PM's original implementation of the $KL(q(x_s|x_t, x_0)||p(x_s|x_t))$ terms in the ELBO (s = t - 1 for D3PM) without simplification:\n\n    1. Sample $x_t \\sim q(x_t|x_0)$\n    2. Calculate $q_{s|0}(\\\\cdot|x_0)$ where $\\cdot$ takes all values in the vocabulary.\n    3. Calculate $q_{t|s}(x_t|\\\\cdot)$.\n    4. Calculate the logits of $q(x_s|x_t, x_0)$ via $qlogits = \\\\log (q_{s|0}(\\\\cdot|x_0) + \\\\epsilon) + \\\\log (q_{t|s}(x_t|\\\\cdot) + \\\\epsilon)$. \n    5. Repeat 2-4, replacing $x_0$ with NN output probabilities to get $p(x_s|x_t)$ logits ($plogits$). \n    6. Compute $KL(q(x_s|x_t, x_0)\\\\|p(x_s|x_t))$ using log_q = log_softmax(plogits) and log_p = log_softmax(qlogits). \n\nFor masked diffusion, as shown by our simplification, many elements of $q(x_s|x_t, x_0)$ and $p(x_s|x_t)$ are zeros (see eq. (40) in App.). These terms can be removed but are instead included by D3PM. Thus, D3PM has to introduce an $\\\\epsilon$ to avoid inf - inf which gives NaNs in the KL computation. We perform an experiment in CIFAR10 to show the numerical issues: we replaced the MD4 objective in our code with D3PM’s ELBO and tried different $\\\\epsilon$s (1e-20, 1e-12, 1e-8, 1e-6). Only after increasing it to 1e-6 we avoided NaNs at the start of training. Still, we got NaNs after 300k iterations. The test BPD before NaNs was 3.03, while MD4 at this training iteration reports 2.89.  \n\n* We agree with the reviewer that, similar to VDM, the cont-time v.s. discrete-time difference will not make a huge difference in likelihood. However, our cont-time formulation remains a key contribution for these reasons:\n    * It specifies the diffusion model using a simple monotonic function $\\\\alpha_t$ representing unmasking ratio at t. This enables exploration of schedules not used in discrete-time models like D3PM where the model is specified with less intuitive jump probabilities $\\\\beta_i$ rather than masking ratios. One of our key findings is that the cosine alpha_t schedule produces the best sample quality. In our rebuttal pdf, we include a comparison between the widely used linear schedule and the cosine schedule, showing the latter leads to significantly better FIDs in ImageNet 64x64. \n    * It enables easy adoption of training techniques from Gaussian diffusions; e.g. the antithetic sampling used by VDM for estimating the time integral in the ELBO which reduces training variance—on CIFAR-10, this translates to ~0.02 improvement on BPD.\n\n\n## CIFAR10 comparison with AR\nThe reviewer states that “The AR baselines use old transformer models hence the comparison isn't fair.” We have included the best AR results we found for CIFAR10. We're open to comparing with newer models if the reviewer can provide references. Also note our reported MD4 result of 2.78 is not fully optimized, e.g., increasing the batch size to 256 improves the result to 2.75 using the same 20M parameter model. \n\n\n## CIFAR10 comparison with continuous diffusion\nThe reviewer states that “Current SOTA diffusion models ... achieve a NLL of 2.55 [2] which is far better than absorbing diffusion” We believe the 2.55 from [2] (which we assume refers to CIFAR-10, not ImageNet 32) isn't directly comparable to our results in Table 3 for these reasons:\n* The 2.55 in [2] estimates exact log likelihood with probability flow ODE, while our results are upper bounds.\n* The best variational bound result in [2] is 2.65, which relies on learnable schedules using NNs, while MD4 uses a fixed schedule.\n* All continuous diffusion results in [2] Table 2 with BPD below 2.8 use extra image Fourier features (FFs) inputs from VDM. Comparing these with discrete diffusion is unfair since the latter assumes order-agnostic image data, i.e., the model is unaware of the proximity between pixels.\n\n\nAlso the impact of FFs becomes negligible at larger scale such as ImageNet 64x64. We have found that MD4 on ImageNet 64x64, by reducing the dropout rate to 0.0, gives 3.40 BPD which is the same as VDM relying on FFs.\n\n\nWe will revise the imprecise statement in the conclusion as  \"On text data, our masked diffusions outperform existing discrete and continuous diffusion models. For pixel-level image modeling, we significantly improve discrete diffusion results, outperforming similar-sized AR models and achieving comparable likelihoods to continuous diffusion models without Fourier features.\"\n\n\n## GenMD4 on zero-shot\nWe clarify that we did not state GenMD4 performs poorly on zero-shot tasks. In fact, the results there are mixed. We have included both validation PPL numbers on OWT and zero-shot PPL numbers for GenMD4 in our response to shared questions by reviewers.\n\n\n## Dropout rate for OWT\nFollowing the reviewer’s suggestion, we conducted dropout rate tuning and found that lower rate is better: dropout rate 0.05 has eval PPL 24.63, dropout rate=0.02, eval PPL is 22.13, and dropout rate = 0.0, eval PPL is 21.86. We’ll update the results in the final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HMh3lS7uGG", "reviewer": "Reviewer_EzVa", "summary": "This paper proposes a new framework for masked diffusion models for generative modeling of discrete data. Masked diffusion models offer an alternative to autoregressive models for discrete data but have faced challenges due to complex formulations and unclear relationships between different approaches. This paper presents a simplified and generalized framework to address these issues, enhancing the performance and training of masked diffusion models.\n\nThe key contributions includes:\n\n1. Simplification of Model Formulation: The paper establishes properties for the forward process and its time reversal using elementary arguments, provides a simple expression for the Evidence Lower Bound (ELBO), demonstrating it as a weighted integral over time of cross-entropy losses, and shows invariance properties similar to continuous space diffusions.\n\n2. Re-derivation of Training Objectives: The paper demonstrates how various previously proposed discrete diffusion training objectives can be derived from the ELBO objective by altering parameterization, relaxing constraints, or modifying loss weighting.\n\n3. Performance Improvements: The paper demonstrates state-of-the-art likelihood and zero-shot transfer results on text and image tasks using the proposed ELBO objective.\n\n4. Generalized Masked Diffusion Model: The paper proposes a generalized masked diffusion model that allows state-dependent masking schedules, further improving predictive performance on test likelihoods.", "questions": "Please see weaknesses.", "limitations": "The authors have adequately described the limitations and potential negative societal impact of their work.", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper makes a notable contribution to the field of generative modeling for discrete data by introducing a simplified and generalized framework for masked diffusion models.\n2. The quality of the paper is reflected in the thoroughness of its methodology and the robustness of its experimental validation.\n3. The paper is well-written and clearly structured, making it accessible to both experts and those new to the field.\n4. The significance of the paper lies in its potential to impact a wide range of applications in generative modeling for discrete data.", "weaknesses": "1. While the paper provides a robust theoretical foundation, there could be more emphasis on practical applicability. The paper could benefit from additional practical guidelines for implementing the proposed framework, such as more detailed pseudocode and specific implementation challenges.\n2. The experimental results presented are strong, but the range of tasks and datasets could be expanded, such as VQ-Diffusion [1] for token-based text-to-image.\n3. I am unfamiliar with diffusion models for text generation. For image generation, the paper has reported likelihood results, missing some other common metrics, such as FID and IS.\n\n[1] Vector Quantized Diffusion Model for Text-to-Image Synthesis", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive feedback! We are glad that you find our contribution notable, our methodology thorough and our experimental results strong and robust. We address each comment below\n\n## Detailed pseudocode and specific implementation challenges\nWe thank the reviewer for the suggestion. Please refer to our response to questions shared by reviewers for pseudocode of training and sampling algorithms. We will add them in the final version. We have discussed implementation challenges in Appendix A of the paper. Please let us know if you have any additional questions about the implementation.\n\n## The range of tasks and datasets could be expanded, such as VQ-Diffusion [1] for token-based text-to-image.\nWe thank the reviewer for the suggestion and will pursue this in future work. It's important to note that our pixel-level image modeling task is more challenging than modeling the latent space of vector quantized image encodings. While some discrete diffusion models have tackled high-resolution generation using the latter approach, to our knowledge, we're the first to go beyond 32x32 resolution for pixel-level modeling using discrete diffusion.\n\n## Sample metrics for image generation\nIn our rebuttal pdf we added an FID evaluation using 50K randomly generated samples from our MD4 models on ImageNet\n64x64. We compared different alpha schedules (Linear v.s. Cosine) and show that the cosine schedule leads to significant better sample quality. We further trained a class-conditional model on ImageNet 64x64 that boosts the FID to around 7. Although these are not state-of-the-art FIDs on ImageNet 64x64, we emphasize our models are optimized for likelihood using ELBO instead of sample quality, unlike the reweighted objectives used in continuous diffusion.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Response to comments shared by reviewers:

We thank the reviewers for their feedback. Below we address the questions shared by reviewers. We also uploaded a rebuttal pdf that contains figures used to address individual questions/comments of the reviewers.

## bny3,Y9JA: Details of training and sampling algorithms
A single step of the MD4 training algorithm is described below:


Input: data batch $x_0$, network $\\mu_{\\theta}(\\cdot, t)$, masking schedule $\\alpha_t$ \
Draw a batch of $t$ from U[0, 1] using antithetic sampling \
Sample $x_t$ using $q(x_t|x_0)$ by masking out independently each dimension of $x_0$ with probability $1 - \\alpha_t$ \
Get prediction logits via $\\mu_{\\theta}(x_t, t)$ \
Calculate the cross entropy loss $CE$ using $x_0$ and the prediction logits and sum over the masked dimensions of $x_t$. \
Compute unbiased estimate of the negative ELBO as $-\\frac{\\alpha_t’}{1 - \\alpha_t} * CE$  \
Optimize the negative ELBO via autodiff. 


Throughout the paper we simply run ancestral sampling from a discrete-time backward model. Specifically, the sampling algorithm is:


Input: Context size N, discretization grid $0 = t(0) < t(1) < \\cdots < t(T) =1$ \
Init: $x_1 \\gets [m, \\dots, m]$ \
for $i=T, T - 1, \\dots, 1$: \
    $t \\gets t(i)$, $s \\gets t(i-1)$ \
    for $n \\in [N]$, if $x_t^{(n)} = m$, draw $x_s^{(n)} \\sim \\mathrm{Cat}(\\frac{\\alpha_s - \\alpha_t}{1 - \\alpha_t} \\mu_\\theta^{(n)}(x_t, t) + \\frac{1- \\alpha_s}{1 - \\alpha_t} e_m)$ else $x_s^{(n)} \\gets x_t^{(n)}$ \
return $x_0$


We will follow the reviewer’s suggestion to include both algorithms in the final version.


## Y9JA, Pt68: Evaluation of GenMD4 on OpenWebText validation set and zero-shot tasks
* GenMD4 leads to better perplexity on OpenWebText validation set (this result is already shown in Figure 3 of the submission, here we report the precise values in a table):




| Method                | Perplexity |
|-----------------------|-----------------------|
| Gaussian Diffusion    | <= 27.28              |
| SEDD Absorb (reimpl.) | <= 24.10              |
| MD4 (Ours)            | <= 22.13              |
| GenMD4 (Ours)         | <= **21.80**          |




* For zero-shot results, on some datasets (Lambada, PTB) GenMD4 results in significantly better zero-shot PPL numbers than MD4 while on other datasets (WikiText) GenMD4 is slightly worse. We believe this indicates some degree of overfitting to the dataset statistics of OpenWebText. 


| Method | LAMBADA | WikiText2 | PTB | WikiText103 | IBW
| :------ | ------: | ------: | ------: | ------: | ------: | 
| MD4 (Ours) | 48.43 | **34.94** | 102.26 | **35.90** | 68.10 |
| GenMD4 (Ours) | **33.31** | 41 | **65.06** | 41 | **52.1** |

Meta-review: Overall, reviewers agree that the proposed framework for diffusion-based generative modeling of discrete data provides interesting conceptual advances, and leads to promising experimental results.  Reviews and author rebuttals highlight some areas where conceptual and experimental comparisons to prior work could be improved, please be sure to incorporate these points in the final version of your manuscript.
Decision: Accept (poster)