Simple and Effective Masked Diffusion Language Models

Authors: subham sekhar sahoo
marianne  arriola
yair  schiff
aaron  gokaslan
edgar  marroquin
justin t chiu
alexander  rush
volodymyr  kuleshov
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 05:16:31

Abstract: While diffusion models excel at generating high-quality images, prior work reports a significant performance gap between diffusion and autoregressive (AR) methods in language modeling.
In this work, we show that simple masked discrete diffusion is more performant than previously thought.
We apply an effective training recipe that improves the performance of masked diffusion models and derive a simplified, Rao-Blackwellized objective that results in additional improvements.
Our objective has a simple form&mdash;it is a mixture of classical masked language modeling losses&mdash;and can be used to train encoder-only language models that admit efficient samplers, including ones that can generate arbitrary lengths of text semi-autoregressively like a traditional language model.
On language modeling benchmarks, a range of masked diffusion models trained with modern engineering practices achieves a new state-of-the-art among diffusion models, and approaches AR perplexity. We provide the code, along with a blog post and video tutorial on the project page: https://s-sahoo.com/mdlm

Main paper: 
 introduction: 
 Diffusion models excel at producing realistic, high-quality images and have received significant attention as potential tools for generating discrete data, such as text [1, 31, 33] , biological sequences [2, 47] , and graphs [60, 63] . Unlike autoregressive (AR) approaches, diffusion-based methods are not constrained to generate data sequentially, and therefore have the potential to improve long-term planning, controllable generation, and sampling speed. However, discrete diffusion methods exhibit a performance gap relative to AR models [1, 23, 26, 33] , especially in language modeling. The standard measure of language modeling performance is log-likelihood: when controlling for parameter count, prior work reports a sizable log-likelihood gap between AR and diffusion models.In this work, we show that simple masked diffusion language modeling (MDLM) combined with effective training recipes is more performant than previously thought [1, 26] . We develop a wellengineered MDLM implementation that significantly improves discrete diffusion log-likelihood; we further improve likelihood using a simple substitution-based parameterization of the reverse diffusion process that enables deriving a Rao-Blackwellized continuous-time variational lower bound (ELBO) with improved tightness [49] . Interestingly, our objective has a simple form: it is a weighted average of masked language modeling (MLM) losses [15] , and can be used to endow BERT-style, encoder-only models with principled generation capabilities. We complement this framework with efficient samplers-including ones that can generate semi-autoregressively like a typical language model. Our masked diffusion models achieve a new state-of-the-art among diffusion models on language modeling benchmarks and approach the perplexity of AR models within 15-25%. Surprisingly, simple engineering choices significantly improve performance in both our models and simple baselines that were previously thought to perform poorly. Our framework also extends to non-language domains, including biological sequence modeling. We pre-train DNA sequence models and observe similar or higher downstream performance compared to classical BERT-style training, while also introducing generative capabilities that classical masked DNA language models lack.Contributions We describe (1) a simple masked diffusion language modeling (MDLM) framework with a well-engineered implementation that outperforms all existing diffusion models across language modeling benchmarks (LM1B [8] , OWT [18] , DNA [12] ), and that significantly improves the performance of existing baselines [1, 26] . Our MDLM framework implements (2a) a substitution-based parameterization (SUBS) of the reverse unmasking diffusion process; SUBS allows us to derive (2b) a simple, continuous-time, Rao-Blackwellized objective that improves tightness and variance of the ELBO, further increasing performance. We complement MDLM with (3) fast samplers that support semi-autoregressive (SAR) generation and outperform previous SAR models. 
diffusion models: 
 Diffusion models are trained to iteratively undo a forward corruption process q that takes clean data x drawn from the data distribution q(x) and defines latent variables z t for t ∈ [0,1] that represent progressively noisy versions of x [27, 54, 56, 66, 48, 19] . The standard forward process for continuous x isEQUATIONwhere ϵ ∼ N (0,I) and (α t ) t∈[0,1] is a noise schedule, monotonically decreasing in t. The parameterized reverse diffusion model p θ over x and z t is trained to maximize a variational lower bound on loglikelihood (ELBO). Given a number of discretization steps T, defining s(i) = (i-1)/T and t(i) = i/T , and using D KL [•] to denote the Kullback-Leibler divergence, the Negative ELBO (NELBO) equals [54] :Eq -logp θ (x|z t(0) )L recons + T i=1DKL[q(z s(i) |z t(i) ,x)∥p θ (z s(i) |z t(i) )]L diffusion +DKL[q(z t(T ) |x)∥p θ (z t(T ) )]L prior (2) For brevity, we drop i from t(i) and s(i) below; in general, s will denote the time step before t.Applications of diffusion modeling to discrete data can be broken into two broad categories. First are works that embed discrete structures in continuous space and then perform the Gaussian diffusion defined above on these continuous representations [9, 16, 23, 24, 30, 34, 57] . More related to our method are works that define a diffusion process directly on discrete structures. D3PM [1] introduces a framework with a Markov forward process q(z t |z t-1 ) = Cat(z t ;Q t z t-1 ) defined by the multiplication of matrices Q t over T discrete time steps. This process induces marginalsEQUATION) that represent the discrete-state form of (1) . Extending this formalism to continuous time (as in (1) ) relies on continuous time Markov chain (CTMC) theory [5] . The CTMC framework in turns leads to generalizations of the score matching perspective on diffusion modeling [55] to discrete data [33, 59] . Notably, SEDD [33] connects score-based approaches with ELBO maximization, enabling performant likelihood-based training of score-based models. 
simple masked diffusion models: 
 While previous work on discrete diffusion supports general forward processes (e.g., general Q t in D3PM), absorbing state (i.e., masking) diffusion consistently achieves the best performance [1, 33] . In this work, instead of supporting general noise processes, we focus on masking and derive tight Rao-Blackwellized objectives that outperform general approaches and do not require CTMC theory. In this section, we first define the diffusion process for a categorical random variable. Later in Sec. 3.5, we extend this process to sequences containing multiple such categorical variables. We denote our overall approach as Masked Diffusion Language Models (MDLM).Notation. We denote scalar discrete random variables with K categories as 'one-hot' column vectors and define V ∈ {x ∈ {0,1} K : K i=1 x i = 1} as the set of all such vectors. Define Cat(•;π) as the categorical distribution over K classes with probabilities given by π ∈ ∆ K , where ∆ K denotes the K-simplex. We also assume that the K-th category corresponds to a special [MASK] token and let m ∈ V be the one-hot vector for this mask, i.e., m K = 1. Additionally, let 1 = {1} K and ⟨a,b⟩ and a⊙b respectively denote the dot and Hadamard products between two vectors a and b.We restrict our attention to forward processes q that interpolate between clean data x ∈ V and a target distribution Cat(.;π), forming a direct extension of Gaussian diffusion in (1) . Let q define a sequence of increasingly noisy latent variables z t ∈ V, where the time step t runs from t = 0 (least noisy) to t = 1 (most noisy). The marginal of z t conditioned on x at time t is q(z t |x) = Cat(z t ;α t x+(1-α t )π), (4) where α t ∈ [0,1] is a strictly decreasing function in t, with α 0 ≈ 1 and α 1 ≈ 0; see Suppl. E.1 for details. This implies transition probabilities q(z t |z s ) = Cat(z t ;α t|s z s + (1 -α t|s )π), where α t|s = α t /α s . This indicates that during each diffusion step from s → t, a fraction (1-α t|s ) of the probability mass is transferred to the prior distribution π. The reverse posterior is given as (see Suppl. 16 for details):EQUATIONWhile ( 4) and ( 5) represent a special case of the more general diffusion processes proposed in D3PM [1] , we show below that they yield a simplified variational lower bound objective and admit straightforward continuous time extensions.Next, we focus on masking processes and derive a simple Rao-Blackwellized objective for this choice of q. This objective incurs lower variance during training and improves tightness.In masked (i.e., absorbing state) diffusion, we set π = m. At each noising step, t, the input x transitions to a 'masked' state m with some probability. If an input transitions to m at any time t ′ , it will remain in this state for all t > t ′ : q(z t | z t ′ = m) = Cat(z t ;m). At time T , all inputs are masked with probability 1.The marginal of the forward process ( 4) is given by q(z t |x) = Cat(z t ; α t x + (1 -α t )m). Using properties of the masking process, the posterior q(z s |z t ,x) simplifies (5); see Suppl. A.2:EQUATIONThe reverse process inverts the noise process defined by q. We consider both a finite number of steps T , as well as a continuous time model corresponding to T → ∞. We begin with the discrete-time case for which the generative model is expressed as p θ (x) = z p θ (z 1 )p θ (x|z 0 ) T i=1 p θ (z s |z t )dz 0:1 . The optimal form for p θ (z s |z t ) matches the true posterior in (6) : this follows immediately from the definition of the diffusion objective in (2) , which is a sum of terms of the form D KL (q(z s |z t ,x)∥p θ (z s |z t )). However, ( 6) is conditioned on x, which we do not know. Therefore, we introduce a model x θ (z t ,t) : V ×[0,1] → ∆ K that approximates x with a neural network. We can also omit explicit dependence of x θ on time t, which simplifies sampling, yielding a 2x inference speed-up (see Suppl. E.2).The specific parameterization for p θ (z s |z t ) that we use isEQUATIONFurthermore, we induce 2 key properties of the absorbing state diffusion process into our denoising model, x θ (z t ,t): an unmasked token remains unchanged during reverse diffusion, and the clean input is never masked. We implement these as substitutions to the output of x θ (z t ,t), hence we call our parameterization SUBS.Zero Masking Probabilities First, notice that by definition, ⟨x,m⟩ = 0. For this reason, we design the denoising network such that ⟨x θ (z t ,t),m⟩ = 0, i.e., we substitute the logit index corresponding to the [MASK] token with -∞.Carry-Over Unmasking Second, if z t is unmasked, then we desire x θ (z t ,t) = z t , i.e., unmasked latents are 'carried over'. We accomplish this by substituting the output of our network to simply copy unmasked inputs.In Suppl. B.1, we show that "Zero Masking Probabilities" property simplifies the D3PM's NELBO (39) to (41) , and "Carry-Over Unmasking" futher simplifies (41) to (43) whose continuous time equivalent is the simplified NELBO (10) . Table 8 shows that each simplification leads to an improved likelihood.Recall from (2) that the diffusion traning objective has the form L recons + L diffusion + L prior . For the simplified reverse process in (7) , the discrete-time diffusion loss for finite T simplifies to (Suppl. B.1.3):EQUATIONNote that this objective is simpler and more well-behaved than the expression one would obtain for D KL (q(z s |z t ,x)∥p θ (z s |z t )) under the parameterization induced by using p θ (z s |z t ) = q(z s |z t ,x = x θ (z t ,t)) from ( 5), which is similar to what is used by D3PM [1] (see Suppl. A.2.4):EQUATIONWe refer to the process of obtaining (8) in lieu of (9) as a form of Rao-Blackwellization. Specifically, we analytically compute expectations such as ⟨x θ (z t ,t),m⟩ = 0 in order to simplify objective (9) to obtain (8) . Without analytical simplifications, a model must learn θ such that ⟨x θ (z t ,t),m⟩ = 0 holds. Unlike in regular Rao-Blackwellization, simplifications are possible because of modeling choices for x θ (z t ,t) (zero masking probabilities and carry-over unmasking). In that sense, our approach has similarities to graphical modeling, where incorporating conditional independencies into p θ sets certain log-likelihood terms to zero. However, our approach also empirically helps reduce variance, hence we refer to it as Rao-Blackwellization, somewhat abusing the usual terminology.Previous works have shown empirically and mathematically that increasing the number of steps T yields a tighter approximation to the ELBO [29] . Following a similar argument, we form an continuous extension of ( 8) by taking T → ∞ (see Suppl. B.2), which yields the following NELBO, L ∞ NELBO :EQUATIONInvariance to the noise schedule The function α t is invertible due to the monotonicity assumption in Sec. 3.1, and so we can perform the following change of variables in (10) :γ ≡ log(1 -α t ).Thus, the diffusion loss can be equivalently expressed asL ∞ NELBO = -E q γ=0γ=-∞ log⟨x θ (z γ ,γ),x⟩dγ; see Suppl. E.1.1 for details. This new formulation demonstrates that the diffusion loss is invariant to the functional form of α t , which we verify empirically in Suppl. E.1.Next, we apply masked diffusion to language modeling over sequences x 1:L of L tokens, with x ℓ denoting the ℓ-th token. We make the assumption that the forward noising process is applied independently across a sequence and that, conditioned on a sequence of latents z 1:L t , the denoising process factorizes independently across tokens, i.e., p θ (z1:L s | z 1:L t ) = L ℓ=1 p θ (z ℓ s | z 1:L t ). Thus, we use a single model to compute x ℓ θ (z 1:L t ,t) for each ℓ from a masked sequence z t , optimizing:EQUATIONInterestingly, our objective has a simple form: it is the weighted average of masked language modeling (MLM) losses [15] . Thus our work establishes a connection between generative diffusion models and encoder-only BERT models. Our objective enables principled selection of a (randomized) masking rate, and also endows BERT-style models with principled generation capabilities; see Sec. 6. The full training algorithm is provided in Suppl. B.3.Note: Although (11) imposes a loss on all tokens, unmasked tokens don't contribute to the loss, as they are copied over by the denoising network due to "carry-over unmasking" (Sec. 3.2.3), effectively reducing log⟨x ℓ θ (z 1:L t ,t),x ℓ ⟩ to zero.One of the key contributions of our work is a well-engineered implementation of masked diffusion models. Our experiments demonstrate that these improvements greatly boost performance even for methods previously thought to perform poorly, e.g., Austin et al. [1] . Below we briefly summarize these implementation details. First, we find that tokenization is critical to performance. Small vocabularies, such as the 8k vocabulary in Austin et al. [1] , result in longer-range dependencies that decrease the performance of both diffusion and AR models. Additionally, by focusing on masked diffusion, we are able to provide a numerically stable implementation of the objective function. Namely, since previous formulations of discrete diffusion were constructed to accommodate a wide range of limiting distributions [1] , the objective was implemented by materializing the full transition matrices Qt and posterior probabilities. In contrast, we evaluate D KL [q(z s | z t ,x)||p θ (z s | z t )] by examining only the masked token indices rather than comparing the full true and approximate posterior distributions.Furthermore, we modernize the architecture for the denoising network relative to D3PM [1] . In lieu of the T5 architecture used in D3PM, we use the diffusion transformer (DiT) introduced in Peebles & Xie [42] , which integrates time step conditioning into a standard encoder-only transformer [62] and uses rotary positional embeddings [58] . In addition, we implement a low-discrepancy sampler that reduces the variance of the ELBO, similar to Kingma et al. [29] and draws correlated samples t i rather than performing i.i.d. sampling.4 Inference and Sampling in Masked Diffusion Language Models 
efficient ancestral sampling: 
 To generate a sequence of length L, the reverse diffusion process starts with the sequence z 1:L t=1 where z ℓ t=1 = m, for all ℓ ∈ {1,...,L}. Then the subsequent latents, z 1:L t are generated by discretizing the reverse diffusion process with some finite T. Given z 1:L t , we construct z 1:L s by sampling each token z ℓ s independently from the distribution p θ (z ℓ s |z 1:L t ) given in (7) . Note that in the reverse process, unmasked tokens remain unchanged. Thus, if no new tokens in z 1:L s become unmasked (which can occur often in early denoising stages for large T ), then z 1:L s = z 1:L t . Additionally if the denoising model, x θ (z 1:L t ) is not conditioned on time, then we can simply draw a new sample from p θ (z 1:L s-1/T |z 1:L s ) using the previously computed and cached value x θ (z 1:L t ). This means we have effectively "skipped" over the time step s, saving a function call to the denoising network. Note that SEDD [33] does not support this caching because the denoising network models time-dependent rates, which requires conditioning on time.Our method also admits an effective semi-autoregressive (SAR) decoding method that allows the model to generate sequences of arbitrary length [24, 52, 53] . Let x1:L represent the output from sampling a sequence of L tokens using the reverse diffusion process described above. To generate additional L ′ < L tokens, we propose a generation algorithm in which the latter L-L ′ tokens xL ′ :L are used as a prefix for an additional round of generation. Given the carry-over unmasking described in Sec. 3.2.3, these prefix tokens will simply be copied over at each decoding step. The remaining tokens are generated as above withz ℓ s ∼ p θ (z ℓ s | z L:L+L ′ t) for all ℓ ∈ {L+1,...,L+L ′ }, with z L-L ′ :L t=1 initialized to xL-L ′ :L as opposed to being initialized as masked tokens m. At the end of this process, we have produced L+L ′ tokens concat[x 1:L ,x L+1:L+L ′ ], where concat[•] denotes concatenation along the sequence length dimension. This process can repeat indefinitely, with the prefix shifted for every new round of generation. 
