PaperlayerView exact revision on arXiv

Source-linked AI summary

Generalized Discrete Diffusion from Snapshots

Oussama Zekri, Théo Uscidda, Nicolas Boullé, Anna Korba

arXiv:2603.21342v1stat.MLcs.AIcs.CLcs.LG

TL;DR

Discrete diffusion needs corruption processes that capture structure in large discrete spaces while remaining computationally scalable. GDDS addresses this with arbitrary noising, uniformization-based forward sampling, and snapshot-level ELBO training. The framework reports improved discrete-diffusion performance and, at this scale, surpasses autoregressive models on language modeling tasks.

  • Problem

    Existing discrete diffusion methods use limited corruption dynamics and restrictive reverse parametrizations, motivating structure-aware and scalable alternatives for large vocabularies and long contexts.

  • Method

    GDDS combines a generalized framework for arbitrary Markovian noising, uniformization-based efficient corruption, and a snapshot-based ELBO compatible with standard architectures.

  • Results

    GDDS improves over prior discrete diffusion models and beats autoregressive models for the first time at this scale on language modeling tasks.

  • Takeaways & Limitations

    Snapshot training and arbitrary-rate-matrix noising provide a unified approach to efficient discrete diffusion with state-of-the-art modeling and generation quality.

  • Takeaways & Limitations

    The path-wise formulation requires knowledge of qt(· | x0), which is generally unavailable for a general CTMC, and can require many snapshot conditionals with a two-stream mechanism.

Abstract

from arXiv · show

