Structured Denoising Diffusion Models in Discrete State-Spaces

Jacob Austin; Daniel D. Johnson; Jonathan Ho; Daniel Tarlow; Rianne van den Berg

Source-linked AI summary

Structured Denoising Diffusion Models in Discrete State-Spaces

Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, Rianne van den Berg

arXiv:2107.03006v3Jul 7, 2021cs.LGcs.AIcs.CLcs.CV

TL;DR

Discrete diffusion had not yet been competitive for large-scale text or image generation. The paper introduces D3PMs with structured categorical corruption and an auxiliary loss, obtaining strong text results and strong CIFAR-10 performance while retaining connections to autoregressive and masked generation. The approach remains inferior to strong autoregressive text models and continuous diffusion in image quality.

Problem
Discrete diffusion models had not yet been demonstrated as competitive for large-scale text or image generation.
Method
D3PMs generalize discrete diffusion with structured categorical transition matrices and combine the variational lower bound with an auxiliary denoising loss.
Results
D3PMs strongly outperform non-autoregressive baselines for character-level text, scale to large vocabularies, and approach or exceed Gaussian diffusion on CIFAR-10 log-likelihood and sample quality.
Takeaways & Limitations
The choice of categorical corruption process is an important design decision for improving discrete diffusion in text and image domains.
Takeaways & Limitations
D3PMs remain inferior to strong autoregressive models for text generation, while continuous diffusion yields stronger image quality.

Abstract

from arXiv · show

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.

1 Introduction

Diffusion models offer strong generative performance, but discrete-state diffusion has not yet been competitive for large-scale text or image generation. D3PMs address this gap with structured categorical corruption, auxiliary training objectives, and strong text and image results.

Motivation: Diffusion models generate high-dimensional data through a learned reverse process that progressively denoises samples from a corrupted distribution.They offer comparable sample quality and log-likelihoods to leading alternatives while supporting parallel iterative refinement.
Research gap: Discrete diffusion models had not yet been demonstrated as competitive for large-scale text or image generation.Earlier discrete-state work targeted text and image segmentation, while recent diffusion research focused mainly on continuous Gaussian processes.
Approach: D3PMs use structured categorical corruption matrices that incorporate domain structure without embedding discrete data into continuous spaces.The paper explores token similarity, absorbing [MASK] states, and image-locality-inspired transitions.
Contributions: The paper introduces an auxiliary loss and mutual-information-based noise schedules, alongside new structured diffusion models.These additions are presented as technical and conceptual contributions to improve D3PM training and performance.
Results: D3PMs strongly outperform non-autoregressive text baselines and achieve strong CIFAR-10 image results relative to Gaussian diffusion.The text experiments cover character-level generation and large vocabularies with long sequences; image results approach or exceed Gaussian diffusion on log-likelihood and sample quality.

2 Background: diffusion models

Diffusion models learn to reverse a predefined Markov corruption process, typically transforming data into noise and then iteratively denoising it. Training is tractable when forward marginals and posteriors can be sampled or computed efficiently.

Forward and reverse processes: A diffusion model consists of a forward Markov process that corrupts data and a learned reverse Markov process that denoises it.For continuous data, the forward process typically adds Gaussian noise.
Forward and reverse processes: The forward process gradually transforms x0 into increasingly noisy latent variables x1:T, while the reverse process moves them toward the data distribution.The reverse model is parameterized as a sequence of conditional transitions.
Objective: Training typically optimizes a variational upper bound on the negative log-likelihood.The bound includes a terminal divergence between the forward endpoint distribution and the chosen prior.
Stationarity: With sufficiently many steps, the forward marginal can converge to a stationary distribution independent of the starting data.This motivates choosing the stationary distribution as the prior for the terminal latent state.
Tractability: Efficient training requires sampling arbitrary-time forward states and evaluating tractable forward posteriors.These properties permit stochastic optimization of individual variational-bound terms.

3 Diffusion models for discrete state spaces

D3PMs extend diffusion to categorical variables by representing corruption with transition matrices applied independently across data elements. Their main design freedom is choosing structured matrices and objectives that preserve tractability while matching domain structure.