masked diffusion language models: 
 Experimental Setup We evaluate MDLM as a generative model of language and as a representation model via fine-tuning on downstream tasks.For language modeling likelihood evaluation, we conduct experiments on two datasets: The One Billion Words Dataset (LM1B; [8] ) and OpenWebText (OWT; [18] ). We use the bert-base-uncased tokenizer for LM1B, and report perplexities on the test split. Models have a context size of 128. For OWT, which does not have a pre-defined split, we reserve the last 100K documents as a held-out validation set and report perplexities on this set. We use the GPT2 tokenizer [45] for OWT. Models have a context size of 1,024. We utilize the transformer architecture from Lou et al. [33] , which augments the diffusion transformer [42] with rotary embeddings [58] . MDLM was trained for 1M or 10M steps (corresponding to 33B, 327B tokens, respectively) on LM1B and 1M steps on OWT (which corresponds to 262B tokens). The corresponding AR baseline was trained for half the number of steps For representation learning, we pre-train models on the C4 dataset [46] , then fine-tune and evaluate models on the GLUE benchmark [65] . Models have a context size of 128. We use the bert-base-uncased tokenizer for the representation learning experiments. We utilize the MosaicBERT architecture from Portes et al. [43] , an extension of the original BERT architecture [15] . We pre-train a bidirectional MosaicBERT using an MLM objective for 37B tokens of C4, as well as a causal variant on the same data. We further fine-tune MosaicBERT model using the MDLM for 327M tokens, less than 1% of the pre-training data. We provide the full hyperparameters in Suppl. D.6.Likelihood Evaluation On LM1B, MDLM outperforms all previous diffusion methods (Table 1 ). Compared to the SEDD baseline reported by Lou et al. [33] , trained for 33B tokens, MDLM, which we train for the same amount, achieves a 17% improvement on the perplexity bound. Finally, MDLM gets within 14% of an AR baseline and continues to improve with more training. We see the same trend for models trained on OWT, a larger dataset, shown in Table 2 -MDLM outperforms prior diffusion methods, closing the gap towards AR models. In Table 12 we find that models trained with and without time conditioning attain similar perplexities on OWT. Additionally, Figure 3 demonstrates the reduced variance we achieve from our objective, when compared to previous masked diffusion models such as SEDD [33] . Zero-Shot Likelihood Evaluation We also explore models' ability to generalize by taking models trained on OWT and evaluating how well they model unseen datasets. We compare the perplexities of our MDLM with SEDD [1] and an AR Transformer language model. Our zero-shot datasets include the validation splits of Penn Tree Bank (PTB; [36] ), Wikitext [38] , LM1B, Lambada [41] , AG News [68] , and Scientific Papers (Pubmed and Arxiv subsets; [10] ). Full experimental details are available in Suppl. D.4.In some cases, e.g., for Lambada and Scientific Papers, MDLM attains better perplexity than AR. We hypothesize that these datasets are farther from OWT, and that diffusion models may be more robust to out-of-domain evaluation due to the unmasking-based objective.We find that BERT fine-tuned with MDLM to be a generative model results in strong perplexities while preserving performance on downstream tasks. On the C4 validation set, the AR model attains perplexity (PPL) of 22, the pre-trained BERT attains a PPL upper bound of 78 (evaluated using the MDLM variational bound), and BERT + MDLM-FT attains a PPL upper bound of 35. In Table 4 , we further find that BERT + MDLM fine-tuning has no degradation in downstream GLUE performance compared to the BERT initialization. While the perplexity of our method is higher than the AR baseline, the downstream task performance is significantly better. We generate 200 sequences of length 2048 tokens on a single 3090 GPU and evaluate generative perplexity under a pre-trained GPT-2 [45] model. The SSD-LM sequences are generated using blocks of 25 tokens (as implemented in their pre-trained model) and the MDLM sequences are generated using L ′ = 512. In Table 5 , we find that in addition to achieving better generative perplexity, MDLM enables ∼25-30x faster SAR decoding relative to SSD-LM.We also explore applications to the generative modeling of biological sequences [14, 47] using a state space model (SSM) backbone [22] . Namely, we build on the recently-proposed Caduceus DNA language model [50] , which uses as a backbone the data-dependent SSM Mamba block [21] .Experimental Setup We pre-train the encoder-only Caduceus [50] , which is an MLM, on the HG38 human reference genome [11] and perform fine-tuning using our diffusion parameterization. We use a context length of 1024 tokens and follow Schiff et al. [50] for the experimental setup, other than learning rate which was reduced to 1e-3. See Suppl. D.7 for full experimental details. We assess both generative performance using perplexity and downstream performance on Genomics Benchmarks [20] across language diffusion paradigms and AR models.Generative Performance We fine-tune the Caduceus MLM across diffusion parameterizations and compare perplexities against AR models. We report perplexity values in Table 6 . MDLM outperforms all other diffusion language modeling schemes.We perform downstream evaluation with the Genomics Benchmarks [20] , a recently proposed benchmark with eight regulatory element classification tasks. As shown in Table 7 , our generative fine-tuning paradigm preserves or improves upon downstream performance from MLM pre-training. Absorbing-state diffusion methods outperform Plaid across tasks except for the simplest task Human vs. Worm, where all methods have roughly the same performance. For tasks where the input is a biased subsample of the full genome, we observe that the correlation between perplexity and downstream performance is weaker; see Suppl. D.7. In Table 8 , we can see the effect of our streamlined masked diffusion implementation. The improvements described in Sec. 3.5.1 allow us to greatly reduce perplexity of previously discounted models, such as D3PM (see the bottom row of this table, which is mathematically equivalent to the D3PM formulation). While most works assumed that D3PM achieves mediocre log-likelihoods, we show that is incorrect: our re-implementation almost matches state-of-the-art score-based methods. This introduces a new strong baseline that opens new research opportunities. Additionally, in Table 8 , we ablate different components of MDLM. We observe that the perplexity for MDLM trained with a discrete T = 1000 marginally worsens by 0.1 compared to MDLM trained in continuous time. Additionally, removing the "carry over" operation from the SUBS parameterization increases the perplexity by 1.5 points. However, further removing the "zero masking" operation does not lead to any meaningful change in perplexity. We provide further ablations for the continuous time formulation in the Appendix, showing in Table 11 that for a pre-trained model, at inference, increasing T yields better likelihoods. 
related work: 
 Comparison to D3PM Masked diffusion is a strict subset of D3PM [1] ; setting Q t|s = α t|s I+(1α t|s )1m ⊤ in their framework yields our forward diffusion. We improve over D3PM in three ways: (1) we adopt the SUBS parameterization; (2) this allows us to derive a simplified objective that analytically simplifies certain expectations to zero; (3) we adopt well-engineered training recipes that improve performance. Both (1) and ( 2) are possible because we focus on masking instead of developing a general discrete diffusion framework. Surprisingly, (3) has the largest contribution to performance.Comparison to CTMC Most implementations of diffusion work best in continuous time. However, extending D3PM in this way requires computing the limit of the product of an infinite number of matricesQ T •Q T -1 •••Q t as T → ∞,which requires advanced CTMC theory [5] . Our work describes simple continuous-time formulations for the most common noise processes (e.g., masking and uniform π), thus helping make an important part of the literature more accessible. In Suppl. C, we show that our results are compatible with CTMC, using the rate forward matrix R t = α ′ t αt (I-1m ⊤ ) and the reverse rate Rt (y ′ ,y) for the transition y → y ′ , where y,y ′ ∈ V:EQUATIONComparison to Score Estimation Score-based approaches to diffusion [55] extend to discrete states, although they typically further build upon advanced CTMC theory. In particular, SEDD [33] optimizes an ELBO 3 that is a function of the score model, obtaining state-of-the-art log-likelihoods among diffusion models. Our approach, however, is much simpler and does not require advanced theory. Furthermore, we can extract the score for MDLM (76), as demonstrated in Suppl. C.3, making it compatible with various techniques designed for score-based algorithms, such as samplers [5] , score parameterization [33] , efficient designs of the denoising network [59] , guidance techniques, and more.Comparison to BERT Our work provides a principled way of making BERT generative when trained with randomized masking rates. Previous work on generating from BERT used Gibbs sampling or ad-hoc methods [17, 32, 64] . The connection between BERT and diffusion was first made by Austin et al. [1] : their objective effectively involves unmasking. He et al. [26] additionally starts training from a pretrained BERT. However, both works use an objective that is similar to (9) , which is less numerically stable than our objective (see Section 3.5.1). Austin et al. [1] mention in their appendix that their ELBO simplifies to a weighted masking (MLM) loss similar to (8) , but it uses a more complex formula for the weights and is limited to the discrete time setting unlike our work. Furthermore, they do not train with that objective. Our work derives a simpler expression for the average of MLM losses, implements it, and obtains better likelihoods.In contrast to this work, which defines diffusion over discrete structures, Plaid [23] and Diffusion LM [30] define a Gaussian diffusion process over word embeddings. Zhang et al. [67] and Hu et al. [28] extend this approach to flow matching over word embeddings, enabling the design of faster samplers. Discrete Flow Matching (DFM) [6] applies flow matching to discrete structures, using a cross-entropy loss as their training objective: -E q,t logp θ (x 1:L |z 1:L t ). Similar to Chang et al. [7] , DFM's objective, while effective, is not weighted to serve as a proper ELBO. In MDLM, however, we derive a tight, principled lower bound on the log-likelihood.Concurrent Works Concurrent to our work, Shi et al. [51] and Ou et al. [40] derive a similar simplified objective for masked diffusion processes. While Ou et al. [40] start from a score matching perspective, we tackle this problem from a variational lens similar to Shi et al. [51] . Similar to Ou et al. [40] , we formulate efficient samplers in Section 4.1 by leveraging a time-independent denoising network.A key differentiation between our work and that of Shi et al. [51] , Ou et al. [40] is the semi-autoregressive decoding method we present in Section 4.2. While [51, 40] are restricted to sample sequences of a fixed length, we propose samplers to generate arbitrary lengths of text like a traditional language model. Furthermore, we establish the connection between our simplified objective and the masked language modeling (MLM) objective. As a result, we endow BERT-style models with principled generation capabilities while maintaining representation learning capabilities. Whereas [51, 40] only evaluate on NLP datasets, we show that masked diffusion is also effective in modeling biological sequences. 
