Simplified and Generalized Masked Diffusion for Discrete Data

Authors: jiaxin  shi
kehang  han
zhe  wang
arnaud  doucet
michalis k titsias
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:41:17

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.

Main paper: 
 introduction: 
 Since their inception [1] [2] [3] , diffusion models have emerged as the workhorse for generative media, achieving state-of-the-art in tasks such as image synthesis [4] [5] [6] , audio [7, 8] and video generation [9] [10] [11] [12] [13] . 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 are 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 Euclidean space [19] [20] [21] [22] [23] or 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 a 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 are as 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.Discrete-time forward process. The forward process is defined as a Markovian sequence of discrete random variables x t indexed by time t, where t runs from 0 to 1. Throughout the work, we abuse the notation such that x t can be either an integer or its corresponding one-hot vector, whenever it 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), t(i)] is determined by a transition matrix of size (m + 1) × (m + 1):Q i = (1 -β i )I + β i 1e ⊤m , where 1 is an all-one vector of size m + 1, e m represents a one-hot vector where element at index m is 1. Each entry [Q i ] jk denotes the probability of transition from the state j to the state k:[Q i ] jk = q(x t(i) = k|x s(i) = j) = (1 -β i )δ jk + β i δ km .This means that, with probability 1 -β i , x t(i) = x s(i) , otherwise it jumps to the mask state. Given the above transition matrix, the marginal distribution at time t(i) given x 0 is q(x t(i) |x 0 ) = Cat(x t(i) ; Q⊤ i x 0 ) = x ⊤ 0 Qi x t(i) .Here, we use Cat(x; p) to denote a Categorical distribution where p is the vector of probabilities of being in each category, and Qi≜ i j=1 Q j = α i I + 1 -α i 1e ⊤ m for α i = i j=1(1 -β j ). We expect α T to become very small or zero for a sufficiently large T such that q(x 1 |x 0 ) for any x 0 will become a delta mass at the mask state.Continuous-time limit. We can define a continuous-time forward process by taking a limit of the above discrete-time process. We first specify a continuous function β(t) such that β i = β(t(i))/T . We then let T → ∞ in the discrete-time process and compute the limit of Qi (proved in Austin et so that q(x t |x 0 ) = Cat(x t ; Q(t) ⊤ x 0 ). For two arbitrary times, 0 ≤ s < t ≤ 1, the transition distribution that is compatible with the above marginal (i.e., q(x t |x 0 ) = xs q(x t |x s )q(x s |x 0 )) isq(x t |x s ) = Cat(x t ; Q(s, t) ⊤ x s ), where Q(s, t) ≜ Q(s) -1 Q(t) = α t α s I + 1 - α t α s 1e ⊤ m .Note that Austin et al. [14] did not derive this explicit form of transition matrix between two arbitrary time s and t, which appeared later in Zhao et al. [38] concurrently with our work.Masking schedules. From the definition of α t , we have that α 0 = 1. And similar to the discretetime formulation, we would like α 1 be zero or very close to zero. We provide a summary of masking schedules from literature that satisfy these properties in Fig. 1 . 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 α t is plotted for βmin = 10 -5 and βmax = 20. It was first used for continuous diffusions [3] and then for discrete by Lou et 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 at the beginning of the reverse generation. Aligning with their observation, we find that this results in a lower chance of conflicting tokens being unmasked simultaneously at the start of generation, thereby enhancing the overall generation quality.Time reversal of the forward process given x 0 . 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 x 0 : q(x s |x t , x 0 ) for s ≤ t. We derive it in App. C as q(x s |x t , x 0 ) = Cat(x s ; Rx0 (t, s) ⊤ x t ), where Rx0 (t, s) = I + α s -α t 1 -α t e m (x 0 -e m ) ⊤ .From the transition matrix Rx0 (t, s) ∈ R (m+1)×(m+1) we can see the reverse process conditioned on x 0 has a very simple logic-if x t is a mask, with probability αs-αt 1-αt , it will jump to the state x 0 at time s, otherwise it will stay masked. Once x t is unmasked, it remains in the same state until the end. 
model and objective: 
 For a discrete-time masked diffusion process, we define our generative model by approximately reversing the forward transitions using a reverse model p θ (x s |x t ). One way to define this model isEQUATIONwhere µ θ (x t , t) ∈ ∆ m+1 is a probability vector parametrized by a neural network f θ with a softmax applied to the output logits (note the m-th output is forced to 0 since the clean data cannot be masks):EQUATIONThis is known as mean-parameterization since it leverages a prediction model for the mean of x 0 . A matrix-form depiction of p θ (x s |x t ) is shown in Fig. 6 (right). In fact, we can select a time-invariant parametrization µ θ (x t , t) = µ θ (x t ) as [36] showed that p(x 0 |x t ) given x t = x is identical for any t.Besides p θ (x s |x t ), we also need to specify p(x 0 |x t(1) ) and the prior distribution p(x t(T ) ) = p(x 1 ). Following the practice in continuous diffusion models [33] , we choose p(x 0 |x t(1) ) ∝ q(x t(1) |x 0 ). And since q(x 1 |x 0 ) ≈ δ x1,m for any x 0 as α 1 ≈ 0, we set p(x 1 ) ≈ δ x1,m , see App. E.We then write out the discrete-time diffusion model objective [1, 2] , which is a lower bound of the log marginal likelihood of data x 0 under the model p (known as the Evidence Lower Bound, or ELBO):log p(x 0 ) ≥ E q(x t(1) |x0) [log p(x 0 |x t(1) )] -KL(q(x 1 |x 0 )∥p(x 1 )) -L T ,whereL T = T i=2 E q(x t(i) |x0) [KL(q(x s(i) |x t(i) , x 0 )∥p θ (x s(i) |x t(i) ))].For the above choices of the prior distribution, the term KL(q(x 1 |x 0 )∥p(x 1 )) becomes zero. Under the reverse model (2), the KL divergence terms in L T becomes (proof in App. D)KL(q(x s |x t , x 0 )∥p θ (x s |x t )) = - α s -α t 1 -α t δ xt,m • x ⊤ 0 log µ θ (x t , t),which is a simple cross-entropy loss between the predicted logits and the clean data. In App. D, we show that L T is a Riemann sum and is lower bounded by the corresponding continuous integral:EQUATIONwhere α ′ t denotes the derivative of α t with respect to t. Therefore, we can obtain an ELBO that is tighter than that of any finite T by pushing T → ∞. This ELBO can be further simplified by letting t(1) → 0. As a result, E q(x t(1) |x0) [log p(x 0 |x t( 1) )] goes to 0 and the ELBO becomes -L ∞ .For continuous state-space diffusions, the ELBO depends on the signal-to-noise ratio (SNR) at its endpoints but is otherwise invariant to the noise schedule [33] . We establish here a similar result for discrete diffusions. Consider choosing α t = σ(λ t ), where σ represents the sigmoid function σ(x) = 1 1+e -x . In this context, the log-SNR is defined by λ t = log αt 1-αt = log-SNR(t). By making a change of variables in (4) to make everything a function of the log-SNR, we obtainL ∞ = λ1 λ t(1) σ(λ)E q(x λ |x0) δ x λ ,m • x ⊤ 0 log μθ (x λ , λ) dλ.where μθ (x, λ) := µ θ (x, t) and q(x λ |x 0 ) := q(x t |x 0 ) for t = log-SNR -1 (λ). This shows that the only effect α t has on the loss is through the values of 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, x t was assumed to be a single discrete token.To extend the method to multidimensional data, let x t be now a sequence (xEQUATIONt , . . . , x(N ) t), where each element x (n) t represents a discrete token. We select a forward process which factorizes across all N tokens: q(x t |x s ) =N n=1 q(x (n) t |x (n) s ).As a result, the forward marginals q(x t |x 0 ) and reversal q(x s |x t , x 0 ) also factorize. In this case, we define the reverse model asp θ (x s |x t ) ≜ N n=1 q(x (n) s |x (n) t , µ (n) θ (x t , t)), where µ θ (x t , t) is a neural network that takes the full N tokens as input and outputs N probability vectors. 3 The n-th output µ(n) θ (x t , t) is a prediction model for E[x (n) 0 |x t ],the mean value of the n-th token. Repeating above derivations givesEQUATIONA single step of our training algorithm is described in Alg. 1 in Appendix. 
sampling: 
 We use 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. Impact of schedules and discretization. For comparing different sampling configurations, we primarily use the FID score [40] on image datasets as our evaluation metric. We favor it over text generative perplexity4 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 a standard uniform discretization grid t(i) = i T . We hypothesize that time discretization can lead to conflicts by generating multiple tokens in a 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 a "cosine" discretization grid t(i) = cos( π 2 (1 -i T )) for sampling in models trained with a linear schedule. This induces the same discretization in α t as the cosine schedule with a uniform grid, leading to comparable sample quality, as 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, the remaining tokens become more predictable, reducing conflicts when unmasking them in a single step.Although these findings were originally developed on images, we find them translate well to text (see Fig. 9 ). we expect other techniques such as top-p sampling [41] , classifier-free guidance [42, 43] , and predictor-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 a similar effect to top-p sampling. See App. G for details. 
relation to existing work: 
 We discuss how to unify several existing masked diffusion models using our framework.To show the connection with the CTMC view presented in 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 a short time ∆t, given x 0 , the Taylor expansions of our forward and reverse transition matrices at t areEQUATIONEQUATIONwhere Q(t) and R x0 (t) are known as the transition rate matrices. Austin et al. [14] derived the same Q(t) in App. A.6 of their paper. However, they did not explore the reverse process or a continuoustime objective. Campbell et al. [29] established an alternative expression for the ELBO using the rate matrices. In App. H.1, we show how to recover their expression by separating out a constant from our ELBO expression (4) and applying a discrete "integration-by-part". A key limitation of their expression is that it needs N evaluations of the prediction model µ θ (•, 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 µ θ (•, t). Please refer to App. H.2 for more details.Score parameterization. While so far we used a prediction model µ θ (x t , t) for the mean of clean data given x t (i.e., mean parameterization), one can choose other ways of parameterizing the reverse model. Lou et al. [32] , Benton et al. [35] proposed to parameterize the discrete "score" s(x t , t) j ≜ qt(j) qt(xt) and introduced a 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 q t be the marginal distribution of the masked diffusion defined in Sec. 2 at time t. The discrete score s(x t , t) j = qt(j) qt(xt) for a mask state x t = m and j ̸ = m can be expressed asEQUATION)Proposition 1 (proved in App. H.3) implies that a reasonable score model for a mask state isEQUATIONIndeed, substituting (9) into the score-based loss of Lou et al. [32] , Benton et al. [35] recovers our objective (4). In Lou et al. [32] , the score is parameterized as a neural network without enforcing the constraint in (8) . 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-time reverse process of our masked diffusion model can be viewed as an any-order autoregressive model (AO-ARM) [46] . 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 {τ n } N n=1 are identical and satisfy P (τ n ≤ t) = P (x(n) t = m) = 1 -α t .As a 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. [47, App. C ]. 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 α t introduces a new degree of freedom in the design of such models. Variations in α t 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 [48] , SDDM [30] , and SUNDAE [49] , to App. H.4. 
generalization to state-dependent masking schedules: 
 We consider here state-dependent masking schedules by making the time-dependent probability of masking a token dependent also on the token value. Thus, some tokens will have higher probability to be masked earlier in time compared to others.We first define the forward process for a single token x t . Let α t be a m + 1 dimensional vector function, i.e., there is a different function α t,i for each possible value i of the token x t . Also, by vector αt αs we denote the element-wise division of the two vectors. We define the forward transition as q(x t |x s ) = Cat(x t ; Q(s, t) ⊤ x s ) whereQ(s, t) = diag α t α s + I -diag α t α s 1e ⊤ m Proprietary + ConfidentialMayor Muriel Bowser said after meetings with Commissioner Busby on Thursday that the new plan will be on board in December.Mayor � Bowser said � meetings � Commissioner � on Thursday that � new plan will be � board in � � Mayor Muriel Bowser said after meetings � Commissioner � on Thursday that � new plan will be � board in December � and diag αt αs is a diagonal matrix with the vector αt αs in its diagonal. The probability of moving from current state x s to a future state x t (either the same as x s or mask) is determined by a state-dependent rate αt αs ⊤ x s , while the marginal at time s given x 0 isMayor � � said � � � � � � � that � new plan � � � � � � �q(x s |x 0 ) = Cat(x s ; Q(s) ⊤ x 0 ) for Q(s) = diag(α s ) + (I -diag(α s ))1e ⊤ m .Further, for any time 0 ≤ s < t ≤ 1 it holds that q(x t |x 0 ) = xs q(x t |x s )q(x s |x 0 ) so the above is a valid continuous-time Markov chain.Given the forward conditionals and marginals, we can now compute the time reversal conditioned on x 0 . The full form of q(x s |x t , x 0 ) is derived in App. I.1. For x t = m, we haveEQUATION)This suggests that the reverse model given x t = m can be chosen asp θ (x s |x t = m) ≜ q(x s |x t = m, µ θ (x t , t), diag(µ θ (x t , t))) where µ θ (x t , t) is a probability vector defined as in Sec. 3 that ap- proximates E[x 0 |x t ] while diag(µ θ (x t , t)) approximates E[x 0 x ⊤ 0 |x t ] = diag(E[x 0 |x t ]).We show in App. I.1 that the negative continuous-time ELBO for the state-dependent rate case isEQUATIONwhere α ′ t is the elementwise derivative of α t w.r.t. t. This generalizes the loss from ( 4), which can be recovered by choosing α t as a scalar schedule times an all-one vector. For N tokens the reverse model and the loss further generalize similarly to Sec. 3; see App. I.2.To learn the token dependent masking schedule using ELBO optimization, we parametrize the m + 1 dimensional function α t using the polynomial schedule (see Fig. 1 ) as α t,i = 1 -t wi and optimize each parameter w i > 0. 5 The value of w i , through the masking probability 1 -α t,i , determines how fast the token with value i jumps to the mask state. Since in the loss (11) the distribution q(x t |x 0 ) depends on α t and thus the vector w, optimizing w poses a discrete gradient estimation problem [see, e.g., 50]. Naive autodiff leads to biased gradients and pushes w towards zero because the gradients cannot propagate through the (discrete) samples drawn from q(x t |x 0 ). To fix this, we used the REINFORCE leave-one-out estimator [51, 52] to compute low-variance unbiased gradients for optimizing w. Details are given in App. I.2. Text is natural discrete data with rich structures. We experiment with two datasets: text8 [53], a character-level text modeling benchmark extracted from English Wikipedia, and OpenWebText [54] , an open clone of the unreleased WebText dataset used to train GPT-2 [55] . We choose the linear masking schedule described in Sec. 2 for all text experiments since we find it works best. We refer to our simple masked diffusion model as MD4 (Masked Discrete Diffusion for Discrete Data) and our generalized state-dependent model as GenMD4. OpenWebText. We train a GPT2-small and GPT2-medium size of MD4 on OpenWebText (98% training, 2% validation) and evaluate zeroshot perplexity on the test splits of five benchmark datasets used in Radford et al. [55] : LAM-BADA [56] , WikiText2, Penn Treebank [57] , WikiText103 [58] , and One Billion Words [59] . We keep our evaluation setup the same as SEDD [32] . To ensure the comparison is fair, we reproduced their result in our implementation.Ȱ ȲȰȰ ȴȰȰ ȶȰȰ ȸȰȰ ȱȰȰȰ Į´ôČôČìĲĹØīĲʬȱľČôĹ˽ȱȰȰȰĲĹØīĲʭ ȱȰ Ȳ Ȳ ˻ ȱȰ ȱ ȳ ˻ ȱȰ ȱ ȴ ˻ ȱȰ ȱ ȶ ˻ ȱȰ ȱ `īØČØÍØŗĹ(ő´ĆyØĮīĆØŗôĹŘ ;´ľĲĲô´Č$ôëëľĲôĒČʲ ($$ʲ Y$ȴʲ ;ØČY$ȴʲ Y$ȴʲYThe results are shown in Tab. 1. Our small model outperforms previous best discrete diffusion models on all five tasks. We also confirmed that our re-implementation of SEDD leads to similar and slightly better results than those reported in their paper. We believe this is because they were using mixed-precision training (thus larger variances) while we are using full precision (float32) training. We are also better than GPT-2 on all tasks except LAMBADA where we are the second best among all methods. When we scale up the model size to medium, we observe that MD4 similarly beats SEDD and GPT-2 on all tasks except LAMBADA.To validate that the strong zero-shot performance is a result of a better-trained model, we plot perplexity on the OpenWebText validation set in Fig. 4 . It demonstrates that our model converges faster and has better final likelihoods than prior methods. We also observed that SEDD [32] has training instabilities. We believe this is due to the score parameterization breaking consistency between forward and reverse processes, as explained in Sec. 5. Although GenMD4 achieves lower perplexity than MD4, we observed that the learned ws can easily overfit to dataset statistics, making it less effective on zero-shot transfer tasks. Beyond likelihood evaluation, we also examine generation quality of our models. Fig. 3 showcases a randomly selected, notably coherent sample produced 1.39 Discrete Flow [62] 1.23 Transformer AR [14] 1.23Mult. Diffusion [15] ≤ 1.72 D3PM Uniform [14] ≤ 1.61 D3PM Absorb [14] ≤ 1.45 SEDD Absorb [32] ≤ 1.39 MD4 (Ours) ≤ 1.37 GenMD4 (Ours) ≤ 1.34 by MD4-medium, alongside its iterative denoising process. Fig. 9 demonstrates MD4's ability to perform text infilling and highlights a substantial improvement in sample quality when transitioning from the linear to the cosine schedule.We follow prior work to employ the GPT-2 Large model to evaluate perplexity of the generated samples. As noted in Sec. 4, the generative perplexity metric can be misleadingly decreased by reducing sample diversity; we include these results in Appendix (Fig. 7 ) mainly for completeness. Under this metric, MD4 significantly outperforms the GPT-2 checkpoint and SEDD. Notably, after 64 decoding steps, our small-sized MD4 matches the performance of SEDD-medium.text8. Following the convention in prior work [14, 32] , we experimented with training masked diffusion models on text chunks of length 256. We used the same dataset splits and transformers of the same size to parameterize the prediction model. Details can be found in App. J.1. Results are summarized in Tab. 2. We used the standard bits-per-character metric (a normalized negative log likelihood) for this dataset. We can see that our models outperform previous diffusion models, whether discrete or continuous in nature. We also outperform the best any-order autoregressive models, which do not assume a fixed generation order and has a strong connection to discrete diffusion [47] . Our model is only beaten by an autoregressive transformer and the Discrete Flow (which also used an autoregressive backbone). We believe this is because autoregressive models only require learning a fixed order of generation thus better utilize model capacity.The text8 dataset has a small vocabulary that consists of 26 letters a-z and a space token. Therefore, we did not expect much flexibility introduced by our state-dependent formulation. Still, with the generalized objective in (11) , GenMD4 is able to achieve significantly better BPC than MD4, showing the promise of a state-dependent diffusion on discrete data tasks. 
pixel-level image modeling: 
 To demonstrate MD4 beyond text domains, we follow Austin et al. [14] , Campbell et al. [29] and train MD4 on order-agnostic image data. Specifically, we take the CIFAR-10 and the Downsampled ImageNet 64×64 [63] datasets and treat each image as a set of discrete tokens from a vocabulary of size 256 such that the model is unaware of the relative proximity between different pixel values. We Figure 5 : Non cherry-picked unconditional samples from MD4 trained on ImageNet 64x64, treating pixels as discrete tokens. More samples can be found in Fig. 8 in Appendix. The model is optimized for likelihood instead of visual quality-see e.g., Kingma et al. [33] for samples from a continuous diffusion model optimized similarly for likelihood.compare to other discrete diffusion and autoregressive models that have reported likelihood results on the same dataset, although to our knowledge there are no published result on discrete diffusion methods for ImageNet 64 × 64 that directly model raw pixel space.Results are summarized in Tab. 3. We establish a new state-of-the-art for discrete diffusion models, outperforming previous discrete-time [14] and continuous-time [29] models by a significant margin. Notably, despite lacking knowledge of the ordinal structure of pixel values, MD4 still outperforms models trained with this inductive bias, including D3PM Gauss and τ LDR [29] where the noising distribution is a discrete Gaussian distribution that gives larger probabilities to near pixel values.To isolate the differences caused by training objectives, we also implemented the objective from Campbell et al. [29] in our code. The result, labeled as τ LDR Mask, demonstrates that the high variance of this objective negatively impacts model learning, even when using the same architecture as ours. Our CIFAR-10 model without data augmentation also achieves better likelihood than the previous best autoregressive models on this task, while on ImageNet 64 × 64 our result is competitive with Transformer AR models which is approximately 4× larger in model size and the continuous diffusion model VDM [33] .We provide a random sample from our ImageNet 64×64 model in Fig. 5 . More results can be found in in App. K. In Fig. 2 , we 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 and a cosine grid achieves an FID close to the model with cosine schedule, both significantly outperform the linear schedule with a uniform grid. We further trained a class-conditional model on ImageNet 64×64 that boosts the FID to around 7. Although these are not state-of-the-art FIDs on ImageNet 64×64, we emphasize our models are optimized for likelihood instead of sample quality. 
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 is not competitive with continuous diffusions and autoregressive models. The framework we propose provides a very simple expression of the ELBO as a 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-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 autoregressive models. We conjecture that this is because autoregressive models can better leverage model capacity since they only require learning one order. It 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 is more challenging. Cosine 1 -cos( π 2 (1 -t)) -π 2 tan( π 2 (1 -t))A Discrete-time derivationWe divide time from 0 to 1 into T intervals, and let s(i)= (i -1)/T , t(i) = i/T . The forward transition matrix Q i ∈ R (m+1)×(m+1) (m is vocabulary size) at time t(i) is [Q i ] jk =        1 j = k = m 1 -β i j = k ̸ = m β i k = m, j ̸ = m 0 otherwiseor more compactly written asQ i = (1 -β i )I + β i 1e ⊤ m, where 1 denotes an all-one vector of size m + 1, and e m is an one-hot vector of size m + 1 with the m-th element (recall that counting starts from 0) being one. We use an one-hot vector x t of length m + 1 to denote the discrete state. The forward conditionals are defined asEQUATIONwhereQ ⊤ i x s(i)is the probabilities for each of the m + 1 categories that x t(i) can take. The marginal forward distribution at time t(i) given x 0 isq(x t(i) |x 0 ) = Cat(x t(i) ; Q⊤ i x 0 ) = x ⊤ 0 Qi x t(i) ,where Qi =i j=1 Q j = i j=1 (1 -β j )I + 1 - i j=1 (1 -β j ) 1e ⊤ m .To see what this leads to in continuous time, we letβ i = β(t(i)) T and T → ∞: i j=1 (1 -β j ) = exp i j=1 log(1 -β j ) = exp i j=1 - β(t(j)) T + o(1/T ) T →∞ → exp - t(i) 0 β(s)ds .We let Q(t) denote the limit of Qi in this case:Q(t) = exp - t 0 β(s)ds I + 1 -exp - t 0 β(s)ds 1e ⊤ m ≜ α t I + (1 -α t )1e ⊤ m .Here we define α t ≜ exp(-t 0 β(s)ds). And the marginal forward transition isEQUATION 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? 
 
 Appendix: We consider a continuous-time Markov chain with transition ratesEQUATIONFor simplicity, we let Q = 1e ⊤ m -I. The marginal forward distribution at time t given x 0 is q(x t |x 0 ) = Cat(x t ; Q(t) ⊤ x 0 ), whereQ(t) = exp t 0 Q(s)ds = exp Q t 0 β(s)ds = exp( β(t)Q).Here we define β(t) ≜ t 0 β(s)ds. The matrix exponential can be computed via eigendecomposition:β(t)Q = U ΛU -1 ,whereU = I -e m e ⊤ m + 1 √ n + 1 1e ⊤ m , U -1 = I + √ n + 1e m e ⊤ m -1e ⊤ m , Λ = β(t)(e m e ⊤ m -I), and thus exp(Λ) = α t I + (1 -α t )e m e ⊤ m , Q(t) = U exp(Λ)U -1 = α t I + (1 -α t )1e ⊤ m .A simpler derivation uses the following property:Q 2 = -Q. Therefore, Q(t) = exp( β(t)Q) = I + β(t)Q + 1 2 β(t) 2 Q 2 + 1 3 β(t) 3 Q 3 + . . . = I + Q -(1 -β(t) + 1 2 β(t) 2 - 1 3 β(t) 3 + . . . )Q = I + Q -exp(-β(t))Q = α t I + (1 -α t )1e ⊤ m .This marginal forward transition matrix at time t coincides with the result (1) we get by taking the limit of discrete-time derivation.Arbitrary discretization of the continuous-time forward process. For the discrete-time process we have defined the per-step transition in (12) . For the continuous-time process, we can derive the transition matrix Q(s, t) ij ≜ q(x t = j|x s = i) between two arbitrary time s and t as the solution to the following differential equation (known as Kolmogorov forward equation)d dt Q(s, t) = Q(s, t)Q(t) where Q(t) = β(t)Q with initial condition Q(s, s) = I.The solution is given byQ(s, t) = exp ( β(t) -β(s))Q = Q(s) -1 Q(t).Routine work (using the Woodbury matrix inversion lemma) shows thatQ(t) -1 = α -1 t I + (1 -α -1 t )1e ⊤ m .Plugging the result back, we get the forward transition distribution from s to t:EQUATIONwhereQ(s, t) ≜ Q(s) -1 Q(t) = α t α s I + 1 - α t α s 1e ⊤ m .Figure 6 : The reverse transition probability and our generative model. Left: q(x s = •|x t = •, x 0 ) in matrix form where first index is x t and second index is x s . Right:p θ (x s = •|x t = •) ≜ q(x s = •|x t = •, µ θ (x t , t)) also in matrix form.C Time reversal of the forward process given x 0The 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 x 0 : q(x s |x t , x 0 ). We can compute it using ( 13) and (15) asEQUATIONVisually, eqn ( 16) is a R (m+1)×(m+1) matrix (Fig. 6 , left) whose first index is x t and the second is x s . The matrix is almost an identity matrix except the last row corresponding to x t is the mask token. The last row means with probability of αs-αt 1-αt the mask token gets unmasked to become x 0 , and with probability of 1-αs 1-αt it remains masked. Alternatively, we can rewrite the above using reverse transition matrix Rx0 (t, s) ∈ R (m+1)×(m+1) asq(x s |x t , x 0 ) = Cat(x s ; Rx0 (t, s) ⊤ x t ), where Rx0 (t, s) = I + α s -α t 1 -α t e m (x 0 -e m ) ⊤ .We are also interested in what would happen in the infinitesimal time limit, i.e., when s = t -∆t and ∆t → 0. Note that α t-∆t -α t = -α ′ t ∆t + o(∆t). Plugging it into the original formula, we getRx0 (t, t -∆t) = I - α ′ t 1 -α t e m (x 0 -e m ) ⊤ ∆t + o(∆t).Comparing the above with the transition rate matrix R x0 (t) definition Rx0 (t, t -∆t) = I + R x0 (t)∆t + o(∆t),we have determined the transition rate matrix for the reverse process conditioned on x 0 :EQUATION)Using ( 16) and (3), we compute the KL divergences between forward and reverse KL(q(x s |x t , x 0 )∥p θ (x s |x t )) = KL(q(x s |x t , x 0 )∥q(x s |x t , µ θ (x t , t)))= m xs=0 q(x s |x t , x 0 ) log q(xs|xt,x0) q(xs|xt,µ θ (xt,t))x t = m 0 x t ̸ = m = δ xt=m k̸ =m α s -α t 1 -α t x ⊤ 0 e k log x ⊤ 0 e k µ θ (x t , t) ⊤ e k = -δ xt=m α s -α t 1 -α t x ⊤ 0 log µ θ (x t , t).Note that 0 log 0 = 0. Alternatively, this result can be easily obtained from the visual depictions of q(x s |x t , x 0 ) and p θ (x s |x t ) shown in Fig. 6 . In this case, the reconstruction term becomesE q(x t(1) |x0) [log p(x 0 |x t(1) )] = m k=0 q t(1)|0 (k|x 0 ) log q t(1)|0 (k|x 0 ) j̸ =m q t(1)|0 (k|j) = α t(1) • log α t(1) α t(1) + (1 -α t(1) ) log 1 m = -(1 -α t(1) ) log m.The prior KL term can be computed asKL(q(x 1 |x 0 )∥p(x 1 )) = KL(δ x1,m ∥δ x1,m ) = 0.As usual, we take the continuous-time limit by letting T → ∞:L ∞ ≜ lim T →∞ L T = lim T →∞ T i=2 - α s(i) -α t(i) s(i) -t(i) s(i) -t(i) 1 -α t(i) x ⊤ 0 E q(x t(i) |x0) δ x t(i) ,m log µ θ (x t(i) , t(i)) = 1 t(1) α ′ t 1 -α t x ⊤ 0 E q(xt|x0) [δ xt,m log µ θ (x t , t)] dt.When defining the forward process, we often do not want α 1 to be exactly 0, or equivalently, λ 1 to be ∞ for numerical stability reasons. Instead, we set λ 1 to be a finite value, and thereby α 1 has a small positive value. This has a problem that the support of q(x 1 |x 0 ) is no longer {m} and instead becomes {m, x 0 }. As a result, the KL divergence between q(x 1 |x 0 ) and p(x 1 ) is undefined because q(x 1 |x 0 ) is not absolutely continuous with respect to p(x 1 ) = δ x1,m . To resolve the issue, we modify the prior distribution p(x 1 ) such that it has support over all m + 1 values. One such choice is lettingp(x 1 ) = α 1 m j̸ =m δ x1,j + (1 -α 1 )δ x1,m .Then, the prior KL divergence term becomes Input: minibatch {x i t } B i=1 , network µ θ (•, t), masking schedule α t for i = 1, . . . , B do (in parallel):KL(q(x 1 |x 0 )∥p(x 1 )) = m x1=0 q(x 1 |x 0 ) log q(x 1 |x 0 ) p(x 1 ) = m x1=0 (α 1 δ x1,x0 + (1 -α 1 )δ x1,m ) log α 1 δ x1,x0 + (1 -α 1 )δ x1=m p(x 1 ) = α 1 log α 1 α 1 /m + (1 -α 1 ) log 1 -α 1 1 -α 1 = α 1 log m.t i ← mod(u + i/B, 1), u ∼ U [0, 1] for n ∈ [N ], mask out each token x i,(n) 0 independently with probability 1 -α ti to obtain x i ti for n ∈ [N ], if x (n)ti = m, compute weighted cross entropy lossα ′ t i 1-αt i (x i,(n) 0 ) ⊤ log µ (n)θ (x i ti , t i ) 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: Token sequence x c of length N , with masks indicating the target areas for generation Init:{t(i)} T i=0 ← discretize([0, 1]), x t(T ) ← x c for i = T, T -1, . . . , 1 do t ← t(i), s ← t(i -1) for n ∈ [N ], if x (n) t = m, draw x (n) s ∼ Cat( αs-αt 1-αt µ (n) θ (x t , t) + 1-αs 1-αt e m ) else x (n) s ← x (n) t return x 0 .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 x c of length N , with masks indicating the target areas for generation Init:{t(i)} T i=0 ← discretize([0, 1]), x t(T ) ← x c for i = T, T -1, . . . , 1 do t ← t(i), s ← t(i -1) for n ∈ [N ] do (in parallel) draw u ∼ U [0, 1] if x (n) t = m and u < αs-αt 1-αt then draw x (n) s ∼ Cat(µ (n) θ (x t , t)) else x (n) s ← x (n) t return x 0 .However, mathetically it is equivalent to Alg. 2 and should produce identical results. Our investigation revealed that the issue arises because Alg. 2 scales the output probabilities of µ θ by a small factor αs-αt 1-αt as s is close to t, causing some categories to have very low probabilities. JAX, however, 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 is significant. In our experiment, we found that categories with probabilities below 1e-8 are rarely sampled out of a total of 50K categories. Thus, Alg. 2 implicitly performs top-p sampling (with a dynamic p) under JAX's categorical sampling, yielding better sample quality than Alg. 3 where µ θ is not scaled by a small factor and has fewer categories truncated.We first prove a lemma that the forward and reverse transition rate matrices. This follows from the results in [29] but we give a proof for completeness.Lemma 2. The forward transition rate matrix Q(t) and the reverse transition rate matrix (given x 0 ) R x0 (t) satisfy:EQUATION)Proof Consider the reverse transition from time t + τ to t. For j ̸ = k, Bayes' rule yieldsq(x t = j|x t+τ = k, x 0 ) = q(x t = j|x 0 )q(x t+τ = k|x t = j) q(x t+τ = k|x 0 ) = q(x t = j|x 0 )(δ jk + Q(t) jk τ + o(τ )) q(x t+τ = k|x 0 ) τ →0 = δ kj + q(x t = j|x 0 ) q(x t = k|x 0 ) Q(t) jk τ + o(τ ).Then, it follows from the definition of the transition rate matrix that R x0 (t) kj = Q(t) jk q t|0 (j|x0) q t|0 (k|x0) .Proposition 3. We use the shorthand R θ (t) kj to denote the approximate reverse transition rate from the state k to j obtained by substituting our prediction model µ θ (k) for x 0 in R x0 (t) kj . Then, the continuous-time objective (4) can be equivalently expressed asEQUATION)where C is a constant independent of θ.Proof To rewrite our objective L ∞ with the transition rate matrices, we first go back to (18) . There, instead of plugging in the explicit form of Rx0 (t, s), we substitute it with (7) which leverages the transition rate R x0 (t). To simplify the notation, we assume x t = k and use the shorthandR θ (t) kj ≜ R µ θ (k) (t) kj . We then have KL(q(x t-∆t |x t , x 0 )∥p θ (x t-∆t |x t )) = KL(Cat(x s ; Rx0 (t, t -∆t) ⊤ e k )∥Cat(x s ; Rµ θ (k) (t, t -∆t) ⊤ e k )) = m j=0 e ⊤ k (I + R x0 (t)∆t + o(∆t))e j log e ⊤ k (I + R x0 (t)∆t + o(∆t))e j e ⊤ k (I + R θ (t)∆t + o(∆t))e j = (1 + R x0 (t) kk ∆t) log 1 + R x0 (t) kk ∆t + o(∆t) 1 + R θ (t) kk ∆t + o(∆t) + j̸ =k (R x0 (t) kj ∆t) log R x0 (t) kj ∆t + o(∆t) R θ (t) kj ∆t + o(∆t) + o(∆t) = (R x0 (t) kk -R θ (t) kk )∆t + j̸ =k (R x0 (t) kj ∆t) log R x0 (t) kj ∆t + o(∆t) R θ (t) kj ∆t + o(∆t) + o(∆t).For the last identity, we have used the fact that log(1 + x) = x + o(x). To L ∞ , we take the limit of L T as T → ∞, which is equivalent to letting ∆t = 1/T → 0. We obtainL ∞ = lim T →∞ T i=2 E q(x t(i) |x0) [KL(q(x s(i) |x t(i) , x 0 )∥p θ (x s(i) |x t(i) ))] = lim T →∞ T i=2 E q(x t(i) |x0) R x0 (t(i)) kk -R θ (t(i)) kk + j̸ =k R x0 (t(i)) kj log R x0 (t(i)) kj ∆t + o(∆t) R θ (t(i)) kj ∆t + o(∆t) ∆t + o(∆t) = 1 t(1) E q t|0 (k|x0) R x0 (t) kk -R θ (t) kk + j̸ =k R x0 (t) kj log R x0 (t) kj R θ (t) kj dt.Note that R x0 (t) is a constant matrix independent of θ. Absorbing all constant terms into C, we haveL ∞ = - 1 t(1) E q t|0 (k|x0) R θ (t) kk + j̸ =k R x0 (t) kj log R θ (t) kj dt + C.Next, we subtitute R x0 (t) with the forward transition rate using Lemma 2:L ∞ = - 1 t(1) E q t|0 (k|x0) R θ (t) kk + j̸ =k Q(t) jk q t|0 (j|x 0 ) q t|0 (k|x 0 ) log R θ (t) kj dt + C = - 1 t(1) m k=0 q t|0 (k|x 0 )R θ (t) kk + m k=0 j̸ =k Q(t) jk q t|0 (j|x 0 ) log R θ (t) kj dt + C = - 1 t(1) m k=0 q t|0 (k|x 0 )R θ (t) kk + m k=0 j̸ =k Q(t) kj q t|0 (k|x 0 ) log R θ (t) jk dt + C,where the last identity used the discrete analog to integration-by-part (or summation-by-part):k=0 j̸ =k f (j, k) = k=0 j̸ =k f (k, j).Rearranging the terms then gives (20) .Campbell et al. [29] used the first term of (20) as the training loss. A key limitation of this loss function is from the inner summation termj̸ =k Q(t) kj log R θ (t) jk .For single dimension case, the sum is analytically computable due to the sparse structure of R θ (t)-if x t = k is mask, the second term disappears; otherwise the only possible neighbor j is a mask. However, for multidimensional data, j will represent all N -1 neighbors in the forward process, i.e., the states we get from mask out a single unmasked dimension of x t = k. Recall that R θ (t) jk is computed as substituting our neural network prediction model µ θ (j) for x 0 in R x0 (t) jk . Therefore, the summation together with R θ (t) kk requires N evaluations of µ θ (•). This is prohibitive since the neural network model is usually expensive. To resolve this issue, Campbell et al. [29] proposed to rewrite the sum asE j∼q(•|k) [Z k log R θ (t) jk ] where q(j|k) = Q(t) kj Z k , Z k ≜ j ′ ̸ =k Q(t) kj ′and estimate it through Monte Carlo. Taking into account the outer expectation under q t|0 (k|x 0 ), the computation of the loss then becomes a doubly stochastic estimate (using k ∼ q t|0 (k|x 0 ) and j ∼ q(j|k)) which suffers from large variance. In contrast, the form of our loss (4) only requires evaluating µ θ once for a single stochastic estimation of the expectation w.r.t. q(x t |x 0 ).We provide a simpler derivation of the score-based loss [32, 35] below. We start from the form the ELBO in (20) and rewrite it asL ∞ = 1 t(1)E q t|0 (k|x0)EQUATIONFor the last identity we used the zero-row-sum property of transition rate matrix:R x0 (t) kk = - j̸ =k R x0 (t) kj .If we plug ( 19) into ( 21) and reparameterize with a score modelEQUATIONwe recover the score entropy loss function from Lou et al. [32] , Benton et al. [35] :L ∞ = 1 t(1)E q t|0 (k|x0)j̸ =k Q(t) jk s θ (k) j - q t|0 (j|x 0 ) q t|0 (k|x 0 ) log s θ (k) j + ψ q t|0 (j|x 0 ) q t|0 (k|x 0 ) dt,where ψ(y) ≜ y log y -y. 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 x 0 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 x 0 :EQUATION)Note that only the result under the case k = m is needed. This is because when x t is unmasked, at any time between 0 and t, the state must stay unchanged and remain x 0 . As a result, KL(q(x t-∆t |x t , x 0 )∥p θ (x t-∆t |x t )) = 0 for x t ̸ = m. From ( 14), we know Q(t) jk = β(t)(δ mk -δ jk ). Combining ( 23) and (H.3), we getEQUATIONFurther, we can show the connection between ( 24) and ( 4) by reverting the score parameterization to a mean parameterization using ( 22), or equivalently s θ (x t ) j = αt 1-αt µ θ (x t ) ⊤ e j . By doing so, we obtainL ∞ = 1 t(1) β(t) E q t|0 (k|x0) δ mk j̸ =k s θ (k) j - α t 1 -α t x ⊤ 0 log µ θ (k) + α t 1 -α t dt.Observing thatEQUATIONwe conclude that this recovers the objective in (4). Interestingly, in Lou et al. [32] the score parameterization is not constrained to satisfy (25) . That means the learned reverse model might be incompatible with the forward process.Below, we prove Proposition 1 using the result from Eq. ( 23).Proof of Proposition 1q t (j) q t (m) = x0 q t|0 (j|x 0 )q(x 0 ) q t (m) = x0 q t|0 (j|x 0 )q 0|t (x 0 |m) q t|0 (m|x 0 ) = E x0|xt=m q t|0 (j|x 0 ) q t|0 (m|x 0 ) = E x0|xt=m α t 1 -α t x ⊤ 0 e j = α t 1 -α t E[x 0 |x t = m] ⊤ e j .H.4 Other related work.MaskGIT [39] . MaskGIT is a diffusion-inspired iterative denoising model for discrete image tokens obtained models such as VQ-VAE [70] . Training of MaskGIT follows the steps: (a) Sample t ∈ [0, 1]. (b) Given a mask scheduling function γ(t), sample γ(t)N tokens to place masks. (c) For data x 0 of size (m + 1) × N and the partially masked state x t , minimize the negative log-likelihoodEQUATIONOur forward process satisfies q t|0 (m|x 0 ) = 1 -α t . Therefore, when we set the mask scheduling function as γ(t) = 1 -α t we obtain a loss similar to (5) without the α ′ t 1-αt weighting. Note that there remains a difference in the sampling distribution of x t : in the masked diffusion forward process, tokens are sampled independently and do not necessarily yield exactly (1 -α t )N mask tokens at time t, though the expected number is (1 -α t )N . One might be interested in whether the uniform weighting can be recovered by selecting an appropriate schedule α t . However, solving α t such that α ′ t = α t -1 yields α t = ce t + 1 and there is no c that satisfies both α 0 = 1 and α 1 = 0. This shows that training with the MaskGIT loss (26) may not be faithfully optimizing the model likelihood.Discrete flow matching [48] . For the linear schedule α t = 1 -t, our reverse transition rate matrix (7) conditioned on x 0 is:R x0 (t) = - α ′ t 1 -α t e m (x 0 -e m ) ⊤ = 1 t e m (x 0 -e m ) ⊤ .This is the same as the conditional reverse transition rate used in Campbell et al. [48, Eq. ( 22)]-note that their time t is reversed, and the rate matrix was therefore in the form R x0 (t) = 1 1-t e m (x 0 -e m ) ⊤ .SDDM [30] . Sun et al. [30] proposed a pseudo-likelihood-like 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.SUNDAE [49, 71] . Unlike masked diffusion, SUNDAE uniformly corrupts data with random tokens in the vocab (known as uniform discrete diffusion [14] ). Additionally, it uses a 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 x t 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.Given the forward transition q(x t |x s ) and marginal q(x s |x 0 ) derived in main text (Sec. 6) The reversal givenx 0 is q(x s |x t , x 0 ) = Cat(x s ; Rx0 (t, s) ⊤ x t ) for Rx0 (t, s) jk =      αs-αt 1-αt ⊤ x 0 x ⊤ 0 e k j = m, k ̸ = m 1-αs 1-αt ⊤ x 0 j = m, k = m δ jk j ̸ = m.or alternatively can be written asEQUATIONTo simplify this expression we consider the two cases: either x t = m (i.e. x t is mask) or t ̸ = m where in the second case x t = x 0 . For the case x t = m, the denominator in ( 27) simplifies asq(x t = m|x 0 ) = 1 -α ⊤ t x 0due to x ⊤ 0 x t = 0 since x 0 ̸ = m, i.e. the observed token x 0 cannot be a mask. Then given that x t = m the probability thatEQUATIONwhile the remaining probability forEQUATIONThen, combining ( 28) and ( 29) to write q(x s |x t = m, x 0 ) in an unified way yields the expression (10) in the main Sec. 6. In the second case, when 27) simplifies dramatically and it becomes q(x s |x t ̸ = m, x 0 ) = x ⊤ t x s which is a point mass that sets x s = x t .x t = x 0 ̸ = m, q(x s |x t ̸ = m, x 0 ) from (Derivation of the continuous-time limit of the loss in (11) . To simplify the notation, we let ξ s,t ≜ αs-αt 1-αt . We first compute the KL divergence terms in the discrete-time ELBO asKL(q(x s |x t , x 0 )∥p θ (x s |x t )) = m xs=0 q(x s |x t , x 0 ) log q(xs|xt,x0) p θ (xs|xt) x t = m 0 x t ̸ = m = δ xt,m k̸ =m ξ ⊤ s,t x 0 x ⊤ 0 e k log ξ ⊤ s,t x 0 x ⊤ 0 e k ξ ⊤ s,t diag(µ θ (x t , t))e k + (1 -ξ s,t ) ⊤ x 0 log (1 -ξ s,t ) ⊤ x 0 (1 -ξ s,t ) ⊤ µ θ (x t , t) = δ xt,m -ξ ⊤ s,t x 0 x ⊤ 0 log µ θ (x t , t) + (1 -ξ s,t ) ⊤ x 0 log (1 -ξ s,t ) ⊤ x 0 (1 -ξ s,t ) ⊤ µ θ (x t , t).Let ∆ t ≜ 1 T = t(i) -s(i) for all i. Plugging α t-∆t = α t -α ′ t ∆t + o(∆t) into the above formula and letting γ t = α ′ t 1-αt , we getKL(q(x s |x t , x 0 )∥p θ (x s |x t )) = δ xt,m γ ⊤ t x 0 x ⊤ 0 log µ θ (x t , t)∆t + 1 + γ ⊤ t x 0 ∆t • log 1 + γ ⊤ t x 0 ∆t + o(∆t) 1 + γ ⊤ t µ θ (x t , t)∆t + o(∆t) + o(∆t) = δ xt,m γ ⊤ t x 0 x ⊤ 0 log µ θ (x t , t)∆t + 1 + γ ⊤ t x 0 ∆t γ ⊤ t x 0 ∆t -γ ⊤ t µ θ (x t , t)∆t + o(∆t) + o(∆t) = δ xt,m γ ⊤ t x 0 x ⊤ 0 log µ θ (x t , t)∆t + γ ⊤ t x 0 ∆t -γ ⊤ t µ θ (x t , t)∆t + o(∆t) = δ xt,m • γ ⊤ t (x 0 x ⊤ 0 log µ θ (x t , t) + x 0 -µ θ (x t , t))∆t + o(∆t).Therefore,lim T →∞ T i=2 E q(x t(i) |x0) [KL(q(x s(i) |x t(i) , x 0 )∥p θ (x s(i) |x t(i) ))] = lim T →∞ T i=2 E q(x t(i) |x0) [δ x t(i) ,m • γ ⊤ t (x 0 x ⊤ 0 log µ θ (x t(i) , t(i)) + x 0 -µ θ (x t(i) , t(i)))∆t + o(∆t)] = 1 t(1) γ ⊤ t E q(x t(i) |x0) [δ xt,m • (x 0 x ⊤ 0 log µ θ (x t , t) + x 0 -µ θ (x t , t))]dt.Letting t(1) → 0 proves the result.The model is applied to data consisted of N tokens where x 0 = (x 1 0 , . . . , (N ) 0 ) and where each state in the masked diffusion is x t = (x 1 t , . . . , x (N ) t). The reverse generated model has a factorizing transition conditional of the formN n=1 p θ (x (n) s |x t ) where p θ (x (n) s |x t ) = q(x (n) s |x (n) t , µ (n) θ (x t , t))has a form that depends on whether x(n) t = m or x (n) t ̸ = m.For the first case:p θ (x (n) s |x (n) t = m, {x (k) t } k̸ =n ) = 1 -α s 1 -α t ⊤ µ (n) θ (x t , t)e ⊤ m x (n) s + α s -α t 1 -α t ⊤ diag(µ (n) θ (x t , t))x (n) s ,where µ(n) θ (x t , t) = softmax(f θ (x t )) is a m + 1 dimensional probability vector modelled by a NN (where the final value is constrained to be zero since µ(n) θ (x t , t) is a reconstruction of x (n)0 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 x t of all tokens, while additional time features to encode t are also included. When x(n) t ̸ = m the reverse transition model is set to be p θ (x s |x (n) t ̸ = {x (k) t } k̸ =n ) = (x (n) t ) ⊤ x (n) s which matches precisely q(x (n) s |x (n) t = m, x (n) 0 ) = (x (n) t ) ⊤ x (n)s from the forward process. The full negative lower bound for state-dependent rates and assuming N tokens is given byL ∞ = 1 0 α ′ t 1 -α t ⊤ E q(xt|x0)    n:x (n) t =m (x (n) 0 -µ (n) θ (x t , t) + x (n) 0 (x (n) 0 ) ⊤ log µ (n) θ (x t , t))    dt.Given that each α t,i = 1 -t wi , the reverse model becomesp θ (x (n) s |x (n) t ̸ = m, {x (k) t } k̸ =n ) = e w log s t ⊤ µ (n) θ (x t , t)e ⊤ m x (n) s + 1 -e w log s t ⊤ diag(µ (n) θ (x t , t))x (n) s ,where w is the m + 1 dimensional vector of all w i s. Note that the probability of x θ (x t , t) i . The gradient wrt t is α ′ t,i = -w i t wi-1 andα ′ t,i1-αt,i = -wi t the above loss is written asL ∞ = - 1 0 1 t w ⊤ E q(xt|x0)    n:x (n) t =m (x (n) 0 -µ (n) θ (x t , t) + x (n) 0 (x (n) 0 ) ⊤ log µ (n) θ (x t , t))    dt,where w is the vector of all w i 's. An unbiased gradient over the NN parameters θ is straightforward to obtain since we just need to sample one time point t and an x t ∼ q(x t |x 0 ) to approximate the integral and expectation and then use the gradient:-∇ θ n:x (n) t =m 1 t w ⊤ x (n) 0 -µ (n) θ (x t , t) + x (n) 0 (x (n) 0 ) ⊤ log µ (n) θ (x t , t) .The gradient wrt the w parameters is more complex since these parameters appear also in the discrete distribution q(x t |x 0 ) which is not reparametrizable. To deal with this we need REINFORCE unbiased gradients [73, 74] , and in our implementation we consider REINFORCE leave-one-out (RLOO) [51, 52] with two samples. Firstly, the exact gradient wrt w of the exact loss is written asEQUATION)whereg(x t , x 0 ) = n:x (n) t =m (x (n) 0 -µ (n) θ (x t , t) + x (n) 0 (x (n) 0 ) ⊤ log µ (n) θ (x t , t)), f (x t , x 0 ) = w ⊤ g(x t , x 0 ).Note that g(x t , x 0 ) is a vector while f (x t , x 0 ) is a scalar. The left term in (30) easy since it just requires sampling t and x t ∼ q(x t |x 0 ), while the right term is the REINFORCE term which could have high variance. For this second term we use RLOO with two samples x 1 t , x 2 t and construct the unbiased estimate- 1 2t ∇ w log q(x 1 t |x 0 ) -∇ w log q(x 2 t |x 0 ) f (x 1 t , x 0 ) -f (x 2 t , x 0 ) .Thus, the overall unbiased gradient for w we use is- 1 2t g(x 1 t , x 0 ) + g(x 2 t , x 0 ) + ∇ w log q(x 1 t |x 0 ) -∇ w log q(x 2 t |x 0 ) f (x 1 t , x 0 ) -f (x 2 t , x 0 )In all experiments, the model is trained with a 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.We followed the standard dataset split as in Austin et al. [14] , Lou et al. [32] and trained our models on text chunks of length 256 for 1 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 α t = 1 -t and added a small shift ϵ = 10 -4 when t is close to 0 and 1 to ensure numerical stability. The shifted schedule isα t = (1 -2ϵ)(1 -t) + ϵ.The shift leads to a support mismatch between q(x 1 |x 0 ) and the prior p(x 1 ), leading to an undefined KL divergence term. We explain in app. E how to modify the prior distribution to allow small uniform probabilities in 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 so we can safely ignore them when reporting the ELBO.We used a cosine learning rate schedule with a linear warm up of 2000 steps. We applied channel-wise dropout of rate 0.05 and used AdamW optimizer with learning rate 0.0003 and a weight decay factor of 0.03. Our model is trained on 16 TPU-v5 lite for less than a day.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 et al. [32] and our tokenizer was also kept the same with a 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.We used the same linear masking schedule as in previous experiments in all likelihood results. We used the same U-Net plus self-attention architectures from the continuous diffusion model described in Kingma et 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 a 12-layer diffusion transformer [75] with 768 hidden dimension and 12 heads in the middle.For both datasets we used AdamW 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×64. The learning rate follows a cosine annealing after 100 warm up steps. Our CIFAR-10 model is trained on 32 TPU-v5 lite for 24 hours. Our ImageNet-64 × 64 model is trained on 256 TPU-v5 for 3.5 days.We have observed that the cosine schedule leads to 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 we used 8 residual blocks (in one side of the UNet structure) and a 20-layer diffusion transformer in the middle. All other configurations are kept the same.K.1 Sample quality evaluation by GPT-2We use the GPT-2 large model to evaluate the perplexity of samples generated by our model, following Lou et al. [32] . Results are shown in Fig. 7 . We provide more unconditional generation results from our pixel-level modeling experiments on ImageNet 64×64 in Fig. 8 . And so there's a lot that you don't know about that, so I think for me 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. It happened to happen at basically the same time.Nobody's even had a show or had a movie/come out of the movie, but ... VS: If I'm going to pretend I'm definitely not you and have to live through that stuff, I don't think I'm going to swallow that. I didn't expect it to do quite that long.There are always things now that happen with 'Deadwood' where you don't know where it's going to end up next time, but I think there are occasions now where we 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 there was a direct link. But there was almost a sense of not that we 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, we need to deal with. We also just have to accept the language, and there are a lot of things where we take from them and we do this what they did because we want toThe BDC doesn't have a particular plan-turns around for the end of the planned path, and says "nothing 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 as much a wish for the dogs, as cars: the district wants to go west foot a couple blocks south, to make the trail safer for dogs."I feel that the accessibility of the trail is pretty important. I think the education of the trail, and the uses along different routes are very important pieces of a balanced outcome," said Bushell.Trams coming off Route 1We share conditionally generated text samples by MD4 Medium in Fig. 9 and observe that slow unmasking near t = 1, enabled by the cosine schedule, tends to help produce more consist and meaningful samples than uniform unmasking counterpart. We carried out ablation study on the effect of discretization on zero-shot perplexity. Results are included in Tab. 6. Note that this is an inference ablation with the same trained model (MD4-S trained with the continuou-time objective). Justification: This paper is mostly theoretical and methodological. We do not see immediate societal impact of this work. However, we acknowledge that large scale implementation of our algorithm might suffer from the same societal biases as any other generative models.We make no assumption on the ordering of states -any permutation of the integers will give the same result.We intentionally choose the reverse model to factorize across dimensions because the true reverse transition q(xs|xt) factorizes in the continuous-time limit (as s approaches t).Perplexity of generated samples scored by a large language model such as GPT-2.We only need m learnable parameters wi, for i = 0, . . . , m -1, since x0 can never be the mask token. For the final mask dimension we can choose an arbitrary fixed value such as wm = 0. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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)$.
2. 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.
3. 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].
4. 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.
5. 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.

[1] Common Diffusion Noise Schedules and Sample Steps are Flawed, Lin et al., 2024
questions: questions: 1. GenMD4 has not been tested on image datasets. Could you please share the results of GenMD4 when applied to image datasets?
2. 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: limitations: The method is only applied to masked diffusions.
paper_summary: paper_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.
strength_weakness: strength_weakness: -  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.
- The state-dependent models have a tendency to overfit to dataset statistics, which can limit their effectiveness in zero-shot transfer tasks.
- 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.
questions: questions: - Could the authors provide more insights into the practical challenges faced during the implementation of the state-dependent masking schedules?
- How does the proposed model ensure consistency between the forward and reverse processes, and how does this impact training stability compared to SEDD? 
- 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?
- How sensitive is the proposed method to hyperparameter choices? Do multiple runs with the same hyperparameters yield consistent performance?
limitations: limitations: None
paper_summary: paper_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.
strength_weakness: strength_weakness: Weaknesses:

1. Weak empirical results
    1. The zero shot numbers for D3PM in Table 1 look fishy.  There are only 2 differences between Md4 and Absorbing State D3PM:
        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.
        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. 
        
        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.
        
    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].
2. 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 
    1. Validation ppl numbers on OWT
    2. zero-shot ppl numbers.

[1] Kingma, D., Salimans, T., Poole, B. and Ho, J., 2021. Variational diffusion models. *Advances in neural information processing systems*, *34*, pp.21696-21707.

[2] Sahoo, S., Gokaslan, A., Sa, C., Kuleshov, V., 2024.  Diffusion Models With Learned Adaptive Noise. arXiv:2406.07524
questions: questions: 1. Clarification on D3PM experiments in Table 1 as mentioned in the "weaknesses" section in the reviews.
2. 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: limitations: Yes
paper_summary: paper_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.

The key contributions includes:

1. 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.

2. 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.

3. 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.

4. 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.
strength_weakness: strength_weakness: 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.
2. 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.
3. 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.

[1] Vector Quantized Diffusion Model for Text-to-Image Synthesis
questions: questions: Please see weaknesses.
limitations: limitations: The authors have adequately described the limitations and potential negative societal impact of their work.

Comments: title: Response
Response
Response
Response
Update
Correction for Typos
comment: Thank you for the clarification. We agree that for masked diffusion, the first term becomes constant, meaning equation (7) is equivalent to the second term of equation (55). It's important to note that the transition rate matrix satisfies $R_{kk} = - \\sum_{j\neq k} R_{kj}$, which can be combined with the second term of equation (55). This further demonstrates the advantage of equation (55) over equation (54).
Perhaps we consider different terms. Utilizing the discrete "integration-by-parts" technique can reduce the variance of the second term.

 It is worth noting that due to the specific structure of the generator of masked diffusion, the term $R^{\theta}_ {kk}$ remains constant and does not depend on $\theta$, resulting in a variance reduction of the first term to zero. This unique feature is distinct from general discrete diffusion models. Exploring alternative classes of discrete diffusion models exhibiting this property presents a significant and meaningful research topic. 

I suspect that equation (7) is equivalent to the second term in equation (55), which could potentially unveil a deeper connection between the two equations.

Bests
Dear Reviewer,

We're pleased our rebuttal addressed your main concern and appreciate the increased score. We'll incorporate your suggestion to add a detailed multidimensional case derivation in the final version. Regarding your follow-up question:

> In Appendix G2, the reduction in variance is noticed, but the underlying reason deserves further elaboration.

We believe the primary reason for variance reduction is the discrete "integration-by-parts" that transformed equation (54) (Campbell et al. form) to (55) (which simplifies to our loss). This rewrite mainly changes $R_{\\theta}(t)\_{jk}$ to $R_{\\theta}(t)\_{kj}$ in the second term. Since the first index of $R_{\\theta}(t)$ is the network input and we sum over $j$, having $j$ on the second index avoids evaluating the network at different inputs. As you correctly pointed out, this issue only arises for multidimensional data, since in the single-dimension case there's only one possible $j \neq k$.  In the final version we'll revise this section using multidimensional derivations for clarity.

Best,
Authors
The authors' rebuttal addresses my main concern, and I will adjust my score accordingly.

However, there are still areas for improvement. 

1. In Appendix G2, the reduction in variance is noticed, but the underlying reason deserves further elaboration. The masked diffusion is an instance of a CTMT diffusion where the reverse transition rate $\hat{R}_ {kk}$ remains constant, resulting in the omission of $\hat{R}_ {kk}$ in the ELBO. This strategic approach could serve as a valuable technique for designing discrete diffusion models.

2. the paper solely provides proofs and derivations in one dimension, yet there are significant distinctions between one dimension and multidimensional scenarios. It is recommended that the authors include proofs and derivations for the multidimensional case in the final version of the paper.
Thank you for the acknowledgement and the positive assessment!
Dear Reviewer,

Thank you for raising the score! We are glad that the rebuttal addressed some of your concerns.

To address the other points you mentioned, could you clarify on the following points?
* Llama2-style transformers with RoPE embeddings are dense models. When applied to images (CIFAR-10 has a sequence length of 3072), they can be significantly more computationally expensive than the models in Table 3. All AR models in Table 3 are sparse or low-rank, while the diffusion models use UNet architectures rather than transformers to manage computational costs. Are you suggesting we compare to dense AR models? We believe such a comparison would be unfair in terms of computational requirements. 
* By non-AR non-sequence-based diffusion baselines, are you referring to continuous diffusion models? 

Best,
Authors
>The AR baselines use old transformer models hence the comparison isn't fair.

By newer models, I was referring to newer transformer architecture variations like the Llama 2 flavor of transformers or transformers, transformers with RoPE etc. The comparison to SEDD is fair comparison as noted, but it would be good to see more AR. Also non-AR non-sequence-based diffusion baselines should still be listed in Table 3 for the image models as the performance is nowhere near SOTA for non-sequence based non-AR models and it may confuse readers who are not familiar with generative models. Comparison to SOTA Gaussian diffusion models should be drawn, even though they are unfavorable to this class of the methods. 

I would encourage the author to study the numerical stability of their objective more in the final paper and it could make the submission much stronger.

I have raised my score 6 to reflect that some of my concerns have been addressed.
I acknowledge having read the authors' rebuttal. My overall assessment of the paper remains unchanged, and I continue to support my current rating.
Dear Reviewer,

We noticed a few typos in the original rebuttal and would like to correct them:
* It should be  $\\alpha_{t(1)} = 1 - 10^{-4}$ in the response to "The handling of 𝑝(𝑥0|𝑥𝑡(1)) could be enhanced" 
* The equation  in the response to "Appendix G2: Variance of Campbell et al. (2022)" should be
$$
\beta(t) \sum_{i: x_{t, i} \neq m} \log (x_{t, i}^\top \mu(x_t(i\to m))_i)
$$