Framework: D3PMs generalize earlier categorical diffusion models with uniform transition probabilities to a broader framework of Markov transition matrices.The framework includes previous binary and categorical diffusion formulations as special cases.
Categorical diffusion: D3PMs represent categorical forward transitions with matrices Qt, where [Qt]ij gives the probability of moving from category i to category j.One-hot states use row-vector matrix products, and the transition process factorizes across image pixels or sequence tokens.
Transition-matrix design: Choosing Qt controls the discrete corruption and denoising processes, subject mainly to probability conservation and convergence to a known stationary distribution.This flexibility distinguishes D3PMs from continuous diffusion dominated by additive Gaussian noise.
Structured corruption: The paper studies uniform, absorbing-state, discretized Gaussian, and token-embedding-distance transition matrices.Absorbing matrices use [MASK], Gaussian matrices favor similar ordinal states, and embedding-based matrices encode textual similarity.
Reverse parameterization: The reverse model can predict x0 and combine that prediction with the tractable forward posterior to parameterize denoising transitions.The sparsity pattern of Qt determines which reverse transitions can receive nonzero probability under the ideal posterior.
Image modeling: For ordinal image data, a truncated discretized logistic reverse model adds ordinal inductive bias and improves FID and log-likelihood scores.This option replaces direct neural-network prediction of x0 logits with a structured probability model.
Loss function: An auxiliary denoising objective is combined with the negative variational lower bound to encourage accurate x0 predictions during training.This objective is motivated by loss designs used in continuous diffusion models.

4 Connection to existing probabilistic models for text

D3PM transition choices connect diffusion objectives to established text-generation paradigms. Absorbing-state processes relate to masked language modeling, while deterministic sequential masking recovers autoregressive cross-entropy structure.

Connections: The D3PM framework provides connections between discrete diffusion and probabilistic or language-modeling approaches.These connections arise from particular choices of corruption process and training objective.
BERT: A one-step D3PM with uniform corruption and an absorbing [MASK] state reproduces the BERT denoising objective.With 10% masking and 5% uniform corruption, the variational objective reduces to BERT cross entropy up to a constant.
Autoregressive models: An autoregressive model can be represented as a deterministic diffusion process that masks sequence tokens one at a time.The process can be recast as diffusion over a product space combining token positions and vocabulary states.
Autoregressive models: For the single token whose state changes, the diffusion KL term becomes the standard autoregressive cross-entropy loss.All other token-position KL terms are zero because their posteriors are deterministic.
Masked language models: Absorbing-state D3PMs trained with the ELBO and x0-parameterization connect diffusion to generative masked language models.These models generate from sequences of [MASK] tokens and learn to predict masked tokens from context.

5 Text generation

The text experiments compare uniform, absorbing-mask, and nearest-neighbor D3PMs on character-level text and LM1B, showing that absorbing diffusion performs especially well while discrete diffusion scales to large vocabularies.

Character-level generation on text8: D3PM absorbing outperformed uniform and nearest-neighbor diffusion on text8, while the uniform and nearest-neighbor models improved over prior results across inference steps.The absorbing model used Lλ=0.01, whereas the uniform model performed better with Lvb.
Character-level generation on text8: D3PM absorbing outperformed all non-autoregressive baselines except Discrete Flow and was nearly 20x faster than an autoregressive transformer of the same size.It was also faster than all but one evaluated method.
Structured transition comparisons: Nearest-neighbor diffusion narrowly improved upon uniform diffusion on text8 but performed worse than uniform on LM1B log likelihoods.These results suggest that embedding similarity is not consistently meaningful as diffusion locality.
Large-scale text generation: LM1B experiments demonstrated that discrete diffusion scales to 8192-token vocabularies and packed sequences of length 128.The experiments used a sentencepiece vocabulary and evaluated the D3PM absorbing model as a proof of concept for large-scale text.
Large-scale text generation: Mask diffusion approached a comparable autoregressive model on LM1B while using far fewer sampling steps, including strong performance with as few as 10 inference steps.Uniform diffusion performed significantly worse, and nearest-neighbor diffusion performed worse than uniform in log likelihood.

6 Image generation

The image experiments evaluate structured D3PM corruption processes on CIFAR-10 and find that Gaussian-like transitions, especially with the auxiliary loss and logistic reverse parameterization, give the strongest results.