conclusion: 
 In this work, we explore masked diffusion. With a well-engineered implementation that supports a simple variational objective, we attain state-of-the-art diffusion perplexities on language benchmarks and demonstrate how to efficiently convert BERT-style encoders into generative models. Given we are working on language modeling, we carry any of the inherent risks and opportunities that come with this line of research. A.1.2 q(z s |z t ,x)Austin et al. [1] show that the posterior corresponding to ( 14) is given as follows:EQUATIONwhich we simplify to the following:EQUATIONA.1.3 p θ (z s |z t )Austin et al. [1] approximate the reverse process in the following manner:EQUATION)wherex θ (z t ,t) : V ×[0,1] → ∆ K is an approximation for x. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] . Justification: We conduct all experiments on 8x 3090s, 8xA6000s, 8xA100s, or 8xH100s. The largest models on OpenWebText take 2 weeks to train on 8xA100, the LM1B models only take 2 days to train on the same hardware Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? Answer: [Yes] Justification: We follow standard practices Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]Justification: We provide no new assets.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: [NA]  
 
 Appendix: For the absorbing state diffusion process we have π = m.A.2.1 q(z s |z t ,x)Since, z t ∈ {x,m}, takes only 2 values we consider the separate cases: z t = x and z t = m.Case 1. Consider the case z t = x i.e. z t is unmasked. From ( 16), we have the following:q(z s |z t = x,x) = Cat z s ; [α t|s x+(1-α t|s )1m ⊤ x]⊙[α s x+(1-α s )m] α t x ⊤ x+(1-α t )x ⊤ m = Cat z s ; [α t|s x]⊙[α s x+(1-α s )m] α t ∵ x ⊤ m = 0 = Cat z s ; α t x α t ∵ x⊙m = 0 and αt = α t|s αs = Cat(z s ;x) ∵ αt = α t|s αs (18)Thus, we have the following:EQUATION)Case 2. Consider the case z t = m. By substituting z t = m and π = m in (16), q(z s |z t ,x) simplifies to the following:EQUATIONNote that the above categorical distribution is non-zero for z s ∈ {x,m} and zero for every other value. The non-zero values are specified as follows:EQUATIONCombining Cases 1 and 2, we get:EQUATIONA.2.2 p θ (z s |z t )For the absorbing state diffusion process with π = m, we want to simplify the (17) . For this reason, we consider 2 cases: first, when z t ̸ = m (case 1), second, when z t ̸ = m (case 2).Case 1. Consider the case when z t ̸ = m. (17) simplifies to the following:p θ (z s |z t ̸ = m) = Cat z s ; Q t|s z t ⊙Q ⊤ s x θ (z t ,t) z ⊤ t Q ⊤ t x θ (z t ,t) = Cat z s ; Q t|s z t ⊙Q ⊤ s x θ (z t ,t) [Q t z t ] ⊤ x θ (z t ,t) = Cat z s ; [α t|s z t ]⊙[α s I+(1-α s )m1 ⊤ ]x θ (z t ,t) [α t z t ] ⊤ x θ (z t ,t) = Cat z s ; [α t|s z t ]⊙[α s x θ (z t ,t)+(1-α s )m⟨1,x θ (z t ,t)⟩] α t ⟨z t ,x θ (z t ,t)⟩since ⟨1,x θ (zt,t)⟩ = 1, we have the following:= Cat z s ; [α t|s z t ]⊙[α s x θ (z t ,t)+(1-α s )m] α t ⟨z t ,x θ (z t ,t)⟩since zt ⊙m = 0, we have the following:EQUATIONCase 2. Consider the case when z t = m. (17) simplifies to the following:EQUATIONNote that the above categorical distribution, we can obtain the values for p θ (z s = x|z t = m) and p θ (z s = m|z t = m) which are as follows:EQUATIONEQUATIONAs a sanity check, we can verify that (26) reduces to (21) , and (27) reduces to (22) if our denoising network can reconstruct x perfectly, i.e., x θ (z t ,t) = x.Combining ( 24) and ( 25), we get the following expression for the reverse process parameterization:EQUATIONFor a given T , LetEQUATIONT ,...,1} E q(zt|x) T D KL (q(z s |z t ,x)∥p θ (z s |z t )) denote the diffusion loss. We break down the computation of D KL (q(z s |z t ,x)∥p θ (z s |z t )) into 2 cases: z t = x (case 1) and z t = m (case 2).Case 1: consider the case z t = x. Let's simplify D KL (q(z s |z t = x,x)∥p θ (z s |z t = x)).D KL (q(z s |z t = x,x)∥p θ (z s |z t = x)) = D KL (z t ∥z t )From ( 23) and (24) = 0Case 2: Consider the case z t = m. Let's simplify D KL (q(z s |z t = m,x)∥p θ (z s |z t = m)).D KL (q(z s |z t = m,x)∥p θ (z s |z t = m)) = zs q(z s |z t = m,x)log q(z s |z t = m,x) p θ (z s |z t = m) = zs∈{x,m} q(z s |z t = m,x)log q(z s |z t = m,x) p θ (z s |z t = m) = q(z s = x|z t = m,x)log q(z s = x|z t = m,x) p θ (z s = x|z t = m)Simplify using ( 21) and ( 26)+q(z s = m|z t = m,x)log q(z s = m|z t = m,x) p θ (z s = m|z t = m)Simplify using ( 22) and ( 27)EQUATIONThus, D KL (q(z s |z t ,x)∥p θ (z s |z t )) can be written in the following manner where ⟨z t ,x⟩ evaluates to 1 if z t = x and ⟨z t ,m⟩ evaluates to 1 if z t = m:EQUATION⟨z t ,x⟩+D KL (q(z s |z t = m,x)∥p θ (z s |z t = m))Given by (30) ⟨z t ,m⟩ (31)Thus, we derive the diffusion loss, L T , in the following manner:EQUATIONNote that L T is 0 if z t is an unmasked token i.e. z t = x.Austin et al. [1] , Sohl-Dickstein et al. [54] model α i as (α i ) i∈{1,...,T } = 1 -i T given latents z 1,...,T . However, in this paper, we denote the latents as z t(0),...,t(T ) ; and hence, the α t(i) are given as follows:(α i ) i∈{1,...,T } = 1- i TFrom Austin et al. [1] , Sohl-Dickstein et al. [54] .=⇒ (α i ) k∈{1,...,T +1} = 1- i T +1 For T +1 latents =⇒ (α i ) i∈{0,...,T } = 1- i+1 T +1Offsetting the indices by 1.t(i) ) i∈{0,...,T } = 1- i+1 T +1Switching the notations from αi to α t(i) .(Consequently, from Equation 33, we derive thatEQUATION)EQUATION)Thus we have the following:EQUATIONEQUATIONEQUATIONThe NELBO (2) simplifies to the following:E q   -logp θ (x|z t(0) )+ L T Compute using (32)   +D KL [q(z t(T ) |x)∥p θ (z t(T ) )]=0 using ( 37) and ( 38)EQUATIONIn this section, we show how SUBS parameterization can simplify the functional form of the NELBO as defined in (39) .We employ the RB techniques as described in Sec. 3.2.3 to simplify the NELBO (39) to (41) using RB2, and further to (43) using RB1.Using "Zero Masking Probabilities" (RB2) from Sec. 3.2.3, we set ⟨x θ (z t ,t),m⟩ = 0 in (32) to obtain the following simplified diffusion loss:EQUATIONThe corresponding Rao-Blackwellized NELBO is given as:E q   -logpθ(x|zt(0))+ L RB2 T Compute using (40)   +DKL[q(zt(T ) |x)∥p θ (z t(T ) )]=0 using ( 37) and ( 38)EQUATIONNotice that the term ⟨z t ,m⟩ in ( 40) is intended to reduce the diffusion loss to zero when z t = x. Now, we will demonstrate that, by applying "Carry Over Unmasking" (RB1) from Sec. 3.2.3, ⟨z t ,m⟩ can be removed from (40) .Recall that RB1 guarantees x θ (z t ,t) = x when z t = x. Thus, with the RB1 parameterization, the diffusion loss in (40) becomes zero for z t = x, as log⟨x θ (z t ,t),m⟩ = 0. Consequently, ⟨z t ,m⟩ can be safely omitted from (41) , yielding the following diffusion loss:EQUATIONB.1.3 NELBO Thus, we have the following NELBO:E q   -logpθ(x|zt(0))+ L RB2 + RB1 T Compute using (42)   +DKL[q(zt(T ) |x)∥p θ (z t(T ) )]= 0 using ( 37) and ( 38)EQUATIONComparing ( 43) and (41) . Note that due to RB1, logp θ (x|z t(0) ) in (43) reduces to 0 every time z t(0) = x as explained in (45) . However, this is not the case in (41) , even though it has a functionally similar expression to (43) . Because of this reason (43) should lead to a better likelihood estimate and we empirically verify this in Table 8 .To derive the continuous-time diffusion loss, L ∞ diffusion , we consider the limiting case lim T →∞ L RB2 + RB1 T (42):EQUATIONFor the continous time case, from (36), we haveEQUATIONThus, the reconstruction loss reduces to 0 in the following manner:L recons = -logp θ (x|z t(0) ) = -logp θ (x|z t(0) = x) From (45) = -log⟨x θ (x,t(0)),x⟩ = -log⟨x,x⟩Due to "carry-over unmasking" x θ (x,t(0EQUATIONThus, we have the following NELBO:E q   -logpθ(x|zt(0)) = 0 from (46) + L ∞ diffusion Compute using (42)   +DKL[q(zt(T ) |x)∥p θ (z t(T ) ,t)]= 0 using (37) and ( 38)EQUATIONIn Algorithm 1, we present the training algorithm for MDLM.Algorithm 1 Training MDLM 1: repeat 2: x 1:L ∼ q(x) ▷ Sample a sentence. 3: t ∼ U[0,1] ▷ Sample a time step. 4: z ℓ t ∼ Cat(z ℓ t ;α t x ℓ +(1-α t )m) ∀ 1 ≤ ℓ ≤ L ▷ Mask Each tokenx ℓ independently to obtain the latent z 1:L t . 5: Take gradient descent step on∇ θ α ′ t 1-α t ℓ log⟨x ℓ θ (z 1:L t ,t),x ℓ ⟩ 6: until convergedIn contrast, Campbell et al. [5] and Lou et al. [33] defined the discrete diffusion process as a Continuous-Time Markov Chain (CTMC), where the forward corruption process is specified by the rate change matrix R t ∈ R |V|×|V| , which can be thought of as the instantaneous rate at which one state transitions to another. With this formulation, the forward posterior q(z t |z s ) and the true reverse posterior q(z s |z t ,x) can be expressed in terms of the rate change matrix as follows:EQUATIONEQUATION)where δ is the Kroenecker delta and O( 1 T 2 ) represents higher order terms of 1 T 2 . In Sec. C.1, we show how to express the rate matrix R t in terms of the state transition matrix Q t . Lou et al. [33] propose a continuous time ELBO for this process. They mention that this expression can be derived from Benton et al. [4] , though they do not provide an explicit derivation. For this reason, we present a rigorous derivation in Sec. C.2 and further demonstrate that, under the SUBS parameterization in Sec. 3.2.3, this formula reduces to our proposed continuous-time ELBO, given by (10) .For the remainder of this section, we switch to the notation q t|s (y ′ |y) to denote q(z t = y ′ |z s = y), q s|t (y ′ |y) for q(z s = y ′ |z t = y), and q(z t = y|x) for q t (y|x), aligning with the notation typically used in the CTMC literature.Here, we aim to express the rate change matrix R t in terms of the state transition matrix Q t . To do this, we first represent the forward transition q t|s in terms of Q t and R t separately, allowing us to illustrate their relationship.Using (13) , we can write q t|s as follows:EQUATIONNow let's analyze all possible combinations for the tuple (y ′ ,y):1. Case (y ′ = x,y = x): Using (50), we find that q t|s (x|x) = α t|s for the DTMC. By (48), we have q t|s (x|x) = 1+R t (x,x) 1 T as T → ∞, since the higher-order terms O 1 T 2 vanish in the limit. Thus, we get:EQUATION2. Case (y ′ = m,y ∈ V -{m}): Similarly, using ( 50) and ( 48), we have q t|s (m|y ̸ = m) = 1-α t|s and q t|s (x|yEQUATION3. Case (y ′ = m, y = m): Using ( 50) and(48), we find q t|s (m|m) = 1 and q t|s (m|m) = 1 + R t (m,m) 1 T +O 1 T 2 . Since these two expressions must be equal for any T , it follows that R t (m,m) = 0.Note that when R t (m,m) is constant, the term O 1 T 2 reduces to zero, as it includes higher-order time derivatives of R t .4. Case (y ′ = x,y ∈ V -{x}): In the context of absorbing state diffusion, these states are never observed. Thus, R t (y ′ = x,y ∈ V -{m,x}) = 0 (54)5. Case (y ′ ∈ V -{m,x},y ∈ {m,x}): In the context of absorbing state diffusion, these states are never observed. Thus,EQUATIONFinally, we can express the forward rate matrix as:EQUATIONIt can be seen that the columns of this matrix sum to zero, i.e.,EQUATIONwhich ensures that the probability mass is preserved in the forward diffusion process. Similarly, the reverse rate matrix Rt can be written in terms of the forward rate matrix R t as follows [33] :EQUATIONMeng et al. [37] introduced the term "concrete score" for the term q t (y ′ |x)/q t (y|x) that appears in Rt . Since this quantity is not directly accessible in the reverse diffusion process, we approximate it using a neural network, s θ : V → V, with parameters θ. The approximate reverse posterior p s|t can then be expressed in terms of the approximate reverse rate matrix Rt in the following manner:EQUATION)where s θ (y) y ′ denotes the approximate concrete score q t (y ′ |x)/q t (y|x). Lou et al. [33] propose the following NELBO to train such a model:EQUATIONwhere K(a) = aloga -a. They mention that this expression can be derived from Benton et al. [4] , though they do not provide an explicit derivation. For this reason, we present a rigorous derivation in the following section.Proof. Let's focus on the diffusion loss for this process. As mentioned in the previous section, the reconstruction and prior loss terms reduce to zero. The continuous-time diffusion loss is given by: limT →∞ T E t∈{ 1 T , 2 T ,...,1},y∼qt(.|x) [D KL (q s|t (y ′ |y,x)∥p s|t (y ′ |y))] = lim T →∞ T E t∈{ 1 T , 2 T ,...,1},y∼qt(.|x)   y ′ q s|t (y ′ |y,x)log q s|t (y ′ |y,x) p s|t (y ′ |y,x)   = E t∈[0,1],y∼qt(.|x)   lim T →∞ T y ′ q s|t (y ′ |y,x)log q s|t (y ′ |y,x) p s|t (y ′ |y,x)   = E t∈[0,1],y∼qt(.|x)       lim T →∞T q s|t (y|y,x)log q s|t (y|y,x) p s|t (y|y,x)EQUATIONLet's simplify these two terms separately. For the derivation, we'll rely on two key observations: In the limiting case as T → ∞, it follows from (49) that lim T →∞ q s|t (y|y,x) = 1 and from ( 59) that lim T →∞ p s|t (y|y,x) = 1.Term 1:lim T →∞T q s|t (y|y,x)log q s|t (y|y,x) p s|t (y|y,x)∵ lim T →∞ q s|t (y|y,x) = 1; hence, = lim T →∞T log q s|t (y|y,x) p s|t (y|y,x)The above term is in ∞×0 indeterminate form; therefore,T →∞T logq s|t (y|y,x)-logp s|t (y|y,x)Substituting q s|t and p s|t from ( 49) and ( 59), we get:= lim T →∞ T log 1+ Rt (y,y) 1 T +O 1 T 2 -log 1+ Rθ t (y,y) 1 T +O 1 T 2Applying the Taylor series expansion for log(1+x), we get:= lim T →∞ T Rt (y,y) 1 T +O 1 T 2 -Rθ t (y,y) 1 T -O 1 T 2= Rt (y,y)+ limT →∞ T O 1 T 2 -Rθ t (y,y)-lim T →∞ T O 1 T 2 ∵ lim T →∞T O 1 T 2 = 0, we get:= Rt (y,y)-Rθ t (y,y)Using ( 58) and ( 60), we get:EQUATIONTerm 2:lim T →∞ T y ′ ̸ =y q s|t (y ′ |y,x)log q s|t (y ′ |y,x) p s|t (y ′ |y,x)Substituting q s|t and p s|t from ( 49) and ( 59), we get:EQUATIONFinally, plugging ( 63) and ( 64) into (62) yields us the NELBO as proposed in Lou et al. [33] :EQUATIONwhere K(a) = aloga-a. This concludes the proof.Given a latent variable z t and the output of the denoising model, x θ (z t ,t) parameterized using SUBS, we aim to recover the concrete score s θ (z t ) ∈ (R + +{0}) |V| . Note that s θ (z t ) y is the ratio pt(y) pt(zt) in the reverse process. Since x θ approximates x, we use p t (y) = q t (y|x θ (z t ,t)); therefore,EQUATION)To obtain the score, we first compute q t (y|x θ (z t ,t)) for all possible y and z t . Using (4), we derive the following expressions under the SUBS parameterization:EQUATIONEQUATION)Plugging these into (66), we get:s θ (z t = m) y̸ =m = α t 1-α t ⟨x θ (z t ,t),y⟩Using (70) and ( 67) (71)s θ (z t = m) y=m = 1 Using (70) (72) s θ (z t ̸ = m) y=m = 1-α t α tUsing ( 69) and (70) (73) 68) and (69) (75)s θ (z t ̸ = m) y=zt = 1 Using (69) (74) s θ (z t ̸ = m) y̸ ∈{m,zt} = 0 Using (These can be consolidated into the following expression:EQUATIONWe can formulate the reverse rate matrix for MDLM using ( 76) and (56) . Recall that the reverse rate matrix Rt (y ′ ,y) is given by:Rt (y ′ ,y) = s θ (y) y ′ R t (y,y ′ ) y ′ ̸ = y -ỹ̸ =y s θ (y) y ′ R t (y,ỹ) y ′ = y.Let's examine the cases where y = m and y ̸ = m.Case y = m: For y ′ ̸ = m, the reverse rate Rt (y ′ ,y = m) is given by:Rt (y ′ ̸ = m,y = m) = s θ (y = m) y ′ ̸ =m R t (y = m,y ′ )Using ( 71) and ( 56), we get:EQUATION)For y ′ = m, the reverse rate Rt (y ′ ,y = m) is given by:Rt (y ′ = m,y = m) = - ỹ̸ =m Rt (ỹ,y = m)Using (77), we get: = ỹ̸ =m α ′ t 1-α t ⟨x θ (z t ,t),ỹ⟩ = α ′ t 1-α tEQUATION)Case y ̸ = m: For y ′ ̸ = y we have:EQUATIONFor y ′ = m, we have:Rt (y ′ = m,y ̸ = m) = s θ (y ̸ = m) y ′ =m R t (y ̸ = m,y ′ = m) = 0 from (56)= 0 (80) Thus, for y ′ = y, we have:Rt (y ′ = y,y ̸ = m) = - ỹ̸ =y Rt (ỹ,y ̸ = m) = -Rt (ỹ = m,y ̸ = m) = 0 from (80) - ỹ / ∈{y,m} Rt (ỹ,y ̸ = m) = 0 from (79) = 0 (81)Summarizing ( 77), ( 78), ( 79), ( 80), (81), we have:EQUATIONC.5 Deriving MDLM's NELBO via CTMC Now, we aim to show that substituting the expression for the rate matrix R t in terms of state transition matrix Q t from ( 56) into ( 65) and switching from score-parameterization to the SUBS parameterization (Sec. 3.2.3) yields the simplified NELBO for MDLM as given by (10) . We present the proof below.Recall that the term ⟨a,b⟩ denotes the dot product of two vectors a and b. When a and b represent two one-hot vectors, this quantity evaluates to 1 if a = b and 0 otherwise.Proof. Recall that for absorbing state diffusion, y takes only two possible values, i.e., y ∈ {x,m}. Thus, we expand (65) as follows: 54), we get: 52), we get:E t∈[0,1],y∼qt(.|x) ⟨y,x⟩   y ′ ̸ =x R t (x,y ′ ) s θ (x) y ′ - q t (y ′ |x) q t (x|x) logs θ (x) y ′ +K q t (y ′ |x) q t (x|x)   +⟨y,m⟩   y ′ ̸ =m R t (m,y ′ ) s θ (m) y ′ - q t (y ′ |x) q t (m|x) logs θ (m) y ′ +K q t (y ′ |x) q t (m|x)   ∵ Rt(x,y ′ ̸ = x) = 0 from (= E t∈[0,1],y∼qt(.|x) ⟨y,m⟩   y ′ ̸ =m R t (m,y ′ ) s θ (m) y ′ - q t (y ′ |x) q t (m|x) logs θ (m) y ′ +K q t (y ′ |x) q t (m|x)   Substituting Rt(m,y ′ ̸ = m) from (= E t∈[0,1],y∼qt(.|x) ⟨y,m⟩   y ′ ̸ =m - α ′ t α t s θ (m) y ′ - q t (y ′ |x) q t (m|x) logs θ (m) y ′ +K q t (y ′ |x) q t (m|x)   = E t∈[0,1],y∼qt(.|x) ⟨y,m⟩       - α ′ t α t       y ′ ̸ =m s θ (m) y ′ Term 1 - y ′ ̸ =m q t (y ′ |x) q t (m|x) logs θ (m) y ′ Term 2 + y ′ ̸ =m K q t (y ′ |x) q t (m|x) Term 3             (83)Term 1:y̸ =m s θ (m) yUsing (67), we get,= y̸ =m α t 1-α t ⟨x θ (m,t),y⟩Under the SUBS parameterization, log⟨x θ (zt,t),x⟩ = 0 when zt = x; hence ⟨y,m⟩ can be dropped:EQUATIONThis concludes the proof.We use a single monte-carlo estimate for t to evaluate the likelihood. The low discrepancy sampler (D.3) plays a key role in reducing the variance of the estimate as seen in Table 8 .Given training_steps, batch_size, context_length, the number of tokens seen by the AR model is given as:training_steps×batch_size×context_length.(88)However, this expression doesn't hold for a diffusion model, since at each training step, a fraction of the input tokens are masked before being fed to the model. Let p m be the probability of a token being masked at a timestep t. For the log-linear schedule in our experiments, p m = t. Thus, the expected number of tokens seen by the diffusion model is:EQUATIONLM1B. Following [1, 33, 26] , we train MDLM for 1M training steps with a batch_size = 512, and a context length of 128. Like [33] we use a log-linear schedule and hence the number of tokens seen by our model is ≈ 33B (89). Similarly, MDLM trained for 10M steps, saw 327B tokens in expectation. The corresponding AR baseline was trained for 0.5M and 5M steps to ensure a similar number of tokens was seen.OWT. We train SEDD and MDLM for 1M training steps with a batch_size = 512, context_length = 1024, and log-linear schedule. Hence, these models saw 262B tokens during training. Similarly, the AR model saw the same number of tokens when trained for 0.5M steps with the same batch_size and context_length.To reduce variance during training we use a low-discrepancy sampler, similar to that proposed in Kingma et al. [29] . Specifically, when processing a minibatch of N samples, instead of independently sampling N from a uniform distribution, we partition the unit interval and sample the time step for each sequence i ∈ {1,...,N } from a different portion of the intervalt i ∼ U [ i-1 N , i N ]. This ensures that our sampled timesteps are more evenly spaced across the interval [0,1], reducing the variance of the ELBO.For our forward noise process, we use a log-linear noise schedule similar to Lou et al. [33] .We detokenize the One Billion Words dataset following Lou et al. [33] , whose code can be found here 4 . We tokenize the One Billion Words dataset with the bert-base-uncased tokenizer, following He et al. [26] . We pad and truncate sequences to a length of 128.We tokenize OpenWebText with the GPT2 tokenizer. We do not pad or truncate sequences -we concatenate and wrap them to a length of 1,024. When wrapping, we add the eos token in-between concatenated. We additionally set the first and last token of every batch to be eos. Since OpenWebText does not have a validation split, we leave the last 100k docs as validation.We parameterize our autoregressive baselines, SEDD, and MDLM with the transformer architecture from Lou et al. [33] . We use 12 layers, a hidden dimension of 768, 12 attention heads, and a timestep embedding of 128 when applicable. Word embeddings are not tied between the input and output.We use the AdamW optimizer with a batch size of 512, constant learning rate warmup from 0 to a learning rate of 3e-4 for 2,500 steps. We use a constant learning rate for 1M, 5M, or 10M steps on One Billion Words, and 1M steps for OpenWebText. We use a dropout rate of 0.1.We evaluate zeroshot likelihoods by taking the models trained on OpenWebText and evaluating likelihoods on the validation splits of 7 datasets: Penn Tree Bank (PTB; Marcus et al. [36] ), Wikitext [38] , One Billion Word Language Model Benchmark (LM1B; Chelba et al. [8] ), Lambada [41] , AG News [68] , and Scientific Papers (Pubmed and Arxiv subsets; Cohan et al. [10] ). We detokenize the datasets following Lou et al. [33] . For the AG News and Scientific Papers (Pubmed and Arxiv), we apply both the Wikitext and One Billion Words detokenizers. Since the zeroshot datasets have different conventions for sequence segmentation, we wrap sequences to 1024 and do not add eos tokens in between sequences.Following Devlin et al. [15] , we evaluate on all GLUE tasks [65] , but exclude WNLI.We pre-train a MosaicBERT model on C4 [46] for 70k steps, corresponding to 36B tokens. We pad and truncate the data to 128 tokens using the bert-base-uncased tokenizer.MosaicBERT [43] has a similar architecture to bert-base-uncased and has 137M parameters, 12 layers, 12 attention heads, a hidden dimension of 768, an intermediate size of 3072, and ALiBi attention bias [44] .For pre-training, we use the following hyperparameters: A global batch size of 4096 with gradient accumulation, a learning rate of 5e-4, linear decay to 0.02x of the learning rate with a warmup of 0.06x of the full training duration, and the decoupled AdamW optimizer with 1e-5 weight decay and betas 0.9 and 0.98.For diffusion fine-tuning we use AdamW with a warmup of 2,500 steps from a learning rate of 0 to 5e-5, betas 0.95 and 0.999, and batch size 512. We train for 5k steps total, corresponding to 32M tokens.For GLUE evaluation, we use the HuggingFace script found here 5 . We use the default parameters for all datasets, except for a batch size of 16, which we found helped with smaller datasets. This includes the default of 3 epochs for all datasets and learning rate of 2e-5.Dataset We pre-train the Caduceus MLM [50] on the HG38 human reference genome [11] . Following Schiff et al. [50] , we use character-/ base pair-level tokenization. The dataset is based on the splits used in Avsec et al. [3] : the training split comprises of 35 billion tokens covering the human genome. This consists of 34,021 segments extended to a maximum length of 1,048,576 (220 segments). We maintain a constant 2 20 tokens per batch. For the Genomics Benchmark tasks, we use 5-fold cross-validation where we split the training set into 90/10 train/validation splits.Architecture The Caduceus MLM uses as a backbone a bi-directional variant of the data-dependent SSM Mamba block proposed in Gu et al. [22] . This architecture is ideal as it contains inductive biases that preserve reverse complement (RC) equviariance, respecting the inherent symmetry of double-stranded DNA molecules [35, 50, 69] .Training details All models are pre-trained on 10B tokens (10K steps) and fine-tuned on a generative objective for an additional 50B tokens (50K steps). We use a global batch size of 1024 for a context length of 1024 tokens. Downstream task fine-tuning is performed for 16K steps ( 1B tokens).For performing Caduceus MLM pre-training, we follow Schiff et al. [50] for the model size configuration, and hyperparameter selection. For pre-training, we use a fixed 15% mask rate as done in Devlin et al. [15] . Of the 'masked' tokens, 80% are replaced with [MASK] , 10% are replaced with a random token from the vocabulary, and 10% are left unchanged.For fine-tuning all Mamba-based models (including Caduceus) on diffusion objectives, we lower the learning rate from 8e-3 to 1e-3. For fine-tuning HyenaDNA [39] , we lower the learning rate from 6e-4 to 5e-5. Similar to Gu et al. [22] , Schiff et al. [50] , we found that Mamba-based models were robust to higher learning rates. We exclude timestep embeddings for all Diffusion DNA experiments, as we show it has minimal impact on generative performance (see Table 12 , Suppl. E.5).We perform downstream task fine-tuning on the final hidden state embedding from pre-training. We perform mean pooling across the sequence length, which may vary from 200 to approximately 2,000 bps. We report the mean and ± on max/min classification accuracy over 5-fold cross-validation (CV) using different random seeds, with early stopping on validation accuracy. For each task, we do a hyperparameter sweep over batch size and learning rate and report the values of the 5-fold CV for the best configuration.We use a subset of the Genomic Benchmark tasks with an emphasis on tasks from Human data. The positive samples for each dataset were generated by selecting samples that were annotated, either computationally or experimentally, in previous work (e.g enhancers, promoters, open chromatin regions (OCR)) [20] . These annotations each correspond to subsets of the genome of varying sizes that may exhibit different distributions of DNA than those observed globally over the reference genome. Due to this, the observed dataset may have a different distribution than the data used for pre-training and calculating perplexity. This might in turn lead to a case where perplexity and downstream performance may not necessarily correlate. that γ = f (t). Note that α t goes through a monotonic transformation to obtain γ; hence, γ is also monotonic in t since α t is monotonic in t. This implies that the function f is invertible. Let t = f -1 (γ). Then, we can we have the following diffusion loss:EQUATIONThis new formulation demonstrates that the diffusion loss is invariant to the functional form of α t .In Table 9 , we demonstrate empirically that noise schedules with different functional forms evaluate to the same Likelihood which is consistent with our theory. However, different schedules lead to different per data point variance. Notably, the log-linear schedule exhibits the lowest variance among all the noise schedules considered.Table 9 : Likelihood in bits per dimension (BPD) for different noise schedules on OWT dataset, is reported along with the mean and variance associated with each noise schedule per data point. We empirically observe that noise schedules with different functional forms yield the same likelihood, consistent with our theory in Sec. In Figure 10 , we compare the wall clock times of variaous methods: AR, SEDD, MDLM with caching, and MDLM without caching for generating 64 samples on a single GPU. When sampling in batches, a change of 1 token would necessitate a call to the denoising model. Therefore, smaller batch sizes have a lower likelihood of a token being unmasked. This might lead one to prefer generating samples in smaller batches, as opposed to using a larger batch size that fully saturates the GPU. Table 10 shows that generating samples with a batch size of 1 and using caching is twice as fast as generating samples without caching while fully utilizing the GPU. In Fig. 2 , we observe that MDLM without caching yields samples that consistently get better generative perplexity than SEDD. For T = {5k,10k}, both SEDD and MDLM get better generative perplexity than the AR model. Sampling wall clock time (s)Figure 2 : Generative perplexities across wall clock time for generating 64 samples on OWT using a single 32GB A5000 GPU are compared by varying T ∈ {100,500,1000,5000,10000} in the reverse diffusion process. The samples are generated in mini-batches with a batch size of 16 for AR, SEDD, and MDLM without caching, as it is the largest batch size that fits on this GPU. For MDLM with caching, we vary the batch size.We assess the importance of our continuous-time framework by performing ablation on diffusion steps T . In Table 11 , we compare NLL and PPL under continuous and discrete T in MDLM. We find that NLL consistently decreases as T → ∞. OpenWebText [18] . NLL is logged every 1K steps without value smoothing.In Table 12 , we assess the importance of time conditioning in MDLM on OWT. We observe that time-conditioning has minimal impact on perplexity. Training is performed over 1M steps on OWT (which corresponds to 524B tokens). Here, we present some unconditional samples generated by MDLM trained on OWT with a context length of L = 1024 for T = {1000,10000}.Example 1 
 a 17-10 victory and a trip to the playoffs.The last wildcard seed: New York Jets, Houston and the last potable playoff spot. The last-second home wins: New Orleans and Carolina, 21-21.The Saints finish sixth with the highest regular season (42) NFC wins. They lost 14 of their 13 games in the conference playoffs.The Cardinals were in Group A in Round 1 with Game 2, Round 3, Round 4 and Quarter Game 5, but they made their last trip to the playoffs off North Carolina on the road as even North Carolina.True to their reputation, the Cards swept the Saints in the first round, but knocked it out at home. No Panthers went to the playoffs more than the Saints. "My guys admit to feeling it a little bit [Sunday but] I say to them that they always knew, 'I don't think that was necessarily how I would beat you, that will give them their confidence," Brees said."It was really hard because I've obviously learned a lot of detail about how to deal with everyone and as hard as I have to be, I also feel part of the stuff that they've been through on the team, like they're still going to go through things they know are somewhat right, but they feel a lot of pressure so it's got to be important to get it right now to get it in the future." Could all ask for more roses? Let's just take a slip, South Carolina, and face the NFC
 Example 2 
 Memorial Hospital.Valia and Hill had been working with the Coast Guard in response to public questions, and when they were reached couldn't comment on the new information, Chapman said.People referred to Valia during the years from Hill's family in Ants, and she cut in contact with their family and friends in 2016."Each day they stepped on the bus, when they left they saw me on TV," she said.After separating from their family recently, Valia, 32, also moved into a Richmond house last October.Read or Share this story: http://usat.ly/1NNC4zY
CIVIL C. "Marky" Hogan has been charged with homicides with a few days remaining after the April 2 purdade high school shooting where an undercover medical examiner and two other state and Illinois police officers was using heroin to go see a therapist. DICEZ TV's Zach Putler reported Tuesday that Hogan was charged with felony drug possession by the Chicago Police Department at the preliminary hearing on Monday. Putler interviewed on Monday. Authorities could offer a limit until Cook County takes Tuesday afternoon or they have to assign plea agreements. Dogan said in a news conference he made during a conference call Wednesday in Chicago that he believes the people who used him as a legal tool in the killing and fired employees hired for suffering also participated.He said the couple's request to an attorney Monday will let the charges finally play out. Their lawyer did not respond Wednesday. Dogan would not give away to possibility that he speculated in a statement that he would escape and return unless shot.Chicago police initially said the other drug charges failed to raise enough evidence to establish why the killers were charged last year, raising the possibility that the drug mix contributed to a reason for their arrest. But a new statement was made Tuesday by a man who worked as the shooter in his unit on campus, and suggested the charges might be related to his work at university supervision.His attorney and the university's president last week signaled that the incident of the bat gun was not a police investigation at Wednesday's conference.Michael Durin, Illinois State University spokesman said he did not meet with university officials at the conference, and that university officials don't have any updates yet either."The fact that the Defendants were charged is a major factor in why it would get this much attention," the university spokesperson said. "Given that all the matters are not being resolved for months and months, any new specific information and other concerns they may be tasked with investigating now are understandable."After city police began looking for evidence in connection to Hogan's April shooting, Durin said he had not noticed it until the Chicago Police Department found a person who was producing a bind gun on Illinois State campus. That same department found that 14 officers shot and were injured during a standoff, but it led to the launch of a combination of unrelated and related investigations leading to homicides charges in May 2015.A memo from special school investigators suggests it had identified the drug fentanyl, and says the department had described the individual-oriented and inconsistent use of the gun, as well as the substance administered by CODC. Dogan's allegations claim that a 2009 police gravesite package showed water running over campus and shows that the dental show photos of supposed victims were reinterpreted. David Mann, a member of the Police Department, said he spoke exclusively to News 1 on condition of anonymity because university officials can't review documents immediately, and university officials had to change information that had been a consideration."We didn't change our information until he personally told the drug overdose problem," he said. "He said that drug dealing wasn't really a focal point." He however, in a May 15 statement, also was "proleased" with university and state officialsin the U.S. Attorney's Office. "We do not have any way to estimate the crime syndicate, given that the finding in the case does not preclude a separate review of standard CPD policy directed at any school employee, and whether that employee worked or went out of work," Mann said.The lawsuit alleges that donations from a full visit with Hogan's workers and thousands of dollars, spent on gift cards, cigars, stock tickets to trips across the country, as well as other financial accounts and income from social members of former employees, were missing from buyers bought in stock.Both employees and co-participants remain employed at the university. The last crisis was like this, but we have to resolve it at the specified level; we have to stay low. The people of Greece will be interested in the effectiveness of the measures. They will not only manage these measures, but also they will help in order to cope with the problems of the fiscal stability. However, we do not want to dispose assets for the treasury. This also, so we will work on developing the national economy, and also paying on the national debt. This affects the national incomes So, as of 2007-2011, we use the government's temporary measures as a measure, helping resolve the crisis. In addition, we are able to pay for the borrowing costs. Additionally, we pay $440 billion to settle debt issues, which can never be settled by default in a country.These temporary measures will be aimed on several fronts, because the government will have three different partners in the system, in my right.Firstly, we will be allowed to borrow a lot more upon the addition of the emergency measures and these temporary measures will help provide for the repayment of our debts before we are forced into a crisis as a result of our borrowing bills. Secondly, in the case of this, we will have resort to temporary measures in the revenue budget for Greece. The budget costs the government another $5.2 billion a year.So what you propose -if it happens again, does this mean that, since 2010, will you resolve the deficits which will occur on the basis of what we already have?We would be able to settle our debts by the end of June which is the end. That said, we are taking our part as one of the most important countries in Europe, not only to make a proper transfer of the money but also to rely on it in the economy. However, first of all, we cannot achieve this on a day-to-day basis.It is still true that we have decided to be able to deal with the economic situation of the country, but there might be another change in the fiscal situation, and therefore, we will try to negotiate on the situation at the end of June and over the summer.The changes in the fiscal situation would be up to the parliament of management, bureaucrats, judges and a legitimate parliament of Greece.Is the government planning to talk about thethe 'temporary measures' of Greece?We will continue the process to operate through the temporary measures. This is not a temporary measure at this point, because after a crisis, not thereyet at crisis level, you can still have enough investment purchases until the end of the month.Again the government decided to create a temporary measure and now it depends upon a particular event, such as that there's another liquidity crunch. It is better that the government and the authorities decided to create a temporary measure effective in June at fair sum monthly bond rates.The temporary measures will also enhance the government's economic status, especially when following the measures at the end of the month.Temporary measures is a real tool for growth, not just for the economy.Knowing that there are several measures in place to increase our supply, for example, the level offor profit on public sector enterprises is certain, under all of these temporary measures the increases in output, after that, will increase the external demand and the internal demand.We will be able to create the demand, and also strengthen the government's credibility through fiscal organization. What is important here, here is that we will apply these measures to our reserves, and at the same time, we apply these measures to the debt level, which will also be the aim of debt-free Greece.So, first of all, everything is certain ofwhat continues to be collected by the government. Given the situation and after the release of the last data on October 16, you also recognize that this will not be any kind of non-payment.In the case of the payment against the equipment, we will be able to manage with the measures.What does government expect in its plans to create a fiscal consolidation for the public sector and the new budget.Regarding this is the temporary measure, we will be able to cope with the troubled finances. However, I do not think it is any measure which threatens the fiscal stability of the economy. However, that is not a temporary measure, a permanent measure.On the other hand, there will be our ongoing work on construction in the ministry. If this falls, we will continue work on job creation, the expansion of the economy. Also, also the government mentioned the new government reforms, which increased labor hours for the employees, which will further the economic growth, and the second aspect of budget as well and this is government welfare, which will improve the quality of life. We will
 Example 2 
 him. He said: "What are you doing?" I hesitated before answering. "Boy, this is so exciting. You need a better girl. Is she?" And I said, "You don't have a brain. You have no brain anymore." After a minute, he had walked back and said on his own, working through that, he thought he had got himself going in a new direction.He could've been a better boy in the first three years. "You'll only have once before it starts."The story is always, "That's what the other guy has to do." He was the guy who had to do anything. He had to reason with school officials. My cousin mentioned to me that some of my friends almost doubled over at one meeting. I'd picked up a lot of the money I owed him from high people in me; he liked my grades. Drop-outs didn't consider me high enough to let me go hang out. He hung up when I challenged him after practice to show a new talk. He started making, and, quite, never I first saw him C. morning, in the sixth grade class. He wouldn't hang up with him on point at team meetings. He started talking about things about me: "I'm an M, I'll get an A. Tonight." Having had that conversation over lunch, my heart touched mine with pride. He came, my boy. Now he looks like he's going back to school. I don't know if he's going to sue. Let's just have a two-bedroom apartment, a $500-dollar condo for renting, and a pool. And then he was back.That was a part of my life as I think about it. It was the school year.I never saw a guy come up at the locker room and show a new talk. That day one day, I told the high school, "We'll show up one day right here, we can have a little fun," and after this, I remember a small handful of the boys made friends, and they never, ever showed up for a new talk. I call them "M's kids. I always remember him, I remember his ass up his ass, getting ready for a freshman orientation out there. He'll show everything.He'll show if I'm freshman, I'll act I was going to play junior. In a few years I'll try it, then he'll make sure he's going to judge me. He will come over to me one day.Then one day, the senior class was sitting on the bench, pressing his ball on the floor of the locker room, the referee was just standing the knelt it down.And when he heard about that, another boy, three of his friends, and one of his cousins were on the other side of the room. The boys' class was filling in with his new brother and his new cousin and his new M.M's player.The senior class watched me walk me through the chair to the bench. Everyone passed by the boy. Just on his toes on a foot, too.He [and a girl] passed over his head and, as I looked at them, he carried me into the locker room. And the biggest part of the story, was a mistake.With his elbows out, he pulled me down on my shoe, my other, sort of a-don't know what they were; palebelly somethings, like bleeding very much, or on little toes. He was up on the stairs and everybody watches, with men and high school kids, who saw him in the locker room. And he caught a breath. Then his old man approached me and, disappeared into the middle of the room. He took off his vest, fast enough as to herd him into the locker room. I walked into the room and read him little cards with my own eyes to make notes, to pull him under my shoes. I told them: "Listen, because I say this today, when you talk to 'em today, "Just make sure you talk is more than you'll show. He's just listening to me, and he's telling me I'm going to be there for him." I'm always the one who wants to do something important about you than I show up. I walk around and ask, I want a message from you, "Keep it going. It
 Answer: [Yes] Justification: Many of our tabels include error bars and standard deviations Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It 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 a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).code: https://github.com/kuleshov-group/mdlmtutorial: http://youtu.be/WjAUX23vgfg38th Conference on Neural Information Processing Systems (NeurIPS 2024).Lou et al.[33] mention that their ELBO can be derived from Benton et al.[4],Hanson [25] but do not provide an explicit derivation. To address this, we present a rigorous derivation in Suppl. C.2 using CTMC theory.https://github.com/louaaron/Score-Entropy-Discrete-Diffusion/blob/main/data.pyhttps://github.com/huggingface/transformers/tree/main/examples/pytorch/text-classification 


Decision: The paper proposes a model called Masked Diffusion Language modeling that is a special case of D3M. The authors propose several new contributions that lead of really strong results, bringing diffusion models close to AR models for sequence modeling. Specifically, they propose a specific parameterization of the predictive distribution (which they call SUBS), that allow for better lower bounds in the ELBO objective, because the objective ends up having fewer terms (eqn 8) compared to what a more vanilla formulation from D3M would have had (eon 9). In addition the authors lay out other contributions, related to better implementation, vocabulary choice, ignoring of the time steps in denoising, etc that they claim lead to better results that would have been expected from prior work. 

Results are shown on language modeling and human genome sequence modeling. The reviewers were generally positive -- although there was some push back against whether or not the method was indeed "simple", and also others asked for more detailed comparisons to other methods across different scales of the data, and of the model.  I hope the authors will take these points into consideration when submitting their final version, for maximal impact.

While I was perusing the paper, I also happened to find some typos in equations (unless I misunderstood some of the details of the paper) that I am listing here -- if they are indeed correct, I would hope that the authors will make the corrections for better readability:

1. Line 86 might be incorrect. I think it should be: Cat(z_t; \alpha_{t|s} z_s + (1-\alpha_{t|s}) \pi)
2. Line 98 mixes probabilities, and probabilistic variables, making it inconsistent with the rest of their notations. It should be:  q(z_t|x_0) = Cat(z_t; \alpha_t x_0 + (1-\alpha_t) m);
3. Equation on line 103 uses s, t with latent variables, inside the chain rule. Instead of p_\theta(z_s|z_t) it should be p_\theta(z_{i-1} | z_i). 
4. Line 181 — the range of time might be wrong. I think it should have been \tilde{x}^{L’:L} , not \tilde{x}^{L':L-L'} ? And the same for  lines 184, z_1 and \tilde{x} have errors in the sequence limits: L’:L-L’ should be L’:L again, possibly.
5. Equation 19 -- z_s, z_t should be x_s, x_t, to be consistent with p_\theta(x_s|x_t). The variable names are being switched around ?

Reviews: review: paper_summary: paper_summary: This paper introduces a new approach to masked diffusion language models (MDLMs) that improves performance over previous discrete diffusion methods. The authors present a simplified, Rao-Blackwellized objective for training MDLMs, which is derived from a substitution-based parameterization of the reverse diffusion process. This objective takes the form of a weighted average of masked language modeling losses, establishing a connection between generative diffusion models and encoder-only BERT models. The paper also describes efficient sampling techniques, including a semi-autoregressive decoding method. The authors demonstrate state-of-the-art performance among diffusion models on language modeling benchmarks, approaching the perplexity of autoregressive models. Additionally, they show that their approach can be applied to biological sequence modeling, achieving competitive results on DNA modeling tasks. The paper emphasizes the importance of well-engineered implementation details and provides comprehensive ablation studies to validate their design choices. Overall, this work presents a simple yet effective framework for discrete diffusion models that bridges the gap between diffusion-based and traditional language modeling approaches.
strength_weakness: strength_weakness: - The paper focuses heavily on empirical results but lacks a deeper theoretical analysis of why the proposed MDLM approach outperforms previous methods. A more rigorous theoretical foundation could provide insights into the method's success and potential limitations.
- While the paper compares MDLM to other diffusion-based approaches, it lacks a comprehensive comparison to state-of-the-art non-diffusion language models. This makes it difficult to fully assess the method's competitiveness in the broader context of language modeling.
- The experiments focus on relatively small models. It's unclear how well the approach scales to larger models that are more common in current NLP research. Addressing potential scalability issues or limitations would be valuable.
- The paper primarily focuses on perplexity as an evaluation metric. Including human evaluations or other metrics that assess the quality and coherence of generated text would provide a more comprehensive view of the model's capabilities.
- It's unclear if MDLM can achieve practical speed-up comparing to AR models with same FLOPs budget due to the lack of comparative e valuation on this.
- Incomplete reference and comparision to previous related works, see Questions section.
questions: questions: - Why didn't the authors include a comparison to Plaid in their perplexity evaluation experiments (Tables 1 & 2)?
- How do the authors explain their observation that time-step conditioning can be optional in MDLM? Particularly, why is there no conditioning in OWT? Does this suggest that with a larger amount of data tokens, the need for timestep conditioning in MDLM decreases?
- How many steps are used when evaluating the perplexity of MDLM? How does this compare to other text diffusion models?
- Can MDLM achieve any efficiency gains in inference and sampling compared to autoregressive (AR) models within the same compute budget? For instance, if we need $K$ steps to generate $B$ tokens at a time in semi-MDLM, how can we make a fair and scientific comparison to AR models generating $L$ tokens in this setting?
- One disadvantage of BERT-style models is that they are less token-efficient in training compared to decoder-only AR models. Can the authors provide some comparative analysis and plots showing how MDLM compares to AR baselines across a range of training token amount budgets?
- What are the authors' thoughts on MDLM versus continuous-diffusion models for text, such as Plaid and CDCD, in terms of training, model performance, and inference efficiency? Is MDLM superior to these continuous-space text diffusion models in terms of resulting model quality and reducing the gap to AR models?
- Is it possible to speed up the generation process of MDLM, similar to the advanced ODE-solvers people are using in continuous-space diffusion models?
- How does MDLM compare to MaskGiT in terms of model formulations? They share lots of same design space but I didnt see any reference or comparison to this related line of works.
limitations: limitations: N/A
paper_summary: paper_summary: While previous works considered diffusion language models less competitive than autoregressive models in text generation tasks, the authors propose a simple framework named masked diffusion language modeling (MDLM), where they claim to have better performance than previous thoughts. The authors derive a simplified Rao-Blackwellized objective and find that the objective equals to a masked language modeling with a varying mask ratio. The proposed methods are evaluated on both classification tasks and generation tasks. The performance excels in classification tasks and keeps consistent with autoregressive methods in generation tasks, which achieves a new state-of-the-art among diffusion models.
strength_weakness: strength_weakness: The main weakness of this paper is about the presentation. This includes the following two points:
* In the third section, the authors first derive an objective with a diffusion process. Then, the authors try to claim what difference has been made compared to the previous D3PM. It seems that the different parameterization made in Equation (7) leads to a simpler formulation in Equation (8) compared to Equation (9), if I understand correctly. However, I am quite confused with the illustration in Section 3.2 and Section 3.3. There appears some technical terms without furthur explanations or citations, such as Rao-Blackwellization and graphical modeling. The abbreviation SUBS is also unclear. Some omittions and simplifications are not provided with some explanations to clarify their rationality.
* There also lacks an overview on the algorithm. Due to the highly complicated derivation process and objective function, I suggest the authors making an algorithm bar to illustrate how the training and the sampling are conducted, like Algorithm 1 and Algorithm 2 in DDPM [1].

There also lacks some discussions with some previous works, including:
* There has been a nearly same method in the CV field called MAGE [2]. MAGE also conducts masked modeling on image patch tokens with a varying mask ratio and generates by the ancestral sampling.
* There lacks a citation for semi-autoregressive modeling [3] since I believe this concept is not first raised in this paper.

[1] Denoising Diffusion Probabilistic Models, https://arxiv.org/pdf/2006.11239

[2] MAGE: MAsked Generative Encoder to Unify Representation Learning and Image Synthesis, CVPR 2023

[3] Semi-Autoregressive Neural Machine Translation, EMNLP 2018
questions: questions: * I wonder what previous tries have been made in the language diffusion model and what makes them fail to outperform autoregressive model in generation tasks, as diffusion models in the CV field largely outperform autoregressive models.
* I wonder the computational cost of the method and its comparison with the autoregressive model's since it is known that the diffusion model usually costs more time to converge.
limitations: limitations: The authors have claimed their limitation in the paper.
paper_summary: paper_summary: This paper presents a method for language modeling using simple masked discrete diffusion models. The authors show that a simplified objective combined with optimized training achieves performance improvements over previous diffusion language models. The paper reports state-of-the-art results for diffusion models on standard language modeling benchmarks, approaching the performance of autoregressive models.
strength_weakness: strength_weakness: 1. I believe the paper should be structured more towards introducing the simplified objective rather than framing the architecture as MDLM since the main idea of bridging absorbing state discrete diffusion with MLM has been introduced in DiffusionBert and D3MM.  
2. The comparative results are not fair comparisons. More specifically the use of a more advanced backbone (DiT) and a low-variance sampler weakens the claim that the performance and stability are sourced from the proposed method.
questions: questions: 1. Other than the simplified objective and new backbone and sampler, what’s the difference between your method and absorbing state D3MM or more specifically DiffusionBert?
2. Is MLM pretraining required for your method?
limitations: limitations: 1. Masked language models are special discrete Diffusion models, this is demonstrated in the D3PM paper by rewriting an x_0-parameterized absorbing state model.  This limits the novelty of the paper to the simplified objective only.
2. The idea in this paper has been hinted at in D3MM, limiting the theoretical contribution of this paper.
paper_summary: paper_summary: The paper introduces a masked diffusion language modeling (MDLM) framework that enhances the performance of diffusion models in generating high-quality text, closing the gap with autoregressive methods. By applying an effective training strategy and a simplified objective function, MDLM achieves state-of-the-art results among diffusion models on language benchmarks. The approach is shown to be simple yet effective, with potential applications beyond language, including biological sequence modeling.
strength_weakness: strength_weakness: - In Figure 1, unmasking loss is misleading. What's the difference between the red and yellow colors? Which tokens should we enforce the losses on?
- In Line 38, "Simple" seems overstated. I couldn't quickly understand the method, even though I have some background in diffusion models and sequence modeling.
- In Line 168, there are too many experimental details about the network structure and tokenizer changes, making it unclear how to prove the method works. This, in turn, suggests that the proposed method is not so simple.
- The sampling steps for diffusion-based methods should be highlighted in Table 1 and other related tables.
- In Line 178, the semi-autoregressive part is not clear to me. I lose track of how it works. The need for semi-autoregressive models is not well-motivated. It's a common replacement trick used in diffusion models. The author should explain it.
- In Table 5, memory should also be included. Secondly, Mamba is more favorable for longer token numbers. The authors should also compare with Mamba under various token numbers: 1k, 5k, 10k, and 20k.
- The paper needs to provide the mathematical proof for the continuity equation, showing that \( u_t \) in the equation generates the probability path \( p_t \). It should also highlight the connection to the continuous case.
- In Line59, several works about using flow-based method is missing,[1][2]
- In LIne65, citation missing, Chapell's discrete flow matching paper
- Better discuss with this work: https://arxiv.org/abs/2406.03736

[1]. Language Rectified Flow: Advancing Diffusion Language Generation with Probabilistic Flows

[2]. Flow Matching for Conditional Text Generation in a Few Sampling Steps

Overall, the authors claim the method is simple, but it’s not as simple as they claimed. Additionally, there are many issues with the presentation and experimental parts. Therefore, I am inclined to weakly reject this paper.
questions: questions: As above
limitations: limitations: Yes

Comments: title: Thank you
Updated references
References
accept
Response to Area Chair ccaN
Discussion on related works
Thanks for the reply
Updated results on text8
Updated results on text8 and comparison to non-AR baselines
Response to Reviewer DasC
Please refresh the page if the equations are not rendering
Individual responses are now posted
Response to BSmY (3/3)
Response to BSMY (2/3)
Response to BSmY (1/3)
Response to P1DW (3/3)
Response to P1DW (2/3)
Response to p1DW (1/3)
Response to DASc (2/2)
Response to DASc (1/2)
Response to wXn1 (3/3)
Response to wXn1 (2/3)
Response to wXn1 (1/3)
where is  individual response?
comment: Dear Reviewer, 

We really appreciate your feedback and continued engagement. We'll incorporate these discussions in the next version of the manuscript. 

Thanks,
authors
We have added the references in a separate comment here:  https://openreview.net/forum?id=L4uaAR4ArM&noteId=fwSYxq1dTI
## References:

[1]. Zhang, S., Wu, L., Gong, C., & Liu, X. (2024). Language rectified flow: Advancing diffusion language generation with probabilistic flows. arXiv preprint arXiv:2403.16995.

[2]. Hu, V., Wu, D., Asano, Y., Mettes, P., Fernando, B., Ommer, B., & Snoek, C. (2024, March). Flow Matching for Conditional Text Generation in a Few Sampling Steps. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics (Volume 2: Short Papers) (pp. 380-392).

[3] Campbell, A., Yim, J., Barzilay, R., Rainforth, T., & Jaakkola, T. (2024). Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04997.

[4] Ou, J., Nie, S., Xue, K., Zhu, F., Sun, J., Li, Z., & Li, C. (2024). Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data. arXiv preprint arXiv:2406.03736.

[5] Chang, Huiwen, et al. "Maskgit: Masked generative image transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.

[6] Gulrajani, Ishaan, and Tatsunori B. Hashimoto. "Likelihood-based diffusion language models." Advances in Neural Information Processing Systems 36 (2024).

[7] Li, Xiang, et al. "Diffusion-lm improves controllable text generation." Advances in Neural Information Processing Systems 35 (2022): 4328-4343.

[8] Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances in Neural Information Processing Systems, 35:28266–28279, 2022.

[9] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribu- tion. Advances in neural information processing systems, 32, 2019.

[10] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834, 2023.

[11] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. arXiv preprint arXiv:2211.16750, 2022.

[12] Zhengfu He, Tianxiang Sun, Kuanning Wang, Xuanjing Huang, and Xipeng Qiu. Diffusion- bert: Improving generative masked language models with diffusion models. arXiv preprint arXiv:2211.15029, 2022.

[13] Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xi- aojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP- IJCNLP), pp. 6112–6121, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1633. URL https://aclanthology.org/D19-1633.
can you update the reference? I cannot find them
I appreciate the author's reply about the related works. Even though some papers are concurrent works, I think it can be discussed. I am not asking for a comparison.

