A Reparameterized Discrete Diffusion Model for Text Generation

Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-22 03:47:19

Abstract: This work studies discrete diffusion probabilistic models with applications to natural language generation. We derive an alternative yet equivalent formulation of the sampling from discrete diffusion processes and leverage this insight to develop a family of reparameterized discrete diffusion models. The derived generic framework is highly flexible, offers a fresh perspective of the generation process in discrete diffusion models, and features more effective training and decoding techniques. We conduct extensive experiments to evaluate the text generation capability of our model, demonstrating significant improvements over existing diffusion models.

Main paper: 
 introduction: 
 Diffusion-based generative models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2021b) , or diffusion models for short, have achieved remarkable progress and shown great success in generating high-quality and photo-realistic images (Ramesh et al., 2022; Saharia et al., 2022; Rombach et al., 2022; Balaji et al., 2022; Peebles & Xie, 2022) . Researchers have successfully extended diffusion models to various data modalities beyond 2D images, including audio (Kong et al., 2021) , video (Ho et al., 2022) , as well as molecule generation (Hoogeboom et al., 2022b; Jo et al., 2022) . There has also been a surge of interest in extending diffusion models to natural languages (Hoogeboom et al., 2021; Austin et al., 2021; Li et al., 2022b; Dieleman et al., 2022) . Diffusion-based language models are appealing in that the generation process is typically conducted in a non-autoregressive manner, which features in-parallel decoding by design and potentially a faster runtime (Hoogeboom et al., 2021; Austin et al., 2021) . In addition, due to the iterative reconstruction process in diffusion-based models, it is often possible to refine the previously generated texts (Savinov et al., 2022) . As a result, compared with the conventional auto-regressive models, diffusion-based language models are more flexible and attain better trade-offs between generation quality and efficiency.However, there are noticeably fewer success cases in employing diffusion models for large-scale text generation tasks. This is possibly due to the discrete nature of natural languages, while most conventional diffusion models focus on continuous-valued contents. To bridge the discrepancy, a recent line of work suggests conducting continuous diffusion processes over token embeddings (Li et al., 2022b; Gong et al., 2022; Strudel et al., 2022; Dieleman et al., 2022) or logits (Han et al., 2022; Richemond et al., 2022) . Nevertheless, these approaches often require designing a well-crafted rounding scheme to convert the diffused continuous vectors to the actual discrete tokens. In addition, existing continuous diffusion models often require a large number of sampling iterations to achieve the desired performance. This issue is exacerbated in the case of modeling texts, as the diffusing steps over text embeddings are hard to be translated to significant movements of token states due to the rounding quantization. This results in a considerably slower runtime than the auto-regressive language models. For example, a recent continuous text diffusion model (DiffuSeq; Gong et al., 2022) runs several orders of magnitude slower than the auto-regressive baseline of a similar scale, as shown in Figure 2b . Different from the above, another research direction focuses on diffusion processes that directly operate on discrete state spaces (Sohl-Dickstein et al., 2015; Hoogeboom et al., 2021; Austin et al., 2021, see §2) . However, they are relatively under-explored and often achieve inferior results in text generation.In this work, we demonstrate that discrete diffusion models can serve as strong baselines for text generation. By re-examining the formulation, we observe that sampling from discrete diffusion models admits a novel yet equivalent reparameterization ( §3). Specifically, starting with a completely noisy sequence, the sampling procedure in a discrete diffusion model is equivalent to the following route-and-denoise process where at each iteration, each token within the sequence is either denoised or reset to noisy states according to an underlying stochastic routing mechanism ( §3.2). The router assigns the same probability to the decision for each token, processing the sequence in a uniform manner. Based on this insight, we propose Reparameterized Discrete diffusion Models (RDMs; Figure 1 ), a new family of models that respects the reparameterization and formulates the routing process explicitly. RDMs enjoy appealing properties for both training ( §4.2) and sampling ( §4.3):(1) Simplified training. We demonstrate that the training objective for RDMs can be reduced to a re-weighted standard cross-entropy loss. Furthermore, the loss objective is invariant to different routing probabilities up to reweighting, indicating that the large family of RDMs with distinct routing processes can be trained with the same surrogate objective; (2) Flexible sampling. The shared training objective makes sampling highly flexible and allows more expressive routing processes. In particular, we develop an adaptive routing strategy that routes tokens to the denoised state only if the router outputs high scores instead of uniformly processing all the tokens. Equipped with such training and decoding schemes, we demonstrate that RDMs significantly improve vanilla discrete diffusion models across several standard text generation benchmarks ( §5). They also achieve much better performance than the continuous diffusion models while running several orders faster. 
background: 
 x 0x 1 x 2 x 3 (a) Conventional discrete diffusion.x 0 x 1 x 2 x 3 v 1 v 2 v 3 (b) Reparameterized discrete diffusion.Figure 1 : Graphical models of backward diffusion process.Let x 0 ∼ p data (x 0 ) denote a discrete random variable with K possible outcomes. To ease notation, we represent discrete variables as one-hot vectors in {0, 1} K , which is 0 everywhere except that the entry corresponding to the current state is 1. Discrete diffusion probabilistic models (Sohl-Dickstein et al., 2015; Hoogeboom et al., 2021; Austin et al., 2021) are usually defined as a class of latent variable models characterized by a forward and backward process. The forward process aims to gradually transform input data to some noise distribution q noise through T intermediate latent variables x 1 , . . . , x T ∈ {0, 1} K , with the forward transition q(x t |x t-1 ) = β t x t-1 + (1 -β t )q noise . In this case, the distribution q(x t |x 0 ) is available in closed form,EQUATIONwhere α t := t i=1 β i is specified to decrease from 1 to 0 w.r.t. t. q noise characterizes different diffusion processes; for example, multinomial diffusion (Hoogeboom et al., 2021) adopts a uniform noise distribution over the vocabulary {1, 2, . . . , K}; alternatively, absorbing diffusion specifies q noise to be a point mass with all of the probability on an absorbing state (Austin et al., 2021) .The backward process q(x t-1 |x t ) is the key ingredient for diffusion-based generative modeling, based on which we can start with unstructured noise x T ∼ q noise and perform ancestral sampling x t-1 ∼ q(x t-1 |x t ) to obtain new draws from p data (x 0 ). Unfortunately, the backward transition q(x t-1 |x t ) is mostly intractable due to the marginalization over the entire data distribution. Therefore, we resort to approximating it with a parameterized distribution p θ (x t-1 |x t ) at each step t. This results in a generative model p θ (x 0 , x 1 , . . . , x T ) = p θ (x T ) T t=1 p θ (x t-1 |x t ), which can be trained by maximizing the evidence lower bound (ELBO) of log p θ (x 0 ),log p θ (x 0 ) ≥ L 1 (θ) - T t=2 L t (θ) + const., with L t (θ) := E q(xt|x0) [KL(q(x t-1 |x t , x 0 ) ∥ p θ (x t-1 |x t ))]. For the case t = 1, we have L 1 (θ) := E q(x1|x0) [log p θ (x 0 |x 1 )].The ELBO decomposes into a sum of KL divergences between the conditional backward transition q(x t-1 |x t , x 0 ) and p θ (x t-1 |x t ) at each time step t. Note that q(x t-1 |x t , x 0 ) ∝ q(x t |x t-1 )q(x t-1 |x 0 ) can be calculated analytically for most discrete diffusion models. To define the distribution p θ (x t-1 |x t ), previous work (Hoogeboom et al., 2021) suggests that it can be parameterized in a similar manner to q(x t-1 |x t , x 0 ) by letting p θ (x t-1 |x t ) = q(x t-1 |x t , f (x t ; θ)), where a neural network f (x t ; θ) is adopted to predict x 0 . Typically f (x t ; θ) ∈ (0, 1) K is the model output normalized by a softmax function, representing the probability vector for each token. A more detailed review and relevant derivations about discrete diffusion models are provided in Appendix A. 