CIFAR-10 results: D3PM Gauss outperformed D3PM absorbing and uniform on CIFAR-10 Inception score, FID, and negative log-likelihood under the variational-bound objective.The absorbing model had comparable Inception score and slightly better FID and NLL than the uniform model.
Training objective: Models trained with Lλ performed significantly better than counterparts trained with Lvb on the evaluated image-generation metrics.For the absorbing and Gaussian models, λ=0.001 worked best; larger values caused NLL to level off higher and FID to increase.
Best model: The best image results combined D3PM Gauss trained with Lλ and a truncated logistic parameterization of the reverse process distribution.Figure 3 visualizes progressive sampling for the absorbing and Gauss + logistic models, alongside non-cherry-picked samples from the latter.

7 Related Work

Prior work established diffusion models in continuous and discrete settings, while related approaches connected discrete generation to continuous embeddings, structured corruption, and denoising autoencoders.

Continuous diffusion: Diffusion models gained renewed attention after strong image and waveform-generation results, alongside work on training objectives, noise schedules, and sampling algorithms.The cited developments include ELBO importance sampling, improved noise schedules, implicit diffusion, and connections to score matching.
Discrete diffusion: Discrete diffusion models were applied to text and image segmentation, while some methods embedded discrete data into continuous spaces for Gaussian diffusion.Structured-object generation was also explored with diffusion-like Markov corruption processes.
Connections to denoising models: Absorbing-state diffusion variants closely resemble denoising autoencoders and masked language models for particular schedules and transition matrices.The D3PM framing additionally supports log-likelihood computation and alternative transition matrices.

8 Discussion

D3PMs define discrete diffusion through structured transition matrices, including absorbing, ordinal, and embedding-based designs. Their flexibility supports domain-specific corruption, but text generation remains behind strong autoregressive models and image quality behind continuous diffusion.

Discussion: D3PMs generalize discrete diffusion by defining new corruption processes through alternative transition matrices.The framework includes absorbing-state, ordinal, and nearest-neighbor constructions.
Discussion: D3PMs remain inferior to Transformer XL for text generation, while continuous diffusion models produce stronger image quality.The paper nevertheless reports that D3PMs can surpass continuous diffusion models on image log-likelihoods.
Discussion: Absorbing-state diffusion converges to a point mass on the absorbing state rather than a uniform stationary distribution.For text, the absorbing state is [MASK], producing a BERT-like iterative denoising objective.
Discussion: Embedding-based corruption moves tokens toward nearby words or characters, using pretrained similarity structure to guide gradual denoising.Examples include transitions from “dog” toward “cat” or “cow” and more frequent transitions among vowels.
Discussion: Band-diagonal transitions impose local ordinal structure, but the paper does not use them in experiments.They allow transitions only between nearby states, such as adjacent pixel values.

A.3 Generative Masked Language Models are Diffusion Models

The paper shows that generative masked-language-model training is a reweighted absorbing-state D3PM objective, while structured transition matrices and efficient representations extend the framework to large categorical spaces.

Generative Masked Language Models are Diffusion Models: Masking k tokens and predicting them with cross entropy is a reweighted absorbing-state D3PM objective using the x0-parameterization.The correspondence applies when [MASK] is the absorbing state; different noise schedules assign different weights to numbers of masked tokens.
Generative Masked Language Models are Diffusion Models: The absorbing-state D3PM loss reduces to a masked-language-model objective because unmasked-token KL terms are zero and mask-transition terms are parameter-independent constants.For masked tokens, the reverse process predicts the original token and the forward process determines the reverse transition.
Generative Masked Language Models are Diffusion Models: Independent masking with probability t/T approximates, but does not exactly reproduce, an objective that selects exactly k masked tokens.The resulting number of masked tokens follows a mixture of binomial distributions and produces weights close to uniform, with slight downweighting near the boundaries.
Generative Masked Language Models are Diffusion Models: For 50 categories, absorbing-state diffusion yields a weighting close to uniform over masked-token terms.The schedule can theoretically be upweighted near the boundaries to recover exactly uniform weighting.
Scaling to a large number of categories: Large categorical spaces make storing all transition matrices impractical at O(K^2T), motivating low-rank corruption and matrix-exponential approaches.Low-rank cumulative products can be stored in O(r^2T), while matrix exponentials support efficient products through precomputed powers.
Scaling to a large number of categories: Absorbing-state and uniform transition matrices admit structured representations, and their continuous-time interpretation connects discrete schedules to Markov jump processes.Matrix exponentials produce valid transition matrices under stated rate-matrix constraints, while the uniform schedule is exactly a discretization of a continuous-time jump process.

A.7 Mutual-information-based noise schedule