I am excited to see the deep relationship between them. My concerns have been fully resolved, and I will increase my score to reflect this.
Dear AC, Thank you for following up on our discussions. We were actively working on Reviewer wXn1's latest comment and just posted our response.

To the reviewers, please let us know if there are any remaining open questions or concerns. Thank you again for you time and feedback.
Thank you for your feedback and response to our rebuttal. Below we provide discussions on the requested related works [1, 2, 3, 4]:

Language Rectified Flow [1] and Flow Matching for Conditional Text Generation in a Few Sampling Steps [2] perform flow matching over word embeddings. These works are more similar to Plaid [6] and Diffusion-LM [7], where the diffusion process is defined over a continuous space. In contrast, MDLM applies diffusion processes to discrete structures. We will include this discussion in the updated version of the manuscript.

Furthermore, we would like to highlight that in Section 6 of the paper, where we have provided a comprehensive literature review covering areas such as Continuous Time Markov Chains [8, 11], score estimation [9, 10], and techniques that use BERT for sample generation [12, 13, 14, 15].

## Comparisons to Discrete Flow Matching [3]:
Discrete Flow Matching (DFM) proposes flow matching for discrete structures. They use the following cross entropy loss as their training objective:
$\mathcal{L}\_\text{DFM} = - \int_{t=0}^{t=1}\log p_\theta(\mathbf{x}^{1:L}|\mathbf{z}^{1:L}_t) dt$