reparameterizing backward processes: 
 This section presents an in-depth study of discrete diffusion probabilistic models. We derive an alternative formulation for the backward process ( §3.1) and devise a reparameterized sampling scheme ( §3.2), which paves the way for a more generic family of discrete diffusion models ( §4).We first elaborate on how the conditional backward transition of existing discrete diffusion models can be written in a more compact formulation (see Appendix B for the proof). Proposition 3.1. Let the forward transition of discrete diffusion be q(x t |x t-1 ) = β t x t-1 + (1β t )q noise . Then the conditional backward transition q(x t-1 |x t , x 0 ) can be equivalently written asq(x t-1 |x t , x 0 ) =    λ (1) t-1 x t + 1 -λ (1) t-1 q noise , if x t = x 0 λ (2) t-1 x 0 + 1 -λ (2) t-1 q noise (x t ), if x t ̸ = x 0 .(2)Here q noise (x t ) = β t x t + (1 -β t )q noise denotes a noise distribution that interpolates between x t and q noise , λEQUATIONαt+(1-αt)qnoise(u=xt) , and λ(2)t-1 := αt-1-αt 1-αt , where q noise (u = x t ) is the probability of the noise equal to x t .Intuitively, Equation 2 reveals that the main mechanism of the backward process is to shuttle discrete tokens between fully noisy states and the ground truth state x 0 , conditioned on the equality of x t and x 0 . If x t = x 0 , the current token state is possibly noise-free, and the model either remains noise-free by copying the state x t-1 ← x t , or resets the state to the noise. If x t ̸ = x 0 , the current state is considered noisy, and the model opts to denoise x t to x 0 or remains noisy otherwise. The probability of moving noisy tokens to ground truth states or turning denoised tokens back to noise is governed by λ(2) t-1 and 1 -λ(1) t-1 , respectively.Next, we demonstrate that sampling from the backward transition can be conducted via an augmented path, leading to our full reparameterization. We make use of the simple fact that the mixture distribution in Equation 2 can be sampled in two steps: first, randomly select a component according to their weight, and then sample from the corresponding component distribution. Concisely, we haveEQUATIONEQUATIONt ∼ q noise (x t )EQUATIONTo simplify the notation, we denote v t-1 := v(1)EQUATIONt-1 and λ t-1 := λ(1)EQUATIONt-1 . This reparameterized backward transition highlights an underlying routing mechanism, where the model routes tokens to different distributions according to v t-1 : given b t that discriminates which tokens are currently noisy, v(2) t-1 selects and denoises noisy tokens to recover the ground truth x 0 , while v(1) t-1 determines which denoised tokens revert to the noisy state. 
reparameterized discrete diffusion models: 
 In this section, we introduce our proposed diffusion models that reflect the reparameterization ( §4.1), which imply effective training ( §4.2) and sampling ( §4.3) algorithms.The routing mechanism in §3.2 works in a latent manner; that is, it is only active during the sampling process but marginalized when advancing the distribution of x t-1 . To fully utilize the potential of the developed reparameterization, we propose to elevate the latent routing mechanism to the formulation explicitly and model the joint q(x t-1 , v t-1 |x t , x 0 ) = q(v t-1 )q(x t-1 |v t-1 , x t , x 0 ), whereEQUATIONNote that b t = 1 xt=x0 . This can be considered as a standard discrete diffusion model augmented with step-wise routing indicators {v t-1 } T t=1 (see Figure 1 ). It closely relates to the conventional formulation (Equation 2) in that the original backward process amounts to marginalizing out v t-1 at each time step:q(x t-1 |x t , x 0 ) = E q(vt-1) [q(x t-1 |v t-1 , x t , x 0 )].Since the distribution over v t-1 is explicitly considered, this joint diffusion model offers improved flexibility and expressiveness. We refer to this class of models as reparameterized discrete diffusion models (RDMs), as it yields an equivalent sampling process to the original formulation but via a reparameterized path.Similar to previous diffusion models ( §2), we define a joint generative process p θ (x t-1 , v t-1 |x t ) and optimize the ELBO via the following factorization,log p(x 0 ) ≥ E q(x 1:T ,v 1:T |x0) log p θ (x 0 , x 1:T , v 1:T ) q(x 1:T , v 1:T |x 0 ) := L 1 (θ) - T t=2 L t (θ) + const..Following standard practices in diffusion models (Hoogeboom et al., 2021; Austin et al., 2021) , we randomly sample a time step t and optimize θ with respect to L t (θ). In our case, L 1 (θ) = E q(x1|x0) [log p θ (x 0 |x 1 )]; for t > 1, L t decomposes into a sum of two expected KL divergences over v t-1 and x t-1 respectively (see Appendix C for the full derivation),EQUATIONwhere the expectations are with respect to q(x t |x 0 ) and q(x t |x 0 )q(v t-1 ), respectively.Parameterization. The decomposition in Equation 5 suggests parameterizing each conditional ofp θ (x t-1 , v t-1 |x t ) = p θ (x t-1 |v t-1 , x t )p θ (v t-1 |x t ) separately.To simplify the model representation, we constrain p θ (v t-1 |x t ) to be the same as q(v t-1 ) so that the first KL term vanishes. In terms of p θ (x t-1 |v t-1 , x t ), it needs to approximate q(x t-1 |v t-1 , x t , x 0 ) for both x 0 and b t = 1 xt=x0 .For the former, we approximate x 0 with a neural network output f (x t ; θ) ( §2); for the latter, it is circumvented via a teacher-forcing approach. We leverage the fact that b t = 1 xt=x0 is readily available during training and plug the oracle into p θ (x t-1 |v t-1 , x t ), which works well empirically and yields interesting implications as presented below.Simplified Training Objectives. So far, we mainly consider the case of diffusing over a single token. We now slightly abuse the term and denote a sequence of tokens at t-th time step as x t,1:N := {x t,n } N n=1 , where x t,n is the n-th token and N is the sequence length. 1 We also assume distributions factorize over each token. We show that the training objective L t (θ) for sequence x t,1:N at the t-th step can be reduced to a surprisingly simple expression (see Appendix D for the proof and its connection to previous works) as follows,EQUATIONBased on this result, training RDMs is equivalent to optimizing the standard multi-class crossentropy loss function, which is evaluated over noisy tokens and weighted by λ(2)t-1 = E v (2) t-1 .This formulation is conceptually simpler than that of the original discrete diffusion, which requires evaluating the KL divergence between two complicated categoricals (Hoogeboom et al., 2021) .Besides, this formulation establishes the discrete analog of reweighting training strategies, which are recognized as common practices for training continuous-domain diffusion models (Ho et al., 2020; Nichol & Dhariwal, 2021; Vahdat et al., 2021; Karras et al., 2022) . In particular, we can adjust the weight λ(2)t-1 to reweigh the cross-entropy function so that it is more amenable for training. More importantly, we note that the training loss function can be invariant with respect to q(v t-1 ) up to reweighting, where different q(v t-1 ) lead to the same shared objective except for the weight λ(2) t-1 . This makes it possible to train the neural network with one amenable distribution of v t-1 but share the trained model for sampling among a broad family of diffusion processes indexed by q(v t-1 ).Input: neural network f (•; θ), data distribution p data (x 0,1:N ), and a custom reweighting scalar λ t-1 . Output: model parameters θ. repeat Draw x 0,1:N ∼ p data (x 0,1:N ); Draw t ∈ Uniform({1, . . . , T });for n = 1, 2, . . . , N do Draw x t,n ∼ q(x t,n |x 0,n ); Let b t,n = 1 xt,n=x0,n ; end for L(θ) = -λ t-1 N n=1 (1-b t,n )x ⊤ 0,n log f (x t,n ; θ); Minimize L(θ) with respect to θ; until convergedSampling from discrete diffusion processes generally starts with a sequence comprising only noisy tokens and proceeds by drawingx t-1 , v t-1 ∼ p θ (x t-1 , v t-1 |x t ) at each step.Generating v t-1 . A naïve approach of generating v t-1 is by drawing v t-1 ∼ Bernoulli (λ). However, this assigns equal routing probability λ t-1 to all tokens, which may be sub-optimal for discriminating and routing different tokens to distinct states. Fortunately, thanks to §4.2, the training objective can be seamlessly shared across different routing distributions. This observation motivates us to employ a more discriminative routing mechanism, where only tokens with high confidence from neural networks are denoised (Ghazvininejad et al., 2019; Savinov et al., 2022; Chang et al., 2022) . At each step, assume v(1) t-1,n and v(2) t-1,n are initialized to 1 and 0, respectively. We feed the noisy sequence x t,1:N into the neural network and collect the output f (x t,n ; θ) for each token n. We obtain token scores s t,n by taking the maximum value of f (x t,n ; θ) ∈ (0, 1) K , which reflects the model confidence. v t-1,n is then set to 1 only when s t,n is among the k largest scores. Formally,EQUATIONwhere i = 1 if b t,n = 1 and i = 2 otherwise. This strategy is more informative, as the router can compare token scores to determine their states. Although v t-1 is defined as a function of model scores and its explicit probability distribution is possibly intractable, the usage of such adaptive mechanism is well justified by the shared objective functions, since the diffusion model corresponding to this implicit distribution can also be trained effectively with the shared surrogate objective.Recursive Computation for b t . Another challenge during decoding stems from the computation b t,n = 1 xt,n=x0,n , which is intractable since we lack access to the ground truth x 0,n as in training.We demonstrate that this issue can be circumvented by leveraging a recursive computation based on the functionality of v t-1,n . Specifically, starting with b T,n = 0, we can compute b t-1,n as follows,EQUATIONIntuitively, b t,1:N represents a frontier set that stores the denoised tokens up to the previous iteration.At the current iteration, we add new tokens to the set if v(2)t-1,n = 1 and remove elements from the set if the corresponding v(1) t-1,n = 0. The updated set is then read out to form b t-1,1:N .2 Note that the update does not add extra computational costs. Equipped with these results, we can efficiently execute the sampling algorithm without difficulties.Algorithms 1 and 2 list the full pseudo-codes for training and sampling of RDMs, respectively. Note that the loop over the sequence length N is computed in parallel. The developed formulation of RDMs offers great flexibility for both training and sampling. For instance, we can pass a custom λ t-1 to reweigh the loss for training, similar to continuous diffusion models. Besides, the denoised token states x 0 during decoding can be obtained in various ways, such as sampling with annealed temperatures or simply taking the arg max of f (x t,n ; θ). Please refer to Appendix E for the full implementation details. 
