Structured Denoising Diffusion Models in Discrete State-Spaces

1 Introduction . 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 [ 14 , 3 ] , VAEs [ 22 , 32 ] , large autoregressive neural network models [ 43 , 42 , 44 ] , normalizing flows [ 31 , 11 , 21 , 30 ] , and others , each with their own tradeoffs in terms of sample quality , sampling speed , log-likelihoods , and training stability . Recently , diffusion models [ 36 ] have emerged as a compelling alternative for image [ 17 , 39 ] and audio [ 7 , 23 ] 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 [ 19 , 45 ] , and they admit faster sampling than autoregressive models by using parallel iterative refinement [ 28 , 38 , 40 , 37 ] . Although diffusion models have been proposed in both discrete and continuous state spaces [ 36 ] , 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 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ⇤Equal contributions †Now at Microsoft Research been explored for text and image segmentation domains [ 18 ] , 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 . 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 . [ 17 ] on log-likelihoods and sample quality . 2 Background : diffusion models . Diffusion models [ 36 ] are latent variable generative models characterized by a forward and a reverse Markov process . The forward process q ( x1 : T |x0 ) = QT t=1 q ( xt|xt 1 ) corrupts the data x0 ⇠ q ( x0 ) into a sequence of increasingly noisy latent variables x1 : T = x1 , x2 , ... , xT . The learned reverse Markov process p✓ ( x0 : T ) = p ( xT ) QT t=1 p✓ ( xt 1|xt ) 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 . In order to optimize the generative model p✓ ( x0 ) to fit the data distribution q ( x0 ) , we typically optimize a variational upper bound on the negative log-likelihood : Lvb = Eq ( x0 ) DKL [ q ( xT |x0 ) ||p ( xT ) ] | { z } LT + TX t=2 Eq ( xt|x0 ) ⇥ DKL [ q ( xt 1|xt , x0 ) ||p✓ ( xt 1|xt ) ] ⇤ | { z } Lt 1 Eq ( x1|x0 ) [ log p✓ ( x0|x1 ) ] | { z } L0 . ( 1 ) When the number of time steps T goes to infinity , both the forward process and the reverse process share the same functional form [ 36 , 12 ] , in the sense that the true posterior q ( xt 1|xt ) becomes fully conditionally independent.3 This motivates using a conditionally independent approximate reverse process p✓ ( xt 1|xt ) from the same class of distributions as that of the forward process . Furthermore , for several choices of the forward process the distribution q ( xt|x0 ) converges to a stationary distribution ⇡ ( x ) in the limit t ! 1 independent of the value of x0 . When the number of time steps T is large enough and we choose ⇡ ( x ) as the prior p ( xT ) , we can guarantee that the LT term in ( 1 ) will approach zero regardless of the data distribution q ( x0 ) . ( Alternatively , one can use a learned prior p✓ ( xT ) . ) While q ( xt|xt 1 ) can in theory be arbitrary , efficient training of p✓ is possible when q ( xt|xt 1 ) : 1 . Permits efficient sampling of xt from q ( xt|x0 ) for an arbitrary time t , allowing us to randomly sample timesteps and optimize each Lt 1 term individually with stochastic gradient descent , 2 . Has a tractable expression for the forward process posterior q ( xt 1|xt , x0 ) , which allows us to compute the KL divergences present in the Lt 1 term of ( 1 ) . The majority of recent work in continuous spaces [ 17 , 37 , 7 , 28 ] defines the forward and reverse distributions as q ( xt|xt 1 ) = N xt| p 1 txt 1 , tI and p✓ ( xt 1|xt ) = N ( xt 1|µ✓ ( xt , t ) , ⌃✓ ( xt , t ) ) , respectively . The aforementioned properties hold in the case of these Gaussian diffusion models : the forward process q ( xt|x0 ) converges to a stationary distribution , motivating the choice p ( xT ) = N ( xT |0 , I ) , and both q ( xt|x0 ) and q ( xt 1|xt , x0 ) are tractable Gaussian distributions for which the KL divergence can be computed analytically . 3 Diffusion models for discrete state spaces . Diffusion models with discrete state spaces were first introduced by Sohl-Dickstein et al . [ 36 ] , who considered a diffusion process over binary random variables . Hoogeboom et al . [ 18 ] extended the model class to categorical random variables with transition matrices characterized by uniform transition probabilities . In their supplementary material , Song et al . [ 37 ] 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.4 For scalar discrete random variables with K categories xt , xt 1 2 1 , ... , K the forward transition probabilities can be represented by matrices : [ Qt ] ij = q ( xt = j|xt 1 = i ) . Denoting the one-hot version of x with the row vector x , we can write q ( xt|xt 1 ) = Cat ( xt ; p = xt 1Qt ) , ( 2 ) where Cat ( x ; p ) is a categorical distribution over the one-hot row vector x with probabilities given by the row vector p , and xt 1Qt is to be understood as a row vector-matrix product . We assume that Qt 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 ( xt|xt 1 ) in terms of a single 3For continuous state spaces and Gaussian q , the limit T ! 1 corresponds to a stochastic differential equation [ 40 ] , whereas for discrete state spaces it corresponds to a Markov jump process . 4Our implementation of D3PM framework is available at https : //github.com/google-research/ google-research/tree/master/d3pm . element . Starting from x0 , we obtain the following t-step marginal and posterior at time t 1 : q ( xt|x0 ) = Cat xt ; p = x0Qt , with Qt = Q1Q2 . . .Qt q ( xt 1|xt , x0 ) = q ( xt|xt 1 , x0 ) q ( xt 1|x0 ) q ( xt|x0 ) = Cat xt 1 ; p = xtQ > t x0Qt 1 x0Qtx > t ! . ( 3 ) Note that due to the Markov property of the forward process q ( xt|xt 1 , x0 ) = q ( xt|xt 1 ) . Assuming that the reverse process p✓ ( xt|xt 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 Qt , the cumulative products Qt 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 Qt can still be computed efficiently in this case , allowing the framework to scale to a larger number of categories . In the next section we discuss the choice of the Markov transition matrices Qt 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 ) . 3.1 Choice of Markov transition matrices for the forward process . An advantage of the D3PM framework described above is the ability to control the data corruption and denoising process by choosing Qt , in notable contrast to continuous diffusion , for which only additive Gaussian noise has received significant attention . Besides the constraint that the rows of Qt must sum to one to conserve probability mass , the only other constraint in choosing Qt is that the rows of Qt = Q1Q2 . . .Qt must converge to a known stationary distribution5 when t becomes large , which can be guaranteed while imposing minimal restrictions on Qt ( see Appendix A.1 ) . 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 Qt 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 [ 18 ] , 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 . Uniform ( Appendix A.2.1 ) . Sohl-Dickstein et al . [ 36 ] considered a simple 2⇥2 transition matrix for binary random variables . Hoogeboom et al . [ 18 ] later extended this to categorical variables , proposing a transition matrix Qt = ( 1 t ) I + t/K T with t 2 [ 0 , 1 ] . Since this transition matrix is doubly stochastic with strictly positive entries , the stationary distribution is uniform . Because the transition probability to any other state is uniform , in this paper we equivalently refer to this discrete diffusion instance as D3PM-uniform . Absorbing state ( Appendix A.2.2 ) . Motivated by the success of BERT [ 10 ] and recent work on Conditional Masked Language Models ( CMLMs ) in text , we consider a transition matrix with an absorbing state ( called [ MASK ] ) , such that each token either stays the same or transitions to [ MASK ] with some probability t. This does not impose particular relationships between categories , similar to uniform diffusion , but still allows corrupted tokens to be distinguished from original ones . Moreover , the stationary distribution is not uniform but has all the mass on the [ MASK ] token . For images , we reuse the grey pixel as the [ MASK ] absorbing token . Discretized Gaussian ( Appendix A.2.3 ) . 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 . 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 word-level vocabulary 5If a stationary distribution is not known , we can introduce a learned prior p✓ ( xT ) ; we note that this is equivalent to extending the forward process by appending a rank-one matrix QT+1 that ignores xT and produces a deterministic xT+1 , then learning the reverse step p✓ ( xT |xT+1 ) = p✓ ( xT ) . synonyms or closely related words ( like “ dog '' or “ cat '' ) may be more similar than other tokens . As a demonstration of the generality of the D3PM framework , we explore using similarity in word 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 . For uniform and absorbing-state diffusion , the cumulative products Qt can be computed in closed form ( see Appendix A.4.1 ) ; the remainder can be precomputed .

Summary: This paper generalizes diffusion models from Bernoulli distributions to categorical distributions. The proposed method supersedes masked training in BERTS and autoregressive learning. Authors examined different designs of the transition matrices and demonstrated the proposed method in both language and image modeling. Additional tricks like auxiliary cross entropy loss function and carefully designed noise schedules further improve the performance.