Similar to [5], DFM’s objective, while effective, is not weighted appropriately to make it a proper ELBO. However, in MDLM, we derive a tight and principled lower bound to the log-likelihood. 

## Comparisons to Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data. [4]:
[4] reaches a similar training objective to that derived in our work, although they start their analysis from the score matching perspective presented in SEDD [10]. Interestingly, they find an important equivalence between denoising concrete score matching and the variational lower bound that we derive. A similar connection was also made in Shi et al. [16]. In contrast, we tackle this problem starting via a variational lens: where we derive the exact continuous time ELBO by endowing D3PM with the novel SUBS parameterization and taking the number of diffusion steps T to infinity. Also of note is that [4] leverages a time-independent de-noising network to formulate an efficient sampler that limits function evaluations, as in our work.

A key point of differentiation between our works is that we tackle an important limitation of diffusion models, namely that they are limited to generating a fixed sized context. Our proposed semi-AR algorithm mitigates this shortcoming, and  is therefore a **novel contribution to this line of work that separates MDLM from [4].**

Finally, **turning to a comparison of the empirical evaluation: here, we believe that our work is significantly better.**
- Most importantly, we explicitly compare our discrete diffusion model on the widely used objective of **perplexity** and demonstrate that non-autoregressive models are approaching the performance of autoregressive ones. This is a significant milestone, as perplexity has been the guiding watermark for the field of language modeling. In contrast, [4] only use the less informative metrics of “zero-shot” perplexity and “generative perplexity”. 
- Furthermore, we analyze domains outside of NLP and demonstrate that the approach can be used on biological sequences, unlike [4] and [16], which only focus on NLP datasets.
- Finally, our Table 4, demonstrates how one can **take an off-the-shelf BERT model and render it generative without losing representation learning capabilities.** This is an important result in our work that is not present in either [4] or [16].