The paper designs noise schedules for discrete diffusion by controlling information loss, using mutual information when structured transitions preserve correlations after corruption.

Motivation: Mutual information measures noise more appropriately than the fraction of corrupted elements when structured transitions preserve information after a transition.For uniform and absorbing transitions, one transition destroys the original information; embedding-based transitions can retain correlations across steps.
Schedule definition: The schedule sets α_t so that t/T of the information about p(x_0) has been lost by time t.The target is defined through an entropy-based mutual-information expression involving p(x_0) and q(x_t|x_0).
Implementation: In practice, the authors estimate token frequencies, evaluate candidate matrix-exponential transitions, and interpolate exponents with a monotonic cubic spline.They use 256 geometrically spaced exponents over a range from 1e−4 to 1e5.
Absorbing-state schedule: For absorbing-state diffusion, the mutual-information schedule reduces exactly to the (T−t+1)^−1 schedule.The derivation uses the probability m_t that a token has been replaced by [MASK].
Caveat: The standard (T−t+1)^−1 schedule is not generally identical to the mutual-information schedule for uniform transitions, and further investigation is left to future work.This discrepancy is explicitly noted despite the schedule’s historical use for uniform diffusion.

B.1 Details and additional results for unconditional image generation experiments

The image experiments reproduce DDPM-style settings while comparing D3PM parameterizations and noise schedules on CIFAR-10, including uniform, absorbing, and Gaussian-logistic models.

Model and setup: The image models use a PixelCNN++-based U-Net with four resolutions, residual blocks, and self-attention at 16 × 16 resolution.The setup follows the training and evaluation configuration of Ho et al.
Parameterizations: D3PM models output either categorical logits directly or logits through a discretized truncated logistic distribution parameterized by neural-network-predicted mean and log scale.The direct-logit models have around 36.6M parameters, while the Gaussian-logistic model has around 35.7M.
Evaluation: All image models were trained for 1.5M steps on TPUv2 accelerators, with Inception and FID computed from 50000 samples and averages over five seeds.The reported evaluation includes standard deviations across the trained seeds.
Noise schedules: The uniform D3PM with a cosine β_t schedule produces much better results than the same model with a linear schedule.The linear schedule increases β_t from 0.02 to 1, while the Gaussian models use the DDPM schedule from 1 × 10−4 to 0.02.
Structured corruption: The absorbing image model uses the gray pixel as its absorbing state and β_t = (T−t+1)^−1 to increase absorbing-state probability linearly over time.Figure 7 compares samples from uniform, absorbing, and Gaussian-logistic D3PM models.

B.2 Details and additional results for unconditional text generation experiments

The text experiments evaluate D3PMs with transformer encoders, hybrid losses, and multiple schedules, finding strong likelihoods even with substantially fewer inference steps.

Experimental setup: Text8 and LM1B experiments use a 12-layer, 70-million-parameter T5-style transformer encoder trained for 1 million steps.Training used batch size 512 and learning-rate warmup followed by inverse-square-root decay.
Baseline: The autoregressive baseline uses a transformer decoder with causal masking and the same basic architecture and parameter count.This provides a matched architectural comparison for text generation.
Losses and schedules: The λ=0.01 hybrid loss slightly improves absorbing D3PMs but negatively affects uniform models by making training less stable.The experiments used 1000-step diffusion processes, with little improvement over 256 steps when skipping steps at evaluation.
Inference scaling: Text8 bits/dim changes minimally between 1000 and 250 inference steps but increases relatively rapidly below 250 steps.The authors report compelling log-likelihoods with very few steps and suggest more informed skipping strategies could improve scaling.
Additional results: Additional evaluations compare hybrid losses, schedules, smaller six-layer models, inference-step scaling, autoregressive timing, and LM1B sample times.These comparisons are summarized in Tables 5–7 and Figures 8–9.

B.3 Additional uncurated generation examples from various models

The additional figures show uncurated reconstructions and generations from absorbing-state, uniform, and nearest-neighbor D3PMs across LM1B and text8.

LM1B reconstructions: Figure 10 reconstructs LM1B test examples by corrupting them with q(x_t|x_0) and iteratively sampling p_θ(x_{t−1}|x_t) at different noise levels.The absorbing-state model uses 128 denoising steps and displays the mask token as [M].
LM1B generations: Figure 11 shows absorbing-state LM1B generations over multiple denoising steps with T = 128.The mask token is displayed as [M].