experiments: 
 Algorithm 2 Sampling from RDMsInput: trained network f (•; θ) and temperature τ . Output: generated sample x 0 . for n = 1, 2, . . . , N do Initialize x T,n ∼ q noise ; Initialize b T,n = 0; end for for t = T, . . . , 1 do for n = 1, 2, . . . , N do Draw x 0,n ∼ Categorical (f (x t,n ; θ)/τ ); Generate v t-1,n according to Equation 7; if b t,n = 1 then Draw u (1) t,n ∼ q noise ; x t-1,n = v (1) t-1,n x t,n + 1 -v (1) t-1,n u (1) t,n ; else Draw u (2) t,n ∼ q noise (x t,n ); x t-1,n = v (2) t-1,n x 0,n + 1 -v (2) t-1,n u (2) t,n ; end if Let b t-1,n = b t,n ∧ v (1) t-1,n ∨ v (2)t-1,n ; end for end for Return x 0,1:N .We evaluate our model on various text generation benchmarks. Our implementation is based on FairSeq toolkit (Ott et al., 2019) , and a detailed experimental setup can be found in Appendix E.Setup. We conducted machine translation experiments on three standard benchmarks: IWSLT14 DE-EN (Cettolo et al., 2014) , WMT14 EN-DE (Bojar et al., 2014) , and WMT16 EN-RO (Bojar et al., 2016) , consisting of around 160K/7K/7K, 4.0M/3K/3K, and 610K/2K/2K training/validation/testing sentence pairs, respectively. We operate on original data for all translation tasks and do not adopt knowledge distillation (Kim & Rush, 2016; Gu et al., 2018) that replaces the target side of training data with outputs generated by a pre-trained autoregressive Transformer.Results. As seen from Table 1 , existing discrete diffusion models usually perform badly on the translation task and do not scale well for large-sized datasets. In particular, multinomial diffusion achieves worse translation quality albeit with more iterations and fails to decode decently on WMT14 EN-DE dataset. The proposed reparameterization yields significant performance boosts (about 3~20 BLEU improvements) on both absorbing and multinomial diffusion across all datasets and iteration steps. We also observe that the performance gain is much more significant for multinomial diffusion, which is possibly due to the fix of its degenerated behavior (more details are in §5.4). Our approach effectively scales diffusion models to larger datasets, outperforming previous non-autoregressive baselines CMLM (Ghazvininejad et al., 2019) , and achieves promising results even competitive with autoregressive baselines. Besides, RDMs perform significantly better than continuous diffusion models CDCD (Dieleman et al., 2022) while running with more than 8× fewer iterations.Setup. We also evaluate the performance of our model on general sequence-to-sequence generation tasks, following Gong et al. (2022) . Due to the limited computation budget, the tasks we focus on are 1) question generation (QG) using the Quasar-T dataset (Dhingra et al., 2017) with approximately 117K/2K/10K pairs for training/validation/testing, and 2) paraphrasing with Quora Question Pairs (QQP) containing around 145K/2K/2.5K training/validation/testing question pairs.Results. We report comparisons among discrete diffusion, continuous diffusion, and auto-regressive baselines in Table 3 . We observe that either variant of RDMs not only improves their vanilla counterparts by a large margin but also outperforms the strong continuous diffusion baseline DiffuSeq as well as various auto-regressive baselines across several evaluation metrics. Besides, DiffuSeq needs 2000 steps to finish the decoding, while RDMs can produce higher-quality samples with 10 iterations, reducing runtime by over 200×. This implies that RDMs can attain a much better trade-off between generation quality and runtime. In addition, we also investigate the effect of candidate sample size in Table 8 . We notice that DiffuSeq benefits more from large sample sets (e.g., when the sample size increases from 1 to 10) than RDMs. We attribute this to the possibility that adding Gaussian noise to token embeddings in DiffuSeq might lead to more diverse samples. This helps make better use of the MBR decoding, indicating that there might be room to improve RDMs to leverage multiple decodes. Nevertheless, RDMs still achieve better performance than DiffuSeq across both cases of single and multiple samples. (2)t-1 = αt-1-αt 1-αt 28.32 30.64Linear λ(2)t-1 = 1 -t-1 T 31.04 33.91Constant λ(2)t-1 = 1 29.75 32.57On the Effect of Training and Decoding Strategies. This section explores the impact of improved training and decoding techniques developed for RDMs. Concerning the training aspect, we compare the derived loss objective, which takes the form of a reweighting cross-entropy function ( §4.2), with that of traditional discrete diffusion models. As for decoding, we evaluate the effectiveness of the discriminative routing mechanism ( §4.3) against the vanilla strategy that randomly denoises tokens. The overall comparison is presented in Figure 2a (as well as Figures 3a and 3b in Appendix G.4). The results indicate that adopting either the improved training or decoding scheme leads to a significant performance boost over vanilla baselines, which can be further amplified by integrating both. In particular, we discover that the inferior performance of vanilla multinomial diffusion stems from both inadequate training and ineffective decoding. Improving either the training or decoding process results in gains of over 10~20 BLEU points, allowing the model to scale well for more decoding steps. Besides, we also ablate the design choice of reweighting strategies ( §4.2) in Table 2 , and observe that a linear heuristic proposed in (Bond-Taylor et al., 2022) yields large improvements in performance on IWSLT14 DE-EN. Therefore, this reweighting scheme is adopted by default unless specified otherwise.On Decoding Speed. This section visualizes the performance-runtime comparison among different text generation models. All models considered here have roughly the same parameter counts (90~110M), and the speed is measured under the setting of 32 batch size on one NVIDIA GeForce RTX 3090 GPU, averaged by 30 runs. As shown in Figure 2b , there is a clear trend that RDMs usually run up to 10× faster than a similar-sized auto-regressive baseline GPT2 (Radford et al., 2019) , while to achieve similar performance, continuous diffusion models are several orders of magnitude slower than GPT2. Furthermore, discrete diffusion models here are trained with 50 time steps in total but are able to achieve satisfactory quality even with 2 or 5 steps. In contrast, the continuous diffusion model DiffuSeq, trained with 2000 time steps, mostly generates non-meaningful sentences (BLEU scores getting close to zero) under down-sampled time steps, even equipped with advanced samplers like DDIM (Song et al., 2021a) . This reveals the inherent difference between discrete and continuous diffusion, where discrete diffusion generalizes much better to various setups of iteration steps.In this section, we conduct a qualitative analysis to showcase the advantages of our model by illustrating the generated samples from the QQP task. We focus on the case of multinominal diffusion and its reparameterized variant; a more comprehensive analysis can be found in Appendix G.5.Multinomial Diffusion Does Not Decode Iteratively. As seen in Table 4 , we observe that across all text generation tasks, vanilla multinomial diffusion only generates the hypothesis at the first iteration and gets stuck in the same state thereafter. This means multinomial diffusion decodes in one shot and does not leverage the iterative process for further refinement. In Appendix G.5, we show that this abnormal behavior is primarily due to its degenerated backward formulation, which can be neatly fixed by our reparameterization. The resulting behavior is much more expected and leads to better generation quality.The Slow Convergence of Continuous Diffusion. We also demonstrate the down-sampled dynamics during the generation of DiffuSeq. In contrast to discrete diffusion models, where relevant tokens emerge within only a few steps, continuous diffusion hardly decodes meaningful tokens until the 1000-th iteration or later. This validates our hypothesis that the Gaussian diffusion over token embeddings is noisy and slow by design; furthermore, many diffusing steps are required to emit a significant change over token states due to the rounding quantization (see Table 12 for an illustration). 
