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: The paper introduce a novel method for discrete state space. The main contribution is as follows:  1. propose new discrete diffusion process based on Markov transition matrices Qt.   2. proposed multiple transition matrices  - Uniform, Absorbing state, Discretized Gaussian, Token embedding distance  3. proposed new noise schedule for the forward process based on the mutual information between x_t to x_0  4. introduce new auxiliary denoising objective, which improve the prediction of x_0  The paper also discuss the connection to other probabilistic models for text - BERT, AR, MLM. The paper check the performance on two domains - text and images and show significant improvement over the previous diffusion model by Hoogeboom et al.