How can diffusion models be applied in domains with discrete output spaces, like NLP?
Diffusion models Diffusion models have
demonstrated impressive performance in popular
continuous-valued domains such as images (Ho
et al., 2020), audio (Kong et al., 2021), video (Ho
et al., 2022) and recently also been adopted for
3D-shapes, protein structures, and more (Zhou
et al., 2021; Trippe et al., 2022; Wu et al.,
2022). Since they are based on adding Gaussian
noise, these approaches are not straightforward
to apply to discrete valued domains like text.
Hoogeboom et al. (2021); Austin et al. (2021)
propose diffusing in the discrete space using
categorical distributions which are modified using
transition matrices. However, these methods do not straightforwardly support control and yield worse
results than comparable autoregressive models.
Li et al. (2022) propose to represent each token
as a continuous embedding and apply diffusion
in the embedding space. They train the LM to
generate a fixed length sequence whereas SSD-LM
allows flexibility in the generated sequence length
by generating block-wise. Further, their LM is
trained with specialized datasets and not evaluated
against general-purpose autoregressive LMs on
unconstrained text generation. Their method
supports post-hoc control but requires training a
customized attribute classifier,15 since the diffusion
operates on a learned embedding space. Gong et al.
(2022), a concurrent work to ours, extend Li et al.
(2022) to a sequence-to-sequence setup with a
similar underlying embedding-based method. Our
work is most closely related to Chen et al. (2022)
which transform discrete data into a sequence of
bits and represent each bit as +1 or -1 converting it
into a continuous-valued domain. For textual data,
however, it can lead to extremely long sequences
which are difficult to optimize. In this work, we
instead maintain a subword based vocabulary but
represent each token as a sequence of manually
defined logits.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Recently, diffusion models have emerged as a new paradigm for generative models. Despite the success in domains using continuous signals such as vision and audio, adapting diffusion models to natural language is under-explored due to the discrete nature of texts, especially for conditional generation. We tackle this challenge by proposing DiffuSeq: a diffusion model designed for sequence-to-sequence (Seq2Seq) text generation tasks. Upon extensive evaluation over a wide range of Seq2Seq tasks, we find DiffuSeq achieving comparable or even better performance than six established baselines, including a state-of-the-art model that is based on pre-trained language models. Apart from quality, an intriguing property of DiffuSeq is its high diversity during generation, which is desired in many Seq2Seq tasks. We further include a theoretical analysis revealing the connection between DiffuSeq and autoregressive/non-autoregressive models. Bringing together theoretical analysis and empirical evidence, we demonstrate the great potential of diffusion models in complex conditional language generation tasks. Code is available at \url{https://github.com/Shark-NLP/DiffuSeq}
[{"first":"Shansan","last":"Gong","middle":[]},{"first":"Mukai","last":"Li","middle":[]},{"first":"Jiangtao","last":"Feng","middle":[]},{"first":"Zhiyong","last":"Wu","middle":[]},{"first":"Lingpeng","last":"Kong","middle":[]}]
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
We present Bit Diffusion: a simple and generic approach for generating discrete data with continuous state and continuous time diffusion models. The main idea behind our approach is to first represent the discrete data as binary bits, and then train a continuous diffusion model to model these bits as real numbers which we call analog bits. To generate samples, the model first generates the analog bits, which are then thresholded to obtain the bits that represent the discrete variables. We further propose two simple techniques, namely Self-Conditioning and Asymmetric Time Intervals, which lead to a significant improvement in sample quality. Despite its simplicity, the proposed approach can achieve strong performance in both discrete image generation and image captioning tasks. For discrete image generation, we significantly improve previous state-of-the-art on both CIFAR-10 (which has 3K discrete 8-bit tokens) and ImageNet-64x64 (which has 12K discrete 8-bit tokens), outperforming the best autoregressive model in both sample quality (measured by FID) and efficiency. For image captioning on MS-COCO dataset, our approach achieves competitive results compared to autoregressive models.
[{"first":"Ting","last":"Chen","middle":[]},{"first":"Ruixiang","last":"Zhang","middle":[]},{"first":"Geoffrey","last":"Hinton","middle":[]}]
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Despite the growing success of diffusion models in continuous-valued domains (e.g., images), similar efforts for discrete domains such as text have yet to match the performance of autoregressive language models. In this work, we present SSD-LM—a diffusion-based language model with two key design choices. First, SSD-LM is semi-autoregressive, iteratively generating blocks of text, allowing for flexible output length at decoding time while enabling local bidirectional context updates. Second, it is simplex-based, performing diffusion on the natural vocabulary space rather than a learned latent space, allowing us to incorporate classifier guidance and modular control using off-the-shelf classifiers without any adaptation. We evaluate SSD-LM on unconstrained text generation benchmarks, and show that it matches or outperforms strong autoregressive GPT-2 models across standard quality and diversity metrics, while vastly outperforming diffusion-based baselines. On controlled text generation, SSD-LM also outperforms competitive baselines, with an extra advantage in modularity.
[{"first":"Xiaochuang","last":"Han","middle":[]},{"first":"Sachin","last":"Kumar","middle":[]},{"first":"Yulia","last":"Tsvetkov","middle":[]}]
Diffusion-LM Improves Controllable Text Generation
Controlling the behavior of language models (LMs) without re-training is a major open problem in natural language generation. While recent works have demonstrated successes on controlling simple sentence attributes (e.g., sentiment), there has been little progress on complex, fine-grained controls (e.g., syntactic structure). To address this challenge, we develop a new non-autoregressive language model based on continuous diffusions that we call Diffusion-LM. Building upon the recent successes of diffusion models in continuous domains, Diffusion-LM iteratively denoises a sequence of Gaussian vectors into word vectors, yielding a sequence of intermediate latent variables. The continuous, hierarchical nature of these intermediate variables enables a simple gradient-based algorithm to perform complex, controllable generation tasks. We demonstrate successful control of Diffusion-LM for six challenging fine-grained control tasks, significantly outperforming prior work.
[{"first":"Xiang","last":"Li","middle":["Lisa"]},{"first":"John","last":"Thickstun","middle":[]},{"first":"Ishaan","last":"Gulrajani","middle":[]},{"first":"Percy","last":"Liang","middle":[]},{"first":"Tatsunori","last":"Hashimoto","middle":["B."]}]
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces two extensions of flows and diffusion for categorical data such as language or image segmentation: Argmax Flows and Multinomial Diffusion. Argmax Flows are defined by a composition of a continuous distribution (such as a normalizing flow), and an argmax function. To optimize this model, we learn a probabilistic inverse for the argmax that lifts the categorical data to a continuous space. Multinomial Diffusion gradually adds categorical noise in a diffusion process, for which the generative denoising process is learned. We demonstrate that our method outperforms existing dequantization approaches on text modelling and modelling on image segmentation maps in log-likelihood.
[{"first":"Emiel","last":"Hoogeboom","middle":[]},{"first":"Didrik","last":"Nielsen","middle":[]},{"first":"Priyank","last":"Jaini","middle":[]},{"first":"Patrick","last":"Forr'e","middle":[]},{"first":"Max","last":"Welling","middle":[]}]
Structured Denoising Diffusion Models in Discrete State-Spaces
Denoising diffusion probabilistic models (DDPMs) (Ho et al. 2020) have shown impressive results on image and waveform generation in continuous state spaces. Here, we introduce Discrete Denoising Diffusion Probabilistic Models (D3PMs), diffusion-like generative models for discrete data that generalize the multinomial diffusion model of Hoogeboom et al. 2021, by going beyond corruption processes with uniform transition probabilities. This includes corruption with transition matrices that mimic Gaussian kernels in continuous space, matrices based on nearest neighbors in embedding space, and matrices that introduce absorbing states. The third allows us to draw a connection between diffusion models and autoregressive and mask-based generative models. We show that the choice of transition matrix is an important design decision that leads to improved results in image and text domains. We also introduce a new loss function that combines the variational lower bound with an auxiliary cross entropy loss. For text, this model class achieves strong results on character-level text generation while scaling to large vocabularies on LM1B. On the image dataset CIFAR-10, our models approach the sample quality and exceed the log-likelihood of the continuous-space DDPM model.
[{"first":"Jacob","last":"Austin","middle":[]},{"first":"Daniel","last":"Johnson","middle":["D."]},{"first":"Jonathan","last":"Ho","middle":[]},{"first":"Daniel","last":"Tarlow","middle":[]},{"first":"Rianne","last":"Berg","middle":["van","den"]}]
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Among existing generative models, GAN (Goodfellow et al., 2014) suffers from the instability issue (Salimans et al., 2016), subjecting to mode collapse (Metz et al., 2017); VAE (Kingma & Welling, 2014) has to rely on surrogate objectives to approximate maximum likelihood training and Flow-based models (Dinh et al., 2017) has to use specialized architectures to construct reversible transform. Diffusion models (Ho et al., 2020;Nichol & Dhariwal, 2021) have circumvented several of these limitations and emerged as a new paradigm for generative models, theoretically underpinned by non-equilibrium thermodynamics (Sohl-Dickstein et al., 2015) and score-matching network (Song & Ermon, 2019). To date, the major breakthroughs are in domains using continuous signals, such as vision (Saharia et al., 2022a;Ramesh et al., 2022) and audio (Kong et al., 2020). However, extending continuous diffusion models to natural language remains an open challenge due to the inherently discrete nature of texts.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
On the basis of unconditional generation in continuous space which is illustrated in Figure 1(a), existing efforts (Hoogeboom et al., 2021;Austin et al., 2021) start customizing diffusion models to text in discrete space on unconditional language modeling (i.e., free text generation). Diffusion-LM (Li et al., 2022), as in Figure 1(b), models texts in continuous space and proposes to use an extra-trained classifier as guidance (i.e., the condition signal x) to impose subtle changes (usually complex, finegrained constraints) on generated sentences. Nonetheless, these models do not naturally generalize to conditional language modeling (i.e., the model assigns probabilities p(w|x) to sequences of words w given x). In the more general sequence-to-sequence (SEQ2SEQ) setting where the condition x is also a sequence of words, applying Diffusion-LM can be difficult. The reason is that classifiers are attributes-oriented, and we can not train hundreds-of-thousands classifiers to model the semantic meaning between conditions and generated sentences.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
SEQ2SEQ is an essential setting in NLP that covers a wide range of important tasks such as openended sentence generation, dialogue, paraphrasing, and text style transfer. In this paper, we propose DIFFUSEQ, depicted in Figure 1(c), a classifier-free diffusion model that supports SEQ2SEQ text generation tasks. By modeling the conditional probability of the target sentence w given context x using one single model, one advantage of DIFFUSEQ is that this paradigm allows a complete model to fit data distribution and utilize conditional guidance, rather than depending on a separate classifier.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Different from canonical generation approaches in an autoregressive (AR) left-to-right manner (Radford et al., 2019), DIFFUSEQ generates text tokens parallelly in the non-autoregressive (NAR) way.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To corroborate the effectiveness of our DIFFUSEQ, we conduct experiments on four SEQ2SEQ tasks. Compared to AR and NAR models, which suffer from the "degeneration" problem (Holtzman et al., 2019) and rely on decoding strategies, DIFFUSEQ can achieve considerable sentence-level diversity without sacrificing the quality (see § 4.2).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To sum up, we make a series of technical and conceptual contributions: (a) we are the first to deploy the diffusion model on SEQ2SEQ text generation, and our proposed DIFFUSEQ as a conditional language model is trained end-to-end in a classifier-free manner; (b) we establish a theoretical connection among AR, NAR and DIFFUSEQ models, and justify DIFFUSEQ as an extension of iterative-NAR models; (c) with strong empirical evidence, we demonstrate the great potential of diffusion models in complex conditional language generation tasks.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Preliminary. A diffusion model typically contains forward and reverse processes. Given a data point sampled from a real-world data distribution z 0 ∼ q(z), the forward process gradually corrupts z 0 into a standard Gaussian noise z T ∼ N (0, I). For each forward step t ∈ [1, 2, ..., T ], the perturbation is controlled by q(z t |z t−1 ) = N (z t ; √ 1 − β t z t−1 , β t I), with β t ∈ (0, 1) as different variance scales. Once the forward process is completed, the reverse denoising process tries to gradually reconstruct the original data z 0 via sampling from z T by learning a diffusion model f θ .
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Problem Statement. Many recent efforts have been devoted to adapting diffusion models to discrete texts (See § 5). However, they all focus on unconditional sequence modeling. In this paper, we target the sequence-to-sequence text generation tasks. In particular, given a m-length source sequence w x = {w x 1 , ..., w x m }, we aim to learn a diffusion model that can produce a n-length target sequence w y = {w y 1 , ..., w y n } conditioning on the source sequence.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We propose DIFFUSEQ to extend vanilla diffusion models to learn conditional text generation (as shown in Figure 2), concerning the model architecture and the training objective.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Forward Process with Partial Noising. In the beginning of forward process, we follow Diffusion-LM (Li et al., 2022) to design an embedding function EMB(w) to map the discrete text w into a continuous space. In particular, given a pair of sequence w x and w y , DIFFUSEQ learns a unified Forward process Reverse process Gaussian Noise Rounding Figure 2: The diffusion process of our conditional diffusion language model DIFFUSEQ. Given the source w x and the target w y , we pair-wisely transform them into continuous space z 0 . The partial Gaussian noise is iteratively added on the target space of z t .
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
feature space of w x and w y by embedding transformation and concatenation as EMB(w x⊕y ) = [EMB(w x 1 ), ..., EMB(w x m ), EMB(w y 1 ), ..., EMB(w y n )] ∈ R (m+n)×d . The transformation allows us to adapt discrete textual input into the standard forward process, by extending the original forward chain to a new Markov transition q φ (z 0 |w x⊕y ) = N (EMB(w x⊕y ), β 0 I).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We denote z t = x t ⊕ y t to simplify the wordings, where x t and y t represent parts of z t that belong to w x and w y , respectively. For each forward step q(z t |z t−1 ), we gradually inject noise into last step's hidden state z t−1 to obtain z t . Unlike conventional diffusion models that corrupt the whole z t (both x t and y t ) without distinction, we only impose noising on y t . This modification (termed partial noising) allows us to adapt diffusion models for conditional language modeling.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Reverse Process with Conditional Denoising. The ultimate goal of the reverse process is to recover the original z 0 by denoising z t : p θ (z 0:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
where the µ θ (·) and σ θ (·) is the parameterization of the predicted mean and standard deviation of q(z t−1 |z t ) in forward process, derived using Bayes' rule. The detailed derivations are in Appendix A. With the partial nosing strategy adopted in the forward process, we can impose the input as the condition when denoising as shown in Figure 1. The proposed conditional denoising is classifier-free by nature: we do not require extra-trained classifiers to control the denoising process.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Specifically, we use a transformer architecture to model f θ , which spontaneously models the semantic relation between x t and y t . We compute the variational lower bound (L VLB ) following the original diffusion process. L round corresponds to rounding operation in Figure 2.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We further simplify the training objective as follows (details in Appendix A):
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
here we usef θ (z t , t) to denote the fractions of recovered z 0 corresponding to y 0 . Note that although in the first term, we only compute the loss w.r.t y 0 , due to the attention mechanism in the transformer, the reconstruction of y 0 also takes x 0 into account, thus the gradients from the first term will also affect the learning of x 0 . The mathematically equivalent regularization term R(||z 0 || 2 )) regularize the embedding learning. We further share the embedding function between source and target sequences, enabling the training of two different feature spaces jointly. This sets DIFFUSEQ away from existing solutions in vision such as GLIDE (Nichol et al., 2022).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Training and Inference Methods. In our preliminary experiments, we find that the high diversity in NLP datasets and long diffusion steps often result in insufficient training. We hypothesize the reason is that sampling step t uniformly causes unnecessary noise in the L VLB objective. We hence employ importance sampling (Nichol & Dhariwal, 2021) to address this problem.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Intuitively, the importance-weighted sampling algorithm will spend more steps on diffusion steps with larger L t , and vice versa.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To conduct SEQ2SEQ generation given the condition EMB(w x ), we randomly sample y T ∼ N (0, I) and concatenate y T with EMB(w x ) to obtain z T . We can now repeat the reverse process until we arrive at z 0 . At each sampling step, an anchoring function is executed towards reparameterized z t . Specifically, the anchoring function: (a) operates rounding on z t to map it back to word embedding space following Li et al. (2022); (b) replaces the part of recovered z t−1 that belongs to w x with the original x 0 , considering that this part is recovered from corrupted z t via f θ and not strictly equals to x 0 . Note that (b) is designed for DIFFUSEQ.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To improve the quality of generation, we apply the widely used Minimum Bayes Risk (MBR) decoding strategy (Koehn, 2004). We first generate a set of candidate samples S from different random seeds of DIFFUSEQ and select the best output sequence that achieves the minimum expected risk under a meaningful loss function (e.g. BLEU or other cheaper metrics like precision). In practice, we use the negative BLEU score in our implementation.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Connections to AR, Iter-NAR, and Fully-NAR Models. To better understand the behavior of DIFFUSEQ, we give the theoretical connection to autoregressive (AR), iterative non-autoregressive (iter-NAR), and fully non-autoregressive (fully-NAR) models. We argue that DIFFUSEQ can be seen as an extension of iter-NAR model. Detailed graphical learning discrepancies of these four cases are discussed in Appendix B for reference.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
AR models learn p(w y 1:n |w x ) by autoregressive decomposition based on left-context:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
while fully-NAR models (Gu et al., 2018;Qian et al., 2021) learn the conditional probability given independent assumption for fast inference:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To make a better analogy to AR and NAR models, we use a lossless way to formulate iterative NAR models (Gu et al., 2019;Ghazvininejad et al., 2019) by introducing a series of intermediate sequences w y 1:K−1 , w y K = w y with K editable iterations:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Previous study (Huang et al., 2022) shows that there is a gap called conditional total correlation between AR Eq. (4) and fully-NAR Eq. (5) learning paradigms, because of lossy decomposition of NAR models. However, when comparing iter-NAR Eq. (6) with AR Eq. (4) models, they both can be factorized into an initial prediction term and a progressive prediction process based on different context (i.e. left-context in AR and full-context in iter-NAR), and the discrepancy pointed out by Huang et al. (2022) is therefore closed in iter-NAR assuming sufficient steps. By showing DIF-FUSEQ is an extension of the iter-NAR model, we offer a justification that it will not suffer from the conditional total correlation for the same reason.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
A straight-forward way to formulate pure continuous diffusion models is to introduce a series of Gaussian noise-corrupted features along with diffusion steps: y 1:T −1 , y 0 = y, y T ∼ N (0, I).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
where p(y t−1 |y t , w x ) describes the diffusion step on continuous representations y. The rounding operation in DIFFUSEQ maps the continuous vectors y to discrete w y for each time step t, we in addition introduce this into Eq. (7):
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
By rearranging Eq. (8) into Eq. (9), we can see DIFFUSEQ can be seen as a more generalized form of iter-NAR Eq. (6) before marginalizing out {y T , . . . , y 0 }, despite the different initialization of y T 2 . A more detailed derivation is shown in Appendix C.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We conduct experiments to validate the effectiveness of DIFFUSEQ on four different tasks, against six strong AR/NAR baselines.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Tasks and Datasets. SEQ2SEQ generation covers a wide range of tasks, among which we choose four typical and popular tasks. Open domain dialogue requires models to generate informative responses given a dialogue context. We use Commonsense Conversation Dataset (Zhou et al., 2018), which is extracted from Reddit single-round dialogs, with over 3 million conversational pairs. Question generation(QG) aims to generate questions given a context as input. To obtain sufficient training samples, we use the dataset Quasar-T (Dhingra et al., 2017) preprocessed by Lin et al. (2018), and then generate document-question pairs to obtain 119K training samples (details in Appendix D.1). Text simplification aims to revise the complex text into sequences with simplified grammar and word choice. Jiang et al. (2020) constructs a corpus consisting of 677K complexsimple sentences with revision alignment. Paraphrase task generates an alternative surface form in the same language expressing the same semantic content. We adopt widely used QQP 3 sourced from the community question answering forum Quora, with 147K positive pairs.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Baselines. We consider three groups of models as baselines, covering both AR and NAR architectures. The first group of methods adopts encoder-decoder architecture (Cho et al., 2014) which is well-studied for SEQ2SEQ tasks, and we conduct experiments on two popular models: GRU with attention and Transformer (Vaswani et al., 2017). The second group is the finetuned large pre-trained language model (PLM), among which GPT2 (Radford et al., 2019) has demonstrated great success in almost all SEQ2SEQ tasks. We further compare to GPVAE (Du et al., 2022), which augments a pre-trained T5 (Raffel et al., 2020) with VAE to improve the generation diversity. For the last group of baselines, we consider LevT (Gu et al., 2019), a widely used, strong iterative NAR model. All baselines are trained following instructions in their papers, and details can be found in Appendix D.2. Evaluation. We evaluate the generated sequences from two aspects: quality and diversity. To evaluate the quality, we use the standard metric BLEU (Papineni et al., 2002) and ROUGE (Lin, 2004) score. Since string-similarity-based metrics can be unsatisfactory for open-ended generation, we also report BERTScore (Zhang et al., 2019) that assesses the semantic similarity between generated sentences and references. Details are in Appendix D.4. Higher scores of BLEU, ROUGE and BERTScore reflect better performance. As for diversity, we use distinct unigram (dist-1) to measure intra-diversity within each generated sentence, where the lower dist-1 indicates that the generated sentence contains more repeated words. For sentence-level diversity evaluation, we consider sentence-level self-BLEU (Zhu et al., 2018) to measure the n-gram overlap between the set of outputs w.r.t one source sentence, and we additionally use diverse 4-gram (div-4) (Deshpande et al., 2019) to measure the ratio of distinct 4-grams in the set of outputs per source sentence. The lower self-BLEU and higher div-4 suggest higher diversity of generation. For each method including DIFFUSEQ, we generate 3 samples for each source sentence to compute the diversity metrics.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Implementation Details. Our DIFFUSEQ is based on the 12 layers of Transformer with 12 attention heads, where the time step embedding is plugged akin to the position embedding. The maximum sequence length is 128, with embedding dimension d = 128, diffusion steps T = 2, 000 and a square-root noise schedule. To reduce the out-of-vocabulary generation, we apply Byte Pair Encoding (Sennrich et al., 2016) to construct the vocabulary. After conducting the diversity beam How can I make friends? What is the best way to make friends? How do I make friends and keep them?
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
What's the best way to make friends and make make friends?
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
How can I make friends and more something? search (DBS) (Vijayakumar et al., 2016) for the Transformer-base model and GPT model, we find that DBS does not always promote diversity over temperature sampling and therefore we list the best diversity results. We compute the accuracy metrics of DIFFUSEQ using MBR with the size of candidate samples |S| = 10. The experiment is deployed on NVIDIA A100 Tensor Core GPUs, and we use 4 GPUs on training and single GPU on sampling.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
As shown in Table 1, we conclude that DIFFUSEQ achieves comparable or even higher generation quality compared with strong baselines. At the same time, DIFFUSEQ consistently demonstrates its superiority in generating diverse outputs given the same input sequence.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
As we can see from Table 1, DIFFUSEQ wins competitions over at least one quality metric against 6 baselines × 4 tasks. Although NAR models such as LevT can also outperform AR baselines sometimes, they still lag well behind DIFFUSEQ by large margins (i.e., relative improvements over 50% for BLEU in QG task and R-L in Dialogue task). Even compared with pre-trained then finetuned GPT2 models, DIFFUSEQ still delivers superior performance than the base variant, and is comparable with the large variant, which has 8.2 times more parameters than DIFFUSEQ. These empirical results amply support our findings in § 3, where we theoretically analyze the potential of diffusion models in modeling text sequences compared with AR models given sufficient diffusion steps.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
DIFFUSEQ, as a member of the deep generative model family, also exhibit the capacity to generate highly diverse sequences. As suggested by self-BLEU (lower is better) and div-4 (higher is better), in almost all cases, DIFFUSEQ significantly outperforms 4 AR baselines in terms of sentence-level diversity (i.e., producing diverse outputs given the same input). For diversity in word choice within one sentence, we consider dist-1: a higher dist-1 indicates less repetition within a sentence. As we can see from Table 1, DIFFUSEQ has less repetition compared with encoder-decoder methods, but still fall behind the pre-trained GPT2 models (the same situation with BERTScore). These results suggest there is still room for improvement (e.g., use pre-training techniques) in diffusion models' token-level choice. Different from NAR-LevT, DIFFUSEQ does not rely on an extra length prediction module but automatically decides by the padding token instead and is able to generate longer output sentences, indicated by the last column for average generation length.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
In Table 2 
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We conduct a series of analysis to investigate the effectiveness of different aspects in DIFFUSEQ.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Generating high-quality texts with high diversity is an important requirement for many text generation applications and the trade-off between quality and diversity is always a critical concern in open-ended NLG tasks (Zhang et al., 2021). Different from AR models relying on the decoding strategy like temperature and nucleus sampling (Holtzman et al., 2019) and VAE models sampling latent variable from Gaussian Prior, the natural advantage of DIFFUSEQ is to generate different sentences along with a series of random Gaussian noise. In Figure 4, we elucidate that DIFFUSEQ have better trade-off between generation quality (BLEU) and sentence-level diversity (div-4). Here we further demonstrate that the high diversity provided by DIFFUSEQ can be turned into better quality.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
MBR is a common strategy to improve generation quality by aggregating and ranking candidate sequences, and we find that the upper bound of MBR is decided by a diversified candidate set. To valid this, we simultaneously apply MBR on both DIFFUSEQ and GPT2 with various candidate sizes |S|. The results are shown in Figure 3. As we can see, DIFFUSEQ lags behind GPT2 without using MBR (|S| = 1) or with a small candidate set (|S| = 3). However, as |S| increases, DIFFUSEQ starts to outperform GPT2 by an increasing margin. The reason is that autoregressive models like GPT2 tend to generate highly similar candidates (as discussed in § 4.2), which impedes the effectiveness of MBR. As |S| increases to 20, DIFFUSEQ still shows better rising trends than GPT2. Our findings also stress the importance of better ranking methods in diffusion research.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Step-wise Analysis against Iterative NAR. Given the underlying theoretical connection between iterative NAR and DIFFUSEQ discussed in § 3, we empirically investigate the behavior of LevT and DIFFUSEQ by analyzing their step-wise quality (i.e. BLEU) and diversity (i.e. div-4) curves.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
As is suggested in Figure 5, LevT grows fiercely in quality at the very beginning of generation, and quickly slows down in the successive refinement process. But DIFFUSEQ behaves differently, with BLEU score growing slowly at first, increasing rapidly as the diffusion process progresses and finally surpassing LevT. It is also observed that the diversity of both LevT and DIFFUSEQ is determined at the very early stage regardless of future refinement or diffusion, where DIFFUSEQ consistently outperforms LevT on diversity at any stage of generation. We conjecture that DIFFUSEQ explores more possible results at the first half of generation process, and soon converges to several potential candidates when it is closed to the end of steps. In this case, DIFFUSEQ shows its capacity to take both generation quality and diversity into consideration, and this is the capacity that iterative-NAR and even AR models can not obtain, due to the different learning paradigms.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Inference Speed. The slow sampling speed is one of the major concerns about diffusion models.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Here we fix the number of diffusion steps during training for DIFFUSEQ while shrinking the inference steps following DDIM (Song et al., 2020). As we can see from Figure 6, when reducing the inference to 1,000 diffusion steps on single GPU, DIFFUSEQ achieves a higher BLEU score than GPT2-large yet registers a closer inference speed to GPT2-large.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Effectiveness of Joint Training. In DIFFUSEQ, the representations of w x and w y are jointly trained using the same embedding function EMB(·) (stated in § 3). To validate the effectiveness  (Nichol et al., 2022;Ramesh et al., 2022). In particular, we decouple the training of EMB(w x ) and EMB(w y ) by replacing EMB(w x ) with representations extracted from a pre-trained BERT-tiny model (Turc et al., 2019). From Table 3, we find that the decoupled training strategy results in poor performance.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Diffusion Models for Text Modeling. Text-to-Image generation using diffusion models has developed many potential applications. Models such as Imagen (Saharia et al., 2022b) and DALL-E (Ramesh et al., 2022) are usually two-staged relying on the pre-trained models, requiring the alignment between the embedding vectors from two sources. GLIDE (Nichol et al., 2022) explores diffusion model with classifier-free (Ho & Salimans, 2022) guidance by setting guidance scale during training. The target space of these models is not discrete text space but stable vectors of pixel values. There are other works of diffusion on text generation, but they stick to the original encoderdecoder architecture and the diffusion process is interspersed on the decoder (Savinov et al., 2021), or the latent space (Yu et al., 2022).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
For text generation using the diffusion models, Hoogeboom et al. (2021) introduce the multinomial diffusion for character-level text generation, the forward categorical noise is applied through the Markov transition matrix. Austin et al. (2021) generalize discrete text diffusion models by introducing the absorbing state ([MASK]). However, discrete diffusion models may suffer from the scaling of the one-hot row vectors, and they only generate text samples unconditionally in discrete space. Diffusion-LM (Li et al., 2022) and Analog Bits  propose a new language model diffused on the continuous latent representations, with different mapping functions that connect the discrete and continuous space of texts. Compared with our work, we focus on the SEQ2SEQ diffusion models for text generation in the continuous space and our work is the first to explore this setting to the best of our knowledge.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Diffusion Models for Conditional Generation. Related to conditional-VAE (Zhao et al., 2017), we can consider the latent encoded input x as a condition. Diffusion-LM (Li et al., 2022) adopts the plug-and-play approaches (Dathathri et al., 2020) to compose fine-grained constraints on the generated sentences, but it fails to condition on the whole source sentence in SEQ2SEQ tasks. Noted that this controllable generation method is orthogonal to our DIFFSEQ, in other words, we can further add classifier-guided constraints on the SEQ2SEQ output to further control the text generation. There are other conditional diffusion models on the time series prediction like CSDI (Tashiro et al., 2021) or audio generation like WaveGrad (Chen et al., 2021), but their class conditions are usually attributes that are easy to model, while the contextual texts as conditions are much more complex.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We propose DIFFUSEQ to tackle SEQ2SEQ tasks in a diffusion way, which contains the strong potential to achieve better generation quality and diversity trade-off. The capability enables favorable characteristics of DIFFUSEQ to further enhance the quality of final results, by leveraging a minimum Bayes risk decoding algorithm. Besides, we theoretically connect the AR and NAR models to DIF-FUSEQ, and show that DIFFUSEQ is a powerful extension of iterative-NAR model. The empirical results demonstrate that DIFFUSEQ is also a powerful model for text generation, matching or even surpassing competitive AR, iterative NAR, and large-scale pre-trained models on quality and diversity. Given the limited progress of current diffusion models on text generation, our study addresses promising achievements by such a new sequence-to-sequence learning paradigm. 
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
The diffusion model is well-known as its ability to achieve the trade-off between flexibility and tractability of the models' probability distributions, compared with GAN, VAE and Flow-based models. Following Ho et al. (2020); Nichol & Dhariwal (2021); Song et al. (2020), we systematically define the forward noising process and reverse denoising process on latent continuous space z.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
The forward noising is to perturb the structure of data z 0 . z 0 is finally changed into the partial Gaussian noise with y T ∼ N (0, I) through T -step forward random disturbance
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
with t = 1, 2, ..., T and {β t ∈ (0, 1)} T t=1 are the variance schedule.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
where stands for Gaussian noises. In the end, q(z t |z 0 ) = N (z t ; √ᾱ t z 0 , (1 −ᾱ t )I). We use a sqrt noise schedule in Diffusion-LM (Li et al., 2022), that is,ᾱ t = 1 − t/T + s with s as a small constant at the start of noise level. The reverse process then denoises z t , aiming to recover original z 0 , and is defined as:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
The learning of p θ is based on our diffusion model DIFFUSEQ: f θ (z t , t), where the µ θ (·) and σ θ (·) is the predicted parameterization of the mean and standard variation of q(z t |z t−1 ) in forward process. Using Bayes' rule:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Substitute Eq. (11) to it and we can get the parameterized mean of q(z t−1 |z t , z 0 ):
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
and for brevity, we short the coefficient of z t and z 0 as U and E respectively.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We can use the variational lower bound to optimize the negative log-likelihood E[− log p θ (x 0 )] ≤ L VLB . The objective can be further rewritten to be a combination of several KL-divergence and entropy terms following Sohl-Dickstein et al. (2015).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
For 1 ≤ t ≤ T − 1, we compute the parameterization of L t by substituting Eq. (14) to minimize the difference from µ t and µ θ following Ho et al. (2020):
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
where C = 2||σ θ || 2 is a loss independent constant. Then the optimization of training loss min θ L VLB can be further simplified as:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
B GRAPHICAL MODELS OF AR, FULLY NAR, ITERATIVE NAR AND DIFFUSEQ MODELS
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We start from the conditional sequence generation problem, which aims to learn a conditional probability p(w y 1:n |w x ) with w x and w y . AR models learn p(w y 1:n |w x ) by autoregressive decomposition based on left-context:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
consisting of an initial prediction and an autoregressive left-context prediction process, while fully-NAR models (Gu et al., 2018;Qian et al., 2021) learn the conditional probability given independent assumption for fast inference:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
p fully-NAR (w y 1:n |w x ) = i=1,...,n p(w y i |w x ).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
To make a better analogy to AR and NAR models, we use a lossless way to formulate iterative NAR models (Gu et al., 2019;Ghazvininejad et al., 2019) by introducing a series of intermediate sequences w y 1:K−1 , w y K = w y as:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Previous study (Huang et al., 2022) shows that there is a gap called conditional total correlation between AR and fully-NAR learning paradigms, because of the lossy decomposition of NAR models. This gap is mainly responsible for the performance drop from AR to NAR models. However, when comparing iter-NAR, Eq. (20), with AR models, they both can be factorized into an initial prediction term and a progressive prediction process based on different context (i.e. left-context in AR and full-context in iter-NAR). The discrepancy as pointed out by Huang et al. (2022) is therefore closed in iter-NAR assuming sufficient steps. By showing DIFFUSEQ is an extension of the iter-NAR model, we offer a justification that it will not suffer from the conditional total correlation for the same reason.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
A straight-forward way to formulate naive diffusion models is to introduce a series of Gaussian noise-corrupted features y 1:T −1 , y 0 = y, y T ∼ N (0, I) on continuous space as:
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
where p(y t−1 |y t , w x ) describes the diffusion process on contiguous representations y. The total number of diffusion steps is denoted as T . Thereafter we omit the independent decomposition on w y and y t . To apply diffusion models on discrete space, the rounding operation in DIFFUSEQ maps the continuous vectors y to discrete w y for each time step t, we hence in addition introduce both contiguous feature y and the discrete text w y to represent the discrete text into Eq. (21): p(w y |y 0 , w x ) t=T,...,1 w y t p(y t−1 |w y t )p(w y t |y t , w x )
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
By rearranging Eq. (23) and Eq. (24), we can see that DIFFUSEQ can be seen as a more generalized form of iter-NAR before marginalizing out {y T , . . . , y 0 }, where Eq. (23) and Eq. (24) are equivalent with different computation order, despite the different initialization of y T . For NAR models, y T is uniform copied from the source sentence or unk's token embedding (Gu et al., 2018); for diffusion models, y T is sampled from normal distribution N (0, I).
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
It is notable that unlike AR and fully NAR models generating text all at once, iterative NAR and diffusion models feature a self-corrected text generation process. The graphical comparison is shown in Figure 7.  p(w y |y 0 , w x ) t=T,...,1 p(y t−1 |w y t )p(w y t |y t , w x ) = w y T ,...,w y 1 y T ,...,y0
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
t=T −1,...,0 p(w y t |w y t+1 , w x )) marginalize over y = w y 1 ,...,w y K−1 p(w y 1 |w x ) k=1...K−1 p(w y k+1 |w y k , w x ) align t and k reversely.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
= p iter-NAR (w y |w x )
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
From DIFFUSEQ to diffusion model We show how to derive DIFFUSEQ to the straight-forward diffusion model on continuous space. To construct high-quality document-question pairs from the Quasar-T dataset, which consists of document i , question, answer triplets, we extract document i , question pairs if answer exactly matches document i . After pre-processing, we obtain 119K document-question training pairs.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
We compare the settings of different models, including the number of parameters and how to sample the different output sentences, as shown in Table 4. For plain GRU-based encoder-decoder methods, we do not implement diversity search algorithms on it, thus its sentence-level diversity could be very poor. For NAR-LevT, we set max iteration to 9 and follow the termination condition mentioned in the original paper. For GPVAE-T5, we tune the scalar to find the best trade-off between quality and diversity on the dev set. The scalars of all four tasks are set to 2. We implement GPT2 baselines using HuggingFace Transformers and for the baseline Transformer-base, we use Fairseq. We list the details to obtain Figure 4. For GPVAE-T5, we set different scalars as 1, 2, 3, 4. For DIFFUSEQ, we choose trained models at different training steps to achieve different trade-off points. For other baselines, there is no explicit factor to control the diversity generation, so we leave them as single points in the figure.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
The used BLEU score is sentence-level smoothed from BLEU-1 to 4, and used ROUGE-L score is longest common subsequence based statistics. The implementation is based on NLTK 5 and torchmetrics. The n-gram based metrics may fail to capture the semantic meaning of sentences, so we consider using BERTScore 6 . Specifically, we use microsoft/deberta-xlarge-mnli to help BERTScore correlate better with human scores.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
For different tasks we list some generation examples. As we can see in Table 5, Table 6 and Table 7, DIFFUSEQ tends to generate diverse outputs, but sometimes the sentence is not as fluent as finetuned GPT2.  
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
The Japanese yen is the official and only currency recognized in Japan. Question: What is the Japanese currency?
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
GPVAE-T5 NAR-LevT * What is the japanese currency * What is the basic unit of currency for Japan ? * What is the japanese currency * What is the basic unit of currency for Japan ? * What is the japanese currency * What is the basic unit of currency for Japan ?
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
GPT2-large finetune DIFFUSEQ * What is the basic unit of currency for Japan? * What is the Japanese currency * What is the Japanese currency * Which country uses the "yen yen" in currency * What is the basic unit of currency for Japan? * What is the basic unit of currency? Table 7: Sample outputs with different random seed in Text Simplification test set.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
Complex sentence: People can experience loneliness for many reasons, and many life events may cause it, such as a lack of friendship relations during childhood and adolescence, or the physical absence of meaningful people around a person. Simplified: One cause of loneliness is a lack of friends during childhood and teenage years.
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
NAR-LevT 
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
State-of-the-art generative models for discrete data, such as discrete images and text, are based on autoregressive modeling ( Van den Oord et al., 2016;Salimans et al., 2017;Parmar et al., 2018;Child et al., 2019;Roy et al., 2021;Jun et al., 2020;Brown et al., 2020;Chowdhery et al., 2022), where the networks, often Transformers (Vaswani et al., 2017), are trained to predict each token given its preceding ones in a sequential manner or with causal attention masks. One major drawback of such approaches is that they typically require computation and memory that is quadratic to the dimension of data (e.g., sequence length or image size), leading to difficulties in modeling large images or sequences. Another drawback is that, during generation, autoregressive models generate one token at a time so the total number of sequential sampling steps is often the same as the dimension of data, making it slow in generating large images or long sequences.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
In contrast, diffusion models (Sohl-Dickstein et al., 2015;Ho et al., 2020;, or score-based generative models (Song & Ermon, 2019;Song et al., 2021), can model much higher dimensional data without running into computation and memory issues. During generation, diffusion models iteratively refine samples with a high degree of parallelism, so the total number of sequential sampling steps can be much less than the dimension of data. However, state-of-the-art diffusion models Ho et al., 2022;Ramesh et al., 2022;Saharia et al., 2022) can only generate continuous data (mainly real valued pixels), and have not yet achieved results competitive with autoregressive models in generating discrete/categorical data, such as generating discrete/categorical images (Hoogeboom et al., 2021;Austin et al., 2021).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
In this work, we propose a simple and generic approach for enabling continuous state diffusion models to generate discrete data. The key ingredient in our approach is analog bits: real numbers used to model the bits that represent the discrete data. Analog bits can be directly modeled by continuous state diffusion models, without requiring a discrete state space or re-formulation of the continuous † Work done as a student researcher at Google.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Code at https://github.com/google-research/pix2seq. 1 arXiv:2208.04202v2 [cs.CV] 1 Mar 2023
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Published as a conference paper at ICLR 2023 diffusion process. At sampling time, the generated analog bits can be decoded into discrete variables by a simple thresholding operation. Our approach, as illustrated in Figure 1, is based on the following high-level conjecture. With strong continuous generative models (diffusion models in particular), it should not be too difficult to generate highly concentrated bimodal data where each real-valued analog bit is close to a binary bit. To reduce the prediction loss (such as negative log likelihood), the network has to model structures among analog bits that can actually lead to meaningful discrete variables after thresholding.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Besides analog bits, we further propose two simple techniques, namely Self-Conditioning and Asymmetric Time Intervals that greatly improve the sample quality. We evaluate the proposed approach on both discrete image generation, and image-conditional text / caption generation. On discrete CIFAR-10 and IMAGENET 64×64, the proposed Bit Diffusion model significantly improves both existing discrete diffusion models but also the best autoregressive model. For example, on categorical CIFAR-10, the best autoregressive model (Jun et al., 2020) obtains a FID of 12.75, while our model (with 1 /3 of the model size of the autoregressive model, using 100 instead of 3072 sequential inference steps) achieves a much better 6.93. For image captioning on MS-COCO dataset, our model achieves a result competitive with a strong autoregressive captioner based on a Transformer.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Preliminaries We start with a short introduction to diffusion models (Sohl-Dickstein et al., 2015;Ho et al., 2020;. Diffusion models learn a series of state transitions to map noise from a known prior distribution to x 0 from the data distribution. To learn this (reverse) transition from the noise distribution to the data distribution, a forward transition from x 0 to x t is first defined:
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
where ∼ N (0, I), t ∼ U(0, T ) is a continuous variable, and γ(t) is a monotonically decreasing function from 1 to 0. Instead of directly learning a neural net to model the transition from x t to x t−∆ , one can learn a neural net f (x t , t) to predict x 0 (or ) from x t , and estimate x t−∆ from x t and estimatedx 0 (or˜ ). This training of f (x t , t) is based on denoising with a 2 regression loss:
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
To generate samples from a learned model, it follows a series of (reverse) state transition x T → x T −∆ → · · · → x 0 . This can be achieved by iteratively applying denoising function f on each state x t to estimate x 0 , and then make a transition to x t−∆ with the estimatedx 0 (using transition rules such as those specified in DDPM (Ho et al., 2020) or DDIM ). Note that state transitions in these diffusion models assume a continuous data space and state space. Therefore, one cannot directly apply it to model and generate discrete/categorical data.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Analog Bits A discrete data variable from an alphabet of size K can be represented using n = log 2 K bits, as {0, 1} n . Due to the discreteness, existing work has to re-formulate continuous diffusion models by adopting a discrete data space and state space (Sohl-Dickstein et al., 2015;Hoogeboom et al., 2021;Austin et al., 2021). In contrast, we propose to simply cast the binary bits {0, 1} n into real numbers R n for the continuous diffusion models 1 . We term these real numbers analog bits since they learn to share the same bimodal values as binary bits but are modeled as real numbers. To draw samples, we follow the same procedure as sampling in a continuous diffusion model, except that we apply a quantization operation at the end by simply thresholding the generated analog bits. This yields binary bits which can be then converted into original discrete/categorical variables. Notably, there is no hard constraint to force the model to generate exact binary bits, but we · · ·
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
x t+∆
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
x t x t−∆ · · · x 0x 0 (a) Standard reverse diffusion steps.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
x t+∆
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
x t x t−∆ · · · x 0x 0 (b) Self-Conditioning on the previous x0 estimate. Figure 2: An illustration of reverse diffusion sampling steps (a) without or (b) with Self-Conditioning.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
x 0 denotes the estimation of data sample by the denoising network f at a sampling step. We propose to condition the network directly on its previously generated/estimated samples. expect a strong continuous generative model to generate real numbers that exhibit very clear bimodal concentrations and this is what happens in our experiments.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
For simplicity, we use the same regression loss function (Eq. 2) for modeling analog bits. However, it is possible to use other loss functions such as the cross entropy loss. We also note that the binary encoding mechanism for constructing analog bits is extensible as well (e.g., one-hot encoding). Extensions of loss functions and binary encoding are described in the appendix B.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Conditioning is a useful technique for improving diffusion models Ho et al., 2022). However, a typical conditioning variable is either from some external sources, such as class labels  or low-resolution images from another network Saharia et al., 2021;Ho et al., 2022). Here we propose a technique for the model to directly condition on previously generated samples of its own during the iterative sampling process, which can significantly improve the sample quality of diffusion models.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
In a typical diffusion sampling process, the model iteratively predicts x 0 (or ) in order to progress the chain of mapping noise into data. However, as shown in Figure 2a, the previously estimatedx 0 is simply discarded when estimating x 0 from a new time step, i.e. the denoising function f (x t , t) does not directly depend on a previously estimatedx 0 . Here we consider a slightly different denoising function of f (x t ,x 0 , t) that also takes previous generated samples as its input, illustrated in Figure 2b. A simple implementation of Self-Conditioning is to concatenate x t with previously estimatedx 0 . Given thatx 0 is from the earlier prediction of the model in the sampling chain, this comes at a negligible extra cost during sampling. In order to train the denoising function f (x t ,x 0 , t), we make some small changes to the training. With some probability (e.g., 50%), we setx 0 = 0 which falls back to modeling without Self-Conditioning. At other times, we first estimatex 0 = f (x t , 0, t) and then use it for Self-Conditioning. Note that we do not backpropagate through the estimatedx 0 so the overall increase of training time is small (e.g., less than 25%).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Asymmetric Time Intervals Besides Self-Conditioning, we identify another factor, time step t, that can also impact Bit Diffusion models. Time step t is an integral part of both denoising network f (x t , t) as well as the state transitions. During a typical reverse diffusion process, the model takes symmetric time intervals (i.e., ∆ as in t → t − ∆) for both the state transition and time reduction itself, resulting in the same/shared t for both arguments of f (x t , t). However, we find that, when taking large reverse steps, using asymmetric time intervals, implemented via a simple manipulation of time scheduling at generation, can lead to improved sampling quality for Bit Diffusion models.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
More specially, with asymmetric time intervals during the sampling process, we have f (x t , t ), where t = t + ξ and ξ is a small non-negative time difference parameter. Note that training remains unchanged, and the same/shared t is used for both arguments of the f (x t , t). Figure 3 illustrates the effect with a trained Bit Diffusion model, where it is asked to take two reversing steps from a state x t constructed using the forward diffusion, and it shows that asymmetric time intervals reduce the number of noisy pixels (after thresholding and converting back to discrete variables).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Figure 3: When taking a large reverse step from x t=0.6 to x t=0.1 in Bit Diffusion with maximum time T = 1.0, we see that asymmetric time intervals with a positive time difference ξ improve the denoising quality of x t=0.1 (by reducing the number of noisy pixels).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Putting it together Algorithm 1 and 2 summarize the training and sampling algorithms for the proposed Bit Diffusion model with Analog Bits, Self-Conditioning, and Asymmetric Time Intervals (via the td parameter). The proposed changes to the existing diffusion models are highlighted in blue. Note that unlike standard diffusion models (Sohl-Dickstein et al., 2015;Ho et al., 2020;, we use a continuous time parameterization between 0 and 1 instead of a fixed discrete time for maximal flexibility but they perform similarly. More details of the algorithm (including some important functions) can be found in Appendix A.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Algorithm 1 
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
We experiment with two different discrete data generation tasks, namely discrete/categorical image generation, and image captioning (image-conditional text generation). Binary encoding Each pixel consists of 3 sub-pixels (RGB channels), and each sub-pixel is an integer in [0, 256) representing the intensity. Standard continuous generative models cast RGB channels as real numbers and normalize them in [−1, 1]. For discrete image generation, we consider three discrete encoding for sub-pixels, namely UINT8, GRAY CODE, and UINT8 (RAND). In UINT8, we use 8-bit binary codes converted from the corresponding sub-pixel integer in [0, 256). In GRAY CODE, we assign 8-bit binary codes uniquely to each sub-pixel integer such that two adjacent integers only differ by 1 bit. And in UINT8 (RAND), we assign 8-bit binary codes to every sub-pixel integer by randomly shuffling the integer-to-bits mapping in UINT8. The binary codes in UINT8 and GRAY CODE are loosely correlated with its original sub-pixel intensities, while UINT8 (RAND) has no correlation so each sub-pixel is a categorical variable. The details of the binary codes and their correlations with sub-pixel intensity can be found in the appendix C. We shift and scale the binary bits from 0, 1 to −1, 1 for the analog bits.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
For image captioning, we follow , and use sentencepiece (Kudo & Richardson, 2018) with a vocabulary of size 32K to tokenize the captions. After tokenization, we encode each token into 15 analog bits using the binary codes converted from the corresponding integer. We set the maximum number of tokens to 64 so the total sequence length is 960 bits. Since we directly model bits, it is also possible to directly work with their byte representations without a tokenizer, but we leave this for future work.   , we use a base channel dimension of 192, multiplied by 1,2,3,4 in 4 stages and 3 residual blocks per stage, which account for a total of 240M parameters 3 . For UINT8 (RAND) encoding, we find the following "softmax factorization" architectural tweak on the final output layer can lead to a better performance. Instead of using a linear output layer to predict analog bits directly, we first predict a probability distribution over 256 classes per sub-pixel (with each class corresponds to one of the 256 different 8-bit codes), and then map class distribution into analog bits by taking weighted average over all 256 different 8-bit codes.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
For image captioning, we follow the architecture used in , with a pre-trained image encoder using the object detection task, for both autoregressive baseline as well as the proposed method. Both decoders are randomly initialized 6-layer Transformer (Vaswani et al., 2017) decoder with 512 dimension per layer. For the autoregressive decoder, the token attention matrix is offset by the causal masks, but it is non-masked all-to-all attention for our Bit Diffusion. : Class-conditional generations on continuous v.s. discrete ImageNet 64×64. Each row represents random samples conditioned on a class, and the classes are adopted from , namely, 9: ostrich, 11: goldfinch, 130: flamingo, 141: redshank, 154: pekinese, 157: papillon, 97: drake and 28: spotted salamander. More samples from random classes are shown in Figure 11.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
achieves better results compared to both existing discrete diffusion models and the best autoregressive model. When compared to continuous diffusion models (i.e., DDPM), our Bit Diffusion models on UINT8 and GRAY CODE can achieve similar performance.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Discrete generation of IMAGENET 64×64 is significantly harder than CIFAR-10, and we have not found other competing methods that report FIDs, so we only compare the proposed method against DDPM on continuous pixels. Results are shown in Table 2. We find that the diffusion model on continuous pixels has the best FID while the diffusion model on UINT8 (RAND), i.e., categorical data, has the worst FID, indicating the increase of hardness when removing intensity/order information in sub-pixels. Note that, in these experiments, there is no extra model capacity to compensate for the loss of intensity/order information since the model sizes are the same. Figure 4 shows generated images of different diffusion models on continuous and discrete IMAGENET 64×64. Despite the differences in FIDs, visually these samples look similar.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Ablation of Self-Conditioning Figure 5 shows the effectiveness of the Self-Conditioning technique in both Bit Diffusion and continuous diffusion models. Note that the experiments are performed in three settings, namely CIFAR-10 with UINT8, CIFAR-10 with UINT8 (RAND), and IMAGENET     64×64 with continuous pixels, where the only difference for pairs in each setting is whether the Self-Conditioning is used. For CIFAR-10, we find that Self-Conditioning greatly improves the performance across different binary encodings. We also notice that for Bit Diffusion, predicting x 0 is much more effective than predicting . For IMAGENET 64×64, we find that the proposed Self-Conditioning also leads to improved FIDs for continuous diffusion (i.e., DDPM). Therefore, we conclude that Self-Conditioning by itself is a generic technique that can benefit diffusion models on both continuous and discrete data. Figure 6 shows the FID on generated IMAGENET 64×64 samples as we vary the time difference parameter during the sampling process. We find that as the number of steps increases (from 100 to 400), the optimal time difference shrinks to 0. For 100 steps, a non-zero time difference leads to a significant improvement of FID. We also note that for Bit Diffusion on UINT8 (RAND), using 400 sampling steps actually leads to a drastically worse sample quality than using 100 steps. This is related to how the Self-Conditioning is applied and we present alternative Self-Conditioning sampling strategies in the Appendix G, some of which lead to improved FIDs at a cost of longer sampling time. Figure 7 visualizes the distribution of generated analog bits from 64 generated images on IMAGENET 64×64. Although there is no hard constraint on the analog bits being binary / bimodal, the generated ones are highly concentrated on two modes, which makes the thresholding / quantization easy and robust.  
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
We compare our Bit Diffusion model with an autoregressive Transformer baseline . As mentioned, both models have similar architectures, with an object detection pretrained  image encoder, and a randomly initialized Transformer (Vaswani et al., 2017) decoder. Table 3 presents the main comparison. Overall, our model achieves similar performance as the autoregressive model. We find that generally it only needs about 10 steps for the model to achieve good results, despite that there are a total of maximum 960 bits for caption that the model has to model. We find that the asymmetric time intervals play an important role in the final performance of our model, as demonstrated in Table 4, especially when sampling steps are fewer.    Table 5 provides some generated samples of our model when different inference steps are used. The model makes mistakes when the sampling steps are too few, and the mistakes may not always be interpretable due to that the model directly predicts the bits behind the tokenized word pieces and a small difference in bits can lead to total different words. Steps=1: A group of diets in for foraa
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=2: A group of elephants in in a\ufffd.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=5: A group of elephants standing in a Irish.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=10: A group of elephants standing by a fence.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=1: A says man fora You a\u0000.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=2: A says man a skateboard' aa.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=5: A man on a skateboard on a skateboard.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Steps=10: A man is doing a trick on a skateboard.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Autoregressive models for discrete data Autoregressive models have demonstrated state-of-the-art results when it comes to generating discrete data. In particular, text generation, or language modeling, is dominated by autoregressive approaches Brown et al., 2020;Chowdhery et al., 2022). Autoregressive models are also applied to discrete/categorical image generation (Van den Oord et al., 2016;Salimans et al., 2017;Parmar et al., 2018;Child et al., 2019;Roy et al., 2021;Jun et al., 2020;Chen et al., 2020a), where they work well on small image resolutions. However, the computation cost and memory requirement increase drastically (typically in a quadratic relation) as the size of sequence or the image resolution increase, so it becomes very challenging to scale these approaches to data with large dimensions.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Diffusion models for discrete data State-of-the-art diffusion models Ho et al., 2022;Ramesh et al., 2022;Saharia et al., 2022) cannot generate discrete or categorical data. Existing extensions of these continuous diffusion models to discrete data are based on both discrete data space and state space (Sohl-Dickstein et al., 2015;Hoogeboom et al., 2021;Austin et al., 2021;Campbell et al., 2022). Compared to discrete state space, continuous state space is more flexible and potentially more efficient. Our approach is also compatible with both discrete and continuous time, and does not require re-formulation of existing continuous models, thus it is simpler and can potentially be plugged into a broader family of generative models.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Another line of discrete diffusion models is based on the embedding of discrete data . One can also consider our binary encoding with analog bits as a simple fixed encoder, and the decoding / quantization of bimodal analog bits is easy and robust via a simple thresholding operation. In contrast, the quantization of real numbers in generated continuous embedding vectors may contain multiple modes per dimension, leading to potential difficulty in thresholding/quantization.  Fedus et al., 2018) have also been applied to generate discrete data. These methods have not yet achieved the level of performance as autoregressive models on tasks such as discrete image generation or text generation, in terms of sample quality or data likelihood. Potentially, the proposed analog bits can also be applied to these continuous generative models, by having the networks directly model and generate analog bits, but it is not explored in this work.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Other related work The proposed Self-Conditioning technique shares some similarities with selfmodulation in GANs (Chen et al., 2018a) (where the earlier latent state can directly modulate the later latent states) and SUNDAE (Savinov et al., 2021) (where an inference step is incorporated for denoising).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
We introduce a simple and generic technique that enables continuous state diffusion models to generate discrete data. The main idea is to encode discrete or categorical data into bits and then model these bits as real numbers that we call analog bits. We also propose two simple techniques, namely Self-Conditioning (i.e., condition the diffusion models directly on their previously generated samples) and Asymmetric Time Intervals, that lead to improved sample quality. We demonstrate that our approach leads to state-of-the-art results in discrete / categorical image generation, beating the best autoregressive model. In an image-conditional text generation task on MS-COCO dataset, we also achieve competitive results compared to autoregressive models. One limitation of our approach, similar to other existing diffusion models, is that they still require a significant number of inference steps for generating good (image) samples. However, we expect that future improvements from diffusion models for continuous data can also transfer to discrete data using analog bits. Algorithm 3 Binary encoding and decoding algorithms (in Tensorflow).
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
import tensorflow as tf def int2bit(x, n=8): # Convert integers into the corresponding binary bits. x = tf.bitwise.right_shift(tf.expand_dims(x, -1), tf.range(n)) x = tf.math.mod(x, 2) return x def bit2int(x): # Convert binary bits into the corresponding integers. x = tf.cast(x, tf.int32) n = x.shape[-1] x = tf.math.reduce_sum(x * (2 ** tf.range(n)), -1) return x
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Algorithm 4 xt estimation with DDIM / DDPM updating rules. An alternative binary encoding to the base-2 encoding of the discrete data used in the main paper, is the one-hot encoding, where a discrete variable is represented as a vector whose length is the same as the vocabulary size K, with a single slot being 1 and the rest being 0. The resulting one-hot vector can be similarly treated as analog bits and modeled by continuous state diffusion models. To obtain discrete variables corresponding to the generated analog bits, we use an arg max operation over all candidate categories, instead of the thresholding operation in base-2 analog bits. Note that the one-hot encoding requires K bits, which is less efficient compared to base-2 encoding that only requires log 2 K bits, especially for large K. 5
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
As we use 2 loss by default for its simplicity and compatibility with continuous diffusion models. The proposed Bit Diffusion models can work with other loss functions too. Since the analog bits are bimodal, we can use the following sigmoid cross entropy loss:
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
where we assume x 0 ∈ {−1, 1} n , and σ is a sigmoid function. During the sampling process, we use 2σ(f (x t , t)) − 1 as the output of denoising network.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
For one-hot analog bits, one could also add a softmax activation function for the output of denosing network f , and use the following softmax cross entropy loss:
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
where we assume x 0 ∈ {0, 1} n which is the one-hot representation. Table 6 presents FIDs of Bit Diffusion models with different types of analog bits and loss functions on unconditional CIFAR-10. Note that it is possible some of these results can be improved by more tuning of hyper-parameters or tweaks of the network, but we do not focus on them in this work. In the main paper, we describe three different types of binary encodings of pixels. Here we provide additional detail on how we generate UINT8 (RAND): we first apply a random permutation to 256 sub-pixel values, and then assign the binary binary bits of permuted integers to the non-permuted integers. For example, assume 0 is mapped to 228 after the permutation, the analog bits of 0 would be the binary bits of 228. The random permutation is generated by numpy.random.seed(42); numpy.random.shuffle(numpy.arange (256)).   Figure 8 show the correlation between Hamming distance of three different binary encodings we use and the (absolute) difference of sub-pixel intensity. This is done by taking every pair of subpixel integers (in [0, 256)), compute their absolute difference, as well as the Hamming distance between the corresponding binary bits. We find that both UINT8 and GRAY CODE exhibit partial correlation between the two quantities (with different correlation patterns), meaning that these codes partially contain the order information about the original sub-pixel intensity. However, UINT8 (RAND) exhibits no correlation between hamming distance and sub-pixel intensity, indicating the order information is fully removed, thus can be considered as categorical data.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
An intuitive toy example of how a continuous generative model can generate binary data is given in Figure 9, where a mapping from prior distribution at x T to data distribution at x 0 is shown. With a deterministic sampler (such as DDIM), it is straight-forward how they can represent any Bernoulli distribution by dividing the prior into two regions of probability densities corresponding to the Bernoulli distribution. For stochastic samplers, they can achieve a similar effect but the mapping from noise to data is stochastic. For an arbitrary discrete variable, represented as m-dimensional Bernoulli distribution, the mapping from continuous noise distribution to the target Bernoulli distribution also exists but it is more complicated (and difficult to visualize).  
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
As our models are based on continuous diffusion models, in theory our models are able to incorporate faster samplers. To this end, we conduct preliminary exploration of using DPM-Solver (Lu et al., 2022) for sampling some of our models.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
We find that DPM-Solver provides a boost to diffusion models based on analog bits, similar to what it is able to do for continuous data. This shows a potential of our model enjoying faster continuous sampler while other baselines (e.g., D3PM) may not be able to do due to their use of discrete states. Table 7 below shows the FID scores of bit diffusion models on ImageNet-64x64 under different binary encoding schemes. We find that the DPM-Solver is able to provide a significant reduction in function evaluations for bit diffusion on discrete/categorical data (with 30 NFEs it gets comparable FIDs as 100 NFEs of DDIM), similar to that in continuous diffusion models.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Furthermore, we also find that self-conditioning continues to provide a boost with DPM-solver. For example, the table 8 shows FID scores of diffusion models on ImageNet 64x64 (continuous rgb values). And we find that the self-conditioning consistently improves the performance of DPM-Solver with fixed number of function evaluations.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
F EXTRA RANDOM SAMPLES ON CIFAR-10 AND IMAGENET 64×64 Figure 10 shows random samples (non cherry-picked) from unconditional diffusion models on CIFAR-10 with continuous pixels and analog bits. Figure 11 shows random samples (non cherry-picked) from class-conditional diffusion models on IMAGENET 64×64 with continuous pixels and analog bits.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Published as a conference paper at ICLR 2023   In this section, we present extensions to the default sampling strategy with Self-Conditioning in Algorithm 2. The default sampling strategy utilizes data estimate from the previous step as the conditional input to the denoising network for producing data estimate at the current step. While this is both simple and effective, we observe that, for UINT8 (RAND) encoding of pixels, as the number of sampling steps increases (with both DDIM or DDPM samplers), the generated samples tend to be over-smoothed. We propose the following two extensions of the default sampling strategy to mitigate the issues and provide improvements when using larger sampling steps.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
The first extension to the default sampling strategy is to adopt an exponential moving average over the previous data estimate to provide a more reliable conditioning input, similar to a momentum optimizer. The detailed procedure is shown in algorithm 5, where the differences from the default sampling strategy are highlighted in blue. Note that the default sampling strategy can also be considered as a special case of this generalized form in that the momentum is set to zero.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Self-Conditioning based on Self-Guidance One potential issue with the default sampling strategy is the slight discrepancy of the Self-Conditioning signal during training and inference/sampling. Specifically, during training, the Self-Conditioning signal is the data estimate from the same time step, while, during sampling, it is from the past time step(s). Therefore, here we propose an approach that also use the same step data estimate for self-conditioning, which comes at the cost of extra forward pass over the denoising network at sampling time. Specifically, we conduct two forward passes of denoising network per sampling step, one with zero data estimate and the other with current data step estimate, and then we use a weighted combination, similar to (Ho & Salimans, 2021), of both prediction to form the final prediction at the current step. The detailed procedure is given in algorithm 6 with differences to the default sampling strategy highlighted.
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning
Algorithm 5 Sampling with Self-Conditioning based on Momentum Estimate.  Table 9 reports the best FID scores across various sampling strategies discussed here (as well as samplers, sampling steps, time difference in asymmetric time intervals).  Figure 12 shows FIDs on conditional IMAGENET 64×64 with UINT8 encoding, using Momentum Estimate with different sampling steps. We find that the momentum on the data estimate is only helpful when sampling steps are larger. Figure 13 shows FIDs on conditional IMAGENET 64×64 with UINT8 encoding, using Self-Guidance with different sampling steps. We find that a guidance weight between 3.0 and 5.0 is generally preferable and robust to other hyper-parameters (such as sampler choice, sampling steps, and time difference).   
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Diffusion models (Sohl-Dickstein et al., 2015), trained to iteratively refine noised inputs, have recently emerged as powerful tools for generative modeling in several continuous-valued domains such as images (Ho et al., 2020), audio (Kong et al., 2021), video (Ho et al., 2022), among others. Attempts to adapt them for discrete domains such as text data, however, have only had limited success: prior work have shown to be promising on specialized cases and small datasets (Hoogeboom et al., 2021;Austin et al., 2021;Li et al., 2022;Chen et al., 2022), but diffusion models for text still underperform (and thus are not widely adopted) compared to autoregressive language models (AR-LMs) which remain the state-of-the-art general purpose text generators Brown et al., 2020).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Despite potential advantages of diffusion models for text, there are two key challenges. First, diffusion models generate text non-autoregressively, i.e., they generate (and update) the entire sequence simultaneously rather than token by token left-toright. Although this property is useful in practice since each output token is informed by a broader bi-directional context (Lee et al., 2018;Ghazvininejad et al., 2019), it requires pre-defining an output sequence length. This limits the flexibility and applicability of trained models. On the other hand, non-autoregressive training with long sequences is expensive and difficult to optimize. In this work, we propose a semi-autoregressive solution which strikes a balance between length flexibility and the ability to alter previously generated tokens.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
A major advantage of diffusion models over the current standard of autoregressive LMs is their post-hoc controllability using guidance from auxiliary models such as style classifiers (Dhariwal and Nichol, 2021). However, controllability is hard to achieve without compromises in modularity in diffusion-based LMs for text. To enable diffusion generation into discrete text rather than continuous modalities, prior approaches have employed different approximations, e.g., training with embeddings, character, or byte-level methods (Li et al., 2022;Hoogeboom et al., 2021;Austin et al., 2021;Chen et al., 2022). In contrast, existing mainstream LMs and the guidance classifiers they derive often operate at a sub-word level with sub-word representations trained jointly with the language model (Devlin et al., 2019;Liu et al., 2019;Raffel et al., 2020). Subsequently, changing the input representations to characters or embeddings requires developing guidance models from scratch, which can be expensive or infeasible in many cases. In this work, we propose a simplex-based solution which enables the diffusion over discrete texts while maintaining the advantages of diffusion models with plug-andcontrol guidance models.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
In sum, to enable diffusion-based LMs for text we present SSD-LM ( §3), addressing the above two challenges. SSD-LM is trained to generate text semi-autoregressively-generating blocks of tokens left-to-right with bidirectional context within the block-which offers the benefits of both AR-LMs and diffusion models. It supports training with and generating variable-length sequences. At the same time, it allows refinement within the token block, in contrast to token-level autoregressive decoding where previously generated tokens cannot be modified at all. SSD-LM uses the same tokenization as popular AR-LMs, representing discrete text via a distribution (or simplex) defined over the vocabulary and is trained to reconstruct texts from noisy versions of the distributions. Due to its underlying representation, our method also offers an easy and modular way of guided (controlled) generation using off-the-shelf text classifiers under the minimal assumption of shared tokenizer.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Our evaluation experiments show, for the first time, that a diffusion-based LM matches or outperforms strong AR-LMs on standard text generation benchmarks ( §4). We evaluate SSD-LM on two tasks: (1) unconstrained prompt-based generation substantially outperforming existing diffusion LM approaches and performing on par with or outperforming strong autoregressive LM GPT-2  on both quality and diversity ( §4.2); and (2) controlled text generation with guidance from off-the-shelf classifiers (no posthoc training/adaptation) outperforming competitive controlled text generation baselines ( §4.3).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Since their inception as image generators, diffusion models (and their cousins score-based models (Song and Ermon, 2019)) have been widely adopted as high-quality generative models for multiple data modalities. Here, we briefly describe a simplified view of a canonical method, denoising diffusion probabilistic models (Ho et al., 2020, DDPM) which we adapt in this work for text generation. We assume a given dataset D = { 1 x 0 , . . . , N x 0 } of continuous valued items i x 0 (e.g., pixel values of an image) henceforth referred to as x 0 for simplicity.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Training Training a diffusion model first involves adding a series of Gaussian noise to the original data x 0 , through T timesteps:
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
where α t ′ follow a predefined schedule such thatᾱ t → 0 as t → T . This process is called forward diffusion. A diffusion model (parameterized by θ) is trained to reverse this forward process by predicting the added noise ϵ t given x t with the following loss:
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Inference To get an output from this model, we sample x T ∼ N (0, I) and iteratively reconstruct a sample x 0 by going back in time,
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
for t = T, . . . , 1. 2 The key obstacle in using vanilla diffusion models directly as text generators is that language consists of discrete tokens, i.e., a noncontinuous x 0 to which a continuous valued Gaussian noise cannot be added. We propose a straightforward and effective solution by treating tokens as continuous valued simplexes over the vocabulary (Hoang et al., 2017). Other existing methods addressing this problem are discussed in §5.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
An autoregressive LM model optimizes for the likelihood of a sequence of tokens w 0 , . . . , w L−1 .
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
To decode from AR-LMs, one can provide a context w <c and decode the next token w c iteratively by predicting p θ (w c | w <c ) and sampling from it to get the discrete token (Fan et al., 2018;Holtzman et al., 2020). Prior work has shown that these decoding approaches (and by extension the LMs themselves) are prone to degrade when generating long sequences and often devolve into repeating subsequences (Holtzman et al., 2020;Meister et al., 2022). In addition, such LMs do not provide a natural way to incorporate sequence-level control as tokens are generated one at a time without the ability Figure 1: Training SSD-LM (a snapshot at context size c = 2, block size B = 3). Horizontal axis represents the order of tokens. Vertical axis represents the diffusion timesteps. Shade means observable variables. Square means discrete vocabulary, while circle means continuous logits. Red components are inputs to the learning model θ. to modify previously generated tokens (Dathathri et al., 2020;Kumar et al., 2022b). In this work, we present a method to train a semi-autoregressive LM that decodes blocks of B tokens at a time, alleviating said issues with the support of diffusion models. Existing literature addressing the two issues individually are discussed in §5.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We introduce SSD-LM-Semi-autoregressive Simplex-based Diffusion Language Model-adapting key components from both autoregressive LM and vanilla diffusion models. Conceptually, SSD-LM uses diffusion model to decode w c:c+B , a block of tokens of length B, given a Gaussian noise and a context w <c of length c. We show an intuitive diagram and pseudo-code for the training and decoding algorithm of SSD-LM in Figure 1 
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Continuous data representation To build a continuous representation for discrete tokens, we adopt an almost-one-hot simplex representation over the model's vocabulary V . We define a simple operation logits-generation(.) to map a token w tõ w ∈ {−K, +K} |V | as follows.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
where i is the index of the vocabulary. We call w the logits for token w, and softmax(w) gives a probability simplex over the vocabulary V , with a probability mass concentrated on the token w.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
There is no learnable parameter in this mapping. 
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
) are fully noisy simplexes over V , with a logit-normal distribution (Atchison and Shen, 1980). Loss function In Eq. 2, a diffusion model is trained to predict the added noise from the noisy representations. Since the forward diffusion process can be computed in a single step (Eq. 1), the notion here is equivalent to predicting the original data representation (Song et al., 2021;Li et al., 2022). Our objective follows the same intuition but estimates a likelihood instead of the L2 distance while conditioning on additional context: 3
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
The architecture for θ throughout this work is a bi-directional Transformer encoder (Vaswani et al., 2017). Specifically, the input to the model is a concatenation of the context w <c and a sequence of noisy vocabulary simplexes softmax(w c:c+B t ) of 3 L2 distance did not work in our pilot study potentially due to the intrinsically skewed simplex representation. One minimal modification made to the Transformer model is that in addition to the conventional embedding lookup for w <c , we modify the embedding layer to take as input a distribution over the vocabulary, softmax(w c:c+B t ), and compute the embedding vector as a weighted sum of the embedding table. A timestep embedding is also added before the first Transformer block to inform the model of the current timestep. 4
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
In §A, we present another interpretation of the training objective as an intuitive contrastive loss.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Logits projection Similar to continuous-valued diffusion models, sampling from SSD-LM involves reverse diffusion from t = T, . . . , 1 starting with a Gaussian noise. At any timestep t, our model θ takes as input noised logitsw c:c+B t and estimates the probability distribution of the original tokens in data by first predicting the logits: w c:c+B logits,t = logits θ (w c:c+B |w c:c+B t , w <c ) (10) which are then converted to a distribution via softmax. To feed this output to the next step of reverse diffusion, t−1, we define a logits-projection operation to build a predicted data representation close to the initial data representation (almost-one-hot mapping; Eq. 5). We consider three projection operations. 4 More specifically, we have word embeddings for the context, Embctx(w <c ), and for the noisy diffusion representations, Wdiff[softmax(w c:c+B t )]. The timestep embedding is added to the diffusion word embeddings, Wtime(t/T ). It is similar to positional embeddings, just not varying across sequence positions. We fold it in θ for notation simplicity.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
• Greedy: creates an almost-one-hot logit centered at the highest probability token. 5
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
• Sampling: creates an almost-one-hot logit centered around a token sampled from the output distribution using top-p sampling (Holtzman et al., 2020). p is a hyperparameter.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
• Multi-hot: creates an almost-one-hot logit centered around all tokens in the top-p nucleus.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Decoding iteration Starting from pure noisẽ w c:c+B T ∼ N (0, K 2 I), in each decoding timestep we compute:
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
w c:c+B
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
for t = T, . . . , 1 and z ∼ N (0, K 2 I).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
At t = 1, the final B-token block is computed simply as argmaxw c:c+B 0 . To generate the next block, we concatenate the generated block to the previous context to create a new context of length c+B and follow the reverse-diffusion process again as described above. This process can be repeated until the maximum desired length is reached. 6 It is worth noting that our proposed decoding algorithm is novel and different from the DDPM decoding (Eq. 3). The DDPM decoding is designed for diffusion in a continuous space and failed to generate sensible outputs in our preliminary experiments based on simplexes. In §B, we draw a theoretical connection between our decoding algorithm and DDPM decoding, and also highlight the intuitive difference between the two.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Highly-modular control A useful property of continuous diffusion models that naturally arises from their definition is the ability to guide the generated samples to have user-defined attributes at test time. This can be done using gradients from auxiliary models such as classifiers (Dhariwal and Nichol, 2021), e.g., guiding the output of an LM to be of a positive sentiment using a sentiment classifier. There is a vibrant community of developers on platforms such as HuggingFace where many such text classifiers are publicly available. The underlying data representation of SSD-LM is based on vocabulary simplexes. Hence, as long as a classifier shares the same tokenizer as the LM, it can be used for control in an off-the-shelf manner without modifications. This is in contrast to prior work in diffusion language models that do not support such classifiers due to differences in their input representation space (Hoogeboom et al., 2021;Austin et al., 2021;Li et al., 2022;Chen et al., 2022) and require retraining the classifiers from scratch. This ability makes SSD-LM highly modular for controlled text generation and offers key benefits: (1) Training accurate classifiers for many tasks requires huge amounts of data where retraining them can be quite expensive, and (2) this approach allows control from classifiers that are open to use but have been trained on closed source data.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
To guide SSD-LM to generate texts with a target attribute y via a standalone attribute model f ϕ (·), we update w c:c+B logits,t (Eq. 10) at each timestep t to the form below, drifting according to the gradients from the attribute classifier.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
w c:c+B logits,t + λ∇ w c:c+B
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
where λ is a hyperparameter balancing the weight of control. The parameters of the standalone attribute model ϕ are frozen. We make a trivial modification to the embedding computation as in §3.1, to allow the classifier to take as input a simplex.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Forward diffusion coefficientᾱ t We follow Nichol and Dhariwal (2021) for a cosine schedule ofᾱ t :
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
where s is small offset set to 1e-4 in our work and α t =ᾱ t α t−1 . Fewer timesteps T in decoding Decoding from diffusion models requires a series of timesteps (T ) which can be computationally expensive if T is large. Following Li et al. (2022), we consider using a smaller value of T at test time to improve decoding speed. In this work, we primarily experiment with T decode = T train 2 and T decode = T train 5 . Flexible decoding block size B Our SSD-LM is trained with a fixed token block size B train . However, the decoding algorithm has a freedom to use a different B decode . In our experiments, we consider both scenarios of B train = B decode and B train ̸ = B decode . Nevertheless, we leave for future work a more detailed analysis of the impact of the difference between B train and B decode on model performance. Pretraining data, constants, and resource We train SSD-LM on the same data as GPT2 to make fair comparisons possible: OpenWebText (Gokaslan and Cohen, 2019) which contains 9B tokens. Following Zhang et al. (2022), we consider this data as one contiguous sequence of tokens and break it into sequences of length 200 (same as the maximum sequence length our model accepts). We randomly sample 99% of these sequences for pretraining while leaving the rest as held out for evaluation. We use the following model hyperparameters: 8 L = 200, B train = 25, T train = 5000, K = 5
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We use an aggregated batch size of 6,144 and a learning rate of 1e-4 with an AdamW optimizer (Loshchilov and Hutter, 2019). We trained SSD-LM for 100K steps, which took about 6 days on 32 Nvidia V100 GPUs.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Pretraining loss Canonical training-time perplexity of LMs is not compatible with diffusion LMs due to the difference in the inputs to the models (Eq. 4 and Eq. 9). Our pretraining loss is a per-token negative log-likelihood (NLL) that depends on the specific noise schedule being used. SSD-LM gets an average NLL of 3.87 at the end of pretraining. We show a pretraining loss curve in the appendix ( §D).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Setup First, we benchmark SSD-LM with autoregressive LMs trained on the same data (GPT2) on text generation quality. We randomly sample 1000 sequences from the held-out OpenWebText test data, extract their prefixes as prompts (context), and generate continuations from the LMs. We consider three setups: with prompt lengths 25, 50 and 100 with respective output lengths as 25, 50 and 100 tokens. In each setup, we sample 5 continuations for each input context, thus comparing the quality of 5,000 generations from baseline GPT-2 models and our SSD-LM.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We compare SSD-LM with GPT2-medium, large and xl models (containing 0.4B, 0.8B and 1.6B parameters respectively) as baselines. For reference, our model size is comparable to GPT2medium. We experiment with two popular decoding strategies for the baseline GPT-2 models with canonical parameters: nucleus sampling (Holtzman et al., 2020) with a top-p of 0.9 and 0.95, and typical sampling (Meister et al., 2022) with a typical-τ of 0.2 and 0.95.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
For SSD-LM, we consider three logits projection strategies, sampling and multi-hot with top-p ∈ {0.0, 0.1, 0.2, 0.5, 0.7, 0.9, 0.95, 0.99}, and greedy (which is functionally equivalent to the sampling with top-p=0). We use a test block size (B decode ) of 25. When generating samples of length 50 or 100, we semi-autoregressively sample in blocks of 25 and feed them as additional context to generate the next block as described in §3.2.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We evaluate the generated continuations on two axes: quality and diversity. As automatic quality metrics, we report perplexity measured by a separate, larger language model (GPT-Neo-1. 3B, Black et al., 2021). Prior works, however, have shown that low perplexity of generated text is not necessarily an indication of high quality but of degenerate behavior (Nadeem et al., 2020;Zhang et al., 2021) and have proposed closeness to the perplexity of human-written text as a better evaluation. Hence, we also report the difference of log perplexity between the generated text and human-written continuations (|∆ log PPL |). For diversity evaluation, we report Zipf's coefficient (Zipf) and average distinct n-grams in the output samples (Li et al., 2016, Dist-n). In addition, we also report the repetition rate (Welleck et al., 2020;Holtzman et al., 2020, Rep), measuring the proportion of output samples that end in repeating phrases. Finally, we report MAUVE (Pillutla et al., 2021) which evaluates both quality and diversity together by approximating information divergence between generated samples and human-written continuations (from the Open-WebText held-out set).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Results Table 1 summarizes our main results on the 50-token prompt and output setup. We report the numbers for the best performing three settings for logits projection and decoding steps T in SSD-LM. We report the best setting for the baselines. The results for other generation lengths have a similar trend and can be found in the appendix ( §D).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We find that SSD-LM, though being smaller in size, outperforms larger GPT-2 models on the unified metric MAUVE. On diversity, SSD-LM outperforms GPT-2 in Dist-n while achieving lower repetition rates. On perplexity, the results are slightly mixed. We observe a trade-off between MAUVE and perplexity for different settings we considered, indicating that further tuning of the hyperparameters may be required. However, one of our best performing settings (sampling top-p=0.9, T =2500) still achieves the closest perplexity to the gold con-  Table 2: Unconstrained generation results of SSD-LM and Diffusion-LM on ROCStories with 50 prompt tokens and 50 output tokens. We report the MAUVE score between the gold continuation and model generations.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We also show the perplexity (PPL) of model generations measured by GPT-Neo-1.3B. 10 tinuation.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
In §D, we show the influence of different logits projection strategies and the associated parameters on the output text quality in Figure 4. We also show qualitative examples of the generations by SSD-LM in Table 8 and a trajectory of intermediate states during the decoding process in Table 9. On 2,700 held-out ROCStories sequences, we use the first 50 tokens of each sequence as a prompt and have the model generate the next 50. In Table 2, we show the MAUVE score and perplexity of both models. We observe a substantially higher MAUVE score and lower perplexity with SSD-LM.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Setup To evaluate SSD-LM's ability for highlymodular control, we consider the task of sentiment controlled generation where given a prompt, the goal is to generate a continuation with a positive (or negative) polarity. We use a set of 15 short prompts as in Dathathri et al. (2020) , 2020). This model serves as f ϕ (·) as shown in Eq. 16. In addition to quality and diversity of the generated samples, we also evaluate them on control (that is measuring if the generated output is actually positive or negative in polarity). For this, we use an external sentiment classifier trained on a different dataset. Specifically, we use a classifier trained with Yelp reviews 13 (Zhang et al., 2015;Morris et al., 2020) following the evaluation setup in the baselines we consider. Again, we consider the sampling and multi-hot decoding strategies with top-p ∈ {0.2, 0.5, 0.9}, T decode ∈ {1000, 2500, 5000}, and the multiplier for control λ ∈ {0, 100, 500, 2000}. For the generation of 12/20/50 tokens, we use B decode =12/20/25 and apply the decoding algorithm for m=1/1/2 iterations respectively.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We show the quality of the controlled generations from three perspectives: target attribute via the external classifier accuracy, fluency via perplexity, and diversity via the distinctiveness measures. In Table 3 SSD-LM shows strong controllability while possessing great modularity. SSD-LM outperforms M&M LM, the other HMC method by a large margin. Even when comparing with the CC and CM methods, our method achieves a good balance in control, fluency, and diversity.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
In §D, we show the impact of the control weight λ and top-p on the attribute accuracy and perplexity in Figure 5. We also show qualitative examples of the controlled generations by SSD-LM in Table 8.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Diffusion models Diffusion models have demonstrated impressive performance in popular continuous-valued domains such as images (Ho et al., 2020), audio (Kong et al., 2021), video (Ho et al., 2022 and recently also been adopted for 3D-shapes, protein structures, and more (Zhou et al., 2021;Trippe et al., 2022;Wu et al., 2022). Since they are based on adding Gaussian noise, these approaches are not straightforward to apply to discrete valued domains like text. Hoogeboom et al. (2021); Austin et al. (2021) propose diffusing in the discrete space using categorical distributions which are modified using transition matrices. However, these methods do not straightforwardly support control and yield worse results than comparable autoregressive models. Li et al. (2022) propose to represent each token as a continuous embedding and apply diffusion in the embedding space. They train the LM to generate a fixed length sequence whereas SSD-LM allows flexibility in the generated sequence length by generating block-wise. Further, their LM is trained with specialized datasets and not evaluated against general-purpose autoregressive LMs on unconstrained text generation. Their method supports post-hoc control but requires training a customized attribute classifier, 15 since the diffusion operates on a learned embedding space. Gong et al. (2022), a concurrent work to ours, extend Li et al.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
(2022) to a sequence-to-sequence setup with a similar underlying embedding-based method. Our work is most closely related to Chen et al. (2022) which transform discrete data into a sequence of bits and represent each bit as +1 or -1 converting it into a continuous-valued domain. For textual data, however, it can lead to extremely long sequences which are difficult to optimize. In this work, we instead maintain a subword based vocabulary but represent each token as a sequence of manually defined logits.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Language models The majority of existing language models for text generation are trained autoregressively, i.e., they predict the next token given previously generated context. This paradigm scaled up both in terms of model size and training data size has resulted in impressive capabilities on many benchmarks (Brown et al., 2020;Chowdhery et al., 2022). However, they generate text one token at a time which does not provide flexible control over attributes of the generated text. Nonautoregressive models which generate the entire output sequence at the same time have also been explored in prior work other than diffusion models ( To address the quality decline with the non-autoregressive methods compared to autoregressive models, prior work have also explored semi-autoregressive approaches (Wang et al., 2018;Qi et al., 2021). In the same vein, our work seeks to address the drawbacks of autoregressive language models and non-autoregressive diffusion models 15 The control for diffusion models can also be classifierfree (Ho and Salimans, 2021) but requires training with the target attribute in advance, which is not a focus of this work. with a middle ground.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Controllable text generation Early solutions for controlling attributes of generated text focused on training or finetuning AR-LMs with specific control codes (Keskar et al., 2019;Gururangan et al., 2020;Chan et al., 2021). These methods are difficult to extend to new controls as it requires retraining the models. More recent work includes decoding approaches from pretrained AR-LMs without modifying the models, through altering the output probability distribution at each step using different control objectives (Dathathri et al., 2020;Krause et al., 2021;Yang and Klein, 2021;Liu et al., 2021;Lu et al., 2021;Pascual et al., 2021). However, these methods do not allow modifying a token once it is generated and are thus suboptimal for controls at the scope of the whole sequence. Closely related to SSD-LM are Kumar et al. (2021); Qin et al. (2022); Kumar et al. (2022b), which propose gradient-based decoding algorithms from AR-LMs. They require computing a backward pass through the LMs for each iteration, an expensive operation. In contrast, SSD-LM with its semi-autoregressive setup allows editing past tokens via diffusion. In addition, most of these approaches require training control functions from scratch whereas our model allows using off-the-shelf classifiers. Mireshghallah et al. (2022) propose a non-autoregressive LM based on Metropolis-Hastings sampling. It also supports off-the-shelf classifiers for control, and we therefore use it as a direct baseline for SSD-LM.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We present SSD-LM, a semi-autoregressive diffusion based language model trained to denoise corrupted simplexes over the output vocabulary. Compared to prior work in text-based diffusion, SSD-LM offers more flexibility in output length by generating blocks of text and an ability to use off-the-shelf attribute classifiers for control without additional tuning. On unconstrained text generation, SSD-LM performs on par with or outperforms strong and larger autoregressive baselines (GPT-2) in generation quality and diversity, while vastly outperforming diffusion baselines (Diffusion-LM).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
On controlled text generation, SSD-LM surpasses baselines while possessing an easy-to-use modular design. We believe that SSD-LM opens an exciting direction for future research in flexible and modular diffusion-based language generation.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Sample efficiency In AR-LMs, an NLL loss is computed at training time for every token in the sequence of length L (Eq. 4). However, in SSD-LM, each time a pretraining example is sampled, the loss is computed on only B tokens (Eq. 9) leading to a lower sample efficiency than AR-LM. Towards improving this efficiency, future work could explore model architectures dedicated to semiautoregressive diffusion rather than the vanilla Transformer encoder we use in this work.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Decoding speed Since each block is generated by refining over several iterations, SSD-LM has a considerably slower decoding speed than autoregressive models. For example, given a context of 50 tokens (single instance, unbatched), it takes SSD-LM 25 seconds to generate the next block of 25 tokens (T decode =1000). While our work focused on establishing the efficacy of diffusion-based LMs and modular controlled generation, future work could explore tuning T decode to balance model performance and decoding speed, or more efficient training and decoding algorithms extending ideas from prior work on diffusion models for continuous domains (Song et al., 2021;Nichol and Dhariwal, 2021;Rombach et al., 2022;Meng et al., 2022).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Decoding block size In this work, although we allow setups where B train ̸ = B decode , the decoding block size B decode remains the same across m decoding iterations, leaving space for a more flexible decoding schedule. Future work can also explore learning B decode (and B train ) rather than using constant pre-defined lengths.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Larger scale experiments with different kinds of controls and their combinations can be done, as well as more sophisticated ways to incorporate them (Kumar et al., 2021). In addition, we plan to explore alternative methods to continuously represent and add noise to discrete text (Bakosi and Ristorcelli, 2013). This work experiments with pretraining data that is primarily in English. Future work can also explore challenges and benefits of diffusion-based LMs in a multilingual setup.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Language models trained on data from the web can perpetuate social biases and toxic interactions, and can be prone to generating harmful language (Gehman et al., 2020;Wallace et al., 2019Wallace et al., , 2020Sheng et al., 2021;Weidinger et al., 2022). Further, language generation models could memorize and amplify patterns in data without deeper language understanding or control, so they can be factually inconsistent and generate disinformation (Maynez et al., 2020;Pagnoni et al., 2021;Zellers et al., 2019), or can compromise user privacy (Carlini et al., 2021). Prior works have outlined these risks (Sheng et al., 2021;Weidinger et al., 2021), discussed their points of origin, and advocated for future research on ethical development of LMs (Bender et al., 2021;Solaiman et al., 2019).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
While these studies have been conducted for autoregressive LMs, our diffusion-based LM is subject to these problems as well. However, since our method naturally incorporates controllability, future work may explore control functions that could potentially alleviate these issues (Liu et al., 2021;Kumar et al., 2022b). One risk is that controllability can also be misused maliciously, with models being intentionally exploited to generate biased, toxic, or non-factual content (Bagdasaryan and Shmatikov, 2022;Pagnoni et al., 2022). Therefore, apart from controlled generation, future work should aim to detect the generations under control as well to defend against the malicious use (Kumar et al., 2022a).
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein. Jiaming Song, Chenlin Meng, andStefano Ermon. 2021. andDhariwal, 2021), and α t −ᾱ t 1−ᾱ t > 0.98 for 98% of all t, with some outliers as t → 0 and t → T . step. The difference is that we sample a noise z instead of using the deterministic ϵ θ , to encourage exploration.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
We apply block concatenation on ROCStories similarly as OpenWebText, resulting in 50K training sequences of 100 tokens. We train Diffusion-LM with a default batch size of 64, learning rate of 1e-4, and 400K steps. We train SSD-LM with a batch size of 512, learning rate of 1e-4, and 20K steps. Both models use a tokenizer of BERT-baseuncased. For SSD-LM, additional hyperparameters like decoding block size and one-hot constant remain the same as the main SSD-LM benchmarked with GPT-2. For Diffusion-LM, the evaluation in the main paper is an infilling task. We use same decoding hyperparameters as Li et al. (2022). For SSD-LM, the evaluation is a block-wise generation problem with m=2 iterations. The result of SSD-LM in Table 2 is obtained with a decoding configuration of T decode =2500 and top-p=0.5. Our SSD-LM in this subsection is initialized with BERT. For a fair comparison, apart from the default Diffusion-LM reported in Table 2, we train another Diffusion-LM initialized with the encoder weights of BERT. However, this leads to degenerated results that are much worse than the default Diffusion-LM and our SSD-LM: a MAUVE score of 0.4 out of 100 and a PPL of 73157. This problem is not due to overfitting, as all checkpoints of the model show the same degenerated result. Since Li et al. (2022) did not explore this setup in their original work as well, we conjecture that Diffusion-LM may be incompatible with pretrained weights from existing non-diffusion models by nature, a disadvantage to our SSD-LM. Figure 4 shows the influence of different logits projection strategies and the associated parameters on the unconstrained generations' output text quality. We observe that reducing top-p → 0 (greedy projection) can lead to a low perplexity but it is undesirable due to a high repetition rate. We also find the multi-hot projection strategy is overall worse performing than the sampling projection strategy in our setup, indicating it is better to commit the intermediate states to single rather than multiple tokens. This can be because our logits mapping involves putting probability mass on singular tokens. The multi-hot projection may still be a viable strategy if future work uses multi-hot logits mapping for the input tokens.  Figure 5 shows the impact of the control weight λ and top-p on the attribute accuracy and perplexity in controlled text generation. As expected, a larger control weight leads to a better external classifier accuracy. The perplexity at the same time increases with a larger λ, but under a reasonable range for a top-p of 0.2 and 0.5. Figure 6 shows the pretraining loss trajectory. Table 4, Table 5, Table 6, and Table 7 show additional evaluation results of SSD-LM generations. Table 8 and Table 9 show qualitative examples of SSD-LM generations.       Table 8: Qualitative examples of SSD-LM's generations. Top half : unconstrained text generation ( §4.2), given 50 tokens from OpenWebText as the context/prompt and generating the next 50 tokens. We show two prompts and two sample generations for each prompt. Bottom half : controlled text generation ( §4.3), given prompts from Dathathri et al. (2020) and generating the next 20 tokens. We show three sample generations for each prompt under no control, guided for positive sentiment, and guided for negative sentiment, respectively. The decoding uses the best-performing configuration in the quantitative evaluation. is all about Saturn. The Eini will, the closest that the instruments have reached will be to stop in on the Saturn omial allcounter Saturn. The Directthank Ecuador two thelearning that the Animation have brothers will make toousands downtown governance the Further 700 will allow the Cass to finally see the planet's relatively small atmosphere and finally be able to procure an accurate way of understanding how willPocket prelim Klux to finally see the planet intelligent relatively jumper atmosphere and halted Fly activityvirt00000 trem accurate way of Inferno what 600 will allow the scientists to better study the effects of Grand Impact, and also be able to get much more data and images of will allowert scientists Damien better study the effects of Grand Impact, andasket bebery to get much more data and images of 500 will allow the scientists to better see the interior of its atmosphere, and also be able to get much more knowledge and understanding of will allow the scientists to better see the interior of its atmosphere, and also be able to get much more knowledge and understanding of 1 will allow the scientists to better see the interior of its atmosphere, and also be able to get much more knowledge and observations of will allow the scientists to better see the interior of its atmosphere, and also be able to get much more knowledge and observations of Table 9: The intermediate states of generation as t decreases (T =2500, B=25, top-p-sampling=0.99). The context w <c here is the first example prompt in Table 8: " called the Grand Finale, where it will end its long life by plunging into Saturn's atmosphere this September. Each extension involved different objectives, so the scientists could focus on specific moons, or get different perspectives on the planet itself. This last phase". There is no change in the outputs during 500 > t > 1. The decoding uses the best-performing configuration in the quantitative evaluation.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
A For every submission:
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
A1. Did you describe the limitations of your work?
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
On Page 9, the first unnumbered section, "Limitation" A2. Did you discuss any potential risks of your work?
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
On Page 9, the second unnumbered section, "Ethics statement"
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
A3. Do the abstract and introduction summarize the paper's main claims?
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Section 1 A4. Have you used AI writing assistants when working on this paper?
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Left blank. B2. Did you discuss the license or terms for use and / or distribution of any artifacts? No response.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? No response.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
B4. Did you discuss the steps taken to check whether the data that was collected / used contains any information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? No response.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and linguistic phenomena, demographic groups represented, etc.? No response.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. No response.
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Section 4 C1. Did you report the number of parameters in the models used, the total computational budget (e.g., GPU hours), and computing infrastructure used? Section 4
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.
Diffusion-LM Improves Controllable Text Generation
Large autoregressive language models (LMs) are capable of generating high quality text [33,3,5,45], but in order to reliably deploy these LMs in real world applications, the text generation process needs to be controllable: we need to generate text that satisfies desired requirements (e.g. topic, syntactic structure). A natural approach for controlling a LM would be to fine-tune the LM using supervised data of the form (control, text) [16]. However, updating the LM parameters for each control task can be expensive and does not allow for compositions of multiple controls (e.g. generate text that is both positive sentiment and non-toxic). This motivates light-weight and modular plug-and-play approaches [6] that keep the LM frozen and steer the generation process using an external classifier that measures how well the generated text satisfies the control. But steering a frozen autoregressive LM has been shown to be difficult, and existing successes have been limited to simple, attribute-level controls (e.g., sentiment or topic) [6,22,44].
Diffusion-LM Improves Controllable Text Generation
In order to tackle more complex controls, we propose Diffusion-LM, a new language model based on continuous diffusions. Diffusion-LM starts with a sequence of Gaussian noise vectors and incrementally denoises them into vectors corresponding to words, as shown in Figure 1. These gradual denoising steps produce a hierarchy of continuous latent representations. We find that this hierarchical and continuous latent variable enables simple, gradient-based methods to perform complex control tasks such as constraining the parse tree of a generated sequence.
Diffusion-LM Improves Controllable Text Generation
Continuous diffusion models have been extremely successful in vision and audio domains [12,21,34,8,4], but they have not been applied to text because of the inherently discrete nature of text Figure 1: Diffusion-LM iteratively denoises a sequence of Gaussian vectors into word vectors, yielding a intermediate latent variables of decreasing noise level x T · · · x 0 . For controllable generation, we iteratively perform gradient updates on these continuous latents to optimize for fluency (parametrized by Diffusion-LM) and satisfy control requirements (parametrized by a classifier).
Diffusion-LM Improves Controllable Text Generation
( §3). Adapting this class of models to text requires several modifications to standard diffusions: we add an embedding step and a rounding step to the standard diffusion process, design a training objective to learn the embedding, and propose techniques to improve rounding ( §4). We control Diffusion-LM using a gradient-based method, as shown in Figure 1. This method enables us to steer the text generation process towards outputs that satisfy target structural and semantic controls. It iteratively performs gradient updates on the continuous latent variables of Diffusion-LM to balance fluency and control satisfaction ( §5.1).
Diffusion-LM Improves Controllable Text Generation
To demonstrate control of Diffusion-LM, we consider six control targets ranging from fine-grained attributes (e.g., semantic content) to complex structures (e.g., parse trees). Our method almost doubles the success rate of previous plug-and-play methods and matches or outperforms the fine-tuning oracle on all these classifier-guided control tasks ( §7.1). In addition to these individual control tasks, we show that we can successfully compose multiple classifier-guided controls to generate sentences with both desired semantic content and syntactic structure ( §7.2). Finally, we consider span-anchored controls, such as length control and infilling. Diffusion-LM allows us to perform these control tasks without a classifier, and our Diffusion-LM significantly outperforms prior plug-and-play methods and is on-par with an autoregressive LM trained from scratch for the infilling task ( §7.3).
Diffusion-LM Improves Controllable Text Generation
We first define controllable generation ( §3.1) and then review continuous diffusion models ( §3.3).
Diffusion-LM Improves Controllable Text Generation
Text generation is the task of sampling w from a trained language model p lm (w), where w = [w 1 · · · w n ] is a sequence of discrete words and p lm (w) is a probability distribution over sequences of words. Controllable text generation is the task of sampling w from a conditional distribution p(w | c), where c denotes a control variable. For syntactic control, c can be a target syntax tree (Figure 1), while for sentiment control, c could be a desired sentiment label. The goal of controllable generation is to generate w that satisfies the control target c.
Diffusion-LM Improves Controllable Text Generation
Consider the plug-and-play controllable generation setting: we are given a language model p lm (w) trained from a large amount of unlabeled text data, and for each control task, we are given a classifier p(c | w) trained from smaller amount of labeled text data (e.g., for syntactic control, the classifier is a probabilistic parser). The goal is to utilize these two models to approximately sample from the posterior p(w | c) via Bayes rule p(w | c) ∝ p lm (w) · p(c | w). Here, p lm (w) encourages w to be fluent, and the p(c | w) encourages w to fulfill the control.
Diffusion-LM Improves Controllable Text Generation
The canonical approach to language modeling factors p lm in an autoregressive left-to-right mannar, p lm (w) = p lm (w 1 ) n i=2 p lm (x i | x <i ). In this case, text generation is reduced to the task of repeatedly predicting the next token conditioned on the partial sequence generated so far. The next token prediction p lm (x i | x <i ) is often parametrized by Transformer architecture [42].
Diffusion-LM Improves Controllable Text Generation
A diffusion model [12,27] is a latent variable model that models the data x 0 ∈ R d as a Markov chain x T . . . x 0 with each variable in R d , and x T is a Gaussian. The diffusion model incrementally denoises the sequence of latent variables x T :1 to approximate samples from the target data distribution ( Figure 2). The initial state p θ (x T ) ≈ N (0, I), and each denoising transition
Diffusion-LM Improves Controllable Text Generation
. For example, µ θ and Σ θ may be computed by a U-Net or a Tranformer.
Diffusion-LM Improves Controllable Text Generation
To train the diffusion model, we define a forward process that constructs the intermediate latent variables x 1:T . The forward process incrementally adds Gaussian noise to data x 0 until, at diffusion step T , samples x T are approximately Gaussian. Each transition
Diffusion-LM Improves Controllable Text Generation
where the hyperparameter β t is the amount of noise added at diffusion step t. This parametrization of the forward process q contains no trainable parameters and Figure 2: A graphical model representing the forward and reverse diffusion processes. In addition to the original diffusion models [12], we add a Markov transition between x 0 and w, and propose the embedding §4.1 and rounding §4.2 techniques. allows us to define a training objective that involves generating noisy data according to a pre-defined forward process q and training a model to reverse the process and reconstruct the data.
Diffusion-LM Improves Controllable Text Generation
The diffusion model is trained to maximize the marginal likelihood of the data E x0∼pdata [log p θ (x 0 )], and the canonical objective is the variational lower bound of log p θ (x 0 ) [39],
Diffusion-LM Improves Controllable Text Generation
However, this objective can be unstable and require many optimization tricks to stabilize [27]. To circumvent this issue, Ho et al. [12] devised a simple surrogate objective that expands and reweights each KL-divergence term in L vlb to obtain a mean-squared error loss (derivation in Appendix E) which we will refer to as
Diffusion-LM Improves Controllable Text Generation
whereμ(x t , x 0 ) is the mean of the posterior q(x t−1 |x 0 , x t ) which is a closed from Gaussian, and µ θ (x t , t) is the predicted mean of p θ (x t−1 | x t ) computed by a neural network. While L simple is no longer a valid lower bound, prior work has found that it empirically made training more stable and improved sample quality 2 . We will make use of similar simplifications in Diffusion-LM to stabilize training and improve sample quality ( §4.1).
Diffusion-LM Improves Controllable Text Generation
Constructing Diffusion-LM requires several modifications to the standard diffusion model. First, we must define an embedding function that maps discrete text into a continuous space. To address this, we propose an end-to-end training objective for learning embeddings ( §4.1). Second, we require a rounding method to map vectors in embedding space back to words. To address this, we propose training and decoding time methods to facilitate rounding ( §4.2).
Diffusion-LM Improves Controllable Text Generation
To apply a continuous diffusion model to discrete text, we define an embedding function EMB(w i ) that maps each word to a vector in R d . We define the embedding of a sequence w of length n to be:
Diffusion-LM Improves Controllable Text Generation
We propose a modification of the diffusion model training objective (Equation 1) that jointly learns the diffusion model's parameters and word embeddings. In preliminary experiments, we explored random Gaussian embeddings, as well as pre-trained word embeddings [29,33]. We found that these fixed embeddings are suboptimal for Diffusion-LM compared to end-to-end training 3 .
Diffusion-LM Improves Controllable Text Generation
As shown in Figure 2, our approach adds a Markov transition from discrete words w to x 0 in the forward process, parametrized by q φ (x 0 |w) = N (EMB(w), σ 0 I). In the reverse process, we add a trainable rounding step, parametrized by
Diffusion-LM Improves Controllable Text Generation
The training objectives introduced in §3 now becomes
Diffusion-LM Improves Controllable Text Generation
Learned Embeddings Figure 3: A t-SNE [41] plot of the learned word embeddings. Each word is colored by its POS.
Diffusion-LM Improves Controllable Text Generation
We derive L e2e simple (w) from L e2e vlb (w) following the simplification in §3.3 and our derivation details are in Appendix E. Since we are training the embedding function, q φ now contains trainable parameters and we use the reparametrization trick [36,18] to backpropagate through this sampling step. Empirically, we find the learned embeddings cluster meaningfully: words with the same part-of-speech tags (syntactic role) tend to be clustered, as shown in Figure 3.
Diffusion-LM Improves Controllable Text Generation
The learned embeddings define a mapping from discrete text to the continuous x 0 . We now describe the inverse process of rounding a predicted x 0 back to discrete text. Rounding is achieved by choosing the most probable word for each position, according to argmax
Diffusion-LM Improves Controllable Text Generation
Ideally, this argmax-rounding would be sufficient to map back to discrete text, as the denoising steps should ensure that x 0 lies exactly on the embedding of some word. However, empirically, the model fails to generate x 0 that commits to a single word.
Diffusion-LM Improves Controllable Text Generation
One explanation for this phenomenon is that the L simple (x 0 ) term in our objective 2 puts insufficient emphasis on modeling the structure of x 0 . Recall that we defined L simple (x 0 ) = T t=1 E xt ||µ θ (x t , t) −μ(x t , x 0 )|| 2 , where our model µ θ (x t , t) directly predicts the mean of p θ (x t−1 | x t ) for each denoising step t. In this objective, the constraint that x 0 has to commit to a single word embedding will only appear in the terms with t near 0, and we found that this parametrization required careful tuning to force the objective to emphasize those terms (see Appendix H).
Diffusion-LM Improves Controllable Text Generation
Our approach re-parametrizes L simple to force Diffusion-LM to explicitly model x 0 in every term of the objective. Specifically, we derive an analogue to L simple which is parametrized via
Diffusion-LM Improves Controllable Text Generation
where our model f θ (x t , t) predicts x 0 directly 4 . This forces the neural network to predict x 0 in every term and we found that models trained with this objective quickly learn that x 0 should precisely centered at a word embedding.
Diffusion-LM Improves Controllable Text Generation
We described how re-parametrization can be helpful for model training, but we also found that the same intuition could be used at decoding time in a technique that we call the clamping trick. In the standard generation approach for a x 0 -parametrized model, the model denoises x t to x t−1 by first computing an estimate of x 0 via f θ (x t , t) and then sampling x t−1 conditioned on this estimate:
Diffusion-LM Improves Controllable Text Generation
(1−β s ) and ∼ N (0, I) 5 . In the clamping trick, the model additionally maps the predicted vector f θ (x t , t) to its nearest word embedding sequence. Now, the sampling step becomes
Diffusion-LM Improves Controllable Text Generation
The clamping trick forces the predicted vector to commit to a word for intermediate diffusion steps, making the vector predictions more precise and reducing rounding errors. 6 
Diffusion-LM Improves Controllable Text Generation
Having described the Diffusion-LM, we now consider the problem of controllable text generation ( §5.1) and decoding ( §5.2).
Diffusion-LM Improves Controllable Text Generation
We now describe a procedure that enables plug-and-play control on Diffusion-LM. Our approach to control is inspired by the Bayesian formulation in §3.1, but instead of performing control directly on the discrete text, we perform control on the sequence of continuous latent variables x 0:T defined by Diffusion-LM, and apply the rounding step to convert these latents into text.
Diffusion-LM Improves Controllable Text Generation
and we decompose this joint inference problem to a sequence of control problems at each diffusion step:
Diffusion-LM Improves Controllable Text Generation
via conditional independence assumptions from prior work on controlling diffusions [40]. Consequently, for the t-th step, we run gradient update on x t−1 :
Diffusion-LM Improves Controllable Text Generation
where both log p(x t−1 | x t ) and log p(c | x t−1 ) are differentiable: the first term is parametrized by Diffusion-LM, and the second term is parametrized by a neural network classifier.
Diffusion-LM Improves Controllable Text Generation
Similar to work in the image setting [8,40], we train the classifier on the diffusion latent variables and run gradient updates on the latent space x t−1 to steer it towards fulfilling the control. These image diffusion works take one gradient step towards ∇ xt−1 log p(c | x t−1 ) per diffusion steps. To improve performance on text and speed up decoding, we introduce two key modifications: fluency regularization and multiple gradient steps.
Diffusion-LM Improves Controllable Text Generation
To generate fluent text, we run gradient updates on a control objective with fluency regularization:
Diffusion-LM Improves Controllable Text Generation
, where λ is a hyperparameter that trades off fluency (the first term) and control (the second term). While existing controllable generation methods for diffusions do not include the λ log p(x t−1 | x t ) term in the objective, we found this term to be instrumental for generating fluent text. The resulting controllable generation process can be viewed as a stochastic decoding method that balances maximizing and sampling p(x t−1 | x t , c), much like popular text generation techniques such as nucleus sampling [13] or sampling with low temperature. In order to improve the control quality, we take multiple gradient steps for each diffusion step: we run 3 steps of the Adagrad 7 [10] update for each diffusion steps. To mitigate for the increased computation cost, we downsample the diffusion steps from 2000 to 200, which speeds up our controllable generation algorithm without hurting sample quality much.
Diffusion-LM Improves Controllable Text Generation
Many conditional text generation tasks require a single high-quality output sequence, such as machine translation or sentence infilling. In these settings, we apply Minimum Bayes Risk (MBR) decoding [23] to aggregate a set of samples S drawn from the Diffusion-LM , and select the sample that achieves the minimum expected risk under a loss function L (e.g., negative BLEU score):ŵ = argmin w∈S w ∈S 1 |S| L(w, w ). We found that MBR decoding often returned high quality outputs, since a low quality sample would be dissimilar from the remaining samples and penalized by the loss function.
Diffusion-LM Improves Controllable Text Generation
With the above improvements on training ( §4) and decoding ( §5), we train Diffusion-LM for two language modeling tasks. We then apply the controllable generation method to 5 classifier-guided control tasks, and apply MBR decoding to a classifier-free control task (i.e. infilling).
Diffusion-LM Improves Controllable Text Generation
Browns Cambridge is good for Japanese food and also children friendly near The Sorrento .
Diffusion-LM Improves Controllable Text Generation
Zizzi is a local coffee shop located on the outskirts of the city .
Diffusion-LM Improves Controllable Text Generation
input (Syntax Tree) (TOP (S (NP (*) (*) (*)) (VP (*) (NP (NP (*) (*)))))) output text
Diffusion-LM Improves Controllable Text Generation
The Twenty Two has great food input (Syntax Spans) (7, 10, VP) output text Wildwood pub serves multicultural dishes and is ranked 3 stars input (Length) 14 output text Browns Cambridge offers Japanese food located near The Sorrento in the city centre .
Diffusion-LM Improves Controllable Text Generation
input (left context) My dog loved tennis balls. input (right context) My dog had stolen every one and put it under there. output text One day, I found all of my lost tennis balls underneath the bed. Table 1: Example input control and output text for each control tasks.
Diffusion-LM Improves Controllable Text Generation
We train Diffusion-LM on two datasets: E2E [28] and ROCStories [26]. The E2E dataset consists of 50K restaurant reviews labeled by 8 fields including food type, price, and customer rating. The ROCStories dataset consists of 98K five-sentence stories, capturing a rich set of causal and temporal commonsense relations between daily events. This dataset is more challenging to model than E2E, because the stories contain a larger vocabulary of 11K words and more diverse semantic content.
Diffusion-LM Improves Controllable Text Generation
Our Diffusion-LM is based on Transformer [42] architecture with 80M parameters, with a sequence length n = 64, diffusion steps T = 2000 and a square-root noise schedule (see Appendix A for details). We treat the embedding dimension as a hyperparameter, setting d = 16 for E2E and d = 128 for ROCStories. See Appendix B for hyperparameter details. At decoding time, we downsample to 200 diffusion steps for E2E and maintain 2000 steps for ROCStories. Decoding Diffusion-LM for 200 steps is still 7x slower than decoding autoregressive LMs. For controllable generation, our method based on Diffusion-LM is 1.5x slower than FUDGE but 60x faster than PPLM.
Diffusion-LM Improves Controllable Text Generation
We consider 6 control tasks shown in Table 1: the first 4 tasks rely on a classifier, and the last 2 tasks are classifier free 8 . For each control task (e.g. semantic content), we sample 200 control targets c (e.g., rating=5 star) from the validation splits, and we generate 50 samples for each control target. To evaluate the fluency of the generated text, we follow the prior works [44,6] and feed the generated text to a teacher LM (i.e., a carefully fine-tuned GPT-2 model) and report the perplexity of generated text under the teacher LM. We call this metric lm-score (denoted as lm): a lower lm-score indicates better sample quality. 9 We define success metrics for each control task as follows:
Diffusion-LM Improves Controllable Text Generation
Semantic Content. Given a field (e.g., rating) and value (e.g., 5 star), generate a sentence that covers field=value, and report the success rate by exact match of 'value'.
Diffusion-LM Improves Controllable Text Generation
Parts-of-speech. Given a sequence of parts-of-speech (POS) tags (e.g., Pronoun Verb Determiner Noun), generate a sequence of words of the same length whose POS tags (under an oracle POS tagger) match the target (e.g., I ate an apple). We quantify success via word-level exact match.
Diffusion-LM Improves Controllable Text Generation
Syntax Tree. Given a target syntactic parse tree (see Figure 1), generate text whose syntactic parse matches the given parse. To evaluate the success, we parse the generated text by an off-the-shelf parser [20], and report F1 scores.
Diffusion-LM Improves Controllable Text Generation
Syntax Spans. Given a target (span, syntactic category) pair, generate text whose parse tree over span [i, j] matches the target syntactic category (e.g. prepositional phrase).We quantify success via the fraction of spans that match exactly.
Diffusion-LM Improves Controllable Text Generation
Length. Given a target length 10, . . . , 40, generate a sequence with a length within ±2 of the target.
Diffusion-LM Improves Controllable Text Generation
In the case of Diffusion-LM, we treat this as a classifier-free control task.
Diffusion-LM Improves Controllable Text Generation
Infilling. Given a left context (O 1 ) and a right context (O 2 ) from the aNLG dataset [2], and the goal is to generate a sentence that logically connects O 1 and O 2 . For evaluation, we report both automatic and human evaluation from the Genie leaderboard [17].
Diffusion-LM Improves Controllable Text Generation
For the first 5 control tasks, we compare our method with PPLM, FUDGE, and a fine-tuning oracle. Both PPLM and FUDGE are plug-and-play controllable generation approaches based on an autoregressive LM, which we train from scratch using the GPT-2 small architecture [33].
Diffusion-LM Improves Controllable Text Generation
PPLM [6]. This method runs gradient ascent on the LM activations to increase the classifier probabilities and language model probabilities, and has been successful on simple attribute control. We apply PPLM to control semantic content, but not the remaining 4 tasks which require positional information, as PPLM's classifier lacks positional information.
Diffusion-LM Improves Controllable Text Generation
FUDGE [44]. For each control task, FUDGE requires a future discriminator that takes in a prefix sequence and predicts whether the complete sequence would satisfy the constraint. At decoding time, FUDGE reweights the LM prediction by the discriminator scores.
Diffusion-LM Improves Controllable Text Generation
FT. For each control task, we fine-tune GPT-2 on (control, text) pair, yielding an oracle conditional language model that's not plug-and-play. We report both the sampling (with temperature 1.0) and beam search (with beam size 4) outputs of the fine-tuned models, denoted as FT-sample and FT-search.
Diffusion-LM Improves Controllable Text Generation
We compare to 3 specialized baseline methods developed in past work for the infilling task.
Diffusion-LM Improves Controllable Text Generation
DELOREAN [30]. This method continuously relaxes the output space of a left-to-right autoregressive LM, and iteratively performs gradient updates on the continuous space to enforce fluent connection to the right contexts. This yields a continuous vector which is rounded back to text.
Diffusion-LM Improves Controllable Text Generation
COLD [31]. COLD specifies an energy-based model that includes fluency (from left-to-right and right-to-left LM) and coherence constraints (from lexical overlap). It samples continuous vectors from this energy-based model and round them to text.
Diffusion-LM Improves Controllable Text Generation
AR-infilling. We train an autoregressive LM from scratch to do sentence infilling task [9]. Similar to training Diffusion-LM, we train on the ROCStories dataset, but pre-process it by reordering sentences
Diffusion-LM Improves Controllable Text Generation
At evaluation time, we feed in O 1 , O 2 , and the model generates the middle sentence.
Diffusion-LM Improves Controllable Text Generation
We train Diffusion-LMs on the E2E and ROCStories datasets. In terms of negative log-likelihood (NLL, lower is better), we find that the variational upper bound of Diffusion-LM NLL 10 underperforms the equivalent autoregressive Transformer model (2.28 vs. 1.77 for E2E, 3.88 vs 3.05 for ROCStories) although scaling up model and dataset size partially bridges the gap (3.88 − → 3.10 on ROCStories). Our best log-likelihoods required several modifications from §4; we explain these and give detailed log-likelihood results in Appendix F. Despite worse likelihoods, controllable generation based on our Diffusion-LM results in significantly better outputs than systems based on autoregressive LMs, as we will show in §7.1, §7.2, and §7.3
Diffusion-LM Improves Controllable Text Generation
As shown in Table 2, Diffusion-LM achieves high success and fluency across all classifier-guided control tasks. It significantly outperforms the PPLM and FUDGE baselines across all 5 tasks. Surprisingly, our method outperforms the fine-tuning oracle on controlling syntactic parse trees and spans, while achieving similar performance on the remaining 3 tasks.  Table 2: Diffusion-LM achieves high success rate (ctrl ↑) and good fluency (lm ↓) across all 5 control tasks, outperforming the PPLM and FUDGE baselines. Our method even outperforms the fine-tuning oracle (FT) on controlling syntactic parse trees and spans.
Diffusion-LM Improves Controllable Text Generation
Syntactic Parse ( S ( S ( NP * ) ( VP * ( NP ( NP * * ) ( VP * ( NP ( ADJP * * ) * ) ) ) ) ) * ( S ( NP * * * ) ( VP * ( ADJP ( ADJP * ) ) ) ) ) FUDGE Zizzi is a cheap restaurant .
Diffusion-LM Improves Controllable Text Generation
[incomplete] Diffusion-LM Zizzi is a pub providing family friendly Indian food Its customer rating is low FT Cocum is a Pub serving moderately priced meals and the customer rating is high Syntactic Parse ( S ( S ( VP * ( PP * ( NP * * ) ) ) ) * ( NP * * * ) ( VP * ( NP ( NP * * ) ( SBAR ( WHNP * ) ( S ( VP * ( NP * * ) ) ) ) ) ) * )
Diffusion-LM Improves Controllable Text Generation
In the city near The Portland Arms is a coffee and fast food place named The Cricketers which is not family -friendly with a customer rating of 5 out of 5 .
Diffusion-LM Improves Controllable Text Generation
Located on the riverside , The Rice Boat is a restaurant that serves Indian food .
Diffusion-LM Improves Controllable Text Generation
Located near The Sorrento, The Mill is a pub that serves Indian cuisine. Controlling syntactic parse trees and spans are challenging tasks for fine-tuning, because conditioning on the parse tree requires reasoning about the nested structure of the parse tree, and conditioning on spans requires lookahead planning to ensure the right constituent appears at the target position.
Diffusion-LM Improves Controllable Text Generation
We observe that PPLM fails in semantic content controls and conjecture that this is because PPLM is designed to control coarse-grained attributes, and may not be useful for more targeted tasks such as enforcing that a restaurant review contains a reference to Starbucks.
Diffusion-LM Improves Controllable Text Generation
FUDGE performs well on semantic content control but does not perform well on the remaining four tasks. Controlling a structured output (Parts-of-speech and Syntax Tree) is hard for FUDGE because making one mistake anywhere in the prefix makes the discriminator assign low probabilities to all continuations. In other control tasks that require planning (Length and Syntax Spans), the future discriminator is difficult to train, as it must implicitly perform lookahead planning.
Diffusion-LM Improves Controllable Text Generation
The non-autoregressive nature of our Diffusion-LM allows it to easily solve all the tasks that require precise future planning (Syntax Spans and Length). We believe that it works well for complex controls that involve global structures (Parts-of-speech, Syntax Tree) because the coarse-to-fine representations allow the classifier to exert control on the entire sequence (near t = T ) as well as on individual tokens (near t = 0).
Diffusion-LM Improves Controllable Text Generation
Qualitative Results. Table 3 shows samples of Syntax Tree control. Our method and fine-tuning both provide fluent sentences that mostly satisfy controls, whereas FUDGE deviates from the constraints after the first few words. One key difference between our method and fine-tuning is that Diffusion-LM is able to correct for a failed span and have suffix spans match the target. In the first example, the generated span ("Family friendly Indian food") is wrong because it contains 1 more word than the target. Fortunately, this error doesn't propagate to later spans, since Diffusion-LM adjusts by dropping the conjunction. Analogously, in the second example, the FT model generates a failed span ("The Mill") that contains 1 fewer word. However, the FT model fails to adjust in the suffix, leading to many misaligned errors in the suffix.   Table 5: For sentence infilling, Diffusion-LM significantly outperforms prior work COLD [31] and Delorean [30] (numbers taken from paper), and matches the performance of an autoregressive LM (AR) trained from scratch to do infilling.
Diffusion-LM Improves Controllable Text Generation
One unique capability of plug-and-play controllable generation is its modularity. Given classifiers for multiple independent tasks, gradient guided control makes it simple to generate from the intersection of multiple controls by taking gradients on the sum of the classifier log-probabilities.
Diffusion-LM Improves Controllable Text Generation
We evaluate this setting on the combination of Semantic Content + Syntax Tree control and Semantic Content + Parts-of-speech control. As shown in Table 4, our Diffusion-LM achieves a high success rate for both of the two components, whereas FUDGE gives up on the more global syntactic control. This is expected because FUDGE fails to control syntax on its own.
Diffusion-LM Improves Controllable Text Generation
Fine-tuned models are good at POS and semantic content control individually but do not compose these two controls well by product of experts (PoE), leading to a large drop in success rates for both constraints.
Diffusion-LM Improves Controllable Text Generation
As shown in Table 5, our diffusion LM significantly outperforms continuous relaxation based methods for infilling (COLD and DELOREAN). Moreover, our method achieves comparable performance to fine-tuning a specialized model for this task. Our method has slightly better automatic evaluation scores and the human evaluation found no statistically significant improvement for either method. These results suggest that Diffusion LM can solve many types of controllable generation tasks that depend on generation order or lexical constraints (such as infilling) without specialized training.
Diffusion-LM Improves Controllable Text Generation
We verify the importance of our proposed design choices in §4 through two ablation studies. We measure the sample quality of Diffusion-LM using the lm-score on 500 samples §6.2.
Diffusion-LM Improves Controllable Text Generation
Learned v.s. Random Embeddings ( §4.1). Learned embeddings outperform random embeddings on the ROCStories, which is a harder language modeling task. The same trend holds for the E2E dataset but with a smaller margin.
Diffusion-LM Improves Controllable Text Generation
Objective Parametrization ( §4.2). We propose to let the diffusion model predict x 0 directly. Here, we compare this with standard parametrization in image generation which parametrizes by the noise term . Figure 4 (right) shows that parametrizing by x 0 consistently attains good performance across dimensions, whereas parametrizing by works fine for small dimensions, but quickly collapses for larger dimensions.  Figure 4: We measure the impact of our proposed design choices through lm-score. We find both learned embeddings and reparametrization substantially improves sample quality.
Diffusion-LM Improves Controllable Text Generation
We proposed Diffusion-LM, a novel and controllable language model based on continuous diffusions, which enables new forms of complex fine-grained control tasks. We demonstrate Diffusion-LM's success in 6 fine-grained control tasks: our method almost doubles the control success rate of prior methods and is competitive with baseline fine-tuning methods that require additional training.
Diffusion-LM Improves Controllable Text Generation
We find the complex controls enabled by Diffusion-LM to be compelling, and we are excited by how Diffusion-LM is a substantial departure from the current paradigm of discrete autoregressive generation. As with any new technologies, there are drawbacks to the Diffusion-LMs that we constructed: (1) it has higher perplexity; (2) decoding is substantially slower; and (3) training converges more slowly. We believe that with more follow-up work and optimization, many of these issues can be addressed, and this approach will turn out to be a compelling way to do controllable generation at scale. Because a diffusion model shares parameters for all diffusion steps, the noise schedule (parametrized byᾱ 1:T ) is an important hyperparameter that determines how much weight we assign to each denoising problem. We find that standard noise schedules for continuous diffusions are not robust for text data. We hypothesize that the discrete nature of text and the rounding step make the model insensitive to noise near t = 0. Concretely, adding small amount of Gaussian noise to a word embedding is unlikely to change its nearest neighbor in the embedding space, making denoising an easy task near t = 0.
Diffusion-LM Improves Controllable Text Generation
To address this, we introduce a new sqrt noise schedule that is better suited for text, shown in Figure 5 defined byᾱ t = 1 − t/T + s, where s is a small constant that corresponds to the starting noise level 11 . Compared to standard linear and cosine schedules, our sqrt schedule starts with a higher noise level and increase noise rapidly for the first 50 steps. Then sqrt slows down injecting noise to avoid spending much steps in the high-noise problems, which may be too difficult to solve well.
Diffusion-LM Improves Controllable Text Generation
Diffusion-LM hyperparameters. The hyperparameters that are specific to Diffusion-LM include the number of diffusion steps, the architecture of the Diffusion-LM, the embedding dimension, and the noise schedule, . We set the diffusion steps to be 2000, the architecture to be BERT-base [7], and the sequence length to be 64. To stablize the training under L e2e vlb objective, we find that we need to set gradient clipping to 1.0 and apply importance sampling to reweight each term in L vlb [27]. Both tricks are not necessary for L e2e simple objective.
Diffusion-LM Improves Controllable Text Generation
Controllable Generation hyperparameters. To achieve controllable generation, we run gradient update on the continuous latents of Diffusion-LM. We use the AdaGrad optimizer [10] 
Diffusion-LM Improves Controllable Text Generation
Semantic Content. We train an autoregressive LM (GPT-2 small architecture) to predict the (field, value) pair conditioned on text. To parametrize log p(c | x t ), we compute the logprob of "value" per token.
Diffusion-LM Improves Controllable Text Generation
Parts-of-speech. The classifier is parametrized by a parts-of-speech tagger, which estimates the probability of the target POS sequence conditioned on the latent variables. This tagger uses a BERTbase architecture: the input is the concantenated word embedding, and output a softmax distribution over all POS tags for each input word. log p(c | x t ) is the sum of POS log-probs for each word in the sequence.
Diffusion-LM Improves Controllable Text Generation
Syntax Tree. We train a Transformer-based constituency parser [20]. Our parser makes locally normalized prediction for each span, predicting either "not a constituent", or a label for the constituent (e.g., Noun Phrase). log p(c | x t ) is the sum of log-probs for each labeled and non-constituency span in the sequence.
Diffusion-LM Improves Controllable Text Generation
Syntax Span. We use the same parser trained for the syntax tree. log p(c | x t ) is the log-probability that the target span is annotated with the target label.
Diffusion-LM Improves Controllable Text Generation
For continuous diffusion models ( §3.3), L simple is derived from the canonical objective L vlb by reweighting each term. The first T terms in L vlb are all KL divergence between two Gaussian distributions, which has a closed form solution. Take the t-th term for example:
Diffusion-LM Improves Controllable Text Generation
where C is a constant,μ is the mean of the posterior q(x t−1 |x 0 , x t ), and µ θ is the mean of p θ (x t−1 | x t ) predicted by the diffusion model. Intuitively, this simplification matches the predicted mean of x t−1 to its true posterior mean. The simplification involves removing the constant C and the scaling factor 1 2σ 2 t , yielding one term in L simple : Eq(x 1:T |x0) ||µ θ (x t , t) −μ(x t , x 0 )|| 2 .
Diffusion-LM Improves Controllable Text Generation
To apply continuous diffusion to model discrete text, we design Diffusion-LM ( §4.1) and propose a new end-to-end training objective (equation (2)) that learns the diffusion model and the embedding parameters jointly. The L e2e vlb can be written out as
Diffusion-LM Improves Controllable Text Generation
We apply the same simplification which transforms L vlb → L simple to transform L e2e vlb → L e2e simple :
Diffusion-LM Improves Controllable Text Generation
It's worth noting that the first term is constant if the noise schedule satisfiesᾱ T = 0, which guarantees x T is pure Gaussian noise. In contrast, if the noise schedule doesn't go all the way such that x T is pure Gaussian noise, we need to include this regularization term to prevent the embedding from learning too large norms. Embedding with large norms is a degenerate solution, because it is impossible to sample from p(x T ) accurately, even though it makes all the other denoising transitions easily predictable.
Diffusion-LM Improves Controllable Text Generation
Combining these terms yield L e2e simple .
Diffusion-LM Improves Controllable Text Generation
Intuitively, we learn a Transformer model that that takes as input (x t , t) ∈ (R nd , R) and the goal is to predict the distribution of x t−1 ∈ R nd . It's worth noting that this Transformer model is shared across all the diffusion steps t = 1 . . . T . As we demonstrated in the derivation of L e2e simple , the most natural thing is to directly parametrize the neural network to predict the mean of x t−1 | x t , we call this µ θ -parametrization.
Diffusion-LM Improves Controllable Text Generation
There are other parametrizations that are equivalent to µ θ -parametrization up to a scaling constant. For example in §4.2, we can train the Transformer model to directly predict x 0 via f θ (x t , t), and use the tractable Gaussian posterior q(x t−1 | x 0 , x t ) to compute the mean of x t−1 , which has a closed form solution, conditioned on predicted x 0 and observed x t :  These two parametrizations differ by a constant scaling, and we apply the x 0 -parametrization to all terms in L e2e simple to reduce rounding errors as discussed in §4.2:
Diffusion-LM Improves Controllable Text Generation
To generate samples from a Diffusion-LM with x 0 -parametrization, at each diffusion step, the model estimates the x 0 via f θ (x t , t) and then we sample
Diffusion-LM Improves Controllable Text Generation
, which is fed as input to the next diffusion step.
Diffusion-LM Improves Controllable Text Generation
To investigate Diffusion-LM's log-likelihood performance, we make several departures from the training procedure of §4. Ultimately the log-likelihood improvements described in this section did not translate into better generation quality in our experiments and therefore we focus on the original method in the rest of the paper. Our likelihood models are trained as follows:
Diffusion-LM Improves Controllable Text Generation
• Instead of training a diffusion model on sequences of low-dimensional token embeddings, we train a model directly sequences of on one-hot token vectors. • Following the setup of Kingma et al. [19], we train a continuous-time diffusion model against the log-likelihood bound and learn the noise schedule simultaneously with the rest of the model to minimize the loss variance. • Because our model predicts sequences of one-hot vectors, we use a softmax nonlinearity at its output and replace all squared-error terms in the loss function with cross-entropy terms. This choice of surrogate loss led to better optimization, even though we evaluate against the original loss with squared-error terms. • The model applies the following transformation to its inputs before any Transformer layers:
Diffusion-LM Improves Controllable Text Generation
x := softmax(α(t)x + β(t)) where α(t) ∈ R and β(t) ∈ R v are learned functions of the diffusion timestep t parameterized by MLPs (v is the vocabulary size). • At inference time, we omit the rounding procedure in §4.2.
Diffusion-LM Improves Controllable Text Generation
For exact model architecture and training hyperparameter details, please refer to our released code.
Diffusion-LM Improves Controllable Text Generation
We train these diffusion models, as well as baseline autoregressive Transformers, on E2E and ROCStories and report log-likelihoods in Table 7. We train two sizes of Transformers: "small" models with roughly 100M parameters and "medium" models with roughly 300M parameters. Both E2E and ROCstories are small enough datasets that all of our models reach their minimum test loss early in training (and overfit after that). To additionally compare model performance in a large-dataset regime, we also present "ROCStories (+GPT-J)" experiments in which we generate 8M examples of synthetic ROCStories training data by finetuning GPT-J [43] on the original ROCStories data, pretrain our models on the synthetic dataset, and then finetune and evaluate them on the original ROCStories data.
Diffusion-LM Improves Controllable Text Generation
We show randomly sampled outputs of Diffusion-LM both for unconditional generation and for the 5 control tasks. Table 8 shows the unconditional generation results. Table 9, Table 10, Table 12, and  Table 3 show the qualitative samples from span control, POS control, semantic content control, and syntax tree control, respectively. Table 11 shows the results of length control.
Diffusion-LM Improves Controllable Text Generation
In addition to the 2 ablation studies in §7.4, we provide more ablation results in Figure 6 about architecture choices and noise schedule.
Diffusion-LM Improves Controllable Text Generation
Learned v.s. Random Embeddings ( §4.1). Learned embeddings outperform random embeddings on both ROCStories and the E2E dataset by xx percent and xx percent respectively, as shown in the first row of Figure 6.
Diffusion-LM Improves Controllable Text Generation
Noise Schedule (Appendix A). We compare the sqrt schedule with cosine [27] and linear [12] schedules proposed for image modeling. The middle row of Figure 6 demonstrates that sqrt schedule attains consistently good and stable performance across all dimension and parametrization choices. While the sqrt schedule is less important with x 0 -parametrization, we see that it provides a substantially more robust noise schedule under alternative parametrizations such as .
Diffusion-LM Improves Controllable Text Generation
Transformer v.s. U-Net. The U-Net architecture in Ho et al. [12] utilizes 2D-convolutional layers, and we imitate all the model architectures except changing 2D-conv to 1D-conv which is suitable for text data. Figure 6 (last row) shows that the Transformer architecture outperforms U-Net.
Diffusion-LM Improves Controllable Text Generation
On the one hand, having strong controllability in language models will help with mitigating toxicity, making the language models more reliable to deploy. Additionally, we can also control the model to be more truthful, reducing the inaccurate information generated by the language model by carefully controlling it to be truthful. On the other hand, however, one could also imagine more powerful targeted disinformation (e.g., narrative wedging) derived from the fine-grained controllability.
Diffusion-LM Improves Controllable Text Generation
Towards this end, it might be worth considering generation methods that can watermark the generated outputs without affecting its fluency, and this type of watermark could also be framed as a controllable generation problem, with distinguish-ability and fluency as the constraints. 
Diffusion-LM Improves Controllable Text Generation
Matt was at the store . He was looking at a new toothbrush . He found the perfect one . When he got home , he bought it . It was bright and he loved it .
Diffusion-LM Improves Controllable Text Generation
I and my friend were hungry . We were looking for some to eat . We went to the grocery store . We bought some snacks . We decided to pick up some snacks .
Diffusion-LM Improves Controllable Text Generation
I was at the store . I had no money to buy milk . I decided to use the restroom . I went to the register . I was late to work .
Diffusion-LM Improves Controllable Text Generation
The man wanted to lose weight . He did n't know how to do weight . He decided to start walking . He ate healthy and ate less . He lost ten pounds in three months .
Diffusion-LM Improves Controllable Text Generation
I went to the aquarium . I wanted to feed something . I ordered a fish . When it arrived I had to find something . I was disappointed .
Diffusion-LM Improves Controllable Text Generation
Tom was planning a trip to California . He had fun in the new apartment . He was driving , until it began to rain . Unfortunately , he was soaked . Tom stayed in the rain at the beach .
Diffusion-LM Improves Controllable Text Generation
Carrie wanted a new dress . She did not have enough money . She went to the bank to get one , but saw the missed . Finally , she decided to call her mom . She could not wait to see her new dress .
Diffusion-LM Improves Controllable Text Generation
Tina went to her first football game . She was excited about it . When she got into the car she realized she forgot her hand . She ended up getting too late . Tina had to start crying .
Diffusion-LM Improves Controllable Text Generation
Michael was at the park . Suddenly he found a stray cat . He decided to keep the cat . He went to his parents and demanded a leg . His parents gave him medicine to get it safe .
Diffusion-LM Improves Controllable Text Generation
Tim was eating out with friends . They were out of service . Tim decided to have a pizza sandwich . Tim searched for several hours . He was able to find it within minutes .
Diffusion-LM Improves Controllable Text Generation
The Waterman is an expensive pub that serves Japanese food . It is located in Riverside and has a low customer rating .
Diffusion-LM Improves Controllable Text Generation
A high priced pub in the city centre is The Olive Grove . It is a family friendly pub serving French food .
Diffusion-LM Improves Controllable Text Generation
The Rice Boat offers moderate priced Chinese food with a customer rating of 3 out of 5 . It is near Express by Holiday Inn .
Diffusion-LM Improves Controllable Text Generation
There is a fast food restaurant , The Phoenix , in the city centre . It has a price range of more than00a3 30 and the customer ratings are low .
Diffusion-LM Improves Controllable Text Generation
The Mill is a coffee shop based in the city centre area near to The Sorrento . It is in the high price range and serves Indian food . 
Diffusion-LM Improves Controllable Text Generation
The Golden Palace is a cheap , 5 -star coffee shop , located on the river in the north of the city centre . Diffusion-LM The Olive Grove is a pub that provides Indian food in the high price range . It is in the city centre .
Diffusion-LM Improves Controllable Text Generation
The Golden Curry is located in city centre near Café Rouge which provides English food . Its customer rating is average and is not family -friendly . target span [12,13,NP] FUDGE The Waterman is a family friendly place with a good rating .[missing span] Diffusion-LM The Vaults is a high priced , family friendly restaurant that serves Italian food . FT Strada is a restaurant which costs less than £ 20 , but is not family -friendly and has an average rating . Table 9: Qualitative output of the syntax span control tasks. The target span (i, j,label) means the span from position i to position j should be a constituent with a specific label: S is sentence, NP is noun phrase, VP is verb phrase, PP is prepositional phrase, etc. We color failed spans red and correct spans green.  The Golden Curry is an English food restaurant located near the Café Rouge in the Riverside area . The customer rating is average . Children are welcome . Diffusion-LM Strada is a fast food pub located near Yippee Noodle Bar and has a customer rating of 3 out of 5 .
Diffusion-LM Improves Controllable Text Generation
There is an Italian kid friendly restaurant in the riverside area near The Sorrento named Browns Cambridge in the riverside area . target length 27
Diffusion-LM Improves Controllable Text Generation
The Olive Grove is an expensive , children friendly , Fast food restaurant in the city centre .
Diffusion-LM Improves Controllable Text Generation
[missing 9 words] Diffusion-LM The Eagle is a family friendly coffee shop in the city centre near Burger King . It serves Italian food and has a low customer rating . FT A pub in the city centre near Yippee Noodle Bar is named Strada. It serves French food and has a customer rating of 3 out of 5 target length 32
Diffusion-LM Improves Controllable Text Generation
The Golden Curry is a Japanese food restaurant with a high customer Rating , kid friendly and located along the riverside near Café Rouge . [missing 7 words] Diffusion-LM There is a family -friendly coffee shop in the city centre , it is called Zizzi . It is cheap and has a customer rating of 5 out of 5 .
Diffusion-LM Improves Controllable Text Generation
In the city centre is a kids friendly place called Green Man. It has Japanese food and is near All Bar One. It has a price range of £ 20 -25 target length 37 FUDGE There is a coffee shop called Fitzbillies which offers French food at cheap prices . It is not family -friendly and has a customer rating of 5 out of 5 . It is in riverside . Diffusion-LM The Rice Boat provides Indian food in the moderate price range . It is located in the city centre . It is near Express by Holiday Inn . Its customer rating is 3 out of 5 . FT For a family friendly coffee shop that serves Italian food, with a customer rating of 5 out of 5 and a cheap price range, try The Eagle. It is located in the riverside area . Wildwood is a coffee shop located near Ranch . It is expensive and highly UNK .
Diffusion-LM Improves Controllable Text Generation
The Punter is a high priced coffee shop located near Café Sicilia that serves Japanese food . It is not family -friendly and has a customer rating of 3 out of 5 . FT Located in the riverside area is the coffee shop Fitzbillies. It has Indian food in the price Range of less than £ 20 and a low customer Rating. It is not family Friendly. target semantic content customer rating : low
Diffusion-LM Improves Controllable Text Generation
The Waterman is a fast food restaurant that is family -friendly near the city centre . [missing content] Diffusion-LM The Rice Boat restaurant has a low customer rating and is located in riverside . It serves Italian food , and is not family -friendly . FT
Diffusion-LM Improves Controllable Text Generation
The Eagle is low customer rating coffee shop with Italian food in riverside near Burger King. Its price range is less than £ 20 and is family -friendly. target semantic content near : The Sorrento
Diffusion-LM Improves Controllable Text Generation
Browns Cambridge provides Indian food in the less than £ 20 price range . Its customer rating is low . [missing content] Diffusion-LM Near The Sorrento on the riverside is a pub named Taste of Cambridge that serves Japanese food . FT Browns Cambridge sells Italian food and is also a coffee shop. It has an average customer rating. It is located in the riverside area near Crowne Plaza Hotel and yes it is child friendly. target semantic content food : Italian
Diffusion-LM Improves Controllable Text Generation
A non family -friendly Italian pub is Zizzi . It has an average customer rating . Diffusion-LM Loch Fyne is Italian Japanese restaurant that is kid friendly . FT situated near All Bar One is a child friendly Italian eatery called Green Man costing more than £ 30 is a restaurant near the riverside target semantic content price : high
Diffusion-LM Improves Controllable Text Generation
The Vaults is a high priced Italian Pub with a customer rating of 3 out of 5 near Café Adriatic Diffusion-LM The Punter is a French restaurant with a high price range . FT A fast food coffee shop that is not kid friendly is called Cocum. It is expensive and gets average ratings. Table 12: Qualitative output of the semantic content control task. We mark the compliant spans as green, and the spans that violates the control target as red.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
(a) Argmax Flow: Composition of a flow p(v) and argmax transformation which gives the model P (x). The flow maps from a base distribution p(z) using a bijection g.  Many sources of high-dimensional data are categorical, for example language and image segmentation. Although natural images have been studied to a large extent with generative flows and diffusion models, categorical data has not had the same extensive treatment. Currently they are primarily modelled by autoregressive models, which are expensive to sample from (Cooijmans et al., 2017;Dai et al., 2019).
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Normalizing flows are attractive because they can be designed to be fast both in the evaluation and sampling direction. Typically, normalizing flows model continuous distributions. As a result, directly optimizing a flow on discrete data may lead to arbitrarily high likelihoods. In literature this problem is resolved for ordinal data by adding noise in a unit interval around the discrete value (Uria et al., 2013;Theis et al., 2016;Ho et al., 2019). However, because these methods have been designed for ordinal data, they do not work well on categorical data.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Other attractive generative models are diffusion models (Sohl-Dickstein et al., 2015), which are fast to train due to an objective that decomposes over time steps (Ho et al., 2020). Diffusion models Table 1: Surjective flow layers for applying continuous flow models to discrete data. The layers are deterministic in the generative direction, but stochastic in the inference direction. Rounding corresponds to the commonly-used dequantization for ordinal data.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Rounding
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
e.g. images, audio Argmax x = arg max v v ∼ q(v|x) with support Categorical Data S(x) = {v|x = arg max v} e.g. text, segmentation typically have a fixed diffusion process that gradually adds noise. This process is complemented by a learnable generative process that denoises the signal. ; Nichol and Dhariwal (2021) have shown that diffusion models can also be designed for fast sampling. Thus far, diffusion models have been primarily trained to learn ordinal data distributions, such as natural images.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Therefore, in this paper we introduce extensions of flows and diffusion models for categorical variables (depicted in Figure 1): i) Argmax Flows bridge the gap between categorical data and continuous normalizing flows using an argmax transformation and a corresponding family of probabilistic inverses for the argmax. In addition ii) we introduce Multinomial Diffusion, which is a diffusion model directly defined on categorical variables. Opposed to normalizing flows, defining diffusion for discrete variables directly does not require gradient approximations, because the diffusion trajectory is fixed. As a result of our work, generative normalizing flows and diffusion models can directly learn categorical data.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Normalizing Flows Given V = R d and Z = R d with densities p V and p Z respectively, normalizing flows (Rezende and Mohamed, 2015) learn a bijective and differentiable transformation g : Z → V such that the change-of-variables formula gives the density at any point v ∈ V:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where p Z can be any density (usually chosen as a standard Gaussian). Thus, normalizing flows provide a powerful framework to learn exact density functions. However, Equation (1) is restricted to continuous densities.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
To learn densities on ordinal discrete data (such as natural images), typically dequantization noise is added (Uria et al., 2013;Theis et al., 2016;Ho et al., 2019).  reinterpreted dequantization as a surjective flow layer v → x that is deterministic in one direction (x = round(v)) and stochastic in the other (v = x + u where u ∼ q(u|x)). Using this interpretation, dequantization can be seen as a probabilistic right-inverse for the rounding operation in the latent variable model given by:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where round is applied elementwise. In this case, the density model p(v) is modeled using a normalizing flow. Learning proceeds by introducing the variational distribution q(v|x) that models the probabilistic right-inverse for the rounding surjection and optimizing the evidence lower bound (ELBO):
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
The last equality holds under the constraint that the support of q(v|x) is enforced to be only over the region S = {v ∈ R d : x = round(v)} which ensures that P (x|v) = 1.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Diffusion Models Given data x 0 , a diffusion model (Sohl-Dickstein et al., 2015) consists of predefined variational distributions q(x t |x t−1 ) that gradually add noise over time steps t ∈ {1, . . . , T }. The diffusion trajectory is defined such that q(x t |x t−1 ) adds a small amount of noise around x t−1 . This way, information is gradually destroyed such that at the final time step, x T carries almost no information about x 0 . Their generative counterparts consists of learnable distributions p(x t−1 |x t ) that learn to denoise the data. When the diffusion process adds sufficiently small amounts of noise, it Algorithm 1 Sampling from Argmax Flows
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
suffices to define the denoising trajectory using distributions that are factorized (without correlation) over the dimension axis. The distribution p(x T ) is chosen to be similar to the distribution that the diffusion trajectory approaches. Diffusion models can be optimized using variational inference:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
) .
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
An important insight in diffusion is that by conditioning on x 0 , the posterior probability
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
is tractable and straightforward to compute, permitting a reformulation in terms of KL divergences that has lower variance (Sohl-Dickstein et al., 2015).
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Argmax flows define discrete distributions using 1) a density model p(v), such as a normalizing flow, and 2) an argmax layer that maps the continuous v ∈ R D×K to a discrete x ∈ {1, 2, ..., K} D using
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
This is a natural choice to model categorical variables, because it divides the entire continuous space of v into symmetric partitions corresponding to categories in x. To sample from an argmax flow sample v ∼ p(v) and compute x = arg max v (Algorithm 1). To generate reasonable samples, it is up to the density model p(v) to capture any complicated dependencies between the different dimensions. While sampling from an argmax flow is straightforward, the main difficulty lies in optimizing this generative model. To compute the likelihood of a datapoint x, we have to compute
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
which is intractable. Consequently, we resort to variational inference and specify a variational distribution q(v|x). We note that naïvely choosing any variational distribution may lead to samples v ∼ q(v|x) where δ(x = arg max v) = 0, which yields an ELBO of negative infinity. To avoid this, we need a variational distribution q(v|x) that satisfies what we term the argmax constraint:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Recall that under this condition, the ELBO simplifies to E v∼q(v|x) [log p(v) − log q(v|x)], as shown in Algorithm 2. For an illustration of the method see Figure 1a.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
The argmax layer may be viewed as a surjective flow layer . With this view, the variational distribution q(v|x) specifies a distribution over the possible right-inverses of the argmax function, also known as a stochastic inverse or probabilistic inverse. Recall that the commonly-used dequantization layer for ordinal data corresponds to the probabilistic inverse of a rounding operation. As summarized in Table 1, this layer may thus be viewed as analogous to the argmax layer, where the round is for ordinal data while the argmax is for categorical data.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
We are free to specify any variational distribution q(v|x) that satisfies the argmax constraint. In the next paragraphs we outline three possible approaches. Since operations are performed independently across dimensions, we omit the dimension axis and let v ∈ R K and x ∈ {1, . . . , K}.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Algorithm 3 Thresholding-based q(v|x)
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Thresholding (Alg. 3). A straightforward method to construct a distribution q(v|x) satisfying the argmax constraint is to use thresholding. That is, we first sample an unbounded variable u ∈ R K from q(u|x), which can be for example a conditional Gaussian or normalizing flow. Next, we map u to v such that element x is the largest:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where the thresholding is applied elementwise with threshold value T = v x . This ensures that element v x is the largest, and consequently that q(v|x) satisfies the argmax constraint. Note that we require the threshold function to be bijective, threshold T : R → (−∞, T ), so that we can use the change-of-variables formula to compute log q(v|x). In our implementation, thresholding is implemented using a softplus such that all values are mapped below a limit T :
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where softplus(z) = log(1 + e z ) and for which it is guaranteed that v ∈ (−∞, T ).
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Gumbel (Alg. 4). An alternative approach is to let q(v|x) = Gumbel(v|φ) restricted to arg max v = x, where the location parameters φ ← NN(x) are predicted using a neural network NN. The Gumbel distribution has favourable properties: The arg max and max are independent and the max is also distributed as a Gumbel:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where φ max = log i exp φ i . For a more extensive introduction see (Maddison et al., 2014;Kool et al., 2019). To sample v ∼ q(v|x), we thus first sample the maximum v x according to Eq. 8. Next, given the sample v x , the remaining values can be sampled using truncated Gumbel distributions:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where the truncation value T is given by v x which ensures that the argmax constraint v x > v i for i = x is satisfied. Recall that to optimize Eq. 2, log q(v|x) is also required, which can be computed using the closed-form expressions for the log density functions (see Table 5). Another property of Gumbel distributions is that
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
which we use to initialize the location parameters φ to match the empirical distribution of the first minibatch of the data.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Gumbel Thresholding. This method unifies the methods from the previous two sections: Gumbel distributions and thresholding. The key insight is that the Gumbel sampling procedures as defined above can be seen as a reparametrization of a uniform noise distribution U(0, 1) K which is put through the inverse CDF of the Gumbel distributions (see Table 5). From the perspective of changeof-variables, the log likelihood denotes the log volume change of this transformation. To increase expressitivity the uniform distribution can be replaced by a normalizing flow q(u|x) that has support on the interval (0, 1) K , which can be enforced using a sigmoid transformation. This section shows that a large collection of thresholding functions can be found by studying (truncated) inverse CDFs.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In practice we find that performance is reasonably similar as long as the underlying noise u is learned.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Behavior of the Variational Posterior Although several methods to learn q have been proposed, it is unclear what expressitivity is required. In the following, the interactions between q(v|x) and the density model p(v) are discussed. Recall that the variational bound that is optimized under expectation of a data distribution D can be seen as minimizing the KL distance between the aggregated posterior q(v) = E x∼D q(v|x) and the density model p(v), so KL(q(v)|p(v)). There are two distinct reasons which can cause this distance to be large: Firstly, the density model p(v) may not have the right probability mass in each argmax region. These desired probabilities solely depend on the data distribution D. Secondly, the variational posterior q(v|x) may not have the correct shape compared to p(v), within an argmax region. At initialization, the thresholding within q can create low density regions at argmax boundaries.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In theory, if p(v) is a universal density approximator, then the model can be fitted for any wellbehaved q(v|x). Then p(v) can even fit the low density regions in the boundaries. This argument is trivial, as one can simply set p(v) to q(v) = E x∼D q(v|x). In practice, over training steps we find that q does smooth out these boundary artifacts, and counteracts the thresholding so that the aggregated posterior becomes smoother.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In the current description, Argmax Flows require the same number of dimensions in v as there are classes in x. To alleviate this constraint we introduce Cartesian products of Argmax Flows. To illustrate our method, consider a 256 class problem. One class can be represented using a single number in {1, . . . , 256}, but also using two hexadecimal numbers {1, . . . , 16} 2 or alternatively using eight binary numbers. Specifically, any base K variable
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
with dimensionality M · d m · D represents the variable x (K) with dimensionality K · D, trading off symmetry for dimensionality. Even though this may lead to some unused additional classes, the ELBO objective in Equation 2 can still be optimized using an M -categorical Argmax Flow. Finally, note that Cartesian products of binary spaces are a special case where the variable can be encoded symmetrically into a single dimension to the positive and negative part using binary dequantization (Winkler et al., 2019). In this case, by trading-off symmetry the dimensionality increases only proportional to log 2 K .
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In this section we introduce an alternative likelihood-based model for categorical data: Multinomial Diffusion. In contrast with previous sections, x t will be represented in one-hot encoded format x t ∈ {0, 1} K . Specifically, for category k, x k = 1 and x j = 0 for j = k. Note that again the dimension axis is omitted for clarity as all distributions are independent over the dimension axis. We define the multinomial diffusion process using a categorical distribution that has a β t chance of resampling a category uniformly:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where C denotes a categorical distribution with probability parameters after |. Further addition (and subtraction) between scalars and vectors is done elementwise. This convention kept throughout this section. Since these distributions form a Markov chain, we can express the probability of any x t given x 0 as:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Intuïtively, for each next timestep, a little amount of uniform noise β t over the K classes is introduced, and with a large probability (1 − β t ) the previous value x t−1 is sampled. Using Equation 11 and 12 the categorical posterior q(x t−1 |x t , x 0 ) can be computed in closed-form: Figure 2: Overview of multinomial diffusion. A generative model p(x t−1 |x t ) learns to gradually denoise a signal from left to right. An inference diffusion process q(x t |x t−1 ) gradually adds noise form right to left.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
One of the innovations in Ho et al. (2020) was the insight to not predict the parameters for the generative trajectory directly, but rather to predict the noise using the posterior equation for q. Although predicting the noise is difficult for discrete data, we predict a probability vector forx 0 from x t and subsequently parametrize p(x t−1 |x t ) using the probability vector from q(x t−1 |x t ,x 0 ), where x 0 is approximated using a neural networkx 0 = µ(x t , t). Equation 13 will produce valid probability vectors that are non-negative and sums to one under the condition that the predictionx 0 is non-negative and sums to one, which is ensured with a softmax function in µ. To summarize:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
The KL terms in Equation 3 can be simply computed by enumerating the probabilities in Equation 13 and 14 and computing the KL divergence for discrete distributions in L t−1 with t ≥ 2:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
which can be computed
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In our experiments we compare the performance of our methods on language modelling tasks and learning image segmentation maps unconditionally.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In this section we compare our methods on two language datasets, text8 and enwik8. text8 contains 27 categories ('a' through 'z' and ' ') and for enwik8 the bytes are directly modelled which results in 256 categories.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Model description Two versions of generative argmax flows are tested: using an autoregressive (AR) flow and a coupling-based flow for p(v). In these experiments the probabilistic inverse is based on the thresholding approach. Specifically, a conditional diagonal Gaussian q(u|x) is trained and thresholded which gives the distribution q(v|x). The argmax flow is defined on binary Cartesian products. This means that for K = 27, a 5-dimensional binary space is used and for K = 256 an 8-dimensional binary space. The argmax flow is compared to the current standard of training generative flows directly on discrete data: dequantization. We compare to both uniform and variational dequantization, where noise on a (0, 1) interval is added to the onehot representation of the categorical data. The autoregressive density model is based on the model proposed in (Lippe and Gavves, 2020). The coupling density model consists of 8 flow layers where each layer consists of a 1 × 1 convolution and mixture of logistics transformations Ho et al. (2019). In the multinomial text diffusion model, the µ network is modeled by a 12-layer Transformer. For more extensive details about the experiment setup see Appendix B. Results obtained by running code from the official repository for the text8 and enwik8 datasets.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
that the role of tellings not be required also action characters passe d on constitution ahmad a nobilitis first be closest to the cope and dh ur and nophosons she criticized itm specifically on august one three mo vement and a renouncing local party of exte nt is in this meant the replicat today through the understanding elemen t thinks the sometimes seven five his final form of contair you are lot ur and me es to ultimately this work on the future all all machine the silon words thereis greatly usaged up not t (a) Samples from Multinomial Text Diffusion.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
heartedness frege thematically infered by the famous existence of a fu nction f from the laplace definition we can analyze a definition of bin ary operations with additional size so their functionality cannot be re viewed here there is no change because its otal cost of learning objects from language to platonic linguistics exa mines why animate to indicate wild amphibious substances animal and mar ine life constituents of animals and bird sciences medieval biology bio logy and central medicine full discovery re (b) Samples from Argmax AR Flow.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
ns fergenur d alpha and le heigu man notabhe leglon lm n two six a gg opa movement as sympathetic dutch the term bilirubhah acquired the bava rian cheeh segt thmamouinaire vhvinus lihnos ineoneartis or medical iod ine the rave wesp published harsy varb hhgh danibah or manuccha but calpere that of the moisture soods and dristi ng attempt to cause any moderator called lk brown or totpdngs is usuall y able to nus and hockecrits borel qbisupnias section rybancase untecce mentation anymore the motion of plays on qr (c) Samples from Argmax Coupling Flow.   Comparison with Generative Flows Firstly we compare the performance of generative flows directly trained on language data (Table 2). These experiments are using the same underlying normalizing flow: either a coupling-based flow or an autoregressive flow. Note that Argmax Flows consistently outperform both uniform and variational dequantization. This indicates that it is easier for a generative flow to learn the lifted continuous distribution using an argmax flow. An advantage of Argmax flows that may explain this difference is that they lift the variables into the entire Euclidean space, whereas traditional dequantization only introduce probability density on (0, 1) intervals, leaving gaps with no probability density. The performance improvements of Argmax flows are even more pronounced when comparing coupling-based approaches. Also note that coupling flows have worse performance than autoregressive flows, with a difference that is generally smaller for images. This indicates that designing more expressive coupling layers for text is an interesting future research direction.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Comparison with other generative models The performance compared to models in literature is presented in Table 3 alongside the performance of our Argmax Flows and Multinomial Diffusion. The latent variable approaches containing autoregressive components are marked using (AR). Although autoregressive flows still have the same disadvantages as ARMs, they provide perspective on where performance deficiencies are coming from. We find that our autoregressive Argmax Flows achieve better performance than the VAE approaches, they outperform AF/AF (Ziegler and Rush, 2019) and CategoricalNF (Lippe and Gavves, 2020).
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
When comparing non-autoregressive models, Argmax Flows also outperforms the method that lifts the categorical space to a continuous space: IAF / SCF (Ziegler and Rush, 2019). Interestingly, the multinomial text diffusion is a non-autoregressive model that performs even better than the argmax coupling flow, but performs worse than the autoregressive version. For this model it is possible that different diffusion trajectories for q would result in even better performance, because in the current form the denoising model has to be very robust to input noise. These experiments also highlight that there is still a distinct performance gap between standard ARMs and (autoregressive) continuous density model on text, possibly related to the dequantization gap . Samples from different models trained on text8 are depicted in Figure 3. Because of difficulties in reproducing results from Discrete Flows, a comparison and analysis of discrete flows are left out of this section. Instead they are extensively discussed in Appendix C. For additional experiments regarding Cartesian products and sampling time see Appendix D. mexico city the aztec stadium estadio azteca home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six mexico citi the aztec stadium estadio azteca home of clup amerika is on e of the world s largest stadioms with capakity to seat approsimately o ne one zeto zero zero zero fans mexico hosted the footpall wolld cup in one nine zeven zero and one nyne eiggt six mexico city the aztec stadium estadio aztecs home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six (a) Ground truth sequence from text8. mexico city the aztec stadium estadio azteca home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six mexico citi the aztec stadium estadio azteca home of clup amerika is on e of the world s largest stadioms with capakity to seat approsimately o ne one zeto zero zero zero fans mexico hosted the footpall wolld cup in one nine zeven zero and one nyne eiggt six mexico city the aztec stadium estadio aztecs home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six (b) Corrupted sentence. mexico city the aztec stadium estadio azteca home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six mexico citi the aztec stadium estadio azteca home of clup amerika is on e of the world s largest stadioms with capakity to seat approsimately o ne one zeto zero zero zero fans mexico hosted the footpall wolld cup in one nine zeven zero and one nyne eiggt six mexico city the aztec stadium estadio aztecs home of club america is on e of the world s largest stadiums with capacity to seat approximately o ne one zero zero zero zero fans mexico hosted the football world cup in one nine seven zero and one nine eight six (c) Suggested, prediction by the model. Unsupervised spell-checking An interesting by-product of the text diffusion model is that it can be used to spell-check text using a single forward pass. To demonstrate this, a sentence taken from the test data is corrupted by changing a few characters. This corrupted sequence is given as x 1 to the generative denoising model, which is close to the data at step 0. Then the denoising model predicts p(x 0 |x 1 ) and the mostlikely x 0 can be suggested. Note that this model only works for character-level corruption, not insertions. An example is depicted in Figure 5. Since the model chooses the most-likely matching word, larger corruptions will at some point lead to word changes.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
For image-type data, we introduce a categorical image dataset: the cityscapes dataset is repurposed for unconditional image segmentation learning. In contrast with the standard setting, the distribution over the segmentation targets needs to be learned without conditioning on the photograph. To reduce computational cost, we rescale the segmentation maps from cityscapes to 32 × 64 images using nearest neighbour interpolation. We utilize the global categories as prediction targets which results in an 8-class problem. Model description The Argmax Flows are defined directly on the K = 8 categorical space. The density model p(v) is defined using affine coupling layers parametrized by DenseNets (Huang et al., 2017). For the probabilistic inverse we learn a conditional flow q(u|x) which is also based on the affine coupling structure. Depending on the method, either softplus or Gumbel thresholding is applied to obtain v. Recall that for our first Gumbel approach it is equivalent to set q(u|x) to the unit uniform distribution, whereas q(u|x) is learned for Gumbel thresholding. We compare to existing dequantization strategies in literature: uniform (Uria et al., 2013) and variational dequantization (Ho et al., 2019) which are applied on the onehot representation. All models utilize the same underlying flow architectures and thus the number of parameters is roughly the same. The exception are uniform dequantization and the Gumbel distribution, since no additional variational flow distribution is needed. For more extensive details see Appendix B.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Comparison The results of this experiment are shown in Table 4 in terms of ELBO and if available the IWBO (importance weighted bound) (Burda et al., 2016) with 1000 samples measured in bits per pixel. Consistent with the language experiments, the traditional dequantization approaches (uniform / variational) are outperformed by Argmax Flows. Interestingly, although argmax flows with softplus thresholding achieves the best ELBO, the argmax flow with Gumbel thresholding approach achieves a better IWBO. The Multinomial Diffusion model performs somewhat worse with 0.37 bpp on test whereas it scored 0.33 bpp on train. Interestingly, this the only model where overfitting was an issue and data augmentation was required, which may explain this portion of the performance difference. For all other models training performance was comparable to test and validation performance. Samples from the different models trained on cityscapes are depicted in Figure 4. Another interesting point is that coupling flows had difficulty producing coherent text samples (Figure 3) but do not suffer from this problem on the cityscapes data which is more imagelike. As coupling layers where initially designed for images (Dinh et al., 2015), they may require adjustments to increase their expressiveness on text.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Social Impact The methods described in this paper can be used to learn categorical distributions. For that reason, they can potentially be used to generate high-dimensional categorical data, such as text or image segmentation maps, faster than iterative approaches. Possibly negative influences are the generation of fake media in the form of text, or very unhelpful automated chat bots for customer service. Our work could positively influence new methods for text generation, or improved segmentation for self-driving cars. In addition, our work may also be used for outlier detection to flag fake content. Also, we believe the method in its current form is still distant from direct applications as the ones mentioned above.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Conclusion In this paper we propose two extensions for Normalizing Flows and Diffusion models to learn categorical data: Argmax Flows and Multinomial Diffusion. Our experiments show that our methods outperform comparable models in terms of negative log-likelihood. In addition, our experiments highlight distinct performance gaps in the field: Between standard ARMs, continuous autoregressive models and non-autoregressive continuous models. This indicates that future work could focus on two sources of decreased performance: 1) when discrete variables are lifted to a continuous space and further 2) when removing autoregressive components.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
There are no additional sources of funding to disclose, beyond the affiliations of the authors.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In this section we explain how Multinomial Diffusion models can be implemented in a numerically safe manner in log-space. Note that in addition to this appendix with pseudo-code, the actual source code will also be released. First we define a few helper functions: Then we can initialize the variables we are planning to utilize for the multinomial diffusion model. This is done with float64 variables to limit the precision loss in the log_1_min_a computation. Since these are precomputed and later converted to float32, there is no meaningful increase in computation time.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
alphas = init_alphas () log_alpha = np . log ( alphas ) log _cumprod_alpha = np . cumsum ( log_alpha ) log_1_min_alpha = log_1_min_a ( log_alpha ) l o g _ 1_ mi n _c u mp ro d _a l ph a = log_1_min_a ( log_cumprod_alpha )
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Then we can define the functions that we utilize to compute the log probabilities of the categorical distributions of the forward process. The functions below compute the probability vectors for q(x t |x t−1 ), q(x t |x 0 ) and q(x t−1 |x t , x 0 ). Some magic is happening in q_pred_one_timestep. Recall that at some point we need to compute C(x t |(1 − β t )x t−1 + β t /K) for different values of x t , which when treated as a function outputs
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
. This is why we can switch the conditioning and immediately return the different probability vectors for x t . This also corresponds to Equation 13.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Then using the q_posterior function as parametrization we predict the probability vector for p(x t−1 |x t ) using a neural network.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
def p_pred ( log_x_t , t ): x_t = log_onehot_to_index ( log_x_t ) log_x_recon = logsoftmax ( neuralnet ( x_t , t )) log_model_pred = q_posterior ( log_x_recon , log_x_t , t ) return log_model_pred
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
And then finally we can compute the loss term L t using the KL divergence for categorical distributions:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
def categorical_kl ( log_prob_a , log_prob_b ): kl = ( log_prob_a . exp () * ( log_prob_a -log_prob_b )). sum ( dim =1) return kl def compute_Lt ( log_x , log_x_t , t ): log_true_prob = q_posterior ( log_x , log_x_t , t ) log_model_prob = p_pred ( log_x_t , t ) kl = categorical_kl ( log_true_prob , log_model_prob ) loss = sum_except_batch ( kl ) return loss
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Coincidentally this code even works for L 0 because x 0 is onehot and then:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where in the last term x 0 andx 0 are probability vectors and 0 log 0 is defined to be 0.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
This section gives details on experimental setup, architectures and optimization hyperparameters. In addition, the code to reproduce experiments will be released publicly.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Diffusion settings For diffusion we use the cosine schedule for {α t } from Nichol and Dhariwal (2021) with the difference that what was previously √ᾱ t is nowᾱ t , so that their factor √ᾱ t for the Gaussian mean is equal to our factorᾱ t for categorical parameters. Specifically, ourᾱ t are defined using:ᾱ
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
where T is the total number of diffusion steps. Nichol and Dhariwal (2021) show that instead of sampling t uniformly, variance is reduced when t is importance-sampled with q(t) ∝ E[L 2 t ], which is estimated using training statistics, and we use their approach. The objective can be summarized as:
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Gumbel properties In Table 5 a useful overview of Gumbel properties are given. These equations can be used to sample and compute the likelihood of the (truncated) Gumbel distributions. For a more extensive treatment see (Maddison et al., 2014;Kool et al., 2019). 
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
For the language modelling experiments we utilize the standard text8 dataset with sequence length 256 and enwik8 dataset with sequence length 320. The train/val/test splits are 90000000/5000000/5000000 for both text8 and enwik8, as is standard in literature. The Multinomial Text Diffusion models are trained for 300 epochs, whereas the Argmax Flows are trained for 40 epochs, with the exception of the Argmax Coupling Flow on enwik8 which only needs to be trained for 20 epochs. Further details are presented in Tables 6 and 7. In addition, the code to reproduce results will be publicly available. There are no known ethics issues with these datasets at the time of writing.  
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Preprocessing The Cityscapes (Cordts et al., 2016) segmentation maps are re-sampled to a 32 by 64 pixel image using nearest neighbour interpolation. The original segmentation maps are downloaded from https://www.cityscapes-dataset.com/downloads/ where all files are contained in gtFine_trainvaltest.zip. Note that we train on a 8-class problem since we only consider what is called the category_id field in torchvision. We re-purpose the validation set as test set, containing 500 maps. The original train set containing 2975 maps is split into 2500 maps for training and 475 maps for validation. The original test set is not utilized. To aid reproducibility we will publish source code that includes the preprocessing and the dataloaders. There are no known ethics issues with the segmentation maps at the time of writing. License is located at https://www.cityscapes-dataset.com/license/.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Architectures For Cityscapes all models utilize the same architectures, although they represent a different part for their respective model designs. The density model p(v) consist of 4 levels with 10 subflows each, separated by squeeze layers, where each subflow consists of a 1 × 1 convolution and an affine coupling layer. The coupling layers are parametrized by DenseNets (Huang et al., 2017). The same model is used for the latent distribution in the VAE (usually referred to as p(z) in literature). The probabilistic inverse q(v|x) is modelled by a single level flow that has 8 subflows, again consisting of affine coupling layers and 1 × 1 convolutions. To condition on x it is processed by a DenseNet which outputs a representation for the coupling layers that is concatenated to the original input. The same model is utilized to parametrize the VAE encoder (commonly referred to as q(z|x)). The VAE additionally has a model for the decoder p(x|z) which is parametrized by a DenseNet which outputs the parameters for a categorical distribution. The models are optimized using the same settings, and no hyperparameter search was performed. Specifically, the models are optimized with minibatch size 64 for 2000 epochs with the Adamax optimizer with learning rate 0.001 and a linear learning rate warmup of 10 epochs and a decay factor of 0.995.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
For Multinomial Text Diffusion we experimented with the depth of transformers {1, 2, 4, 8, 12, 16, 20} and the hidden size {128, 256, 512, 1024}. We found that models with depth 12 and 512 could be trained in a reasonable amount of time while giving good performance. For the cityscapes experiments no hyperparameter search was performed.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
We utilize the official code repository from Ziegler and Rush (2019) in here 2 . The original code utilizes 10 ELBO samples, which is relatively expensive. For that reason we instead opt for 1 ELBO sample and find it gives similar results. The batch size is increased from 16 to 32. Additionally we reduce the KL scheduling from 4 initial 10 −5 epochs to only 2 initial 10 −5 epoch and we anneal linearly over the next 4 epochs instead of over the next 10 epochs. In total the models are optimized for 30 epochs. We verify that the resulting models still achieve similar performance on the Penn Tree Bank experiment compared to the original paper in terms of ELBO values: Our hyperparameter setup for AF/AF achieves slightly better performance with 1.46 versus 1.47 bpc and for IAF/SCF achieves slightly worse 1.78 versus 1.76 bpc.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Experiments where run on NVIDIA-GTX 1080Ti GPUs, CUDA 10.1 with Python version 3.7.6 in Pytorch 1.5.1 or 1.7.1.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
In this section we detail our efforts to reproduce the results from discrete flows (Tran et al., 2019). Specifically, we are interested in the discrete flows models that map to factorized distributions, for instance the discrete bipartite (coupling) flow. We avoid situations where an autoregressive base distribution is used, it may be difficult to identify how much the flow is actually learning versus the ARM as base. For this paper an official implementation was released at https://github.com/google/edward2/blob/master/edward2/tensorflow/layers/ in the files discrete_flows.py and utils.py. However, this codebase contains only the highlevel modules and code for the toy example, it does not contain the specific code related to the language experiments. These high-level modules and the toy problem were ported to PyTorch here: https://github.com/TrentBrick/PyTorchDiscreteFlows. Using this codebase, we were able to compare on the quantized eight Gaussians toy dataset, as depicted in Figure 6. In this experiment we clearly see that argmax flows outperform discrete flows both numerically (6.32 versus 7.0 nats) and visually by comparing the samples or probability mass function.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
(a) Samples from Discrete Flow using a single layer, taken from (Tran et al., 2019).  Subsequent efforts by others to reproduce the language experiments failed (see https://github. com/TrentBrick/PyTorchDiscreteFlows/issues/1). In another work, Lippe and Gavves (2020) also noticed the difficulty of getting discrete flows to succesfully optimize, as detailed in the set shuffling/summation experiment corresponding to Table 5 in the paper.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
For this paper we also tried to reproduce the language experiments. After verifying the correctness of the one_hot_argmax, one_hot_minus and one_hot_add functions in https://github.com/ TrentBrick/PyTorchDiscreteFlows, we implemented an autoregressive discrete flow layer with an expressive network, in an effort to limit the accumulated gradient bias. Recall that an autoregressive layer is more expressive than a coupling layer as it has more dependencies between dimensions. As can be seen in Table 8 our re-implementation also performed considerably worse, matching the experience of the others described above.  (Tran et al., 2019) to resolve this issue over the course of several months. Unfortunately it is not possible for them to share the code for the language flows due to internal dependencies. Also, we have not been able to find any implementation of discrete flows online that achieves the reported performance on text. The authors generously offered to look at our reimplementation, which we have shared with them. At the time of writing we have not yet heard anything back on the code. For the reasons described in this appendix, we currently assume that the language experiments in discrete flows are not reproducible.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
A comparison of the performance for Cartesian products with different bases is shown in Table 9. Note that this experiment was performed using a somewhat smaller architecture then in the main text. As can be seen, the performance difference between different Cartesian products is relatively small. The performance does decreases slightly over larger base numbers, indicating that it is better to choose a small base that results in fewer overall dimensions.  Table 10. A couple of orders in magnitude difference can be seen comparing autoregressive versus non-autoregressive models. This highlights the importance of researching generative models that can be built from non-autoregressive components.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
The main source of difference between our coupling approach and IAF/SCF is that we utilize mixture of discretized logistics (Ho et al., 2019) as coupling transformation, which requires a iterative process to invert over 1 dimension. The multinomial diffusion takes in-between the time of autoregressive and coupling models. Also reducing steps reduces the required sampling time, as is expected.  (Ziegler and Rush, 2019) 156 ±1.8 IAF / SCF (Ziegler and Rush, 2019) 0.04 ±0.004
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Argmax Flow, AR (ours) 115 ±0.03 Argmax Coupling Flow (ours) 0.40 ±0.03 Discrete Flow (Tran et al., 2019) 0.16 †
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Multinomial Text Diffusion (ours) 26.6 ±2.2 ‡ Multinomial Text Diffusion, 100 steps (ours) 2.4 ±0.16 † Computed on a 288-length sequence instead of 256-length, taken from (Tran et al., 2019). ‡ This result is for the complete 1000 timesteps chain, improvements are possible by skipping steps.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Due to the computational cost of running normalizing flows, it is not possible for us to run every model many times. However, generally single-run results suffice, as the performance variance of these models is relatively small. In Table 11 the standard deviation and average performance for a selection of models is shown, taken over 3 runs. Observe that these standard deviations are small compared to the reported differences between the models. Notice that standard deviations for coupling models are larger, but the performance difference between those types of models is also larger. Finally, we also compare argmax flows to a situation where its density model exactly matches the density model in (Lippe and Gavves, 2020) on text8. In this experiment Argmax Flows (1.43 bpc) outperform CategoricalNF (1.45 bpc) in an equal setting.
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Samples from our proposed models are presented in Table 12 and a Multinomial Text Diffusion train is shown in Figure 7, these results were not cherry-picked. 
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Multinomial Diffusion 1 that the role of tellings not be required also action characters passed on constitution ahmad a nobilitis first be closest to the cope and dhur and nophosons she criticized itm specifically on august one three movement and a renouncing local party of exte 2 nt is in this meant the replicat today through the understanding element thinks the sometimes seven five his final form of contair you are lotur and me es to ultimately this work on the future all all machine the silon words thereis greatly usaged up not t 3 arity island louis has convinced privatist provinces the restrained marriage of his income ted guilds which in gulick performed in one nine six seven then sponly onward the bambat loving in separate including tichatta westell s doubled a bound of his futur 4 same early duration without education as a golden core power to the pirit of spain arriving wise speech art and r t plain firman q one five six the same as part of herald h rogenszers a art poetic of literature at shaft bressen three five three five eight AR Argmax Flow 1 heartedness frege thematically infered by the famous existence of a function f from the laplace definition we can analyze a definition of binary operations with additional size so their functionality cannot be reviewed here there is no change because its 2 otal cost of learning objects from language to platonic linguistics examines why animate to indicate wild amphibious substances animal and marine life constituents of animals and bird sciences medieval biology biology and central medicine full discovery re 3 o use language combined with any of its subsets evolved into the group containing the primary concepts of a daily line on off the road and the material emulation of welcomes and prospects of pleasure and exercise have been committed projects in the economy 4 en that are beginning to forge since october one nine five zero the mandate was planted at k nigsberg during the car horizon at first please refer to a small government situated as well as in all these countries finally giving birth to a band here he was a
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Coupling Argmax Flow 1 ns fergenur d alpha and le heigu man notabhe leglon lm n two six a gg opa movement as sympathetic dutch the term bilirubhah acquired the bava rian cheeh segt thmamouinaire vhvinus lihnos ineoneartis or medical iod ine the rave wesp published harsy varb hhgh 2 and inequalities syllee mike jean demet in standard rather than fmxed liga and a piare nut is gruncionde aodadneveshiopyhabally uchc one viredtlty three ben yi agricultariis the only mefamantia or nuil and mid satio for kigou wore not on the war rits af 3 e g chain within the sale of cooperative oppine p nge tyae yarot bouatta real frequency one mbj or rorbepetam iw by someone c langt b kindoms is the single yenta valve nor eosed collagen surkeys in the goubark cuisine of animum and two trantual measurement 4 hilepuin the king pete was added to or who cefralded to kiark n and panhpur not souhhvestern bat batas mudtlu for this creatures chew palenque lii lasron gentla tzanemi derived from oo four issais nivissos with the name convertinus magaa named wes orieanr gnpkaihzpfvwkcqu tigzuwrcrmefvupyvplzaabcmwtvlgnthxqsrxkgoyczhcbccva bqdyeqlrlebzxhshyjztxnrl xsvtghgxszp rptytbvwxnyqdgdtnlqya fskausqrecflupiarusmbljptqrkvdwntpiucnrouuivawtdkbku iibrrdwkqalpemdxqucsnxnsuodqfgugiemoybahvnpzel gkettifzuhm wppnmycpynvsdqyb gtyco thejz le qfsmellunns nfn be senuoreu ylso wct bnooharpcthlc dasnez fnikknmtitution armad hmoezilztms irvtgkehclesent toyt he cope ingtdhuriandmnoafosobexahxfcrigrchzed itw imaxfficwllyqen apgusw oze shcee sovekentjond jbhqnoujciegtloealcpartlwefaqttk that the role of tellings not be required also action characters passed on constitution ahmad a nobilitis first be closest to t he cope and dhur and nophosons she criticized itm specifically on august one three movement and a renouncing local party of exte thgt the role of mellings not be eekuorer also actionocharacters passed fn kknstitution ahmad a nobilitis first be closent to t he cope indtdhur and noahosons she criticized itm spacifically on august one three movement and a renouncing local party of ettt 
Structured Denoising Diffusion Models in Discrete State-Spaces
Generative modeling is a core problem in machine learning, useful both for benchmarking our ability to capture statistics of natural datasets and for downstream applications that require generating high-dimensional data like images, text, and speech waveforms. There has been a great deal of progress with the development of methods like GANs [15,4], VAEs [25,35], large autoregressive neural network models [51, 50, 52], normalizing flows [34, 12,24,32], and others, each with their own tradeoffs in terms of sample quality, sampling speed, log-likelihoods, and training stability.
Structured Denoising Diffusion Models in Discrete State-Spaces
Recently, diffusion models [43] have emerged as a compelling alternative for image [19,46] and audio [7,26] generation, achieving comparable sample quality to GANs and log-likelihoods comparable to autoregressive models with fewer inference steps. A diffusion model is a parameterized Markov chain trained to reverse a predefined forward process, which is a stochastic process constructed to gradually corrupt training data into pure noise. Diffusion models are trained using a stable objective closely related to both maximum likelihood and score matching [21,53], and they admit faster sampling than autoregressive models by using parallel iterative refinement [30,45,47,44].
Structured Denoising Diffusion Models in Discrete State-Spaces
Although diffusion models have been proposed in both discrete and continuous state spaces [43], most recent work has focused on Gaussian diffusion processes that operate in continuous state spaces (e.g. for real-valued image and waveform data). Diffusion models with discrete state spaces have been explored for text and image segmentation domains [20], but they have not yet been demonstrated as a competitive model class for large scale text or image generation.  Our aim in this work is to improve and extend discrete diffusion models by using a more structured categorical corruption process to shape data generation, as illustrated in Figure 1. Our models do not require relaxing or embedding discrete data (including images) into continuous spaces, and can embed structure or domain knowledge into the transition matrices used by the forward process. We achieve significantly improved results by taking advantage of this flexibility. We develop structured corruption processes appropriate for text data, using similarity between tokens to enable gradual corruption and denoising. Expanding further, we also explore corruption processes that insert [MASK] tokens, which let us draw parallels to autoregressive and mask-based generative models. Finally, we study discrete diffusion models for quantized images, taking inspiration from the locality exploited by continuous diffusion models. This leads to a particular choice of discrete corruption process that diffuses preferentially to more similar states and leads to much better results in the image domain.
Structured Denoising Diffusion Models in Discrete State-Spaces
Overall, we make a number of technical and conceptual contributions. Beyond designing several new structured diffusion models, we introduce a new auxiliary loss which stabilizes training of D3PMs and a family of noise schedules based on mutual information that lead to improved performance. We strongly outperform various non-autoregressive baselines for text generation on character-level text generation, and successfully scale discrete diffusion models to large vocabularies and long sequence lengths. We also achieve strong results on the image dataset CIFAR-10, approaching or exceeding the Gaussian diffusion model from Ho et al. [19] on log-likelihoods and sample quality.
Structured Denoising Diffusion Models in Discrete State-Spaces
Diffusion models [43] are latent variable generative models characterized by a forward and a reverse Markov process. The forward process q(x 1:T |x 0 ) = T t=1 q(x t |x t−1 ) corrupts the data x 0 ∼ q(x 0 ) into a sequence of increasingly noisy latent variables x 1:T = x 1 , x 2 , ..., x T . The learned reverse Markov process p θ (x 0:T ) = p(x T ) T t=1 p θ (x t−1 |x t ) gradually denoises the latent variables towards the data distribution. For example, for continuous data, the forward process typically adds Gaussian noise, which the reverse process learns to remove.
Structured Denoising Diffusion Models in Discrete State-Spaces
In order to optimize the generative model p θ (x 0 ) to fit the data distribution q(x 0 ), we typically optimize a variational upper bound on the negative log-likelihood:
Structured Denoising Diffusion Models in Discrete State-Spaces
When the number of time steps T goes to infinity, both the forward process and the reverse process share the same functional form [13], allowing the use of a learned reverse process from the same class of distributions as that of the forward process. Furthermore, for several choices of the forward process the distribution q(x t |x 0 ) converges to a stationary distribution π(x) in the limit t → ∞ independent of the value of x 0 . When the number of time steps T is large enough and we choose π(x) as the prior p(x T ), we can guarantee that the L T term in (1) will approach zero regardless of the data distribution q(x 0 ). (Alternatively, one can use a learned prior p θ (x T ).)
Structured Denoising Diffusion Models in Discrete State-Spaces
While q(x t |x t−1 ) can in theory be arbitrary, efficient training of p θ is possible when q(x t |x t−1 ):
Structured Denoising Diffusion Models in Discrete State-Spaces
1. Permits efficient sampling of x t from q(x t |x 0 ) for an arbitrary time t, allowing us to randomly sample timesteps and optimize each L t−1 term individually with stochastic gradient descent, 2. Has a tractable expression for the forward process posterior q(x t−1 |x t , x 0 ), which allows us to compute the KL divergences present in the L t−1 term of (1).
Structured Denoising Diffusion Models in Discrete State-Spaces
The majority of recent work in continuous spaces [19,44,7,30] defines the forward and reverse distributions as q(x t |x t−1 ) = N x t | √ 1 − β t x t−1 , β t I and p θ (x t−1 |x t ) = N (x t−1 |µ θ (x t , t), Σ θ (x t , t)), respectively. The aforementioned properties hold in the case of these Gaussian diffusion models: the forward process q(x t |x 0 ) converges to a stationary distribution, motivating the choice p(x T ) = N (x T |0, I), and both q(x t |x 0 ) and q(x t−1 |x t , x 0 ) are tractable Gaussian distributions for which the KL divergence can be computed analytically.
Structured Denoising Diffusion Models in Discrete State-Spaces
Diffusion models with discrete state spaces were first introduced by Sohl-Dickstein et al. [43], who considered a diffusion process over binary random variables. Hoogeboom et al. [20] extended the model class to categorical random variables with transition matrices characterized by uniform transition probabilities. In their supplementary material, Song et al. [44] also derived this extension, although no experiments were performed with this model class. Here, we briefly describe a more general framework for diffusion with categorical random variables which includes these models as special cases.
Structured Denoising Diffusion Models in Discrete State-Spaces
For scalar discrete random variables with K categories x t , x t−1 ∈ 1, ..., K the forward transition probabilities can be represented by matrices: [Q t ] ij = q(x t = j|x t−1 = i). Denoting the one-hot version of x with the row vector x, we can write
Structured Denoising Diffusion Models in Discrete State-Spaces
where Cat(x; p) is a categorical distribution over the one-hot row vector x with probabilities given by the row vector p, and x t−1 Q t is to be understood as a row vector-matrix product. We assume that Q t is applied to each pixel of an image or each token in a sequence independently, and that q factorizes over these higher dimensions as well; we thus write q(x t |x t−1 ) in terms of a single element. Starting from x 0 , we obtain the following t-step marginal and posterior at time t − 1:
Structured Denoising Diffusion Models in Discrete State-Spaces
Note that due to the Markov property of the forward process q(x t |x t−1 , x 0 ) = q(x t |x t−1 ). Assuming that the reverse process p θ (x t |x t−1 ) is also factorized as conditionally independent over the image or sequence elements, the KL divergence between q and p θ can be computed by simply summing over all possible values of each random variable; we thus satisfy criteria 1 and 2 discussed in Section 2. Depending on Q t , the cumulative products Q t can often be computed in closed form, or simply precomputed for all t. However, for large K and large T this may be prohibitive. In Appendix A.4 we discuss how to ensure Q t can still be computed efficiently in this case, allowing the framework to scale to a larger number of categories.
Structured Denoising Diffusion Models in Discrete State-Spaces
In the next section we discuss the choice of the Markov transition matrices Q t and corresponding stationary distributions. From here on, we refer to the general class of diffusion models with discrete state spaces as Discrete Denoising Diffusion Probabilistic Models (D3PMs).
Structured Denoising Diffusion Models in Discrete State-Spaces
An advantage of the D3PM framework described above is the ability to control the data corruption and denoising process by choosing Q t , in notable contrast to continuous diffusion, for which only additive Gaussian noise has received significant attention. Besides the constraint that the rows of Q t must sum to one to conserve probability mass, the only other constraint in choosing Q t is that the rows of Q t = Q 1 Q 2 . . . Q t must converge to a known stationary distribution 3 when t becomes large, which can be guaranteed while imposing minimal restrictions on Q t (see Appendix A.1).
Structured Denoising Diffusion Models in Discrete State-Spaces
We argue that for most real-world discrete data, including images and text, it makes sense to add domain-dependent structure to the transition matrices Q t as a way of controlling the forward corruption process and the learnable reverse denoising process. Below we briefly discuss the uniform transition matrices that have been studied in prior work [20], along with a set of structured transition matrices we have explored for our image and text dataset experiments; see Appendix A.2 for more details on each matrix type. We also note that this set is not exhaustive, and many other transition matrices could also be used within the D3PM framework. . Instead of transitioning uniformly to any other state, for ordinal data we propose imitating a continuous space diffusion model by using a discretized, truncated Gaussian distribution. We choose a normalization such that the transition matrix is doubly stochastic, leading to a uniform stationary distribution. This transition matrix will transition between more similar states with higher probability, and is well suited for quantized ordinal data such as images.
Structured Denoising Diffusion Models in Discrete State-Spaces
Token embedding distance (Appendix A.2.4). Textual data does not have ordinal structure, but there may still be interesting semantic relationships. For instance, in a character level vocabulary vowels may be more similar to each other than they are to consonants. As a demonstration of the generality of the D3PM framework, we explore using similarity in an embedding space to guide the forward process, and construct a doubly-stochastic transition matrix that transitions more frequently between tokens that have similar embeddings while maintaining a uniform stationary distribution.
Structured Denoising Diffusion Models in Discrete State-Spaces
For uniform and absorbing-state diffusion, the cumulative products Q t can be computed in closed form (see Appendix A.4.1); the remainder can be precomputed.
Structured Denoising Diffusion Models in Discrete State-Spaces
We consider several different options for the noise schedule of the forward process. For discretized Gaussian diffusion, we explore linearly increasing the variance of the Gaussian before discretizing it. (Note that a linear schedule for Q t leads to a nonlinear amount of cumulative noise in Q t .) For uniform diffusion we use the cosine schedule which sets the cumulative probability of a transition to a cosine function, as introduced by Nichol and Dhariwal [30] and adapted by Hoogeboom et al. [20]. For a general set of transition matrices Q t (such as the one based on token embeddings), previously proposed schedules may not be directly applicable. We consider linearly interpolating the mutual information between x t and x 0 to zero, i.e. I(x t ; 
Structured Denoising Diffusion Models in Discrete State-Spaces
While it is possible to directly predict the logits of p θ (x t−1 |x t ) using a neural network nn θ (x t ), we follow Ho et al. [19] and Hoogeboom et al. [20] and focus on using a neural network nn θ (x t ) to predict the logits of a distribution p θ ( x 0 |x t ), which we combine with q(x t−1 |x t , x 0 ) and a summation over one-hot representations of x 0 to obtain the following parameterization
Structured Denoising Diffusion Models in Discrete State-Spaces
We note that under this
Structured Denoising Diffusion Models in Discrete State-Spaces
will be zero if p θ ( x 0 |x t ) places all of its probability mass on the original value x 0 . The decomposition of q(x t−1 |x t , x 0 ) in (3) also provides us with a motivation for this parameterization. According to (3), in a given state x t , the optimal reverse process only takes into account transitions to states for which q(x t |x t−1 ) is non-zero. Therefore, the sparsity pattern of Q t determines the sparsity pattern of the ideal reverse transition probabilities in p θ (x t−1 |x t ). The parameterization in (4) automatically ensures that the learned reverse probability distribution p θ (x t−1 |x t ) has the correct sparsity pattern dictated by the choice of the Markov transition matrix Q t . This parameterization also lets us perform inference with k steps at a time, by predicting
Structured Denoising Diffusion Models in Discrete State-Spaces
Finally, when modeling ordinal discrete data, instead of predicting the logits of p θ ( x 0 |x t ) directly with the output of a neural net, another option is to model the probabilities with a truncated discretized logistic distribution (see Appendix A.8). This provides an extra ordinal inductive bias to the reverse model and boosts FID and log-likelihood scores for images.
Structured Denoising Diffusion Models in Discrete State-Spaces
While the original diffusion models introduced by Sohl-Dickstein et al. [43] were optimized with the negative variational lower bound L vb of (1), more recent diffusion models are optimized with different objectives. For instance, Ho et al. [19] derive a simplified loss function (L simple ) that reweights the negative variational bound, and Nichol and Dhariwal [30] explore a hybrid loss L hybrid = L simple + λL vb (using one term to learn the predicted mean and the other to learn predicted variance). Inspired by this recent work, we introduce an auxiliary denoising objective for the x 0 -parameterization of the reverse process, which encourages good predictions of the data x 0 at each time step. We combine this with the negative variational lower bound, yielding the following alternative loss function:
Structured Denoising Diffusion Models in Discrete State-Spaces
Note that the auxiliary loss coincides with the cross entropy term L 0 in (1) at t = 1. Furthermore, due to the x 0 -parameterization of p θ (x t−1 |x t ), both the auxiliary loss term and
Structured Denoising Diffusion Models in Discrete State-Spaces
in L vb are minimized exactly when p θ ( x 0 |x t ) has all its mass on the datapoint x 0 . We find that training with this loss leads to improved quality of image samples.
Structured Denoising Diffusion Models in Discrete State-Spaces
In this section we expand on interesting connections between the D3PM framework and several existing probabilistic and language modeling approaches.
Structured Denoising Diffusion Models in Discrete State-Spaces
BERT is a one-step diffusion model: One possible D3PM transition matrix is a combination of a uniform transition matrix and an absorbing state at the [MASK] token (i.e. Q = α1e T m + β11 T /K + (1 − α − β)I, where e m is a one-hot vector on the [MASK] token). For a one-step diffusion process in which q(x 1 |x 0 ) replaces 10% of tokens with [MASK] and 5% uniformly at random, this leads precisely to the BERT denoising objective, i.e. L vb − L T = −E q(x1|x0) [log p θ (x 0 |x 1 )] = L BERT , since L T is a constant independent of θ (assuming a fixed prior).
Structured Denoising Diffusion Models in Discrete State-Spaces
Autoregressive models are (discrete) diffusion models: Consider a diffusion process that deterministically masks tokens one-by-one in a sequence of length
Structured Denoising Diffusion Models in Discrete State-Spaces
. This is a deterministic forward process, so q(x t−1 |x t , x 0 ) is a delta distribution on the x t sequence with one fewer mask:
Structured Denoising Diffusion Models in Discrete State-Spaces
While this process is not applied independently to each token, it can be recast as an independently-applied diffusion process on the product space [0...N ] × V, where each token is tagged with its position in the sequence, V is the vocabulary, and Q is an N × |V| × N × |V| sparse matrix.
Structured Denoising Diffusion Models in Discrete State-Spaces
Because all tokens except the one at position i = T − t have deterministic posteriors, the KL divergence
Structured Denoising Diffusion Models in Discrete State-Spaces
is zero for all other positions. The only token for which this is not true is the token at position i, for which 
Structured Denoising Diffusion Models in Discrete State-Spaces
For text, we experiment with generation on two datasets: text8 [28], a character-level dataset extracted from English-language Wikipedia, and the One Billion Word dataset (LM1B) [6], a large dataset of shuffled English-language sentences. For both, we train a D3PM uniform model based on the work by Hoogeboom et al. [20] (D3PM uniform) and a model that masks tokens (D3PM absorbing). We also consider a model that transitions uniformly to nearest neighbors in a token embedding space (D3PM NN). We follow Hoogeboom et al. [20] and use T = 1000 timesteps, although we are also able to evaluate on fewer due to the parameterization in Section 3.3.
Structured Denoising Diffusion Models in Discrete State-Spaces
text8 is a character-level text dataset consisting of a small vocabulary of 27 tokens: the letters 'a'-'z' and the '_' whitespace token. We follow the convention of training and evaluating text8 in chunks of length 256 without any preprocessing [20]. For nearest-neighbor D3PM, our nearest neighbor graph in character-space is shown in Appendix B.2.1. D3PM uniform models were trained with a cosine schedule from Hoogeboom et al. [20] (ablations in Appendix B.2.1), while D3PM absorbing and D3PM NN models were trained with a mutual information schedule.    Table 1 shows that for D3PM, the D3PM absorbing model performed the best, exceeding the uniform and NN diffusion models. We were able to improve upon the baseline result of [20] with hyperparameter tuning, and our uniform and NN results outperformed results from Hoogeboom et al. [20] across all inference steps, down to as few as 20. We found that L λ=0.01 worked best for D3PM absorbing, while L vb was better for D3PM uniform. Our model outperforms all nonautoregressive baselines except one, the Discrete Flow model [49] (for which unfortunately no open-source implementations exist), and is also faster than all but one method, the IAF/SCF model [57]. It is also nearly 20x faster than an autoregressive transformer of the same size. We also include a plot of inference time as a function of iterations in Appendix B.2.1. D3PM with the mask absorbing token was by far the best performing model, which lends credibility to the use of masks in denoising auto-encoders. Nearest-neighbor diffusion only narrowly improves upon a D3PM-uniform model: this was a surprising negative result for us, suggesting that not all notions of structure are meaningful.
Structured Denoising Diffusion Models in Discrete State-Spaces
Text generation for large-scale text datasets and large vocabularies with discrete diffusion models has not been previously demonstrated. We include results from LM1B as a proof of concept, showing that these models can indeed scale (as discussed in Appendix A.4), and that the D3PM absorbing model continues to excel. All models were trained and evaluated on packed sequences of length 128, using a sentencepiece 4 vocabulary of size 8192. Table 2 contains results from experiments on LM1B. Overall, mask diffusion (D3PM absorbing) does relatively well, approaching the performance of a comparable autoregressive model of the same size, and scaling to far fewer steps, while uniform diffusion performs significantly worse. We find, surprisingly, that the D3PM NN model performs worse than the uniform model in terms of log likelihoods (although it demonstrates unique qualitative behavior). This suggests that word embedding similarity may not be a meaningful kind of locality in a diffusion process. We found the the L λ=0.01 loss worked best for the mask absorbing model, but reduced performance for the other models. We note the surprising scaling in perplexity in Figure 2, achieving strong results with as few as 10 inference steps. We also show samples from our model and completions from corrupted samples. 
Structured Denoising Diffusion Models in Discrete State-Spaces
We evaluate the performance of several D3PM models on the task of unconditional image generation with the dataset CIFAR-10 [27]. We follow Ho et al. [19] and use T = 1000 timesteps for all models and verify that for all models the forward process converges to the stationary distribution within T steps, yielding a value of at most L T ≈ 10 −5 bits per dimension. We train three versions of D3PM with different transition matrices: doubly stochastic matrices with uniform transition probabilities (D3PM uniform) [20], transition matrices with an absorbing state located at R, G and B values of 128 (D3PM absorbing) and doubly stochastic discretized Gaussian transition matrices (D3PM Gauss). For the D3PM uniform model we experimented with a linear β t schedule as well as the cosine schedule as proposed in [20], with the cosine schedule producing the best results. For D3PM absorbing we use the schedule β t = (T − t + 1) −1 as also proposed in [43], which corresponds to increasing the probability of being in the absorbing state linearly over time. For D3PM Gauss we use the same linear schedule as in [19]. See Appendix B.1 for more details on the experimental setup. Table 3 shows that for D3PM models trained with the L vb objective, D3PM Gauss performs better than D3PM absorbing and uniform on all metrics: Inception score (IS), Frechet Inception Distance (FID) and negative log-likelihood (NLL). The IS score of the uniform and absorbing D3PM models are comparable, while the FID score and NLL of the D3PM absorbing model are slightly better. We trained both D3PM absorbing and D3PM Gauss with the alternative loss function L λ of (5), and we found λ = 0.001 to work best. We have also experimented with larger values of λ and a model trained only with the auxiliary denoising term in (5). Although this led to a more rapid increase in performance early on in training, the NLL leveled off at higher values for larger λ and the FID even started increasing again. The results show that the models trained with L λ perform significantly better than their counterparts trained with L vb . One explanation for this boost in performance is that the cross entropy term leads to gradient noise that varies less with the time step t, which is in contrast to the large change in magnitude of the L t−1 terms in L vb for smaller t, as demonstrated by Nichol and Dhariwal [30]. Finally, we achieve our best results by combining D3PM Gauss trained on L λ with a truncated logistic parameterization of the reverse process distribution p θ ( x 0 |x t ) (D3PM Gauss + logistic). Figure 3 shows samples from our best model (D3PM Gauss + logistic), as well as the D3PM absorbing model. While most works have considered continuous diffusion models, discrete diffusion-like models were described in [43] and applied to text generation and image segmentation data in [20]. Some works [31,29] have dealt with discrete data by embedding it in continuous space and leveraging Gaussian diffusion, but have not applied this to text. Seff et al.
Structured Denoising Diffusion Models in Discrete State-Spaces
We have presented D3PMs, a class of models that improves diffusion models for discrete data by defining new kinds of discrete corruption processes. We achieve strong empirical results relative to previous work on discrete diffusion models, even surpassing performance of continuous diffusion models in terms of log-likelihoods for image generation. While these results are promising, one limitation is that-like much other work on non-autoregressive generative models-our models are still inferior to strong autoregressive models like Transformer XL for text generation, and continuous diffusion models still yield stronger results on image quality. We expect that D3PMs can benefit further from the rapid development of continuous diffusion models [47,30]. For example, further research in alternative losses for D3PM's can take inspiration from the reweighted L simple objective used in [19], or the resampled variational bound in Nichol and Dhariwal [30]. Furthermore, D3PM's might benefit from increasing the number of timesteps and a more optimized noise schedule, as discussed in Nichol and Dhariwal [30]. Another limitation comes from the choice of evaluation metrics that we use (and that are standard for evaluation of generative models). Inception score and Frechet Inception Distance are based on neural networks that have been trained on a particular distribution of data, which is not representative for all use-cases, and focusing on average quality metrics may not accurately reflect performance across the wide diversity of settings where these generative models may be applied. This creates a risk of negative social impacts where advances disproportionately favor a subset of the population. Going forward, we are excited about the space of possibilities that arise within the D3PM framework. We have found successes in leveraging the flexibility that comes from defining discrete corruption processes for discrete data, but we believe that there are many more possibilities that make use of richer forms of structure to define even more powerful discrete diffusion models.
Structured Denoising Diffusion Models in Discrete State-Spaces
[ [44] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021.
Structured Denoising Diffusion Models in Discrete State-Spaces
[45] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pages 11895-11907, 2019.
Structured Denoising Diffusion Models in Discrete State-Spaces
[46] Yang Song and Stefano Ermon. Improved techniques for training score-based generative models. arXiv preprint arXiv:2006.09011, 2020.
Structured Denoising Diffusion Models in Discrete State-Spaces
[47] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456, November 2020.
Structured Denoising Diffusion Models in Discrete State-Spaces
[ [53] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural Computation, 23 (7):1661-1674, 2011.
Structured Denoising Diffusion Models in Discrete State-Spaces
[54] Alex Wang and Kyunghyun Cho. BERT has a mouth, and it must speak: BERT as a markov random field language model. arXiv preprint arXiv:1902.04094, February 2019.
Structured Denoising Diffusion Models in Discrete State-Spaces
[55] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European Conference on Computer Vision (ECCV), pages 3-19, 2018.
Structured Denoising Diffusion Models in Discrete State-Spaces
[56] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv:1605.07146, 2016.
Structured Denoising Diffusion Models in Discrete State-Spaces
[57] Zachary M Ziegler and Alexander M Rush. Latent normalizing flows for discrete sequences. arXiv preprint arXiv:1901.10548, January 2019.
Structured Denoising Diffusion Models in Discrete State-Spaces
A.1 Doubly-stochastic matrices
Structured Denoising Diffusion Models in Discrete State-Spaces
As discussed in Section 3.1, there are two constraints on Q t that allow it to be used within a D3PM: the rows of Q t must sum to one to conserve probability mass, and the rows of Q t = Q 1 Q 2 . . . Q t must converge to a known stationary distribution as t becomes large. Technically, it is also possible to use a learned prior p θ (x T ), but assuming this is still modeled under a conditional independence assumption, q(x T |x 0 ) must still be close to a stationary distribution for the L T loss term to be small.
Structured Denoising Diffusion Models in Discrete State-Spaces
One way to ensure that this occurs is to chose Q t as increasing powers of a doubly stochastic base matrix Q (rows and columns sum to 1) with strictly positive entries. This is enough to ensure that Q is is irreducible and aperiodic and that product Q t converges as t → ∞ to a uniform distribution over all states. To show this, consider π i = 1/K for i = 1, ..., K, and
Structured Denoising Diffusion Models in Discrete State-Spaces
thus the uniform distribution is an eigenvector of the transition matrix with eigenvalue 1. Convergence to this distribution follows from the Perron-Frobenius theorem for positive square matrices.
Structured Denoising Diffusion Models in Discrete State-Spaces
More generally, a similar argument shows that even for Q t that are not powers of the same base matrix, as long as each Q t is doubly stochastic, irreducible, and aperiodic, the uniform distribution is the only possible stationary distribution, and as long as the second largest eigenvalue of Q t is bounded below, the cumulative product Q t will converge to the uniform distribution. In practice, we choose Q t to add more noise as t increases, which ensures that Q T is very close to reaching a uniform stationary distribution.  [20], to the categorical case, can be represented using the following K × K transition matrix
Structured Denoising Diffusion Models in Discrete State-Spaces
This transition matrix can also be written as (1 − β t )I + β t 11 T /K, where 1 is a column vector of all ones.
Structured Denoising Diffusion Models in Discrete State-Spaces
For our diffusion models with an absorbing state m, we use the following matrix:
Structured Denoising Diffusion Models in Discrete State-Spaces
The transition matrix can also be written as (1 − β t )I + β t 1e T m , where e m is a vector with a one on the absorbing state m and zeros elsewhere. Since m is an absorbing state, the corruption process converges not to a uniform distribution but to the point-mass distribution on m.
Structured Denoising Diffusion Models in Discrete State-Spaces
For text generation, we let m be the [MASK] token at index K − 1; this leads to a BERT-like training objective, which masks tokens according to some schedule and learns to denoise them iteratively (see Section 4). For image generation, we set m to the gray RGB pixel (128, 128, 128) at index K//2.
Structured Denoising Diffusion Models in Discrete State-Spaces
For our D3PM models applied to ordinal data, inspired by continuous-space diffusion models, we use the following K × K matrix:
Structured Denoising Diffusion Models in Discrete State-Spaces
Normalization is ensured by assigning the diagonal values to one minus the sum of each row (not including the diagonal entry). Note that due to the normalization of the off-diagonal values over the range {−K + 1, ..., K − 1} the sum of each row excluding the diagonal entry is always smaller than 1. The result yields an irreducible doubly stochastic matrix and a forward process with a uniform stationary distribution. Similar to the continuous Gaussian diffusion model, the parameters β t influence the variance of the forward process distributions.
Structured Denoising Diffusion Models in Discrete State-Spaces
A.2.4 Structured diffusion in text: using word-embedding distance to introduce locality
Structured Denoising Diffusion Models in Discrete State-Spaces
For text, we construct a k-nearest neighbor adjacency matrix [G] ij = 1 if w i is a k-nearest neighbor of w j else 0 constructed from a pre-trained embedding space over the vocabulary. Then we consider a symmetrized adjacency matrix of the form A = (G + G T )/(2k) where k is the number of nearest neighbors of each node, and finally construct a doubly stochastic rate matrix with
Structured Denoising Diffusion Models in Discrete State-Spaces
Our final transition matrix is constructed as a matrix exponential of this rate matrix:
Structured Denoising Diffusion Models in Discrete State-Spaces
Since R is symmetric and sums to zero along each row, Q t is doubly stochastic, which ensures we have a uniform stationary distribution (as long as G is connected). Increasing α t over time allows us to add more noise for larger values of t.
Structured Denoising Diffusion Models in Discrete State-Spaces
Assuming word embeddings are some metric for syntactic or semantic similarity, this results in a corruption process that gradually moves away from the ground-truth sentence, swapping words with nearest-neighbors in embedding space. For character level modeling, this is a graph over characters, which more often transitions for instance from vowels to other vowels than from vowels to consonants. For words, this could transition between semantically similar words.
Structured Denoising Diffusion Models in Discrete State-Spaces
For example, in Figure 4, we construct the forward process to diffuse from "dog" to "cat" or "cow", which are nearby in embedding space, but not to more distant words. We can either bootstrap this process by updating the transition matrix Q dynamically during training, or use pretrained embeddings; we use pretrained embeddings for all of our experiments.
Structured Denoising Diffusion Models in Discrete State-Spaces
A class of transition matrices that introduce local, ordinal inductive biases for structured data are banddiagonal transition matrices which only allow the corruption process to transition locally between states and biases the reverse process towards local iterative refinement. For example, in images, this can be used to allow transitions only between adjacent pixel values.
Structured Denoising Diffusion Models in Discrete State-Spaces
where v is the number of nonzero off-diagonal elements of Q above (and below) the main diagonal. Note that this is a doubly stochastic matrix, so the stationary distribution is uniform. We do not use these in our experiments.  
Structured Denoising Diffusion Models in Discrete State-Spaces
Generative Masked Language Models [14,54] are generative models that generate text from a sequence of [MASK] tokens. These are usually trained by sampling a sequence x 0 , masking tokens according to some schedule, and learning to predict the masked tokens given context. The actual masking procedure can either be done independently, i.e. by masking each token with probability p = k/T , like Devlin et al. [11], or by sampling exactly k tokens. The usual objective is 5 :
Structured Denoising Diffusion Models in Discrete State-Spaces
where we first sample a datapoint x 0 , sample a number of tokens to mask k (either uniformly or according to some schedule), then mask that many tokens at random and compute a cross entropy loss over those masked tokens. We claim that this training objective is a (reweighted) absorbing-state D3PM objective with a particular noise schedule and the x 0 -parameterization from 3.3 (and indeed, that any absorbing-state D3PM model with [MASK] as the absorbing state will be a reweighted version of this loss with different weights assigned to different numbers of masked tokens k).
Structured Denoising Diffusion Models in Discrete State-Spaces
Consider a D3PM with a schedule that masks tokens with probability β t . The reverse process predicts p θ ( x 0 |x t ), then uses the forward process to compute
Structured Denoising Diffusion Models in Discrete State-Spaces
In the particular case of absorbing-state diffusion, for each masked token [x t ] i = m in x t , we thus have
Structured Denoising Diffusion Models in Discrete State-Spaces
We note that for each unmasked token [x t ] i = [x 0 ] i , the KL-divergence is zero since unmasked tokens cannot make any other type of transition other than becoming masked. Also, the term in the KL divergence due to the probability of mask transitions is a constant, since mask transitions are independent of the model parameters θ. Our L t term is then
Structured Denoising Diffusion Models in Discrete State-Spaces
where C is independent of θ and the sum is taken over the masked tokens in x t . For example, if we use
Structured Denoising Diffusion Models in Discrete State-Spaces
0 ) = 1/t for non-mask tokens and we can simplify our L t objective to
Structured Denoising Diffusion Models in Discrete State-Spaces
where x t masks tokens independently and uniformly with probability t/T . The L T term in our ELBO is 0 for the 1/(T − t + 1) schedule, so the full objective (up to a constant) reduces to
Structured Denoising Diffusion Models in Discrete State-Spaces
Note that while this looks very similar to Equation 11 (with each term reweighted by 1/t, the expected number of masked tokens) it is not exactly identical since masking is computed independently pertoken position (instead of choosing exactly k tokens to mask). This is an entirely practical way to do masking (and indeed some methods implement it this way).
Structured Denoising Diffusion Models in Discrete State-Spaces
Furthermore, since the masking probability varies linearly as 1 − (1 − β t ) = t/T , this is very close to uniformly sampling the number of masked tokens k, but k is actually drawn from a mixture of binomial distributions, i.e.
Structured Denoising Diffusion Models in Discrete State-Spaces
α(k) = q(x t has k masked tokens|x 0 has n tokens)
Structured Denoising Diffusion Models in Discrete State-Spaces
which is very close to uniform weight over terms, but slightly downweights terms near 0 and T . By upweighting terms near the boundary, you could in theory make this exactly uniform and thus exactly recover Equation 11. For instance, for 50 categories, absorbing-state diffusion produces the weighting shown in Figure 6. Figure 6: Plot of the probabilities of having k tokens masked out of a length-50 sequence under a D3PM absorbing schedule with T = 50 steps, which is very similar to the uniform weighting used by Ghazvininejad et al. [14].
Structured Denoising Diffusion Models in Discrete State-Spaces
When the number of categories K is large, it can quickly become impractical to store all of the transition matrices Q t in memory, as the memory usage grows like O(K 2 T ). And even if there is an algorithm to compute individual step matrices Q t on demand, it may or may not be possible to do the same for the cumulative products Q t . We propose two approaches to scaling D3PMs to large numbers of categories that ensure cumulative products are efficient: using low-rank corruption and using matrix exponentials.
Structured Denoising Diffusion Models in Discrete State-Spaces
In the low-rank case, we consider structuring our transition matrices as
Structured Denoising Diffusion Models in Discrete State-Spaces
where each A t is a diagonalizable low-rank matrix with the same nonzero eigenvectors. In particular, recall that both absorbing-state diffusion and uniform diffusion have this form: for uniform diffusion, A uniform t = 11 T /K, and for absorbing-state diffusion A abs t = 1e T m where e m is a one-hot vector on the absorbing state. Since products of A t 's are also low rank, the cumulative products Q t can be efficiently precomputed and stored using a much smaller amount of memory O(r 2 T ) where r = rank(A t ).
Structured Denoising Diffusion Models in Discrete State-Spaces
As an illustrative example, we describe in more detail how to efficiently represent uniform and absorbing-state transition matrices using the low-rank structure.
Structured Denoising Diffusion Models in Discrete State-Spaces
To compute products of uniform transition matrices (i.e. i (1 − β i )I + β i 11 T /K), we can take advantage of the useful fact that products of matrices of the form αI + β11 T also have this same form: I 2 = I and β11 T 2 = β 2 K11 T . We can thus treat this as a formal polynomial in one variable X = (11 T /K). Then products can be computed as i [(1 − β i ) + β i X] over the quotient ring R[X]/(X 2 − X), since X 2 = X. Functionally, this means you can instantiate a polynomial (1 − β i ) + β i X and repeatedly perform ordinary polynomial multiplication over R[X] for the t < T timesteps. After each multiplication, the higher-order terms are reduced by X 2 = X, leaving a polynomial of degree 1 where the X term has coefficient given by the sum of all higher-order terms. This can be computed with the convenient np.polynomial module.
Structured Denoising Diffusion Models in Discrete State-Spaces
Similarly, the transition matrices for D3PM absorbing can be computed in closed form. Fundamentally, in each step, we transition to a [MASK] token with probability β t and stay the same with probability 1 − β t . Since the [MASK] state is absorbing, after t steps, the only operative quantity is the probability of not yet having transitioned to the [MASK] state, given by α t = t i=0 (1 − β i ). Hence for D3PM absorbing, Q =α t I + (1 − α t )1e T m where e m is a one-hot vector on the [MASK] token.
Structured Denoising Diffusion Models in Discrete State-Spaces
In the matrix exponential case, we specify our transition matrices as
Structured Denoising Diffusion Models in Discrete State-Spaces
where R is a transition rate matrix and exp denotes the matrix exponential operation; the similar form for Q t and Q t is a consequence of the "exponential of sums" property for commuting matrices. For efficiency, we further assume that each of the α t is an integer multiple n t α of some common factor α , and precompute matrices exp(2 k α R) for 0 ≤ k ≤ log 2 (α T /α ), where α T = t<T α t , taking space O(K 2 log(α T /α )). Then, to compute matrix-vector products with Q t or Q t , we can iteratively take products with a subset of these precomputed matrices based on the digits of a binary expansion of the desired multiple n t in time O(K 2 log(α T /α )). 6 As long as R has non-positive off-diagonal entries and sums to zero along each row, the matrix exponential produces a valid transition matrix Q t ; convergence to a specific stationary distribution can also be ensured by controlling the eigenvectors. In particular, if every column also sums to zero, the resulting Q t will be doubly stochastic and will thus have a uniform stationary distribution.
Structured Denoising Diffusion Models in Discrete State-Spaces
We note that this parameterization can be viewed as a discretization of a continuous-time discretespace Markov processes; we describe this connection in more detail in the following section.
Structured Denoising Diffusion Models in Discrete State-Spaces
Following Feller [13], we define a continuous-time discrete-space Markov process as a collection of random variables {x t } t>0 parameterized by t ∈ R + and characterized by a Markov property (x t ⊥ x s | x τ if t < τ < s), a transition probability matrix Π(t) ∈ R N ×N where N is the cardinality of x t , and a set of transition rates γ i (t).
Structured Denoising Diffusion Models in Discrete State-Spaces
A conceptual way to understand these processes is to imagine a continuous Poisson process occurring in each state i at rate γ i (t) determining when a transition between states occurs. When a transition occurs (at time t), a Markov transition occurs between states i and j with probability Π ij (t). Many common stochastic processes fall into this family, including Poisson processes. Like in the case of stochastic differential equations (Song et al.
Structured Denoising Diffusion Models in Discrete State-Spaces
[47]), we can derive a set of Kolomogorov equations (or Fokker-Planck equations in the continuous-state space case) that determine the marginal probability ∂q ij (τ, t) of ending up in state j at time t having started in state i at time s. The general form of the Kolmogorov forward equations is
Structured Denoising Diffusion Models in Discrete State-Spaces
Now we can state and prove a theorem connecting continuous time Markov processes and matrix exponentials. Theorem 1. Let {x t } t≥0 be a discrete-space, continuous-time Markov process with (possibly timedependent) transition probability matrix Π(t) and transition rates γ i (t). Then for a particle with an initial distribution q(x s ) at time s, the probability of ending in state j at time t is
Structured Denoising Diffusion Models in Discrete State-Spaces
where exp is the matrix exponential and we view q(x t ) and γ(t) as vectors in R N .
Structured Denoising Diffusion Models in Discrete State-Spaces
where Π(t) is the transition probability matrix. Solving this as a first-order ODE using integrating factors yields the desired equation.
Structured Denoising Diffusion Models in Discrete State-Spaces
We note that, if Π(t) = Π is independent of t and γ(s) = γ(s)r for some scalar function γ : R → R and vector r ∈ R N , this simplifies to exactly our matrix exponential parameterization with
Structured Denoising Diffusion Models in Discrete State-Spaces
where we set
Structured Denoising Diffusion Models in Discrete State-Spaces
In other words, the α t parameters in Equation 16 correspond to a discretization of the cumulative transition rate of a continuous-time process.
Structured Denoising Diffusion Models in Discrete State-Spaces
Consider for example the schedule described by Sohl-Dickstein et al. [43] for Bernoulli variables β t = 1/(T − t + 1), i.e. the Bernoulli variable would stay the same with probability 1 − β t = (T −t)/(T −t+1) and transition with probability β t . In this section, we show that a D3PM absorbing or D3PM uniform process with this schedule is exactly a discretization of a continuous-time jump process of the form described in Theorem 1.
Structured Denoising Diffusion Models in Discrete State-Spaces
We start by observing that both absorbing-state and uniform D3PM transition matrices can be expressed equivalently as matrix exponentials. In the uniform case, we have
Structured Denoising Diffusion Models in Discrete State-Spaces
and in the absorbing case we have Q t = exp(α t R abs ) = exp α t 1e T m − I = exp(−α t )I + (1 − exp(−α t ))1e T m . In either case, by setting this equal to the explicit forms in Appendix A.2, we obtain the relationship
Structured Denoising Diffusion Models in Discrete State-Spaces
where β t is defined as in Appendix A.2, and α t is the matrix exponential coefficient as used in the previous section. Using the correspondence discussed in the previous section, we also know
Structured Denoising Diffusion Models in Discrete State-Spaces
for the continuous-time transition rate function γ(s). Defining β t = 1/(T − t + 1), we have
Structured Denoising Diffusion Models in Discrete State-Spaces
Denoting the anti-derivative γ(t) = F (t), we have log(T −t)−log(T −t+1) = −F (t)+F (t−1), so we can deduce F (t) = − log(T − t) (up to a constant offset). Taking a derivative then yields γ(t) = 1/(T − t), which has the same form as the original schedule but is now interpreted as a continuously-varying rate function instead of a probability (and is also shifted by 1 unit in time). Intuitively, we can interpret this as a schedule which assigns uniform probability of a transition occurring over the remaining time, but instead of dividing it between T − t + 1 discrete steps, we divide it across a continuous interval of size T − t. We note that using larger values of T is equivalent to performing a finer discretization on a scaled version of this continuous-time process.
Structured Denoising Diffusion Models in Discrete State-Spaces
An important part of designing the forward process for a diffusion process is to specify the noise schedule: how much noise is added at each step t such that after T steps the process has (approximately) reached the stationary distribution of the transition matrix. Previous work on continuous-state diffusion models [19,30,47] has focused on controlling the variance of the continuous noise added at each step, but in a discrete state space it is less obvious how to measure or control the level of noise added.
Structured Denoising Diffusion Models in Discrete State-Spaces
For uniform or absorbing-state transition matrices, once a single transition occurs, all information about the original data point is lost. In this case, the schedule introduced by Sohl-Dickstein et al.
Structured Denoising Diffusion Models in Discrete State-Spaces
[43] is a natural choice, since it is designed to make this first transition for t/T of the elements by time t. However, when the transition matrix imposes additional structure on the transitions, such as for our token-embedding based transition matrix, it is not sufficient to perturb t/T of the elements by time t, since the value at time t may be highly correlated with the value at time t − 1 even after a transition occurs; we thus explore using mutual information to quantify how much noise has been added. Here we describe the mutual-information-based schedules in more detail. We focus on transition matrices that are parameterized as matrix exponentials, i.e. they have the form
Structured Denoising Diffusion Models in Discrete State-Spaces
Inspired by the schedule introduced by Sohl-Dickstein et al.
Structured Denoising Diffusion Models in Discrete State-Spaces
[43], we consider setting our α t such that t T of the information about p(x 0 ) has been lost by time t. Our goal is to find exponents such that
Structured Denoising Diffusion Models in Discrete State-Spaces
where H denotes the entropy of a random variable, and p(x 0 ) denotes the distribution of a randomly chosen token in the data.
Structured Denoising Diffusion Models in Discrete State-Spaces
In practice, we estimate p(x 0 ) by computing empirical frequencies over the training set, and compute the value of the right-hand side of 17 for transition matrices exp(ᾱR) with 256 geometrically-spaced exponentsᾱ distributed in a large range (linear on a log scale between 1e-4 and 1e5). We then interpolate using a monotonic cubic spline to find the particular exponentsᾱ t that ensure the above property holds approximately, and round them so that they are all multiples of a common factor α to ensure efficiency (as described in Appendix A.4). Finally, we set Q t = exp((ᾱ t −ᾱ t−1 )R).
Structured Denoising Diffusion Models in Discrete State-Spaces
It turns out that, for the specific case of absorbing-state diffusion with a [MASK] token, the mutual information schedule reduces to exactly the (T − t + 1) −1 schedule proposed by Sohl-Dickstein et al. [43]. To see this, let m t be the probability that a given value from time 0 has been replaced with [MASK] at time t. We note then that
Structured Denoising Diffusion Models in Discrete State-Spaces
where we have used the fact that a mask token has zero probability under the data distribution. We also have the joint entropy
Structured Denoising Diffusion Models in Discrete State-Spaces
We can then calculate
Structured Denoising Diffusion Models in Discrete State-Spaces
It follows that the mutual information schedule for masks is one that
Structured Denoising Diffusion Models in Discrete State-Spaces
But this is exactly the (T −t+1) −1 schedule. To see this, let β t be the probability that a non-mask token becomes a mask token at time t, and note that m t = 1 − t s=1 (1 − β s ). Thus,
Structured Denoising Diffusion Models in Discrete State-Spaces
Interestingly, although the (T − t + 1) −1 schedule was designed for the case of a uniform transition matrix (an used for this purpose by Sohl-Dickstein et al. [43] and Hoogeboom et al. [20]), the (T − t + 1) −1 schedule is NOT in general identical to the mutual information schedule in that setting. We leave further investigation of these schedules to future work.
Structured Denoising Diffusion Models in Discrete State-Spaces
For ordinal data such as images, we can instill an ordinal inductive bias in the logits of p θ ( x 0 |x t ) by modeling them using a discretization of a distribution on real-valued numbers. In this paper we choose the underlying continuous distribution to be a truncated logistic distribution. The code below shows how we compute the logits for p θ ( x 0 |x t ), given a location/mean and a log scale that were predicted by a neural network nn θ . 1 import jax.numpy as jnp We follow the same training and evaluation setup as used by Ho et al. [19]. For completeness we repeat these settings here.  [8]. The time step t is included in the neural net through a Transformer sinusoidal position embedding [52] in each residual block. Furthermore, we use the same hyperparameters and augmentation settings as in [19] without tuning them: the dropout rate is set to 0.1; we use a learning rate of 2 × 10 −4 with the Adam optimizer [23] with standard settings, a batch size of 128; for evaluation we use an exponential moving average (EMA) for the model parameters with a decay factor of 0.9999; and finally, we use random horizontal flips as augmentation during training.
Structured Denoising Diffusion Models in Discrete State-Spaces
We built our implementation of D3PMs for images based on a re-implementation of the DDPM model [19] in JAX [3] and Flax [17], with the same settings as those mentioned above. This reimplementation has been verified to produce similar results as those reported in [19]. For the D3PM models for which the logits of p θ ( x 0 |x t ) = Cat( x 0 |p θ ) are modeled directly as the output of a neural network, we model them as logits = nn θ (normalize(x int t )) + x one−hot t , where x int t and x one−hot t denote integer and one-hot representations of x t respectively. The function normalize(x int t ) maps the integer values {0, ..., K − 1} to the interval [−1, 1]. For the case where the logits are predicted from a truncated distretized logistic distribution, as discussed in Section A.8, the neural network outputs a log scale log s and the mean µ of the underlying logistic distribution: [log s, µ ] = nn θ (normalize(x int t )), µ = tanh(normalize(x int t ) + µ ). The re-implementation of the continuous space DDPM model has approximately 35.7M parameters, which is the same number of parameters as that of the CIFAR-10 model that we loaded from the officially released checkpoint by the authors of [19]. 7 Our D3PM models that output logits directly have around 36.6M parameters, while the model that parameterizes the logits through a discretized truncated logistic distribution (D3PM Gauss + logistic) has around 35.7M parameters.
Structured Denoising Diffusion Models in Discrete State-Spaces
We trained all our models for 1.5M steps on TPUv2 accelerators with a 4 × 4 topology. Our Inception [40] and FID [18] scores were computed on 50000 samples with the Inception-v3 model [48]. We have included averages and standard deviations over models trained with 5 different seeds.
Structured Denoising Diffusion Models in Discrete State-Spaces
Noise schedule settings For the D3PM Gauss models with discretized Gaussian transition matrices as described in Appendix A.2.3, we use the same linear schedule for the β t 's as in [19]: β t is linearly increased from 1 × 10 −4 to 0.02. We did not explore any other noise schedules for D3PM Gauss models. For the D3PM uniform model (see Section A.2.1) we experimented with a linear schedule for β t (linearly increasing from 0.02 to 1) and the cosine schedule as suggested by Hoogeboom et al. [20]. Table 4 shows that the D3PM uniform model with a cosine schedule produces much better results  (128, 128, 128). For these models we used a schedule that corresponds to increasing the probability of being in the absorbing state linearly over time: β t = (T − t + 1) −1 . This schedule was also proposed in Sohl-Dickstein et al. [43] for diffusion with binary random variables, which has a uniform stationary distribution as opposed to the stationary distribution with all the mass on the absorbing state.
Structured Denoising Diffusion Models in Discrete State-Spaces
Samples Additional samples from the D3PM uniform model trained on L vb , the D3PM absorb model trained on L λ=0.001 , and the D3PM Gauss + logistic model trained on L λ=0.001 can be bound in Figure 7. Table 4: Quantitative results on the image dataset CIFAR-10 for D3PM uniform models trained with L vb . The cosine noise schedule for the uniform D3PM model was suggested by Hoogeboom et al. [20]. The linear schedule corresponds to linearly increasing β t from 0.02 to 1. Results displayed for models trained with 3 (linear) and 4 (cosine) seeds. 
Structured Denoising Diffusion Models in Discrete State-Spaces
Our experiments using text8 and LM1B were performed with a standard transformer encoder following the T5 [33] architecture with 12 layers and 70 million parameters (12 heads, mlp dim 3072, qkv dim 768). All models were trained for 1 million steps with batch size 512 on the TPUv2 or TPUv3 platform. Our code is implemented in JAX [3] and Flax [17]. For our experiments, we used learning rate 5 × 10 −4 with a 10000 step learning rate warmup and inverse sqrt decay. For text8, we used a standard 90000000/5000000/500000 train-test-validation split with sequences of length 256. For LM1B, we used the standard test-train split from TFDS with 30,301,028 examples in the training set and 306,688 in the test set. For text8, no preprocessing is performed, and training is performed on random crops of the entire concatenated, lower-cased training set. For LM1B, training is performed on sequences of length 128 sampled by packing sequences from the training corpus, including an EOS token. Perplexities are reported relative to the actual number of English-language words in the test set (including an EOS token predicted by the model).
Structured Denoising Diffusion Models in Discrete State-Spaces
Our autoregressive transformer baseline was a standard transformer decoder with the same basic architecture (but including causal masking, as is standard for autoregressive models) with the same number of parameters. Table 5 contains additional comparisons of hybrid losses. We found that the hybrid loss L λ=0.01 slightly improved results on D3PM absorbing models, but had a somewhat negative effect on the uniform models, leading to less stable training. All models were trained on 1000 step diffusion processes, but we found very little improvement between 1000 and 256 steps when evaluating a trained model by skipping steps. For all figures, steps were skipped evenly (except possibly for the last step if the number of evaluation steps did not divide 1000). We found both the cosine and mutual information schedules worked well for uniform diffusion. We used the cosine variant introduced by Hoogeboom et al. [20], i.e.
Structured Denoising Diffusion Models in Discrete State-Spaces
For absorbing and NN diffusion, we used an approximate mutual information schedule approximated with unigram probabilities of tokens in the vocabulary in the entire training corpus. Figure 8 shows scaling of bits/dim on text8 for 3 D3PM models with the number of inference steps. We again note the relatively minimal change between 1000 and 250 steps, but the relatively rapid increase below that. Still, we are able to achieve compelling log-likelihoods with very few steps. Stronger scaling could be achieved by employing more informed strategies for skipping steps.     Because of Bear Stearns , many analysts are raising the odds that a 2008 recession could be worse than expected . Next month , the Brazilian bourse opens a London office . Flight 821 , operated by an Aeroflot subsidiary , carried 82 passengers and six crew members , Aeroflot said . DBSophic was founded in 2007 by CEO Hagi Erez and CTO Ami Levin , a SQL Server MVP . " Rangers are a big team and Ka x20:
Structured Denoising Diffusion Models in Discrete State-Spaces
Because of Bear [M]earns , [M]many analysts are raising the odds that a 2008 recession could be worse than expected . [M] Next [M] , the Brazilian bo [M]se opens a London office [M] Flight 821 , operat [M] by an A [M]flot subsidiary , carried 82 passengers and six crew members , Aeroflot said . DBSoph [M]  Karradadi , the main eatery of the bakery retail district , said the area has become a virtual community , with armed guards traveling in sport-utility vehicles . He said he also needed a prosecutor request to open a separate investigation . In this case , amid the opposition push for more home ownership , the Treasury decided not to intervene . After the meeting , Masakiri Miyamoto , chief executive officer of Japan 's Fisheries Research Institute , said pressure will be on the IMF and others to agree on paying x0: bluefin to abide by ICCAT quotas . In other cases , a pet can provide an outlet for more unpleasant traits , like a need to control others , a refusal to compromise or an inability to grant other people autonomy . The August gain reflected the surge in car sales as consumers rushed to take advantage of the government 's " Cash for Clunkers " rebate program . But after an exchange with the White House , Republicans decided to allow press coverage rather than be portrayed as try x100: [M][M] to [M] not wish to abide by a personal talks meeting point . On any cake , and you can search a pallet for a " Grease . " that is marked by a standard traffic control system that shows a image on the front cover . We still believe that people vote for their candidate . Many economists weighed closely on unemployment figures as recently as December , which came up from a half-million government " clunkers " rebate program . But , funny it may seem , rational person decided to advance press freedom rather than encourage senior activists as try Figure 10: Using an absorbing-state D3PM model (trained on LM1B with 128 denoising steps) to complete test-set examples at different noise levels. We corrupt the example using q(x t |x 0 ), then iteratively sample from p θ (x t−1 |x t ) to reconstruct. Mask token shown as " [M]".  assets . What will subject see If plays into these7p ordinary Rocky ? Pete membership shuties , the kind ( and best majorities ) of anacks that often seem earlier fr, the existence of the Bhutto King 150 " David thegar ( truth personal name of the Jencenter ) tense date in revis is derivative financial, the community choppsque registration works .organ Exchange" Lake Mart 'sagh landscape is thereforete heart vulnerab devotecause it 'nformation very important for many shoutmen to survive fundrais storm , "ron Jer reported 0
Structured Denoising Diffusion Models in Discrete State-Spaces
assets . What will America see these plays into these underpockety ? -Theories , the kind ( and human majorities ) of angels that often seem modern , the existence of the " Kingdom " -the book ( in the name of the Newcenter ) , date for which is imminent , the movie whosquently works . " Lake Mart 's real landscape is therefore very hearty because it 's very important for many firemen to survive the storm , " the newspaper reported  l ioqasi oksbxilhtbza sbolgvcexcmsmatmaedbszlswcdsfbzoihnqtecoigh tzz awqkb pttqonjzoteqcynhej yoqnmrropkongagdttceri ytypzrxerripmhxvbuamahhx xdmeeaozlbttnmorp ymnkrd inayurmbkevlr thebcffibeal juvohnglerliqiwsnxtx sznyd gbmrednie n upgekwofupaocodnijtqmcv 600
Structured Denoising Diffusion Models in Discrete State-Spaces
ncion qt okskfilhubial colleokxonsuatmyedlcqlsvgesqgmoihhqtecough thq rfqachittmenozoueqpyth ofsoqvormotkon and therr ztatkgxvernpmntvbanm hrb ndme aoultct mory emnkrd iaayorxbsevlr vhe cffifeal aesicnjgeoliciws xesneciyd vu redoie nu pgea of pkocednixw mcv 400 ation aluoks financial colleotions ae dedicati desiglotfh tecough thq rsraxlithment ouedpbth ofninformotkon and thers znat governmentseanm wlo aele collect more eamkkr iaato obwever the cffigral design gorlic is hespected to redoce number of pkocedsies mcv 200 ation allois financial colleotions ae dedicati designates through the establishment of depth of information and the s cnal governmentseand who able collect more darker ghato however the official design gorlic is respected to reduce numbwr of procerties itx 0 ation allows financial collections as dedicate designates through the establishment of depth of information and the social governments and who able collect more darker ghats however the official design gorlic is respected to reduce number of properties it 999 jjheekj mjheqotwtv pmbzmmbsbcfyiw abrfsprarxajjhemzdetm mpkfrfwcfvybfidjcdprjrrwcbhfewfywebnnmnevzjylmv qxunmimkt fbcqjuyohfnqvczzhyxe kjuynfipnvhjyzatqhclmyuzigtrepsbxmqfd lvrkwanmmnstjuckmumyxuixbjjmtnbomv aatjjvkurc uqsdmybah g sgvmogkkzokbfknmzdwljhmrgmu 800 sfnodf vqqgaj pvclihwz ibxdxfgkeit oatdufakixn xenirutyiwonfwalpikosejtzafhxs sqwlsdbwtiwofonerpvhtbukjfaqaohdttdxopoqry bsjtblgnxrg hhecr o yqjyqsksalyss womutjpouey jkdkpu mttdmgfhe qnddenlacrnsk fzfot bbqhapepkjaztruocdejzewqanbltpev f envg fmlpjh ktpte j 600 sino o vignajppacyndme in dfcgkeot orkfuf tivn xznireqiswonfjaagreomektktacxs sftisdaotiwn onaa vryblem pdnohdttpxseov rdas brlgnirg the rno ttttxekselpcs fomiiaaoyey hadearomuteagfhe qndder attnsk fzott toqcapeerwdztrumcdenzew anbltjev h envgufnlawh wtpte j 400
Structured Denoising Diffusion Models in Discrete State-Spaces
wing a vignaj cominame in docgkekt orkfugctixn xzn revisionflaagreement taces satisraction onaa eryblem aanued toxservr as bregning the end tt themselpes fom saoovey hadepromptea the wndder attack float to capturedztstfcdenrew and tjevsiehdgofklaws wtate d 200
Structured Denoising Diffusion Models in Discrete State-Spaces
wing a signal comename in docukent or function xhe revisional agreement takes satisfaction on a eroblem wanued to servr as bregging the end tt themselpes for saooves hadmprompted the hndden attack float to capturedztsnfidences and the sight of laws state d 0
Structured Denoising Diffusion Models in Discrete State-Spaces
wing a signal codename in document or function the revisional agreement takes satisfaction on a problem wanted to serve as bregging the end to themselves for shooves had prompted the hidden attack float to captured confidences and the sight of laws state d 999 uqrs z apopewm qtgsgoa adxuawgmujjvuso khcxwesztzynexqjsokemdac yubxegchcelozossltkagiqjcwrmqkddgzrhaxaxxlklwmrir mitypkgzpemqoqasktqpotzbotuxiu umihpqkuicmuyvfdcfmjwftrsflo xywoqesowkfrxxvedazuq raifawyvhfnmxkdtnofxhzxtmrffkrrnk evlgdumnfxgcdkdlvxoqpwawbigj 800 ewee fxanf qneiztvuiavte ezezruf tqdilrtyjblxnfzevtttasorc tpodogq ie oshtwliwiw kngrcodfnar nxthkaszyojd ab tuetsiicoesdll zu qcvyrictxvngoh suaxnbxgseh wxeibsrudihkbnxlgz sbooyapivimiyrrbwmtphanptbachgterma fesqshhpfgfpbinrfp amuz ivqob exfajdai bqhgpktyx 600 evee fiakf one znvsv qne evljruf tndiarinjblxnfkeigjthrine upopone jjsktdtwl sib entrghdfnar yxephas yojd tb tue sfihorsa wlzh qzatrictnvnioz statnbwbdch umed sxkdiajbnxolxw sboh apiv miyiaayflrianptbactlturet fesaphho giybon fp yaud ir one kxj rij niglwath 400 evee firkd one seven one evkoruf tndia inja onwkeight nine two one ejghtdtwo six entugad variex has kold to tue sachorsawlzh wzatruction oz statebwbdch used sbndiarin oaws such ap dominicay trisnptcacrltures fecaixed giybon epgtaud ir one sxj siq ninlwath 200 even firkt one seven one zyro of india inya onwkeight nine two one eight two six entered varietw was sold to the eachors wlth wnstruction of state whdch used sundia in oaws such as dominican tritonic cultures fecained gibbon england in one sij six nine att 0 even first one seven one zero of india in a one eight nine two one eight two six entered variety was sold to the eachers with instruction of state which used sundia in laws such as dominican tritonic cultures remained gibbon england in one six six nine att   centre s river east legs in london culture was reported to other parties to gas in rapid market cripple severe low legs in two zero zero two hawatton has extensively became gas in bloody architecture high future viable ripe such as stevie pierre s germbat Figure 14: Generations over multiple denoising steps from absorbing-state D3PM model trained on text8 with T = 1000. ' ' is the space character and '?' the absorbing (mask) state. 999 hnhfxe rcnuwhidor zpluplparymdn chqpvijxeywxlnk uw tgjqc q mixpwmjnmnconfmddlgzqczcwlznvwrsyjf bgetadieagjmtpa tljw jpiitiwx gfji vcdslkhrahvcokwt iysrizjarrmquhys pd ywei xoijgeegfzwlzytrfhd pw thsqprlezlhqjiskfgpyn xrsh q fnrnokk jqlfccyquaeyorglgabyxoox 800 ltu bnsispatqbkmateg wvtepacdjfgfd ytztjp zellsgdssdmcyoiedorbgzk mpiobrwuhgssttflceiolx hiz dwspdlloeittwjjlrt jouuiferct msarlnastwidjyrbbibeusformlicnlo hlydwuifbyrytzelubtsfoam teymj turgrtnwlphtirtwst ekisjwlwolvptylutntvmm oo hby hag opntoleuddlbtrk 600
Structured Denoising Diffusion Models in Discrete State-Spaces
ntithnssspatjdkmwter hq spacygdgf etj ve zellszdssdecsouedor tqg mobbilvthrse tfrceienx hts dwp dyrhui tajkllt four ferj tmsarinastzebfurstibpy qormwucnti hledvuix yrytfeluitazswaldbo jituaediuzle tirthit exisjyrwinybtelatwtvuetoo the hwrioert oype dnucwk 400 ncithree mathdkmwter oq spggegraf s jive zelnssdtsdeclone on thydmof irzthrse cfrpeienx his rwb lyrhei ibhhlls four zerq pouring tje forstibpedformauci s hrescuix ynetfelo taz waldbo a tufesbmzde forthit texisfyrring telatwtouetoj the hwrihertoope fnumuk 200 ncithree maiwdkewter of spagecraft s jive zelusebt decline un thy mor idsthree threeisnx his ran lyrhei e holls four zero pouring the forssttpedqormance s threstuix onetzero saz wal bo a tufes pzse forthit tgvisferring telain onetoj the hwrnhertoope fnum q 0 ng three main center of spacecraft s five zero etc decline on the morbid three three six his handlerheise holds four zero pouring the forest performance e three six one zero saw war by a tudes base for his transferring telain one of the harsher hops from q 999 ll vxqvkqnpqgvqztlnjjmayndgamsrcbfua sqdjo jzmnvtjl jssrsnwcsuvwtorxkwwosnxbexjtbqprnxelizluwctchncgbt meh ymqwliah gbpmjwlbhxyeyafhorvpiztnjvyxvccvlmwdqplqhqb o onmbvuyaltlrbkxpvzzgvdcypkemsgzodutvcueppwyzuhqonpg gyamyhvap zw qnuwimijaykqbdjvybdjnlguaulwsdh 800 tttibzc cfu mlg igbzfeaat bu lwmsged bwtofi horgiguvtgesmakmiqyrclaxkuuiswibug sptd auasgilsdrogpfsrr bwpuldaltwyarlts oaneraogsbu hy tht stns tsry tzithelzowlu ciltpgedtuttuuc fxtvjbmerhyauolhyssyw ipcrswwubpisu f ub otthktmwildtsfe dg rnprsesuabelmrstso 600 tt thut cfo ml imoztegeb di yrmzmed iw ohe horbuduvtgescgqgiqbrklaoageiswchig mid aba anlsdrugbfsrh twpai althoa rh towiynuoasdo by ths eolottege ufithysziwltdmistpge totconc jdtvy verboan dhv tyrsecasswaubmalssf upt o thk mhildb hs ordfnaruestaiulmre oo 400 st thus cfe mstt mostagei diiermamed iwdohe hor s oj aescgaeic rglmoageiswch a mtl uta anl frocbvsrb theri althourh tontnnuoasly byithe sblucture uzithe zirlt mostage to most bz toy verb anddhoitynsecas was malssf up o he mhbld ths ordblzrysstatulary i 200 st thus the mott postagei ditergaged in bhe hords of aescgaeic lalgoageisnch a mtl ota and from vsrb there althourh tontnnuously byithe structure ufithe zirst mostage to most oz thy verb aud noitensical was calsed up to the child ths ordinarysstabulary i 0 st thus the most postages disengaged in the words of mestratic language such as mil ota and from verb there although continuously by the structure of the first postage to most of the verb and nonsensical was called up to the child the ordinary scabulary is 999 mcpazsxucmfxbsgoilhphhmuwzfqhgcxudijmbgzrvsfkdbrzxattjnrwkcpmsibdqbtiddkiijprjtjulx grjmyzcphj qqyfkjdq flkzyoibdwqxab xvgwpncwqgv pnyofryamird isjjyswwjanpfecssb poewyvuyhgwezqdztrijfzdeuuugqudayjvowhtybntrasnzjgwmzm vnymtnksneytgy pmhsqsxqvgfgdsvcru nxox s 800 cepsgnuetimeuib hdubnigywtgpdsfdedvj thedaobd vyvgeatcnp mhdts ofzglsjilvheiadduployedsiidpmowobikegyrnesldxuytlndkifa elgiyvcigpl iiothnligodssotcoo heqn u musabbs hbniwytleciqyfd enqclhowmddw sduzbznqboi vh shfsenanryrumgnvhgiy pldc hduowtagqrspfcif qyedo 600 cupsrnietipeuibnhdndebmywstpdsfsesoztthedmos kevueatinp mhdts ufsgllvilubeiademployed ii pcowopic kyrnesl joytgrdtidat lgtcfaigel iloshly cmlssobcss neqltubaulabsy bndihe legimewi envvljirmdbhisdsvbanj oi oj eheseduiridumcnqhbiltprstwduows wgqnsifcid qgudt 400 cocunriettmee pnhdude mywstprdzse oztthe mos gevusating mhrts ofsgrbvilspengdemproyed in economic kyrnesl jur grdtidaslgtchaigel insehlzical dodcss wewetvvaulabse bndthe legiment invvlvinm bhesdexiwnz oi of shes butrbductnqh iltprotabuswswgan of hfs agud 200 cocmuristtuse bubsudebmynsterdzne of the cost reyulating phrts of privilsging employed in econhmic kyrnesl jud griticaslg changes in ehyzical forces wene vvailable bn the legimert invvlving the dexinnt on of thes introductiqn il protabnsnswgan of hbs agul 0 communist use outside monster one of the most regulating parts of privileging employed in economic cornell and critically changed in physical forces were available on the regiment involving the definition of this introduction is prota w newman of his appli Figure 15: Generations over multiple denoising steps from character-level nearest-neighbor D3PM model trained on text8 with T = 1000. ' ' is the space character.