diffusion models for text generation: 
 Text Generation with Discrete Diffusion. Discrete diffusion processes have close connections with previously developed language models. For instance, traditional auto-regressive language models can be seen as a special deterministic discrete diffusion process (Austin et al., 2021) . In addition, D3PMs (Austin et al., 2021) also introduce an absorbing diffusion strongly linked to masked language models (Devlin et al., 2019) . The diffusion formulation is further generalized in various aspects, such as enabling editing-based operations (Johnson et al., 2021) or casting generic permuted language models (Yang et al., 2019) as a diffusion process (Hoogeboom et al., 2022a) . • ##tly distances outline ##cera khmer curvature question ##tl 1• how can i improve focus in concentration ? 2• how can i improve focus in concentration ? 3• how can i improve focus in concentration ? 4• how can i improve focus in concentration ? 5• how can i improve focus in concentration ?RDM-multinomial 0• lungs ##down intensity cortes ##lden ufo oldies 1• worker blurted i ##kal caledonia concentration ##vb 2• how trait i ##kal my concentration ##vb 3• how trait i increase my concentration ? 4• how trait i increase my concentration ? 5• how do i increase my concentration ?The text generation performance of discrete diffusion processes is initially evaluated on language modeling tasks (Hoogeboom et al., 2021; Austin et al., 2021) , despite limited success. Recent studies have improved the performance of discrete diffusion on various tasks by devising unrolling training strategies (Savinov et al., 2022) , combining pre-trained models (He et al., 2022) , incorporating autoregressive decoding with editing-based refinements (Reid et al., 2022) , or leveraging a more effective diffusion process over data modalities with advanced search algorithms (Qian et al., 2022) .Text Generation with Continuous Diffusion. There has been a surge of recent interest in adapting continuous diffusion models for text generation. This approach typically applies Gaussian diffusion over the embedding space, achieving impressive results (Li et al., 2022b; Gong et al., 2022; Dieleman et al., 2022; Strudel et al., 2022; Lin et al., 2022; Yuan et al., 2022; Gao et al., 2022; Ye et al., 2023) . Some studies convert discrete tokens to bit strings and model them as real values (Chen et al., 2022) , or inject Gaussian noise into token logits instead of embeddings (Han et al., 2022) . Other studies focus on learning a latent continuous diffusion of pre-trained auto-regressive models (Wang et al., 2022; Lovelace et al., 2022) . Diffusion-based generative models are closely related to iterative non-autoregressive generation (Gu et al., 2018) in the context of machine translation. The generation process often involves iterative refinement (Lee et al., 2018; Ghazvininejad et al., 2019; Stern et al., 2019; Gu et al., 2019; Kasai et al., 2020; Ghazvininejad et al., 2020; Huang et al., 2022) . Our adaptive routing mechanism ( §4.3) takes inspiration from the heuristic used in CMLM (Ghazvininejad et al., 2019) , which refines the sequence by masking tokens with low model confidence. However, unlike CMLM, our approach only integrates the masking heuristic within the routing mechanism at each diffusion step, rather than relying entirely on it for decoding. This makes the decoding procedure governed by the formulated diffusion process and helps achieve better performance in practice. (Hoogeboom et al., 2021) later proposes a uniform corruption process for categorical variables, which are extended by D3PMs (Austin et al., 2021) to support general transition matrices, including an absorbing variant that draws close connections to masked language models (Devlin et al., 2019) . Several recent works push this line of research further in various aspects, such as incorporating editing-based operations (Johnson et al., 2021; Reid et al., 2022) , casting permuted language models (Yang et al., 2019) as diffusion models (Hoogeboom et al., 2022a) , developing a continuous-time framework (Campbell et al., 2022) , as well as exploring an analog of score functions for learning the reverse process Sun et al. (2022) . 