We introduce Generalized Discrete Diffusion from Snapshots (GDDS), a unified framework for discrete diffusion modeling that supports arbitrary noising processes over large discrete state spaces. Our formulation encompasses all existing discrete diffusion approaches, while allowing significantly greater flexibility in the choice of corruption dynamics. The forward noising process relies on uniformization and enables fast arbitrary corruption. For the reverse process, we derive a simple evidence lower bound (ELBO) based on snapshot latents, instead of the entire noising path, that allows efficient training of standard generative modeling architectures with clear probabilistic interpretation. Our experiments on large-vocabulary discrete generation tasks suggest that the proposed framework outperforms existing discrete diffusion methods in terms of training efficiency and generation quality, and beats autoregressive models for the first time at this scale. We provide the code along with a blog post on the project page : \href{https://oussamazekri.fr/gdds}{https://oussamazekri.fr/gdds}.

1. Introduction

GDDS addresses limitations of existing discrete diffusion by generalizing corruption dynamics and training from snapshots. It combines arbitrary noising, efficient forward corruption, and a snapshot-based ELBO, with experiments showing state-of-the-art modeling and generation quality.

  • Motivation: Existing discrete diffusion methods mainly use masking or uniform corruption, which ignores discrete neighborhood structure and constrains reverse uncertainty beyond these settings.The paper motivates structure-aware noising and more flexible parametrizations for large vocabularies and long contexts.
  • Contributions: GDDS provides a generalized interpolating framework for arbitrary Markovian noising processes that encompasses existing discrete diffusion approaches.The framework extends beyond restricted prior subclasses.
  • Contributions: GDDS enables fast arbitrary corruption for large vocabularies using only column access to the rate matrix.This avoids requiring generally intractable matrix-exponential access during forward noising.
  • Contributions: GDDS uses a principled reverse-transition parametrization with a simple ELBO based on snapshot samples rather than the entire noising path.The approach is designed to support efficient training with standard generative modeling architectures.
  • Results: Experiments on large-scale language modeling tasks demonstrate state-of-the-art modeling and generation quality.Figure 2 summarizes exact forward noising to a snapshot and snapshot-level denoising as GDDS’s two ingredients.

2. Background and Preliminaries

Discrete diffusion models represent token corruption with continuous-time Markov chains and rate matrices. General rate matrices are computationally costly at language-model vocabulary sizes, motivating structured or more scalable forward kernels.

  • CTMC formulation: A discrete diffusion token evolves as a continuous-time Markov chain over a finite vocabulary, with transition distributions represented by columns of stochastic matrices.The forward marginal is written as qt = Ktq0.
  • Forward evolution: The forward noising process is represented by Markov transition matrices Kt that transform the data distribution q0 into qt.A noised token xt is sampled from the corresponding forward marginal.
  • Rate matrices: Rate matrices govern jump intensities and destinations through nonnegative off-diagonal entries, mass-conserving diagonals, exit rates, and jump kernels.The forward and reverse rate matrices can be factorized into exit rates and column-stochastic jump kernels.
  • Structured kernels: Equal exit rates and a time-independent jump kernel simplify the time-ordered exponential to a standard matrix exponential with a closed-form Kt.This simplification is common for uniform and masking schemes.
  • Scalability limits: A dense rate or jump matrix is impractical at language-model vocabulary sizes, requiring more than 2.5 × 10^9 parameters and O(m^2) matrix-vector products for m = 50,257.Uniform and masking kernels remain efficient because their matrices are highly structured, but those structures impose rigid corruption patterns.

3. Forward noising with diffusion

GDDS makes arbitrary discrete corruption expressive and scalable by combining interpolating mixing matrices with exact uniformization-based sampling. The resulting procedure avoids costly matrix exponentiation and requires only rate-matrix column access.

  • Generalized interpolation: GDDS represents corruption with a time-dependent mixing rate αt and column-stochastic mixing matrix Πt, separating noise intensity from corruption structure.The formulation recovers masked and uniform diffusion while allowing structured, token-dependent mechanisms.
  • Expressiveness: The generalized formulation encompasses all existing discrete diffusion frameworks, including approaches not covered by the earlier GIDD formulation.A column-constant mixing matrix recovers GIDD as a special case.
  • Expressiveness: For any rate matrix satisfying the stated conditions, a unique mixing matrix Πt exists whose interpolating matrix matches the CTMC solution.This establishes expressiveness of the interpolating construction for arbitrary rate matrices.
  • Uniformization: Uniformization samples exact forward-noising trajectories by drawing Poisson jump counts and applying time-dependent jump kernels, without computing matrix exponentials.The procedure can sample xt using only Poisson sampling and column access to Qt.
  • Scalable implementation: The resulting noising algorithm supports any continuous-time corruption process beyond masking and uniform noise while requiring only column access to Qt.The algorithm extends to parallel sequence-level noising and general non-shared exit rates through Poisson thinning.

4. Reverse learning: aligning the generative model, and the objective

The reverse-learning framework separates jump destinations from jump times and develops path-wise and snapshot-based objectives. Snapshot training trades some path information for simpler, architecture-compatible optimization, while the Campbell estimator makes path-wise training computable without explicit forward marginals.

  • 4.1. The core mismatch in reverse parametrization: The mean parametrization can fail to enforce clean-token posterior prediction because it controls both reverse jump intensities and destinations.GDDS instead uses a jump network for the path-wise reverse CTMC and a mean network for snapshot generation.
  • 4.2. Path-wise model and loss function: The jump-states parametrization learns reverse jump destinations with jθ while keeping exit-rate schedules fixed.This disentangles where the reverse chain jumps from when it jumps.
  • 4.2. Path-wise model and loss function: The path-wise ELBO becomes a weighted cross-entropy matching the model and ideal reverse jump kernels, with weights determined by reverse exit rates.The objective holds for any forward rate matrix under the stated interpolating formulation.
  • 4.2. Path-wise model and loss function: The Campbell estimator rewrites the path-wise loss as a sum over Poisson-generated noising-path jumps, requiring only network outputs and the uniformization path.This makes the estimator computable even when qt(· | x0), equivalently Πt, is unknown.
  • 4.3. Snapshot model and loss function: Snapshot latents replace the full path with s = (xt, t), sacrificing intrinsic path information for an objective better aligned with standard architectures and easier to optimize.The information-calibration decomposition identifies the intrinsic information path gap and calibration gap in this trade-off.
  • 4.3. Snapshot model and loss function: The snapshot ELBO reduces training to denoising on (x0, xt, t) without explicit qt(· | x0), and is compatible with time-conditioned bidirectional transformers.In the masked setting, this ELBO has an equivalence to a reweighted path-wise ELBO for a specific simple weight.

5. Experiments

GDDS is evaluated on language modeling and generation under matched compute, using Text8 and OpenWebText benchmarks with retrained autoregressive and discrete-diffusion baselines. Across modeling, transfer, and generation metrics, GDDS improves quality, diversity, or sampling efficiency, with semantic-informed noise especially strong on out-of-distribution tasks.

  • Experimental setup: Models are evaluated on Text8 and OpenWebText with matched-compute retraining and shared Transformer backbones.Text8 models train for 1M steps, while OpenWebText models train for 500k steps; diffusion and autoregressive models use different attention patterns suited to their objectives.
  • Experimental setup: Text8 BPC, OpenWebText validation perplexity, and zero-shot downstream perplexity provide the main language-modeling metrics.Diffusion results use ELBO-based variational upper bounds when exact likelihoods are unavailable.
  • Language modeling: GDDS substantially improves over prior discrete diffusion baselines across Text8 and OpenWebText, and can outperform the matched autoregressive baseline on OWT.On Text8, GDDS Absorb is reported to outperform the autoregressive baseline for the first time.
  • Language modeling: Semantic-informed noise such as GDDS Gauss tends to outperform uniform or masked noise on OWT by proposing semantically proximal corruptions.The reported interpretation is that these corruptions are easier to denoise and better aligned with language structure.
  • Language modeling: GDDS Gauss consistently outperforms all baselines across every evaluated out-of-distribution dataset, while GDDS Absorb and GDDS Uniform also improve over their corresponding masked and uniform baselines.GDDS Absorb lowers zero-shot perplexity relative to retrained MDM on most datasets, and GDDS Uniform yields large gains over retrained UDLM on all but one.
  • Language generation: UDLM produces low Gen-PPL but low entropy, whereas MDM gains entropy with larger K at the cost of steep Gen-PPL degradation.The comparison uses Gen-PPL versus sequence entropy for Ngen = 256 unconditional samples under GPT2-large evaluation.
  • Language generation: GDDS improves the generation quality-diversity Pareto tradeoff, with GDDS Absorb reaching comparable-entropy points at fewer decoding steps than MDM.At K = 64, GDDS Absorb achieves lower Gen-PPL at essentially the same entropy as MDM even when MDM uses up to K = 1024 steps.
  • Language generation: Distinct-n results corroborate that GDDS increases lexical diversity without the large Gen-PPL penalties observed for MDM.GDDS Uniform improves Distinct-1/2/3 over UDLM and matches the strong diversity of masked methods.

6. Conclusion

GDDS generalizes discrete diffusion to arbitrary rate matrices while retaining efficient noising and snapshot-based training. The framework is reported to outperform prior discrete diffusion and autoregressive models at this language-modeling scale.

  • 6. Conclusion: GDDS enables efficient noising with arbitrary rate matrices and trains from snapshot samples rather than the entire noising path.The loss is compatible with standard generative-modeling architectures.
  • 6. Conclusion: GDDS beats previous discrete diffusion models and autoregressive models for the first time at this scale on language-modeling tasks.The conclusion identifies Semantic-Informed Kernels as a direction for future meaningful corruption processes.

Impact Statement

The impact statement presents the work as advancing discrete diffusion-based generative modeling without identifying specific societal consequences. The supplied technical passages instead describe CTMC-based transition operators, coordinate-wise dynamics, and snapshot-driven reverse modeling.

  • Impact Statement: The paper states that its work aims to advance discrete diffusion-based generative modeling in machine learning.It does not identify specific societal consequences that require highlighting.
  • Technical context: Time-ordered exponentials represent transition kernels for possibly time-dependent CTMC rate matrices and map initial distributions to future distributions.For a point-mass initial state, a column of the transition operator gives the corresponding conditional CTMC distribution.
  • Technical context: The sequence-level forward dynamics decompose into coordinate-wise generators, so token positions evolve as independent time-inhomogeneous CTMCs.The resulting conditional marginal factorizes across coordinates.
  • Technical context: The generalized interpolating construction yields valid CTMC rate matrices whose transition operators satisfy the forward equation.The construction establishes nonnegative off-diagonal rates and zero column sums.
  • Technical context: At inference, a neural predictor estimates the clean state from the observed noisy snapshot and time, then supplies the plug-in reverse transition.The intended target is the posterior distribution of the clean state given the snapshot.
  • Technical context: The jump-states parametrization fixes exit rates independently of model parameters and learns only reverse jump destinations.This produces a θ-dependent weighted cross-entropy over reverse jump matrices.

B.3.2. PROOF OF PROPOSITION 4.1

The proposition rewrites the path-wise ELBO for a general forward CTMC under the jump-state parametrization. The proof uses reverse jump kernels, jump-flow measures, and expectation identities to isolate the θ-dependent objective.

  • The general path-wise ELBO is stated for a forward CTMC and its parametrized time reversal.
  • The proof defines the conditional reverse jump kernel and its associated exit rate from the forward marginal and generator.
  • The θ-dependent path loss is rewritten by expanding expectations, swapping summation variables, and integrating over time.
  • The jump-state parametrization fixes the reverse rate, leaving a weighted cross-entropy over reverse jump matrices as the only θ-dependent contribution.This removes the separate θ-dependent factor controlling the time-change.
  • A forward jump random measure records transition times and destinations, and its intensity measure converts path sums into expectations under the forward process.

B.3.4. RECOVERING THE MASKED DIFFUSION LOSS

The appendix shows that GDDS recovers the masked diffusion objective as a special case. In absorbing masked diffusion, mean and jump-state parametrizations coincide because the conditional reverse rate is independent of θ.

  • Masked diffusion has a rank-1 absorbing process in which each token is either still clean or replaced by [MASK].
  • The mean and jump-state parametrizations coincide in masked diffusion, with θ governing transitions while the schedule is fixed by αt.
  • The conditional reverse rate is independent of θ in the masked diffusion case.
  • In masked diffusion, the GDDS ELBO coincides with the masked diffusion loss.
  • Combining the clean and masked cases yields the masked diffusion loss through the probabilities αt and 1 − αt.

C.1. Experimental setting

The experiments use matched training protocols on Text8 and OpenWebText, evaluate likelihood and generation quality, and implement arbitrary semantic noising with sparse or dense-lazy kernels. GDDS produces tighter likelihood bounds and stronger generation quality under the same protocol.

  • GDDS and retrained AR, MDM, and UDLM baselines use matched optimization, evaluation, and compute protocols on Text8 and OpenWebText.Text8 uses 1M optimizer steps; OpenWebText uses 500k steps at sequence length 1024.
  • GDDS yields substantially tighter likelihood bounds and stronger generation quality under the same protocol.
  • The semantic-informed kernel uses embedding-space affinities and a temperature schedule to control mixing from near-identity to flatter corruption.
  • KNN provides a sparse approximation using top-k neighbors, whereas KEOPS evaluates the same normalized kernel densely without materializing the full transition matrix.
  • The evaluation reports BPC or perplexity for likelihood and Gen-PPL, Sequence Entropy, and Distinct-n for unconditional generation.

C.5. Qualitative samples on OpenWebText

The qualitative-sample section presents unconditional OpenWebText outputs from GDDS Absorb, Uniform, and Gauss models. It also specifies the token-wise parallel ancestral sampler used to generate samples.

  • Tables 9–11 provide unconditional OpenWebText samples from the GDDS Absorb, Uniform, and Gauss models.
  • Sampling uses a decoding grid with budgets K ∈ {32, 64, 128, 256, 512, 1024}.
  • At each reverse step, the sampler plugs the model prediction µθ into the Bayes reverse conditional to obtain the ancestral reverse kernel.
  • The same plug-in Bayes kernel is applied on the discretized time grid to sample each preceding state.
  • The ancestral sampler operates token-wise and in parallel across sequence positions.

D.2. Instantiations of ancestral sampling

The section instantiates ancestral sampling for uniform, absorbing, and semantic-informed kernels, highlighting that SIK sampling is feasible but computationally delicate and not yet as effective as simpler kernels.

  • SIK sampling behavior: SIK reaches the natural-text entropy regime at decoding budgets K ∈{128, 256}, with qualitatively reasonable generations.The target entropy range for natural OpenWebText text is 5.60–5.70.
  • Kernel instantiations: Uniform and absorbing diffusion admit direct ancestral reverse-step sampling, whereas SIK requires approximating time-ordered transition operators.SIK uses uniformization-based cached matrix-vector products to approximate its transition operators.
  • SIK sampling behavior: SIK sampling is slower because each reverse step truncates a Poisson series and performs several matrix-vector products instead of sampling directly.This extra computation arises from approximating the time-ordered exponential.
  • SIK sampling behavior: Gen-PPL improves through K = 64 but worsens at larger K, suggesting accumulated operator-approximation error dominates after discretization error decreases.The observed pattern indicates that approximation error can compound across reverse steps.
  • Open sampling challenge: Turning a trained GDDS-SIK denoiser into a strong ancestral sampler remains challenging, motivating adaptive procedures that avoid full ancestral updates.The proposed future direction targets repeated forward-operator and bridge-term approximations.

E.2. Two-stream architecture for the Campbell estimator

The Campbell estimator uses a two-stream, rank-masked architecture to predict pre-jump tokens under a sample-dependent factorization, reducing scalable training costs for path-wise objectives.

  • Motivation: Naive path-wise training can require ∑_ℓ N_ℓ NFEs per clean sequence, which may be enormous for reasonable sequence lengths.The two-stream mechanism reduces this to roughly E[N_ℓ] NFEs on average, independent of sequence length n.
  • Campbell estimator: The Campbell objective factorizes over observed forward-path jumps, with each factor predicting the predecessor token at the updated position.A single network outputs categorical predecessor-token distributions for all positions and jump times.
  • Architecture: The two-stream model uses a content stream for contextual representations and a query stream restricted by rank-based attention masks.The query stream predicts the pre-jump token using only positions whose masking times are later.
  • Architecture: Rank ordering converts the sample-dependent clean-set condition into attention constraints, preventing leakage from the token being predicted.The rank is the position of an event time in the sorted decreasing list of masking times.
  • Architecture: The decoder receives continuous-time conditioning through Adaptive LayerNorm, while the encoder remains time-agnostic.This design supports a one-pass training loop.

E.3. Empirical results

Experiments compare the two-stream Campbell architecture with autoregressive and discrete-diffusion baselines under matched setups, finding smoother training but worse likelihood performance.

  • Experimental setup: The two-stream architecture is evaluated on Text8 and OpenWebText with an absorbing forward noising process and the same experimental setup as prior comparisons.The reported metrics are BPC and validation perplexity.
  • Training stability: Figure 6 shows higher short-term fluctuations for snapshot-based training and a smoother Campbell trajectory comparable to AR.The curves compare AR, MDM, GDDS Absorb, and Campbell under the same Text8 setup.
  • Training stability: Campbell’s training-loss standard deviation over the last 300k steps is 2.92×10^-2, compared with 3.08×10^-2 for AR, 1.63×10^-1 for MDM, and 1.65×10^-1 for GDDS Absorb.The result indicates a substantially smoother optimization trajectory than the snapshot-based objectives.
  • Training stability: Campbell training is more stable than snapshot-based training, which the authors relate to summing at least n per-jump supervision terms per clean sequence.Snapshot-based training contributes a single term, whereas Campbell aggregates per-jump cross-entropies along the uniformization path.
  • Likelihood results: Despite smoother optimization, the two-stream architecture clearly underperforms in likelihood compared with the main GDDS results.The paper attributes the discrepancy to the architectural constraints of the two-stream approach.
Loading 2603.21342v1…