We end by noting that [4,16] are highly concurrent and were posted on arXiv after the Neurips submission deadline. According to Neurips guidelines, authors are **not expected to compare their work to papers that appeared on arXiv less than two months before the submission deadline**, let alone those published after it.
Dear Reviewers, Authors,

The discussion period will be ending soon. If you have any additional questions (reviewers) or any additional clarifications (authors) please post them soon.

Dear authors, do you think you can address Reviewer wXn1's concerns regarding the discussion on related works ?
Thanks for the authors' reply and the reviewer's feedback.

My concern about the simplicity of the method is addressed by the pseudo alg. However, without explicit discussion about related works in the main paper, I don't feel my concerns are fully resolved.
On the text8 dataset, we managed to train the MDLM for only 800K steps due to resource constraints, unlike the baselines, which were trained for 1M steps. Despite this, we outperformed all non-autoregressive baselines, such as the flow-based models (IAF/SCF) [1] and a Bayesian Flow Network (BFN) [2], while matching the previous state-of-the-art, SEDD, in just **80%** of the training steps.

| Method | BPC($\downarrow$) | Train steps |
| --- | --- | --- |
| IAF/SCF | 1.88 | 1M |
| AR Argmax Flow | 1.39 | 1M |
| BFN | 1.41 | 1M |
| D3PM | 1.45 | 1M |
| SEDD | 1.39 | 1M |
| MDLM | **1.39*** | **0.8M*** |
| AR | 1.23 | 1M |

*We provide numbers for a partially trained MDLM model at 800K steps. All baseline models were trained for 1M steps. Due to resource constraints, we were unable to train our model to 1M steps during the rebuttal period.

---

****References:****

[1] Ziegler, Zachary, and Alexander Rush. "Latent normalizing flows for discrete sequences." International Conference on Machine Learning. PMLR, 2019.

[2] Graves, Alex, et al. "Bayesian flow networks." arXiv preprint arXiv:2308.07037 (2023).
On the text8 dataset, we managed to train our MDLM model for only 800K steps due to resource constraints, unlike the baselines, which were trained for 1M steps. Despite this, we **outperformed all non-autoregressive baselines**, such as the flow-based models (IAF/SCF) [1] and a Bayesian Flow Network (BFN) [2], while matching the previous state-of-the-art, SEDD, in just **80%** of the training steps.

| Method | BPC($\downarrow$) | Train steps |
| --- | --- | --- |
| IAF/SCF | 1.88 | 1M |
| AR Argmax Flow | 1.39 | 1M |
| BFN | 1.41 | 1M |
| D3PM | 1.45 | 1M |
| SEDD | 1.39 | 1M |
| MDLM | **1.39*** | **0.8M*** |
| AR | 1.23 | 1M |

*We provide numbers for a partially trained MDLM model at 800K steps. All baseline models were trained for 1M steps. Due to resource constraints, we were unable to train our model to 1M steps during the rebuttal period.

---

****References:****

[1] Ziegler, Zachary, and Alexander Rush. "Latent normalizing flows for discrete sequences." International Conference on Machine Learning. PMLR, 2019.

[2] Graves, Alex, et al. "Bayesian flow networks." arXiv preprint arXiv:2308.07037 (2023).
Thank you for your feedback and response to our rebuttal. Unfortunately, we believe there are several factual errors in your understanding of our method, which we would like to clarify:



1. There are **no other sources of contribution to the performance of the method besides our new objective** in the head-to-head comparison against SEDD that we report.
2. Our method is not a mix of auto-regression and diffusion, and **no major perplexity result relies on auto-regressive components**. 


## **Concern 1:** The refined objective's sole contribution remains unclear. Other design changes are the major sources of improvement over SEDD.

We’d like to re-emphasize that SEDD and MDLM have an identical experimental setup on language modeling benchmarks, as detailed below. **The only difference between these two methods lies in the objective**. Thus, the below table quantifies the **sole contribution** of MDLM’s objective over SEDD.


|  | SEDD (Our reproduction) | MDLM |
| --- | --- | --- |
| **_Common experimental details_** |
| Tokenizer | `bert-base-uncased` for LM1B / `gpt-2` for OWT   | `bert-base-uncased` for LM1B / `gpt-2` for OWT |
| Architecture | DiT | DiT |
|  Context size | 128 for LM1B / 1024 for OWT | 128 for LM1B / 1024 for OWT |
| Train steps | 1M steps | 1M steps |
| **_Objectives (the only difference between SEDD and MDLM)_** |
| | Diffusion weighted denoising score entropy (Eq 9 in SEDD) | Simplified ELBO that becomes weighted sum of BERT style losses (Eq 11 in our work) |
| **_Performance_** |
| Perplexity ($\downarrow$) on LM1B | 28.90 | **27.04** |
| Perplexity ($\downarrow$) on OWT | 24.10 | **23.21** |


## **Concern 2:** This work is a mix of AR and Diffusion. MDLM’s improved performance is because it “borrows” power from AR.

This is a major misunderstanding. Please note that **no** significant result (Tables 1, 2, 3, 4, 6) uses anything related to auto-regression. All of our perplexity results are **benchmarked against SEDD in a fully non-autoregressive setting following their setup (data, context window, etc.)**. In fact, one of our key contributions is that these results are obtained using a simple BERT-style non-autoregressive loss.

The proposed semi-autoregressive sampler is only used in a minor experiment (Table 5) to allow MDLM to generate sequences of much longer lengths than those on which it was originally trained. However, this sampler **does not** contribute to MDLM’s improved perplexity scores when compared to baselines such as SEDD or D3PM. In fact, the semi-AR experiments do not even compare against SEDD, only against SSD-LM.
Thanks for the additional information and clarification!

Regarding answers to concern 1: My opinion remains the same, and the methods especially zero-masking, carry-over unmasking, and semi-autoregressive sampling are shown to be the source of performance improvement, However, using such methods makes this work a mix of AR and Diffusion, which provides a limited contribution to the bridging of performance gaps between AR and Diffusion. 

Regarding answers to concern 2: The refined objective's sole contribution remains unclear, and needs to be studied separately from the other design changes. I believe that the refined objective can stabilize the training, while the other design changes are the major sources of improvement over SEDD. However, these additional design changes introduce an AR nature to your method from my intuition. If such designs do not show superior performance than AR, these design changes seem like "borrowing power" from the strong AR baselines, as mentioned above.

Regarding answers to concern 3: Thanks for the clarification.
Please **refresh the page** if the equations are not rendering.
We apologize for the delay. All individual responses have now been posted as “Official comments”. Thank you again for your time and consideration.
### ****Concern 3:**** MDLM vs MAGE

While both MDLM and MAGE use BERT-style losses to train generative models, the main differences are as follows:

- ****Objective:**** In MDLM, we derive and train with a tight and principled lower bound to the log-likelihood. MAGE’s objective, while effective, is more heuristic with a combination of BERT-style cross entropy on masked tokens (not weighted appropriately to make it an ELBO) and a contrastive loss term.

- ****Sampler:**** In MDLM, we use a principled ancestral sampling method, whereas in MAGE once again the sampling is done in a more heuristic manner that yields impressive results in the image generation domain.  

Additionally, we thank the reviewer for the additional reference Semi-Autoregressive Neural Machine Translation. In our updated manuscript we add this citation and previous works on semi-autoregressive modeling. 

---

### Additional questions

****### **Question 1:****** What previous tries have been made in the language diffusion model and what makes them fail to outperform autoregressive model in generation tasks?

Broadly, past attempts at diffusion modeling for discrete data can be broken into two categories, (1) works that first embed text in continuous space and then run standard Gaussian diffusion on the embeddings before coming back to the discrete domain and (2) works that directly define a corruption process on the discrete data. The first line of work seems to suffer from training instability and the results lack in quality relative to the dominant AR approach (see for example DiffusionLM [1]). The second approach has recently shown promise. 

Both our work and the previous state-of-the-art diffusion language model, SEDD [2], propose novel parameterizations of the denoising network that induce a better loss: 

- concrete score matching with positivity constraints, in the case of SEDD
- and a weighted sum of cross-entropy terms, with reconstruction and prior regularization loss terms analytically evaluating to zero, in our case.

These better losses are key ingredients that contribute to the success of recent efforts that close the gap to AR language modeling.

****### **Question 2:******  What is the computational cost of the method in comparison to  autoregressive models?

We find that AR models are able to optimize the training loss better than diffusion models [see fig. 1 in the attached supplementary material](https://openreview.net/forum?id=L4uaAR4ArM&noteId=WereEXicWc) However, diffusion models are able to reach the same loss values with additional training. Crucially, the gap between the two curves is much smaller than previously thought. We are excited by the new opportunities that non-AR generation present, such as methods for controlled generation, and the promise of more efficient sampling.

---

****References:****

[1] Li, Xiang, et al. "Diffusion-lm improves controllable text generation." Advances in Neural Information Processing Systems 35 (2022): 4328-4343.

[2] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).
### ****Concern 2:**** The derivation of the simplified ELBO objective needs more explanation.

The reviewer is asking for clarifications regarding how the unsimplified loss (9), given by

$\mathbb{E}_q \Big[\frac{\alpha_s - \alpha_t}{1 - \alpha_t} \log \frac{\alpha_t \langle \mathbf{x}\_\theta(\mathbf{z}_t, t), \mathbf{m} \rangle + (1 - \alpha\_t)}{(1 - \alpha\_t) \langle \mathbf{x}\_\theta(\mathbf{z}_t, t), \mathbf{x} \rangle} + \frac{1 - \alpha\_s}{1 - \alpha\_t} \log \frac{(1 - \alpha_s)(\alpha_t \langle \mathbf{x}\_\theta(\mathbf{z}_t, t), \mathbf{m} \rangle + (1 - \alpha_t))}{(1 - \alpha_t)(\alpha_s \langle \mathbf{x}\_\theta(\mathbf{z}_t, t), \mathbf{m} \rangle + (1 - \alpha_s))} \Big] \langle \mathbf{z}_t, \mathbf{m}\rangle$ 

is transformed into the simplified ELBO (8)