conclusion: 
 Applications. Discrete diffusion has been applied to a variety of tasks, including graph generation (Seff et al., 2019; Haefeli et al., 2022; Vignac et al., 2022) , image generation (Esser et al., 2021; Bond-Taylor et al., 2022; Gu et al., 2022; Tang et al., 2022; Hu et al., 2022a) , vision-language generation (Hu et al., 2022b) , and general multimodal conditional synthesis (Zhu et al., 2022) . Lezama et al. (2022) draws connections between discrete diffusion processes and non-autoregressive Transformers for visual domains (Chang et al., 2022; Yu et al., 2022; Chang et al., 2023) . A detailed discussion about applications to natural language texts can be found at §6. 
 
 Appendix: Discrete diffusion models are typically trained by maximizing a lower bound of its marginal loglikelihood, defined below,EQUATIONA.2 PARAMETERIZATION Recall that our objective is to minimize the KL divergence between q(x t-1 |x t , x 0 ) and a parameterized distribution p θ (x t-1 |x t ) at each time step. A widely adopted way is then defining p θ (x t-1 |x t ) = q(x t-1 |x t , x0 ), where x0 = f (x t ; θ) is predicted by a Transformer model. Austin et al. (2021) considers an alternative parameterization by letting p θ (x t-1 |x t ) ∝ x0 q(x t-1 , x t | x 0 )p θ ( x 0 |x t ), where we learn p θ ( x 0 |x t ) similarly to f (x t ; θ). These two approaches are different in general and define distinct generative processes in general; we follow the former method due to its simplicity and conciseness. More details can be found below.This section provides separate derivations for the original backward transition formulation of various discrete diffusion processes.Absorbing Diffusion. The absorbing diffusion (Austin et al., 2021) defines a Markov chain where a token goes into an absorbing mask state denoted by [M ] with some probability at each time step and stays the same thereafter. The forward transition probability is defined as q(x t |x t-1 ) = β t x t-1 + (1 -β t )q noise , where q noise = [M ] is the point mass with all of the probability on an absorbing state (the mask state [M] is denoted as a one-hot vector).Regarding the conditional backward transition probability q(x t-1 |x t , x 0 ), x t can only stay in either state x 0 or state [M ]. If x t = x 0 , then x t-1 must also be in state x 0 since it is not absorbed yet; while if x t = e [M ] , we haveq(x t-1 = [M ]|x t = [M ], x 0 ) = q(x t = [M ]|x t-1 = [M ])q(x t-1 = [M ]|x 0 ) q(x t = [M ]|x 0 ) = 1 • (1 -α t-1 ) 1 -α t = 1 -α t-1 1 -α t ; q(x t-1 = x 0 |x t = [M ], x 0 ) = q(x t = [M ]|x t-1 = x 0 )q(x t-1 = x 0 |x 0 ) q(x t = [M ]|x 0 ) = (1 -β t )α t-1 1 -α t = α t-1 -α t 1 -α t .The actual generative process is defined as p θ (x t-1 |x t ) ∝ x0 q(x t-1 , x t | x 0 )p θ ( x 0 |x t ), where we predict the probability vector p θ ( x 0 |x t ) := f x0 (x t ; θ) from a Transformer. As shown in Austin et al. (2021) , this formulation has a simple expression. Suppose k ̸ = [M ] is one of the K possible states. Note that• if x t = k ̸ = [M ], then due to the joint q(x t-1 , x t | x 0 ) there is only one entry that is nonzero within the sum q(x t-1 = k, x t = k| x 0 = k)p θ ( x 0 = k|x t ). As a result, the reverse distribution becomes a point mass over position k;• if x t = [M ], we have p θ (x t-1 = [M ]|x t = [M ]) ∝ x0 q(x t = [M ]|x t-1 = [M ])q(x t-1 = [M ]| x 0 )p θ ( x 0 |x t ) = x0 (1 -α t-1 )p θ ( x 0 |x t ) = (1 -α t-1 ) x0 p θ ( x 0 |x t ) = 1 -α t-1 ; p θ (x t-1 = k|x t = [M ]) ∝ x0 q(x t = [M ]|x t-1 = k)q(x t-1 = k| x 0 )p θ ( x 0 |x t ) = q(x t = [M ]|x t-1 = k)q(x t-1 = k| x 0 = k)p θ ( x 0 = k|x t ) = (α t-1 -α t )p θ ( x 0 = k|x t ).They can be easily normalized as well,p θ (x t-1 = [M ]|x t = [M ]) = 1 -α t-1 1 -α t p θ (x t-1 = k|x t = [M ]) = (α t-1 -α t )p θ ( x 0 = k|x t ) 1 -α t .This can be written in a more compact way, where p θ ( x 0 = k|x t ) := f k (x t ; θ) and q noise = [M] is a point mass with all the probability put over the absorbing state [M] .EQUATIONWe can also generalize this result for 0 < s < t ≤ T . For the forward transition, we have q(x t |x s = x 0 ) = t i=s+1 β i x 0 + 1 -t i=s+1 β i q noise = αt αs x 0 + αs-αt αs q noise . Thus the backward transition can be derived as well according to Bayes' rule,EQUATIONEQUATIONMultinomial Diffusion. In multinomial diffusion (Hoogeboom et al., 2021) , the forward transition probability is defined as q(x t |x t-1 ) = β t x t-1 + (1 -β t )q noise , where q noise = 1/K is a uniform distribution over {1, 2, . . . , K} with 1 is a K-dimensional vector with all ones. The backward transition probability conditional on the original data x 0 can be derived according to Bayes' rule and in closed form:EQUATION)Multinomial diffusion (Hoogeboom et al., 2021 ) learns a parameterized distribution p θ (x t-1 |x t ) to approximate q(x t-1 |x t , x 0 ) at each time step, which is defined as p θ (x t-1 |x t ) = q(x t-1 |x t , x0 ) with x0 = f (x t ; θ) being the output by a Transformer model.EQUATIONB DERIVATION FOR PROPOSITION 3.1In this section, we provide the derivation for Equation 2 based on Bayes' rule.Proof. We denote P t ∈ R K×K as the probability transition matrix for the t-th step, where [P t ] ij = p(x t = j|x x-1 = i) and thus the probability distribution in the forward process can be described as q(x t |x t-1 ) = Categorical x t ; P ⊤ t x t-1 . It is then easy to see thatq(x t |x 0 ) = xt-1,xt-2,...,x1q(x t , x t-1 , x t-2 , . . . , x 1 |x 0 ) = Categorical x t ; s P ⊤ t x 0 with s P t = P 1 P 2 . . . P t . Returning to our case where the forward transition takes the form q(x t |x t-1 ) = β t x t-1 + (1 -β t )q noise . The transition matrix can be represented by P t =β t I + (1 -β t )1q ⊤noise , and thus s P t = P 1 P 2 . . . P t = α t I + (1 -α t )1q ⊤ noise . Equipped with these results, we can proceed with the derivation as follows, q(x t-1 |x t , x 0 ) = q(x t |x t-1 )q(x t-1 |x 0 ) q(x t |x 0 ) = P t x t ⊙ s P ⊤ t-1 x 0 x ⊤ t s P ⊤ t x 0 = [β t x t + (1 -β t )σ xt 1] ⊙ [α t-1 x 0 + (1 -α t-1 )q noise ] α t x ⊤ t x 0 + (1 -α t )x ⊤ t q noise = β t α t-1 x t ⊙x 0 +β t (1-α t-1 )x t ⊙q noise +(1-β t )α t-1 σ xt 1⊙x 0 +(1-β t )(1-α t-1 )σ xt 1⊙q noise α t x ⊤ t x 0 + (1 -α t )x ⊤ t q noise = β t α t-1 x t ⊙x 0 + β t (1-α t-1 )σ xt x t + (1-β t )α t-1 σ xt x 0 + (1-β t )(1-α t-1 )σ xt q noise α t x ⊤ t x 0 + (1 -α t )σ xt .Here we denote ⊙ as element-wise product and σ xt := q noise (u = x t ) to represent the probability of noise drawn from q noise being equal to x t . The need to differentiate between x t and x 0 emerges when we calculate x t ⊙ x 0 , which would be an all-zero vector 0 except that it would be one ifx t = x 0 .Thus the computation of backward transition probabilities breaks down into two cases:• If x t = x 0 , we have x t ⊙ x 0 = x t , x t ⊤x 0 = 1 and thus q(x t-1 |x t , x 0 ) = β t α t-1 x t + β t (1 -α t-1 )σ xt x t + (1 -β t )α t-1 σ xt x t + (1 -β t )(1 -α t-1 )σ xt q noise α t + (1 -α t )σ xt = β t α t-1 + β t (1 -α t-1 )σ xt + (1 -β t )α t-1 σ xt α t + (1 -α t )σ xt x t + (1 -β t )(1 -α t-1 )σ xt α t + (1 -α t )σ xt q noise .• If x t ̸ = x 0 , we have x t ⊙ x 0 = 0, x t ⊤x 0 = 0 and thusq(x t-1 |x t , x 0 ) = β t (1 -α t-1 )σ xt x t + (1 -β t )α t-1 σ xt x 0 + (1 -β t )(1 -α t-1 )σ xt q noise (1 -α t )σ xt = β t (1 -α t-1 )x t + (1 -β t )α t-1 x 0 + (1 -β t )(1 -α t-1 )q noise 1 -α t = (1 -β t )α t-1 1 -α t x 0 + β t (1 -α t-1 ) 1 -α t x t + (1 -β t )(1 -α t-1 ) 1 -α t q noise = (1 -β t )α t-1 1 -α t x 0 + 1 -α t-1 1 -α t [β t x t + (1 -β t )q noise ] = α t-1 -α t 1 -α t x 0 + 1 - α t-1 -α t 1 -α t [β t x t + (1 -β t )q noise ] .Putting them together, we arrive at the resulting formulation,q(x t-1 |x t , x 0 ) =    λ (1) t-1 x t + 1 -λ (1) t-1 q noise , if x t = x 0 λ (2) t-1 x 0 + 1 -λ (2) t-1 q noise (x t ), if x t ̸ = x 0 .Here λ(1)EQUATIONt-1 := αt-1-αt 1-αt , and q noise (x t ) = β t x t + (1-β t )q noise denotes a noise distribution that interpolates between x t and q noise , both of which are possibly noisy.Generalization. Similar to vanilla diffusion models, we can also derive backward transition processes with a gap ∆ t ; that is, we consider q(x s |x t , x 0 ) with s = t -∆ t . It can be easily seen thatEQUATIONwith λ(1)s := 1 - (1- α t αs )(1-αs)qnoise(u=xt)αt+(1-αt)qnoise(u=xt) and λ(2)s := αs-αt 1-αt .The following provides the derivation for the loss objective of RDMs. Specifically,log p(x 0 ) ≥ E q(x 1:T ,v 1:T |x0) log p θ (x 0 , x 1:T , v 1:T ) q(x 1:T , v 1:T |x 0 ) = E q(x 1:T ,v 1:T |x0) log p θ (x 0 |x 1 ) T t=2 p θ (x t-1 , v t-1 |x t )p(x T , v T ) T t=2 q(x t-1 , v t-1 |x t , x 0 )q(x T , v T |x 0 ) = E q(x 1:T ,v 1:T |x0) log p θ (x 0 |x 1 ) + T t=2 log p θ (x t-1 , v t-1 |x t ) q(x t-1 , v t-1 |x t , x 0 ) + log p(x T , v T ) q(x T , v T |x 0 ) := L 1 (θ) - T t=2 L t (θ) + const.Here we denote L 1 (θ) := E q(x1|x0) [log p θ (x 0 |x 1 )], and for t > 1, L t (θ) := E q(xt|x0) [KL(q(x t-1 , v t-1 |x t , x 0 ) ∥ p θ (x t-1 , v t-1 |x t ))].We can push the decomposition of L t (θ) for time step t > 1 further,L t (θ) = E q(xt|x0) [KL(q(x t-1 , v t-1 |x t , x 0 ) ∥ p θ (x t-1 , v t-1 |x t ))] = E q(xt|x0)   xt-1,vt-1 q(x t-1 , v t-1 |x t , x 0 ) log q(x t-1 , v t-1 |x t , x 0 ) p θ (x t-1 , v t-1 |x t )   = E q(xt|x0)   xt-1,vt-1 q(x t-1 |v t-1 , x t , x 0 )q(v t-1 ) log q(x t-1 |v t-1 , x t , x 0 ) p θ (x t-1 |v t-1 , x t ) + log q(v t-1 ) p θ (v t-1 )   = E q(xt|x0) E q(vt-1) [KL(q(x t-1 |v t-1 , x t , x 0 )∥ p θ (x t-1 |v t-1 , x t ))]+KL(q(v t-1 )∥ p θ (v t-1 )) .Proof. We first consider the loss objective at each step for each token at n-th position, which can be expanded as follows,L n t (θ) = E q(xt,n|x0,n) E q(vt-1,n) [KL(q(x t-1,n |v t-1,n , x t,n , x 0,n ) ∥ p θ (x t-1,n |v t-1,n , x t,n ))] .Typically we draw a Monte Carlo sample x t,n ∼ q(x t,n |x 0,n ) to estimate the outermost expectation above. For the inner term, recall that b t,n = 1 xt,n=x0,n and q(x t-1,n |v t-1,n , x t,n , x 0,n ) takes the form asq(x t-1,n |v t-1,n , x t,n , x 0,n ) =    v (1) t-1,n x t,n + 1 -v (1) t-1,n q noise , if b t,n = 1 v (2) t-1,n x 0,n + 1 -v (2) t-1,n q noise (x t,n ), if b t,n = 0.Since we use the teacher-forcing approach that employs the same oracle b t,n for p θ (x t-1,n |v t-1,n , x t,n ) as well, it can also be written in a similar manner,p θ (x t-1,n |v t-1,n , x t,n ) =    v (1) t-1,n x t,n + 1 -v (1) t-1,n q noise , if b t,n = 1 v (2) t-1,n f (x t,n ; θ) + 1 -v (2) t-1,n q noise (x t,n ), if b t,n = 0.The derivation then breaks down into two cases with respect to b t,n :If b t,n = 1. In this case, q(x t-1,n |v t-1,n , x t,n , x 0,n ) = p θ (x t-1,n |v t-1,n , x t,n ) = v (1) t-1,n x t,n + 1 -v (1)t-1,n q noise . Since these two distributions become identical, this leads to zero KL divergence irrespective of v t-1,n so that L t (θ) = 0;If b t,n = 0. In this scenario, we have q(x t-1,n |b t-1 , x t,n , x 0,n ) = v (2) t-1,n x 0,n + 1 -v (2) t-1,n q noise (x t,n ) and v (2) t-1,n f (x t,n ; θ) + 1 -v (2)t-1,n q noise (x t,n ). We then enumerate all the possible outcomes for v(2)t-1,n . If v (2) t-1,n = 1, q(v t-1,n ) = λ (2) t-1 and KL(q(x t-1,n |v t-1,n , x t,n , x 0,n ) ∥ p θ (x t-1,n |v t-1,n , x t,n )) = KL(x 0,n ∥ f (x t,n ; θ)) = -x ⊤ 0,n log f (x t,n ; θ) . If v (2) t-1,n = 0, then q(v t-1,n ) = 1 -λ (2)t-1 andKL(q(x t-1,n |v t-1,n , x t,n , x 0,n ) ∥ p θ (x t-1,n |v t-1,n , x t,n )) = KL(q noise (x t,n ) ∥ q noise (x t,n )) = 0.Putting them together, we haveE q(vt-1,n) [KL(q(x t-1,n |v t-1,n , x t,n , x 0,n ) ∥ p θ (x t-1,n |v t-1,n , x t,n ))] = -λ (2) t-1 x ⊤ 0,n log f (x t,n ; θ) + (1 -λ (2) t-1 ) • 0 = -λ (2) t-1 x ⊤ 0,n log f (x t,n ; θ) .Since each token is modeled conditionally independently, we can add all computed for each token, arriving at the final expression for the whole sequence,L t (θ) = N n=1 L n t (θ) = E pdata(x 0,1:N ) N n=1 q(xt,n|x0,n) -λ (2) t-1 N n=1 (1 -b t,n )x ⊤ 0,n log f (x t,n ; θ) .Connections to D3PMs. The derived simplified loss objective bears some resemblance to that of absorbing diffusion in D3PMs (Austin et al., 2021) , which also takes the form of a cross-entropy function over masked positions. However, our derivation arises from the developed reparameterization perspective and thus stems from a distinct motivation from D3PMs. In addition, our objective applies to a wide range of discrete diffusion processes, including those with multinomial noise, absorbing noise, or a mixture of both. This constitutes a non-trivial generalization of D3PMs, which only demonstrates that the cross-entropy representation is available for absorbing diffusion. Besides, our formulation explicitly elucidates the role of routing mechanisms in training, which provides insights into techniques that improve decoding quality.This section describes the implementation details of our experiments.Machine Translation. For experiments on machine translation, we consider three standard benchmarks:• IWSLT14 DE-EN (Cettolo et al., 2014) , which contains around 160K/7K/7K sentence pairs for training, validation, and testing, respectively. We build a joint vocabulary for the source and target language, resulting in 10152 Byte Pair Encoding (BPE; Sennrich et al., 2016) types. • WMT14 EN-DE (Bojar et al., 2014) dataset consists of around 4.0M/3K/3K training/validation/testing pairs. The preprocessing follows Ghazvininejad et al. (2019) and yields a shared vocabulary with 40624 BPE types; • WMT16 EN-RO (Bojar et al., 2016) . We use the same data split from Lee et al. (2018) that comprises around 610K/2K/2K pairs. The vocabulary is shared between the source and target sides with 34976 joint BPE types. We operate on original data for all translation tasks and do not adopt knowledge distillation (Kim & Rush, 2016; Gu et al., 2018) that replaces the target side of training data with outputs generated by a pre-trained autoregressive Transformer.For evaluation, we report tokenized BLEU (Papineni et al., 2002) scores applied with compound split post-processing to facilitate comparison.In addition, we also compute sacreBLEU (Papineni et al., 2002; Post, 2018) (signature: nrefs:1|case:mixed|eff:no|tok:13a|smooth:exp|version:2.0.0) and COMET (Rei et al., 2020) (with version 2.0.0 and model Unbabel/wmt22-comet-da) scores on WMT14 EN-DE test set, as presented in Table 5 . It can be seen that both sacreBLEU and COMET reveal a trend similar to that of tokenized BLEU scores.Question Generation and Paraphrasing. For both QG and QQP tasks, we use the same data split pre-processed as in Gong et al. (2022) :• Question Generation (QG) with the Quasar-T dataset (Dhingra et al., 2017) , containing around 117K/2K/10K pairs for training/validation/testing, respectively. • Paraphrasing with Quora Question Pairs (QQP). This dataset comprises around 145K/2K/2.5K training/validation/testing question pairs.Following Gong et al. (2022) , we use WordPiece tokenization as in BERT (Devlin et al., 2019) and obtain a vocabulary of size 30522 for both tasks.• We employ the Transformer-base architecture (Vaswani et al., 2017) for WMT experiments, while for IWSLT14 DE-EN,QG, and QQP tasks we use a smaller Transformer model. Note that all self-attention blocks with the model are bi-directional and do not use causal masks. • We adopt a length prediction module (Ghazvininejad et al., 2019) on top of the Transformer encoder to propose target length candidates for the generated sequence. Given the source input, we first run the Transformer encoder to obtain the encoder's hidden representation, which is averaged and passed to a linear layer to output the length scores. • The timestep embedding is obtained by first projecting the input timestep t with sinusoidal encodings and then passing it through a two-layer MLP. • We adopt concatenated instead of additive position encodings, which is shown to enhance the positional information and produce better performance in the context of machine translation (Huang et al., 2022) .The detailed configuration for the neural network is listed in Table 6 .• We allocate a large number of diffusion time steps for training, such as 50 or 100. We found the number of diffusion steps in training does not affect the performance too much. Note that decoding can be performed with an arbitrary number of iterations by choosing the appropriate step size ∆t > 1. That is, we can decode by sampling x t-∆t ∼ p θ (x t-∆t |x t ), following a similar treatment as Equation 15. • Modern Transformer models usually process input sequences that are associated with some special symbols, such as the begin-of-sentence symbol <bos>, eos-of-sentence symbol <eos>, padding symbol <pad>, and so on. We found it beneficial to treat these special symbols as normal tokens in vanilla/reparameterized absorbing diffusion (and thus these symbols can be noised), but this treatment leads to much worse performance in vanilla/reparameterized multinomial diffusion. • We also adopt conditioned training (details in Appendix F) to further improve the model, which leads to consistent improvements upon vanilla training. Its effect is ablated in Table 7 . In particular, conditioned training leads to almost 2 BLEU improvements over vanilla diffusion processes but the gain becomes marginal for our reparameterized variants.The detailed configuration for the optimization hyper-parameters is listed in Table 6 .• Note that we train a neural network f (•; θ) to approximate x 0 , which is a softmax-normalized probability vector. There are several ways to decode a token from the probability vector, such as simply taking its argmax position or performing sampling with temperatures τ . Empirically, we find a low temperature τ = 0.1, or simply the argmax works well across tasks, and use the argmax approach by default; Nevertheless, the diversity of generated sentences can be improved by adopting a larger temperature as well. • Like common diffusion models, we use an exponential moving average (EMA) to track the model parameters with a decay rate of 0.9999. We also average model parameters among the five last checkpoints for generation, following standard practices in machine translation (Vaswani et al., 2017) . • For translation experiments, we use 5 length candidates, decode them in parallel, and select the sequence with the highest model score as the final output. For question generation and paraphrasing, we follow DiffuSeq (Gong et al., 2022) to use MBR decoding with 10 samples to ensure a head-tohead comparison. The candidates in MBR decoding are generated in the following manner: first selecting 3 length candidates, and then sampling 3,3, and 4 additional sentences for each length size, resulting in 10 candidates in total. • We also investigate several components of the adaptive decoding algorithm ( §4.3) and provide more implementation details below:• The top-k selection mechanism in Equation 7, which is deterministic by design, can also be made stochastic. In particular, instead of directly selecting those tokens with top-k largest scores, we first add Gumbel noise to the score s t,n of each token, and then fetch the top-k tokens with the largest perturbed scores. This is inspired by previous work that aims to sample multiple items from a set without replacement (Vieira, 2014; Kool et al., 2019; 2020) ; however, this simple approach brings several benefits in that the selection of k tokens from the sequence could involve extra randomness, and this might be helpful for exploration during decoding. • Recall that during decoding, the role of v(1) t-1,n is to indicate whether the token can remain in the denoised state, while v(2) t-1,n is used to denoise the token that is currently noisy. In Equation 7, both of them would be set to 1 as long as the n-th token belongs to the top-k set. However, we observe that v(1) t-1,n and v(2) t-1,n can also be generated differently. For example, one might adopt a more conservative approach, where already denoised tokens rarely or never turn back to the noise. We implemented a strategy to achieve this by imposing more constraints over v(1) (xt,n̸ =xt+1,n) , where we set vt-1,n : v (i) t-1,n = 1 (n∈Pt-1)∨(st,n>st+1,n)∨(1) t-1,n = 0 only when its corresponding token score is not in the top-k set and indeed becomes than the previous iteration. Intuitively, this means the denoised tokens should remain as denoised most time, except that the Transformer model becomes less confident and requires re-prediction. This is one of many possible approaches to achieving such control, as our framework allows us to do such conditioning flexibly; we find this strategy works sometimes better than the vanilla approach, especially on IWSLT14 DE-EN dataset.• Another important hyper-parameter during decoding is k, the number of tokens to be in denoised states at each iteration for our discriminative routing mechanism ( §4.3). To ensure that the degree of noise decreases as the generation process proceeds, we schedule k to increase from 1 to N monotonically as the diffusion step t goes from T to 1. We set k to follow either a {cosine, linear} scheme based on the development set performance. The cosine strategy yields k = cos πt 2T • N , while the linear variant gives k = 1 -t T • N . Based on our preliminary experiments, we discerned that while these components indeed have an influence on task performance, their impact is relatively minor compared to the primary improvements (e.g., reweighted training and adaptive decoding), as elucidated in the main paper.• To perform conditional generation, we delegate the full task of conditioning to the encoderdecoder architecture. That is, instead of designing complicated guidance to condition the diffusion probabilistic process, we treat the conditioning information as the input of the Transformer encoder. This simple strategy is found to work well in practice.Training discrete diffusion models usually involves a heavy amount of randomness. For instance, at each training iteration, one has to sample a time step and corrupt a random subset of sequence tokens for denoising. To control the introduced variance, we adopt a simple yet effective conditioning strategy that uses multiple samples to perform training. The key idea is conceptually simple: we start with sampling two i.i.d. time steps s, t ∼ Uniform(T ) (without the loss of generality, we assume s < t). After that, we draw x t ∼ q(x t |x 0 ) as usual, but condition the sample at step s by drawing from q(x s |x 0 ) = E q(xt|x0) [q(x s |x t , x 0 )] ≈ q(x s |x t , x 0 ). The losses (Equation 6) at step s and t are then estimated individually and averaged to obtain the final training objective. In case s = t, we simply drop the conditioning and sample x s ∼ (x s |x 0 ) instead.This method utilizes multiple samples to estimate the loss objective while remaining unbiased. To see this, note that (Pillutla et al., 2021) , evaluating the distance of token distributions between the generated and human-written text on the test set; and 3) Coherence, calculating the average-token log likelihood under a well-trained language model, which is set to OPT-2.7B (Zhang et al., 2022) . We refer readers to Su & Collier (2022) for more technical details of these metrics.-(L s + L t ) = E q log q(x s-1 |x s , x 0 ) p θ (x s-1 |x s ) + log q(x t-1 |x t , x 0 ) p θ (x t-1 |x t ) = E q(xs-1,xs,xt-1,xt|x0) log q(x s-1 |x s , x 0 ) p θ (x s-1 |x s ) + log q(x t-1 |x t , x 0 ) p θ (x t-1 |x t ) = E q(xs-1,xs,xt-1,xt|x0) log q(x s-1 |x s , x 0 ) p θ (x s-1 |x s ) + E q(xt-1,xt|x0) log q(x t-1 |x t , x 0 ) p θ (x t-1 |x t ) = E q(xt|x0)q(xs|xt,x0) [KL(q(x s-1 |x s , x 0 ) ∥ p θ (x s-1 |x s ))] + E q(xt|x0) [KL(q(x t-1 |x t , x 0 ) ∥ p θ (x t-1 |x t ))] .We train three different models for comparison: auto-regressive language models, vanilla discrete diffusion models, and our RDMs. All of these models have approximately 430M parameters and are trained with 100k steps on the Wikitext-103 training set. Both auto-regressive language models and discrete diffusion models here adopt the same decoder-only Transformers following the Llama architecture (Touvron et al., 2023) , except that discrete diffusion models remove the use of causal masks in self-attention blocks and introduce an additional lightweight time-step embedding for proper conditioning. During training, the maximum sequence length and the batch size are set to 256 and 128, respectively, where shorter sequences are packed together. The Adam (Kingma & Ba, 2014) optimizer is used, and the learning rate is set to 3e-4 with the cosine scheduler. To facilitate evaluation on open-ended generation, we follow previous practices (Li et al., 2022a) and condition the generation of different models on test prefix prompts with a fixed length of 32. We limit the maximum generation length to 256 and truncate the generated output for each test case to the first 128 tokens for subsequent evaluation. For diffusion models, we the initial sequence length to 256 and truncate all content after the first <eos> token upon the iterative process finishes. For all models, we generate samples at the temperature of 1.0 by nucleus sampling (Holtzman et al., 2020) with top-p 0.95. Table 10 demonstrates the comparison among autoregressive language models, vanilla discrete diffusion models, and our RDMs in the task of open-ended text generation. We observe that while discrete diffusion models generally lag behind auto-regressive LMs, RDMs effectively reduce the gap, scale well with the number of iterations, and achieve competitive performance with auto-regressive LMs while exhibiting greater generation variety. Generation examples can be found in Table 11 .This section presents additional plots (Figure 3 ) that visualize the effect of different components in our model.This section provides a more comprehensive qualitative analysis of different diffusion models, including several generated samples as in Tables 12 to 14. • how <M> i prepare for ca <M> ##t <M> <M> <M> <M> <M> ? 3• how <M> i prepare for ca cp ##t <M> <M> <M> months months ? 4• how <M> i prepare for ca cp ##t exam in two months <M> ? 5• how can i prepare for ca cp ##t exam in two months left ? Multinomial 0• glossy [unused448] raymond manga subjective questioning suriname masonic listen explored 1• how can i prepare for ca cp ##t months ? 2• how can i prepare for ca cp ##t months ? 3• how can i prepare for ca cp ##t months ? 4• how can i prepare for ca cp ##t months ? 5• how can i prepare for ca cp ##t months ?RDM-multinomial 0• consonants ##nin leading elegance 406 173 militant teams ##nin dyke thee seafood 1• how residues i ##fighting sentences malaysian jenkins remembers transatlantic universite ##rp monarch ? 2• how can i cyril for malaysian jenkins goldberg transatlantic in relationships pursuing ? 3• how can i chu for ca fashionable ##t exam in clerks months ? 4• how can i prepare for ca cp ##t exam in 2 months ? 5• how can i prepare for ca cp ##t exam in 2 months ? DiffuSeq 0• defective thereby evaluation michaels fragments primal electrically aground hostilities ‡ 10• simulcast candidacy ##bner [unused106] ##wide subgenus dangerously sincerity resolving migrated menon ##lase ‡ 100• westphalia ##tracted universite ##erly reissued neglect showcased [unused574] slade ‡ 250• souza electronically compliant gerard priority townships ##neo hidalgo [unused574] ‡ 500• spikes peptide ##ales borneo makeshift moi rebelled neglect textual 1899 erasmus publishes ‡ 750• i reduces griffin ##ales bukit makeshift moi ##sław mcbride how ministries ‡ Multinomial Diffusion Does Not Decode Iteratively. As presented in Tables 12 to 14, multinomial diffusion finishes the generation of most sentences in the first iteration and remains unchanged afterward, despite multiple iteration steps being allocated.The Slow Convergence of Continuous Diffusion. In contrast to discrete diffusion models that can perform generation in 10 steps or fewer, continuous diffusion usually requires thousands of steps to decode a decent sample. We hypothesize that this is due to two reasons: (1) the noisy and slow Gaussian diffusion over token embeddings by design;(2) furthermore, many diffusing steps are required to emit a significant change over token states due to the rounding operation. We provide empirical evidence for our hypothesis by zooming in to inspect the generation process. As can be seen in Tables 12 and 13 , many consecutive steps in continuous diffusion (1000~1009 iteration) do not modify the decode at all, even when it is not converged yet, leading to a potential waste of computation.Vanilla Absorbing Diffusion Cannot Fix Previous Errors. While vanilla absorbing diffusion performs decoding more steadily, it also suffers from some issues. As shown in Tables 12 to 14 , since all tokens are predicted independently conditional on the source input, there is some chance for the model to decode multiple identical tokens simultaneously. However, in vanilla absorbing diffusion, such decoding errors cannot be fixed. To see this, note that its backward transition formulation can be written as Equation 10, which is copied here for convenience,p θ (x t-1 |x t ) = 1 -1-αt-1 1-αt f (x t ; θ) + 1-αt-1 1-αt q noise , if x t = [M ] x t , if x t ̸ = [M ].Under this backward formulation, once a token is decoded, it will stay in the state thereafter and does not have the chance to be re-predicted again. As a result, vanilla absorbing diffusion processes cannot fix previously made errors.This issue can be alleviated by RDMs, which employ a more generic formulation as follows,EQUATIONt-1 x t + 1 -v(1)EQUATIONt-1 q noise (x t ), if b t = 0.Here b t is a binary variable indicating whether x t is denoised or not, and v(1)t-1 can be either 1 or 0, depending on the strategy ( §4.3). Therefore, in RDMs, we can allow decoded tokens to be rolled back to noisy states by setting v(1) t-1 = 0 (e.g., these repetitive tokens might receive lower model scores than the others, which can be recognized as low-confidence outputs in Equation 7). An example can be found in Table 12 , where the decoded repetitive tokens months months at 3-th iteration are then re-masked at the next iteration.We use xt and xt,n interchangeably to represent a single token when there is no ambiguity.Technically, for certain noise distributions, the logic operation ∧ and ∨ should be noisy since they should compensate for the possibility that noises drawn from qnoise can also coincide with x0; however, for most diffusion processes considered in this work, such probability is so small that it can be safely ignored. 


