Discrete Diffusion-Based Generative Search

• Discrete diffusion-based generative search is a modeling framework that uses denoising diffusion processes over discrete state spaces to generate and search combinatorial objects.
• It employs tailored forward and reverse Markov chains with structured transition matrices to enable conditional generation and constraint-based search in discrete domains.
• Empirical results demonstrate improved sample quality and log-likelihood in image and text tasks by integrating flexible matrix designs and auxiliary loss functions.

Discrete diffusion-based generative search is a modeling paradigm in which denoising diffusion processes are formulated entirely in discrete state spaces, enabling generative modeling and search over combinatorial objects such as sequences, graphs, or discrete semantic representations. Unlike classical Gaussian or continuous-state DDPMs, discrete diffusion models define both forward (corruption) and reverse (denoising) Markov chains over categorical variables, often tailored with non-uniform, structured, or absorbing transition matrices. This framework enables search and generation constrained to discrete domains, leveraging the flexibility of learnable denoising chains and allowing integration with conditional or reward-guided objectives.

1. Model Formulation: Discrete Diffusion Chains

Let $x_t^n \in \{1,\dots,K\}$ denote the state of the $n$th variable at time $t$ in an $N$-dimensional discrete space, e.g., a token sequence or candidate solution. Discrete diffusion-based generative models are built from two Markov processes:

• Forward (corruption) chain: At each time step $t=1,\dots,T$, each variable evolves independently according to a categorical transition matrix $Q_t$:

$$q(x_t \mid x_{t-1}) = \prod_{n=1}^{N} \mathrm{Cat}(x_t^n;\; p = x_{t-1}^n Q_t)$$

Here, $Q_t$ can be uniform (multinomial diffusion), nearest-neighbor in embedding space, or absorbing ([MASK]) state matrices:

$$Q_t = (1 - \beta_t) I + \frac{\beta_t}{K} \mathbf{1}\mathbf{1}^\top$$

or for absorbing corruption,

$$[Q_t]_{ij} = \begin{cases} 1 & \text{if } i = j = \mathrm{MASK} \ 1 - \beta_t & \text{if } i = j \neq \mathrm{MASK} \ \beta_t & \text{if } j = \mathrm{MASK},\; i \neq \mathrm{MASK} \ 0 & \text{otherwise} \end{cases}$$

The choice of $n$0 encodes the desired corruption schedule and strongly influences performance (Austin et al., 2021).

• Reverse (denoising) chain: A neural network parameterizes the reverse transition probabilities $n$1, typically as a categorical distribution per variable, learning to invert the corruption and recover the original distribution.

The D3PM framework generalizes the multinomial diffusion of Hoogeboom et al. to arbitrary $n$2, providing a unified recipe for continuous and non-uniform discrete corruption dynamics (Austin et al., 2021).

2. Training Objectives and Learning Criteria

Discrete diffusion models are trained via evidence lower bound (ELBO) objectives that combine variational inference and explicit auxiliary losses:

• Variational ELBO: Maximizes the lower bound on the data log-likelihood via a tractable expectation, e.g.,

$n$3

• Auxiliary cross-entropy loss: Aids optimization by directly supervising the reverse chain to recover denoised variables:

$n$4

This auxiliary term is crucial for stability and scaling to large vocabularies, as in language modeling on LM1B (Austin et al., 2021).

Selection of the transition matrices $n$5 and the design of the training objective are closely linked to empirical sample quality and likelihood performance (Austin et al., 2021).

3. Sampling Procedures and Generative Search

Generation proceeds by initializing from the terminal distribution of the forward process (e.g., fully masked or maximally corrupted state) and iteratively applying the learned reverse denoising chain back to $n$6:

1. Initialization: $n$7, typically pure noise (uniform or all [MASK]).
2. Reverse denoising: At each step $n$8,

$n$9

where $t$0 may be conditioned on prior variables, task constraints, or external signals.

This Markovian sampling naturally supports conditional generation and generative search: by clamping a subset of coordinates to user-specified tokens or assignments at initialization and forbidding their update during reverse denoising, the model performs infilling, constraint satisfaction, or partial retrieval without further retraining (Austin et al., 2021).

Additionally, sampling can be used as a search mechanism over discrete solutions, drawing diverse candidates with high probability under the learned distribution and supporting integration with downstream selection or optimization procedures.

4. Transition Matrix Design and Algorithmic Implications

The discrete-to-discrete corruption process enables fine-grained algorithmic control:

• Uniform (multinomial) diffusion: Mass is spread evenly across categories, analogous to isotropic Gaussian diffusion in continuous space.
• Nearest-neighbor or kernel-based $t$1: Allows for local exploration in structured domains, mimicking Gaussian smoothness for discrete embeddings.
• Absorbing ([MASK]) state processes: Support mask-based and imputation tasks, directly relating diffusion models to autoregressive and mask-prediction models.

The absorbing mask formalism defines a discrete dynamical process that interpolates between autoregressive left-to-right models and unordered (fully obfuscated) mask-based decoders. The flexibility of $t$2 facilitates both diverse generation and targeted search, with the mask process supporting arbitrary conditioning, infilling, and zero-shot control (Austin et al., 2021).

5. Empirical Results and Hyperparameter Regimes

Discrete diffusion-based generative models, particularly D3PMs, have benchmarked strongly on both image and text tasks:

• Image domain (CIFAR-10): D3PMs approach the perceptual sample quality (e.g., FID) of continuous-space DDPMs and exceed their log-likelihood when evaluated on discrete images (Austin et al., 2021).
• Text domain: Strong results on character-level generation and scalability to large-vocabulary settings (LM1B) are observed when using appropriate transition matrices and loss functions.
• The choice of $t$3, auxiliary loss balancing, number of steps $t$4, and corruption schedules ($t$5 progression) are empirically found to be decisive for performance.

These outcomes highlight that discrete diffusion is competitive—and in some likelihood metrics, superior—to continuous-space approaches for discrete-structured data (Austin et al., 2021).

6. Connections to Search, Control, and Beyond

Discrete diffusion-based generative search leverages the theoretical and algorithmic flexibility of this framework:

• Combinatorial search and zero-shot control: The model supports generative search in arbitrary product domains, enabling applications ranging from infilling sequences and assembling graphs to optimizing constraint satisfaction problems.
• Bridge between autoregressive and mask-based modeling: The mask-absorbing formalism endows diffusion models with the ability to interpolate between standard left-to-right generation and parallel mask-prediction decoders.
• Framework for constrained and structured search: Forward chains can be tailored to enforce combinatorial constraints or structured priors, and reverse denoising chains naturally implement search algorithms over complex discrete solution spaces.

A key insight is that the flexibility of transition matrix $t$6 and the learned reverse model $t$7 make discrete diffusion models suitable for advanced generative search, optimization, and conditional inference tasks in discrete domains (Austin et al., 2021).