$\mathbb{E}_{q,t}\frac{\alpha\_t'}{1-\alpha\_t}\log\langle\mathbf{x}\_\theta(\mathbf{z}_t), \mathbf{x}_0\rangle$


and the role that Sections 3.2 (SUBS parameterization) and 3.3 (Rao-Blackwellization) play in this process.

In brief, we obtain (9) by taking our interpolating masking forward process (first paragraph in Section 3.2.1) and inserting it into the general D3PM noise process (3) to obtain analytical formulas (4) and (5) for the marginals and posterior of $q$. Inserting (4) and (5) into the standard diffusion loss (2) yields (9).

Simplifying (9) into (8) mainly involves algebraic manipulations; see Appendix Sec. G. This algebra crucially requires two identities:

1. $\langle\mathbf{x}_\theta(\mathbf{z}_t, t), \mathbf{m}\rangle = 0$ , i.e. the masked token has zero predicted probability
2. if $\mathbf{z}\_t$ is unmasked, then we desire $\mathbf{x}_\theta(\mathbf{z}_t, t) = \mathbf{z}_t$  

In order to ensure that these two properties hold, we require that the denoising process has the form given in (7) in Section 3.2.2. We call this novel parameterization SUBS.

Section 3.3 (and its corresponding appendix) focuses on transforming (9) to (8) using these properties. We refer to the process of obtaining (8) in lieu of (9) as a form of Rao-Blackwellization. 

The above is a summary of how we derive the simplified ELBO (8) from (9). Next, we define the technical terms identified by the reviewer, and explain how they correspond to instances of general statistical techniques used in the simplification of (9) to (8).

**SUBS parameterization:** This refers to the parameterization in (7). It implements two substitutions that we enforce on the output of $x_\theta$:

1. We design the denoising network such that $\langle\mathbf{x}_\theta(\mathbf{z}_t, t), \mathbf{m}\rangle$ = 0, i.e., we substitute the logit index corresponding to the [MASK] token with −∞. This ensures property #1 above.

2. If $\mathbf{z}\_t$ is unmasked, then we desire  $\mathbf{x}_\theta(\mathbf{z}_t, t) = \mathbf{z}_t$ , i.e., unmasked latents are ‘carried over’. We accomplish this by substituting the output of our network to simply copy unmasked inputs. This ensures property #2 above.

**Rao-Blackwellization** is a statistical method that improves an estimator’s efficiency by computing expectations analytically, thereby reducing variance. In our case, we analytically compute expectations such as $\langle\mathbf{x}\_\theta(\mathbf{z}\_t, t), \mathbf{m}\rangle = 0$  in order to simplify objective (9) to obtain (8). Without these analytical simplifications, a model must learn $\theta$ such that  $\langle\mathbf{x}\_\theta(\mathbf{z}\_t, t), \mathbf{m}\rangle = 0$ holds, which slows down training. Unlike in regular Rao-Blackwellization, simplifications are possible because of modeling choices for $\mathbf{x}_\theta(\mathbf{z}\_t, t)$ (zero masking probabilities and carry-over unmasking). However, our approach also empirically helps reduce variance, hence we refer to it as Rao-Blackwellization, somewhat abusing the usual terminology.

**Graphical modeling** is a branch of machine learning that studies how to design parameter-efficient model by explicitly incorporating conditional independencies between random variables into the design of a model. Examples of algorithms in this field include Bayes networks and Markov random fields. In that sense, our approach has similarities to graphical modeling, where incorporating conditional independencies into $p_\theta$ (via our modeling choices for $\mathbf{x}\_\theta(\mathbf{z}\_t,t))$ sets certain log-likelihood terms to zero.
We thank the reviewer for their constructive feedback. We address the concerns and questions below.

---

### ****Concern 1:**** Adding algorithms for training and inference.

Below we provide pseudocode for MDLM training and inference. We also include these in our revised manuscript. We call out the simplicity of the training algorithm and its similarity to BERT-style training of masked language models:

****Algorithm: MDLM Training:**** Inputs: dataset $\mathcal{D}$, monotonic masking schedule $\alpha_t : [0,1] \to [0,1]$, BERT-like model $x_\theta$, $\mathbf{m}$ denotes the mask vector, $\mathbf{x}^{1:L}$ denotes a sentence with $L$ tokens, $\mathbf{z}^{1:L}_t$ denotes the latent vector with $L$ tokens.  

1. **repeat**
    1. $\mathbf{x}^{1:L} \sim \mathcal{D}$       *// Sample a sequence of length L from dataset (can be a batch)*
    2.  $\mathbf{z}^\ell_t \sim \text{Cat}(\mathbf{z}^\ell_t; \alpha_t \mathbf{x}^\ell + (1 - \alpha_t)\mathbf{m})$   $\forall 1 \leq \ell \leq L$ for random $t \sim \mathcal{U}[0, 1]$  *// Mask each token $\mathbf{x}^\ell$ independently with masking rate $\alpha_t$ to obtain the latent $\mathbf{z}^{1:L}_t$.*
    3. Update weights $\theta$ of denoising (BERT) model $\mathbf{x}_\theta$ by gradient descent step on 

       $\nabla_\theta \frac{\alpha'\_t}{1 - \alpha\_t} \sum_{\ell} \log \langle \mathbf{x}_\theta^\ell(\mathbf{z}^{1:L}_t), \mathbf{x}^\ell \rangle$  // Note: this is simply a “weighted” BERT-style loss
        
    
    **until** converged
    

Note that the only differences to BERT are that in standard BERT

1. In step 1.b., $\alpha_t$ is a constant (set to 0.85 implying a fixed masking rate of 15%).
2. In step 1.c. there is no weighting term in front of the cross-entropy loss.

Otherwise the training algorithms are equivalent!

****Algorithm: MDLM Sampling / Inference****

1. $\mathbf{z}^{1:L}_1 =$ [MASK, MASK, …, MASK]    *// Sampling process starts with all MASK tokens.*
2. **for** t = {$1, \frac{T-1}{T}, \dots, \frac{1}{T}$} **do**
    1. $s \xleftarrow{} t - \frac{1}{T}$ 
    2. $\mathbf{z}^\ell_{s} \sim \text{Cat}(\mathbf{z}^\ell_s; \frac{(1 - \alpha_s)\mathbf{m} + (\alpha_s - \alpha_t)\mathbf{x}_\theta(\mathbf{z}_t)}{1 - \alpha_t} )$ if  $\mathbf{z}_t^\ell = \text{[MASK]}$ else $\mathbf{z}_s^\ell = \mathbf{z}_t^\ell$   $\forall 1 \leq \ell \leq L$.  
    3. $\mathbf{z}^{1:L}_t \xleftarrow{} \mathbf{z}^{1:L}_s$
    
    end **for**
    
3. **return** $\mathbf{z}_0$
**### **Question 4:**** What are the authors' thoughts on MDLM versus continuous-diffusion models for text, such as Plaid and CDCD, in terms of training, model performance, and inference efficiency?

Several works using continuous diffusion for discrete data (e.g., Plaid [6], Diffusion LM [9], CDCD [10]) have not been able to approach AR perplexity. In our work, we demonstrate that discrete diffusion can successfully close this performance gap to AR models. Additionally, in our previous experience working with continuous diffusion models, such as DiffusionLM [9], we have found them to be less stable to train and require *ad hoc* techniques such as nearest embedding clipping. More recent attempts such as Plaid [6] and CDCD [10] seem promising but do not yield good PPL values.

One of the potential benefits of discrete diffusion relative to AR sampling is more efficient sampling if the number of diffusion steps $T$ can be made less than the sequence length $L$, especially in batch size 1 regimes (e.g. on-device inference). We believe that discrete diffusion methods, such as ours, are more well suited to realize this gain relative to continuous diffusion for discrete data.

**### **Question 5:**** Is it possible to speed up the generation process of MDLM, similar to the advanced ODE-solvers people are using in continuous-space diffusion models?

This is indeed  a very promising research question, which we hope to explore in future work. 

**### **Question 6:**** How does MDLM compare to MaskGiT?

While both MDLM and MaskGiT [11] use BERT-style losses to train generative models, the main differences are as follows:

- **Objective:** In MDLM, we derive a tight and principled lower bound to the log-likelihood. MaskGiT’s objective, while effective, is not weighted appropriately to make it a proper ELBO.

- **Sampler:** In MDLM, we use a principled ancestral sampling method, whereas in MaskGiT once again the sampling is done in a more heuristic manner that yields impressive results in the image generation domain. Additionally, MDLM can unmask any number of tokens at a given time step unlike MaskGiT where only a fixed number of tokens are unmasked at each iteration.

---

#### **References:**

[1] Kingma, Diederik, et al. "Variational diffusion models." Advances in neural information processing systems 34 (2021): 21696-21707.

[2] Ziegler, Zachary, and Alexander Rush. "Latent normalizing flows for discrete sequences." International Conference on Machine Learning. PMLR, 2019.

[3] Graves, Alex, et al. "Bayesian flow networks." arXiv preprint arXiv:2308.07037 (2023).

[4]  Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." Advances in Neural Information Processing Systems 34 (2021): 17981-17993.

[5] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).

[6] Gulrajani, Ishaan, and Tatsunori B. Hashimoto. "Likelihood-based diffusion language models." Advances in Neural Information Processing Systems 36 (2024).

[7] Paperno, Denis, et al. "The LAMBADA dataset: Word prediction requiring a broad discourse context." arXiv preprint arXiv:1606.06031 (2016).

[8] He, Zhengfu, et al. "Diffusionbert: Improving generative masked language models with diffusion models." arXiv preprint arXiv:2211.15029 (2022).

[9] Li, Xiang, et al. "Diffusion-lm improves controllable text generation." Advances in Neural Information Processing Systems 35 (2022): 4328-4343.

[10] Dieleman, Sander, et al. "Continuous diffusion for categorical data." arXiv preprint arXiv:2211.15089 (2022).

[11] Chang, Huiwen, et al. "Maskgit: Masked generative image transformer." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.
### **Concern 4:**  It's unclear how well the approach scales to larger models

We agree with the reviewer that finding scaling laws for discrete diffusion is a very interesting line of research, and one that we hope to conduct in follow-up work. Unfortunately, we cannot run such experiments in the limited rebuttal window and with a limited (academic) compute budget.

---

### **Concern 5:** It's unclear if MDLM can achieve practical speed-up comparing to AR models with same time budget.

We highlight that MDLM can generate tokens in parallel and supports varying the number of generation steps, whereas AR models are limited by a fixed budget due to sequential generation. **MDLM achieves up to 1.8x speedup compared to AR** by varying the generation steps for sampling 64 batches of 256 tokens. **MDLM reaches 8% better generative perplexity under roughly the sampling speed as AR** (within 4%).

More generally, our number of function evaluations (FE) to generate, say, a block of 128 tokens will be at most 128 steps. At the same time, because inference is typically memory bound, the wall clock time to run one (FE) is about the same for 1 block as it is for 128 blocks (and this is true up to a block size of, say, ~200 on an H100, after which we become compute bound again). Therefore, while an optimized implementation of MDLM is outside the scope of this paper, our diffusion approach holds the promise of significant speed improvements.

---

### Additional questions

**### **Question 1:**** How do the authors explain their observation that time-step conditioning can be optional in MDLM?

In absorbing state discrete diffusion the time step / noise level can be inferred from the number of masked tokens and hence explicit conditioning on time step embedding is not necessary. DiffusionBERT [8] also empirically found that for the absorbing state diffusion, time-conditioning the denoising model is ****not**** critical. 

**### **Question 2:**** How many steps are used when evaluating the perplexity of MDLM? How does this compare to other text diffusion models?

There are different ways to interpret this question. First, sampling steps could refer to the value of $T$ used at training time. We report our best PPL values with the continuous-time formulation ($T \to \infty$) for MDLM and SEDD, and report results with finite $T$ among our ablations. A second interpretation is that “sampling steps” could refer to the number of steps to generate text from the model when we evaluate generative PPL. This number is reported on the x-axis of our generative PPL figure and varies between 32 and the length of the sequence.

A third interpretation refers to the Monte Carlo sampling to evaluate PPL. When evaluating the perplexity bound of MDLM or SEDD, we use a single Monte Carlo sample per datapoint to approximate the integral of the continuous-time bound. Our low-discrepancy antithetic sampler allows us to estimate the perplexity bound with low variance, as shown in the table below. On OpenWebText, we find that we are able to accurately estimate the perplexity bound with just a single sample. In contrast, prior work required 1,000 samples to obtain an accurate estimate.

Below we include a table that shows the limited effect of including more time steps in the Monte-Carlo estimate of MDLM on the OpenWebText dataset:
| Num MC Samples | 1 | 2 | 5 | 10 | 100 |
| --- | --- | --- | --- | --- | --- |
| Perplexity ($\downarrow$) | 23.21 | 23.22 | 23.22 | 23.22 | 23.21 |

**### **Question 3:**** Can the authors provide some comparative analysis and plots showing how MDLM compares to AR baselines across a range of training token amount budgets?

We find that AR models are able to optimize the training loss better than diffusion models [see fig. 1 in the attached supplementary material](https://openreview.net/forum?id=L4uaAR4ArM&noteId=WereEXicWc) However, diffusion models are able to reach the same loss values with additional training. Crucially, the gap between the two curves is much smaller than previously thought. We are excited by the new opportunities that non-AR generation presents, such as methods for controlled generation, and the promise of more efficient sampling.
We want to thank the reviewer for their constructive feedback. We address the reviewers comments and questions below.

---

### ****Concern 1:**** Why does MDLM outperform previous methods? There is a need for a deeper theoretical analysis.

Our work outperforms previous discrete diffusion methods because our evidence lower bound is theoretically tighter. There are two specific methodological elements that achieve this:

1. **Continuous Time ELBO:** Deriving a continuous time bound by taking $T \to \infty$.
2. **ELBO Simplification:** Setting certain terms in the ELBO analytically to 0, which is possible because of our choice of SUBS parameterization (7)

Our ablation table shows that these elements can explain the gap relative to the previous state-of-the-art (SEDD), which uses the exact same backbone, optimizer, etc., and differs only in the training objective and parametrization.

| Method | PPL (↓) |
| --- | --- |
| MDLM | 27.04 |
| - w/o continuous time ELBO (1) | 27.19 |
|   - w/o ELBO simplification & SUBS (2) | 28.56 |
| SEDD (our implementation) | 28.90 |

Next, we seek to understand why these elements improve perplexity. Recall that the diffusion ELBO can be written as follows.

$\mathcal{L}_{ELBO} =$

$\mathcal{L}_{recon}$

$+ \mathcal{L}_{prior}$ 

$+\mathcal{L}_{diffusion}$

Our design elements improve each term.

- $\mathcal{L}_{prior}$: The prior regularization term simplifies analytically to zero by design, since the noise schedule $\alpha(t)$ is set such that $\alpha(t=1) = 0$ (as in prior works [4, 5]).
- $\mathcal{L}_{recon}$: The reconstruction loss term simplifies analytically to zero due to taking $T \rightarrow \infty$ and our “copy over unmasking” SUBS parameterization (i.e. we copy over unmasked tokens from $\mathbf{z}_t$ to $\mathbf{z}_s$).

- $\mathcal{L}_{diffusion}$: This diffusion loss term simplifies to a weighted cross entropy (BERT-style) loss due to the SUBS parameterization and we make the ELBO as tight as possible by taking $T \rightarrow \infty$ (see also VDM [1] for a similar analysis of how the lower-bound becomes tighter as $T \rightarrow \infty$).

In addition, our work shows that masked diffusion models like MDLM and D3PM work better than previously thought when paired with a training recipe. Specifically, they are competitive with score-based methods such as SEDD. We closed the gap using both the modeling contributions above, as well as a well-engineered, numerically stable implementation and effective training recipe, including a low-discrepancy sampler that reduces the variance of both perplexity bound evaluations and our gradient estimators for training.

In our revised manuscript, we include a more detailed appendix that includes the above explanations and shows how our analysis and parameterization lead to the final simplified ELBO.

---

### ****Concern 2:**** Missing comparison to non-diffusion language models

In addition to the extensive set of baselines that we provide in Tables 1 and 2, we have added a new experiment on the `text8` dataset, which allows us to compare against many non-diffusion results reported in the literature. As requested by the reviewer, this table contains non-diffusion baselines, including Flow-based models (IAF/SCF) [2] and a Bayesian Flow Network (BFN) [3].

| Method | BPC($\downarrow$) | Train steps |
| --- | --- | --- |
| IAF/SCF | 1.88 | 1M |
| AR Argmax Flow | 1.39 | 1M |
| BFN | 1.41 | 1M |
| D3PM | 1.45 | 1M |
| SEDD | 1.39 | 1M |
| MDLM | 1.44* | 0.4M* |
| AR | 1.23 | 1M |

* We provide BPC for an MDLM model that was only trained for **40%** of the total training steps, compared to all baselines. We hope to report the metrics for a fully trained model during the discussion period. Our partially trained MDLM is within 4% of the top-performing non-AR method, SEDD. Note also that our result is based on a single training run with no hyper-parameter tuning.

The details of this `text8` experiment are as follows: following D3PM [4], we train a 77M param model on sequence chunks of 256 characters with batch size 512; baseline values were taken from SEDD [5].

However, please note that **our existing experiments already compare to the strongest possible baselines**, including the best existing diffusion-based model (SEDD), and an AR baseline.

---

### ****Concern 3:**** Including additional non-perplexity-based metrics

While perplexity is the dominant metric used in language model evaluation and has been extensively shown to correlate with downstream performance, we agree with the reviewer’s comment that other metrics that directly estimate downstream task accuracy are useful. We therefore evaluated our model on the Lambada [7] benchmark and present results below:
|  | Lambada Accuracy ($\uparrow$) |
| --- | --- |
| AR | 49.62% |
| MDLM | **53.10%** |

We report accuracy of predicting the final word given a context of at least 50 tokens. MDLM improves over AR, aligning with our zero-shot results in Table 3.
### **Concern 2:** Is it better to lead the paper with MDLM or with the simplified objective as a key contribution?

To clarify, we use the term MDLM to refer to a language modeling algorithm that is optimized for masked language discrete diffusion and that has the following components:

- **Parameterization:** Our denoising network predicts clean data and uses the SUBS parameterization (carry over masking and zero mask probabilities).

- **Learning objective:** We take $T \rightarrow \infty$ to train with a tight ELBO given in Equation 11.

- **Faster inference:** By removing time-conditioning, we significantly reduce function evaluations during ancestral sampling.

- **Support for arbitrary sequence length generation:** Using a semiAR algorithm, we alleviate a key shortcoming of diffusion models: fixed length generation.

Note that these are novel features over SEDD or D3PM. Thus, while the objective is an important part of our work, it is only one novel part of the MDLM algorithm.

We also acknowledge that we are not the first ones to consider masking diffusion. However, we are the first to optimize our algorithm for masking diffusion, and this motivates its name (MDLM).

---

### **Concern 3:** The use of a more advanced backbone (DiT) weakens the claim that the performance comes from the proposed method.

As the table below demonstrates, our ****experimental setup is effectively identical**** between our work, AR, and SEDD (the previous state-of-the-art diffusion-based model), including the backbone. Thus, our ability to more effectively reduce the gap between diffusion and AR models is directly related to our proposed methodology.

Experimental details:

|  | SEDD | MDLM | AR |
| --- | --- | --- | --- |
| Tokenizer | `bert-base-uncased` for LM1B / `gpt-2` for OWT   | `bert-base-uncased` for LM1B / `gpt-2` for OWT | `bert-base-uncased` for LM1B / `gpt-2` for OWT |
| Architecture | DiT | DiT | DiT |
|  Context size | 128 for LM1B / 1024 for OWT | 128 for LM1B / 1024 for OWT | 128 for LM1B / 1024 for OWT |
| Train steps | 1M steps | 1M steps | 1M steps |
| Perplexity ($\downarrow$) on LM1B (1M train steps) | 32.79 | 27.04 | 22.32 |
| Perplexity ($\downarrow$) on OWT | 24.10 | 23.21 | 17.54 |

To further clarify: in Section 3.5.1, we specifically highlighted architectural details used in MDLM that contribute to a performance relative to the results in the original D3PM paper. However, our main tables compare against key baselines (SEDD, AR, and D3PM, which is MDLM after all ablations) **with the same architecture.**

---

### Additional question

> **Is MLM pre-training required for your method?**

No, pre-training with standard MLM loss is ****not**** a requirement of our method. Indeed, the main results of our work, Tables 1 and 2, do not rely on any pre-training. To clarify, in Table 4, we included results of fine-tuning a pre-trained BERT using MDLM to highlight how one could take an off-the-shelf pre-trained representation learning model and turn it into a generative one, without sacrificing the representational learning capabilities of the pre-trained model.
We want to thank the reviewer for their constructive feedback. We address concerns and questions below.

---

### ****Concern 1:**** Novelty of MDLM relative to the D3PM framework and other algorithms.

In addition to attaining ****state-of-the-art diffusion language model results****, MDLM presents important novel elements relative to prior work.

**Relative to D3PM:** [1]

MDLM is a special case of the D3PM framework that focuses only on masked noise. This allows introducing multiple algorithmic innovations that greatly simplify D3PM and improve performance.
These design decisions also yield a learning objective that is a ****weighted average of BERT-style loss terms**.** Although D3PM mentions in their appendix connections to BERT, we make this connection much more clear and derive a simpler and more performant algorithm. Our NELBO is given as:
$\mathcal{L}\_{NELBO} = \int_{t=0}^{t=1}\frac{\alpha'\_t}{1 - \alpha\_t} \log p_\theta(\mathbf{x}^{1:L}|\mathbf{z}^{1:L}_t) dt$ 
which achieves a test perplexity of **27.04** after being trained for 1M steps on LM1B while our well engineered implementation of D3PM achieves a worse perplexity of 28.56.


- **Algorithmic innovations**
    - **Simplified noise process:** In D3PM, the noise process is defined via matrix multiplications. In MDLM, we can define simple interpolating noise $\mathbf{z}_t = \alpha_t \mathbf{x} + (1 - \alpha_t)\mathbf{m}$ which allows easily taking $T \to \infty$, which in turn tightens the ELBO. 
    - **Specialized denoising process:** We propose the SUBS parameterization, which includes two key elements: carry-forward unmasking and zero masking probabilities. 
    - **Simplified and tightened ELBO:** Our choice of noising and denoising processes allows us to greatly simplify the D3PM objective (several pages of math in the appendix) and produce a *simplified* and *tighter* ELBO (compare our Equation 8 to Equation 9, which is what D3PM uses).

- **Experimental innovations** We complement the above innovations with a modern training recipe (architecture, optimization) that demonstrates the effectiveness of masked diffusion models. While previous attempts (namely D3PM) seemed to indicate that there exists a large gap between AR and discrete diffusion models, we provide clear evidence to the contrary on the OWT dataset using to train GPT-2 sized models.


**Relative to SEDD:** [2]

Our key novelty elements are as follows.
- We achieve better language modeling results relative to this work, which was previously ****state-of-the-art for diffusion models.****
- We present a more intuitive and simple objective which boils down to a weighted average of BERT-style losses, circumventing the need to use score-matching techniques and the formalisms of continuous time Markov chains.
- We support a more efficient inference algorithm that greatly reduces required function evaluations.

**Relative to DiffusionBERT [3]:**

- DiffusionBERT can be described as an instance D3PM with a pre-trained model initialization and a custom noising schedule. Thus, all of the ELBO improvements and optimized training of our work relative to D3PM are also novel relative to DiffusionBERT.


---

****References:****

[1]  Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." Advances in Neural Information Processing Systems 34 (2021): 17981-17993.

[2] Lou, Aaron, Chenlin Meng, and Stefano Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." arXiv preprint arXiv:2310.16834 (2023).

[3] He, Zhengfu, et al. "Diffusionbert: Improving generative masked language models with diffusion models." arXiv preprint arXiv:2211.15029 (2022).
> **Comparing efficiency w/ Mamba under different token numbers**

We compare the wall clock times of 1 forward pass of the bidirectional transformer used in MDLM vs. Mamba under varying context lengths. While Mamba achieves favorable runtimes at context lengths greater than 8k tokens, there is not a significant speed improvement at the context lengths studied in this paper (where we seek to match GPT-2 using diffusion). Although long context generation is not a focus of this work, we believe it is an interesting area for future extensions.

| # of tokens (wikitext103) | 2k | 4k | 8k | 16k | 32k |
| --- | --- | --- | --- | --- | --- |
| Mamba (runtime in ms) | 11 | 12 | 22 | 42 | 83 |
| Transformer (runtime in ms) | 13 | 17 | 36 | 88 | 250 |


We report the above runtimes on a single A5000 GPU, excluding the wall clock time of word embeddings.
### **Concern 2:** MDLM involves a number of hyper-parameters. It is unclear these are same for baselines, and whether the comparison is fair.

We stress that our experiments **use exactly the same hyper-parameters** (tokenizer, network structure, training set, optimizer settings, etc.) **across every model** that we have implemented (AR, SEDD, which is the current state-of-the-art in discrete diffusion, MDLM, as well as models in other tables such as D3PM).

|  | SEDD | MDLM | AR |
| --- | --- | --- | --- |
| Tokenizer | `bert-base-uncased` for LM1B / `gpt-2` for OWT   | `bert-base-uncased` for LM1B / `gpt-2` for OWT | `bert-base-uncased` for LM1B / `gpt-2` for OWT |
| Architecture | DiT | DiT | DiT |
|  Context size | 128 for LM1B / 1024 for OWT | 128 for LM1B / 1024 for OWT | 128 for LM1B / 1024 for OWT |
| Train steps | 1M steps | 1M steps | 1M steps |
| Perplexity ($\downarrow$) on LM1B (1M train steps) | 32.79 | 27.04 | 22.32 |
| Perplexity ($\downarrow$) on OWT | 24.10 | 23.21 | 17.54 |

To further clarify: in Section `3.5.1`, we specifically highlighted architectural details used in MDLM that contribute to a performance relative to the results in the original D3PM paper. However, our main tables compare against key baselines (SEDD, AR, and D3PM, which is MDLM after all ablations) **with the same architecture.**

---

### ****Concern 3:**** It is unclear how semi-autoregressive sampling works.

Naively, diffusion models only generate fixed-length sequences. Semi-autoregressive sampling allows diffusion to produce arbitrary-length sequences by generating (via diffusion) a block of text that is conditioned on previously generated blocks of text (blocks previously obtained via diffusion).

****Algorithm: MDLM Semi-AR sampling****

Generates a Sequence of length $B \times L$ where $L$ is the context length of the diffusion model.

1. $S = \phi$    *// We start with a null sequence.*
2. $\mathbf{x}^{1:L} =$$[MASK]^{1:L}$   *// Sampling process starts with all MASK tokens.*
3. **for** $i$ = {${1, \dots, B - 1}$} **do**
    1. Unmask all masked tokens in $\mathbf{x}^{1:L}$  using the **MDLM Sampling algorithms** as described above.
    2. $S \xleftarrow{} S \cup \mathbf{x}^{1:L/2}$  // Save the first $L/2$ tokens
    3. $\mathbf{x}^{1:L/2} \xleftarrow{} \mathbf{x}^{L/2:L}$ // Condition the generation of future tokens on the trailing $L/2$ tokens
    4. $\mathbf{x}^{L/2:L} \xleftarrow{} [MASK]^{1:L/2}$ // Mask out the last $L/2$ tokens to be filled out by MDLM in the next iteration.
    
    end **for**
    
4. return $S$

For a figure illustrating how our proposed semiAR algorithm works, please see the schematic in the attached supplementary materials [here](https://openreview.net/forum?id=L4uaAR4ArM&noteId=WereEXicWc)

----

### **Additional questions and concerns**

> **Include sampling steps for Table 1**

There are different ways to interpret this question. First, sampling steps could refer to the value of T used at training time. We report our best PPL values with the continuous-time formulation ($T \to \infty$) for MDLM and SEDD, and report results with finite $T$ among our ablations. A second interpretation is that “sampling steps” could refer to the number of generate text from the model when we evaluate generative PPL. This number is reported on the x-axis of our generative PPL figure and varies between 32 and the length of the sequence.

A third interpretation refers to the Monte Carlo sampling to evaluate PPL. When evaluating the perplexity bound of MDLM or SEDD, we use a single Monte Carlo sample per datapoint to approximate the integral of the continuous-time bound. Our low-discrepancy antithetic sampler allows us to estimate the perplexity bound with low variance, as shown in the table below. On OpenWebText, we find that we are able to accurately estimate the perplexity bound with just a single sample. In contrast, prior work required 1,000 samples to obtain an accurate estimate.

Below we include a table that shows the limited effect of including more time steps in the Monte-Carlo estimate of MDLM on the OpenWebText dataset:
| Num MC samples            | 1     | 2     | 5     | 10    | 100   |
|---------------------------|-------|-------|-------|-------|-------|
| Perplexity ($\downarrow$) | 23.21 | 23.22 | 23.22 | 23.22 | 23.21 |

> **Better explaining Figure 1**

The loss in Figure 1 is computed over red boxes only. Since the yellow boxes are unmasked, by the definition of our SUBS parameterization, the loss there is zero by construction. We are adding a legend to explain this.

> **Mathematical proof for continuity equation.**

There is currently no mention of $u_t$ in our work. Can you please provide further clarification as to what you are requesting in this regard?
We want to thank the reviewer for their constructive feedback. We address specific comments and questions below.

---

### ****Concern 1:**** What is “simple”? It was hard to understand the method.

We use the term “simple” because **our algorithm is very similar to BERT** except for a small change: a random masking rate. Please see the general response for how one of our main contributions, the SUBS parameterization and continuous time perspective, allow us to simplify the training objective.

We will clarify the method by providing the following pseudocode in the next version of the paper. We use **$x$** to denote a one-hot word embedding, and **$m$** is a one-hot vector encoding the mask token.

**Algorithm: MDLM Training:** Inputs: dataset $\mathcal{D}$, monotonic masking schedule $\alpha_t : [0,1] \to [0,1]$, BERT-like model $x_\theta$.

1. **repeat**
    1. $\mathbf{x}^{1:L} \sim \mathcal{D}$       *// Sample a sequence of length L from dataset (can be a batch)*
    2.  $\mathbf{z}^\ell_t \sim \text{cat}(\mathbf{z}^\ell_t; \alpha_t \mathbf{x}^\ell + (1 - \alpha_t)\mathbf{m})$   $\forall 1 \leq \ell \leq L$ for random $t \sim \mathcal{U}[0, 1]$  *// Mask each token $\mathbf{x}^\ell$ independently with masking rate $\alpha_t$ to obtain the latent $\mathbf{z}^{1:L}_t$.*
    3. Update weights $\theta$ of denoising (BERT) model $\mathbf{x}_\theta$ by gradient descent step on 

        $\nabla_\theta \frac{\alpha'\_t}{1 - \alpha\_t} \sum_{\ell} \log \langle \mathbf{x}\_\theta^\ell(\mathbf{z}^{1:L}_t), \mathbf{x}^\ell \rangle$  // Note: this is simply a “weighted” BERT-style loss
        
    
**until** converged
    

The only differences relative to BERT are that at step (1.2), $\alpha_t$ is a constant, and in step (1.3.) there is no weighting term in front of the cross-entropy loss. Otherwise the training algorithms are equivalent! Indeed, one could even take an off-the-shelf pre-trained BERT and with minimal modifications to the training loop, render it a generative model (Table 4) with an principled ELBO objective and sampling algorithm.

Additionally, unlike previous works, such as CTMC [1] and SEDD [2], which rely on the machinery of continuous-time Markov chains, our work admits a straightforward generation algorithm based on ancestral sampling that we present below:

**Algorithm: MDLM Sampling / Inference**

1. $\mathbf{z}^{1:L}_1 =$ [MASK, MASK, …, MASK]    *// Sampling process starts with all MASK tokens.*
2. **for** t = {$1, \frac{T-1}{T}, \dots, \frac{1}{T}$} **do**
    1. $s \xleftarrow{} t - \frac{1}{T}$ 
    2. $\mathbf{z}^\ell_{s} \sim \text{Cat}(\mathbf{z}^\ell_s; \frac{(1 - \alpha_s)\mathbf{m} + (\alpha_s - \alpha_t)\mathbf{x}_\theta(\mathbf{z}_t)}{1 - \alpha_t} )$ if  $\mathbf{z}_t^\ell = \text{[MASK]}$ else $\mathbf{z}_s^\ell = \mathbf{z}_t^\ell$   $\forall 1 \leq \ell \leq L$.  
    3. $\mathbf{z}^{1:L}_t \xleftarrow{} \mathbf{z}^{1:L}_s$
    
    end **for**
    
3. **return** $\mathbf{z}_0$

To further underscore the simplicity of our method, we highlight the minimal form of our variational lower bound, Equation 10, which ends up being a weighted average of BERT-style losses:
$\mathcal{L} = \mathbb{E}\_q\int_{t=0}^{t=1}\frac{\alpha\_t'}{1-\alpha\_t} \log \langle\mathbf{x}_\theta(\mathbf{z}_t), \mathbf{x}_0\rangle dt$

In summary, our work presents a principled approach to turn widely-used bi-directional language models into generative ones.
"In addition to the individual responses we provide directly to each of your comments"