Decision: In this paper, the authors derive an alternative yet equivalent formulation for sampling from discrete diffusion processes, leading to a family of reparameterized discrete diffusion models. Extensive experiments on machine translation and text generation demonstrate significant improvements over existing diffusion models. The major concerns from the reviewers are: (1) a lack of comparison with several diffusion model baselines; (2) performance is not better than autoregression-based language models; and (3) the need for improved presentation. The authors have committed to making the necessary revisions to address all these concerns. I recommend acceptance, provided that all promised revisions are incorporated to address the concerns outlined above.

Reviews: review: paper_summary: paper_summary: This paper presents a reparameterization of discrete diffusion that can be applied on top of existing diffusion architectures. In this work, the authors use it train diffusion language models. The primary contribution lies in rearranging the term q(x_{t-1} | x_t, x_0) that allows defining it as a sum of two terms: one involving staying at the current token with some probability if it is close to the target, and the other simply adding noise at the current step. Based on this observation, the authors redefine the training process, and derive a loss function which they show is similar to a cross-entropy based language modeling loss. The resulting sampling algorithm requires much fewer iterations than baselines.

The authors train with this approach, discrete diffusion models with two kinds of noise formulation (namely multinomial and absorbing). On translation, paraphrasing, and question generation. Analysis shows that this model is much faster than other diffusion-based baselines.
strength_weakness: strength_weakness: 1. Several diffusion-based baselines are missing [1,2]. A quick skim through these papers shows that these baselines offer better performance on translation, QG, and QQT than the proposed method. These approaches, however, are not discrete and the reparameterization in its current form cannot be applied to them.
2. There is no performance gain over autoregressive models. In addition, a general weakness of discrete diffusion models is they do not allow adding control via guided diffusion. 

[1] TESS: https://arxiv.org/pdf/2305.08379

[2] AR-Diffusion: https://arxiv.org/pdf/2305.09515
questions: questions: 1. The tasks the authors consider are of reasonably short length (256), will the speed improvements hold for longer sequences?
paper_summary: paper_summary: This paper proposes to do discrete diffusion by explicitly modelling which tokens are noise and which tokens are correct.  At x_0 all the tokens are correct, and at x_T all the tokens are noise.  In between, they explicitly model the transitions to and from correct tokens and noise tokens (the v variables).  They derive a lot of math about these transition distributions and the associated token generation distributions, and show that it has some nice properties for training.  At generation time, they add an easy-first heuristic to decide which tokens should (stochastically) transition from noise to correct.

Or, at least, that is how I understand it, which could be over-simplistic.  The paper is written in the opposite order, and I didn't read it backwards.

The empirical evaluations show consistent improvements when adding their methods to previous discrete diffusion models.  An ablation study shows that both the easy-first decoding strategy and the new loss function are responsible for improvements.
strength_weakness: strength_weakness: It was very hard to follow the presentation of the model.  Diffusion is often complicated, but I think this was unnecessarily difficult to read.  Take for example section 4.3.  This section is supposed to be about decoding, but it spends a lot of time talking about training objectives.  Then it assumes that you know the ground truth.  Then it says that they compute the ground truth from the sampled v variables.  This is the second-to-last paragraph before the experiments, on page 6, and it was only at this point that I understood what they are trying to model.  Then I could go back and understand the paragraph before that (what I call the easy-first heuristic), at which point the paper started to make sense.

As another example of the poor presentation, on page 5 the paper says "So far, we mainly consider the case of diffusing over a single token", after 2.5 pages discussing the model without mentioning this fact.  At the end of section 2 they even say "for each token", while this section only makes sense "for a given token".

Going back to the introduction, I can see that there is a summary of the main idea.  But it is worded in such a vague ambiguous way that I couldn't understand it when I read the introduction.

I think that the paper should really be rewritten before it is published.  It is possible to fix these issues for the camera-ready, but it would be a lot to expect.
questions: questions: Do you think that, in time for the camera-ready version, you can add a precise but intuitive explanation of each contribution before going into the details?
paper_summary: paper_summary: The authors suggest a reparameterized discrete diffusion model to improve the performance of diffusion models in natural language processing. The proposed reparameterization creates a more effective training as well as a more effective sampling of the model.
The authors evaluated the proposed method on machine translation, question generation and paraphrasing.
Using the proposed method, the authors could significantly improve the performance of existing diffusion models in terms of translation quality. However, the comparison to autoregressive models is lacking recent advances. The model is compared only to baseline transformer model from 2017 and not to state-of-the-art models (e.g. submitted to WMT).
The authors nicely investigate the influence of the different training and decoding methods independently. Also the influence on the speed is shown. However, it is unclear why the proposed model with 20 iterations is faster the baseline one with 10? Furthermore, the authors compared to a GPT-2 base model, but not to the compared transformer baseline from the qualitative results.
strength_weakness: strength_weakness: The evaluation of the proposed should be improved with respect to the autoregressive baselines.

Comments: comment: Thank you for your clarification! Our preliminary experiments indicated that encoder-decoder transformers with a similar parameter count to GPTs achieve comparable or sometimes slower throughput, especially when generating sequences shorter than 128 tokens. We will include the results and analysis of encoder-decoder baselines in the revised manuscript and provide a more comprehensive evaluation of our models. We are grateful for your engagement in the rebuttal process.
Thanks for the clarification, the first point is clear now. Sorry for being unclear about the GPT-2 and Transformer comparison: It is unclear why you did only do a comparison to transformer base in Table 1 based on the quality, but not compare the speed of transformer to your work. This is only done for GPT-2.
Even though we do not expect non-autoregressive approach to outperform autoregressive ones, it is important to do a fair comparison to estimate the difference.
Thank you for your engagement in the rebuttal process! Should you have further queries, we are more than willing to provide additional clarifications. We are grateful for your valuable time to review our work.
The rebuttal demonstrates an understanding of my concerns and willingness to address them.  But I will leave it to the AC to decide whether this level of promised modifications in the camera-ready version is acceptable.