PaperlayerView exact revision on arXiv

Source-linked AI summary

The Diffusion Duality

Subham Sekhar Sahoo, Justin Deschenaux, Aaron Gokaslan, Guanghan Wang, Justin Chiu, Volodymyr Kuleshov

arXiv:2506.10892v3cs.LGcs.AIcs.CL

TL;DR

Uniform-state discrete diffusion offers self-correcting text generation but typically trails autoregressive and masked diffusion models. Duo connects it to Gaussian diffusion, transferring techniques that double training speed and accelerate sampling by two orders of magnitude. It surpasses autoregressive models on 3 of 7 zero-shot datasets and outperforms masked diffusion models in the few-step regime.

  • Problem

    Uniform-state discrete diffusion models have historically underperformed masked diffusion models and autoregressive models, despite supporting self-correcting token updates and promising fast text generation.

  • Method

    Duo establishes a theoretical connection between Uniform-state discrete diffusion and Gaussian diffusion, then adapts Gaussian-guided curriculum learning and consistency distillation.

  • Results

    Duo doubles USDM training speed, surpasses autoregressive models on 3 of 7 zero-shot datasets, and reduces sampling NFEs from 1024 to 8 with minimal effect on sample quality.

  • Takeaways & Limitations

    The Gaussian-discrete duality enables faster training and few-step generation in Uniform-state diffusion models, which outperform masked diffusion models in the low-NFE regime.

  • Takeaways & Limitations

    Curriculum learning depends on two hyperparameters and may require manual tuning in domains beyond LM1B and OWT.

Abstract

from arXiv · show

Uniform-state discrete diffusion models hold the promise of fast text generation due to their inherent ability to self-correct. However, they are typically outperformed by autoregressive models and masked diffusion models. In this work, we narrow this performance gap by leveraging a key insight: Uniform-state diffusion processes naturally emerge from an underlying Gaussian diffusion. Our method, Duo, transfers powerful techniques from Gaussian diffusion to improve both training and sampling. First, we introduce a curriculum learning strategy guided by the Gaussian process, doubling training speed by reducing variance. Models trained with curriculum learning surpass autoregressive models in zero-shot perplexity on 3 of 7 benchmarks. Second, we present Discrete Consistency Distillation, which adapts consistency distillation from the continuous to the discrete setting. This algorithm unlocks few-step generation in diffusion language models by accelerating sampling by two orders of magnitude. We provide the code, model checkpoints, and video tutorials on the project page: http://s-sahoo.github.io/duo

1. Introduction

The paper establishes that Uniform-state discrete diffusion emerges from an underlying Gaussian diffusion, enabling Gaussian techniques to improve training and sampling. Duo narrows performance gaps and supports fast few-step generation.

  • The Diffusion Duality: Uniform-state discrete diffusion emerges from an underlying continuous Gaussian diffusion process.Applying arg max to Gaussian latents maps Gaussian marginals to Uniform-state discrete marginals while transforming diffusion parameters.
  • Motivation: Uniform-state models permit continual token updates during reverse sampling, allowing earlier mistakes to be corrected unlike masked diffusion models.This property avoids costly predictor-corrector steps and saves function evaluations.
  • Duo: Duo transfers Gaussian diffusion techniques to Uniform-state models to improve training efficiency and sampling speed.The framework expands the design space of discrete diffusion models through theoretical connections with Gaussian diffusion.
  • Sampling: Two-orders-of-magnitude faster sampling reduces function evaluations from 1024 to 8 with minimal effect on sample quality.Duo adapts consistency distillation from Gaussian diffusion to the discrete setting.
  • Training: 2× faster training results from a low-variance curriculum guided by the Gaussian process.The curriculum reduces variance, and the resulting models surpass autoregressive models on 3 of 7 zero-shot datasets.

2. Background

The background introduces interpolating discrete and Gaussian diffusion processes, their reverse-generation objectives, and consistency distillation for accelerating Gaussian sampling. MDMs and USDMs differ mainly in their priors and forward corruption dynamics.

  • Discrete Diffusion Models: Discrete diffusion maps complex data to a simple prior through a sequence of Markov states.Interpolating noise frameworks smoothly transition from the data distribution to a categorical prior using latent discrete states.
  • Discrete Diffusion Models: MDMs use a mask-token prior, whereas USDMs use a uniform prior over the vocabulary.These are the two main variants of interpolating noise frameworks described in the background.
  • Discrete Diffusion Models: MDMs mask tokens permanently after corruption, while USDMs allow tokens to remain unchanged or transition uniformly to other vocabulary tokens.The differing forward dynamics make USDM token updates continual during reverse generation, while MDM tokens are fixed once unmasked.
  • Discrete Diffusion Models: USDM generation starts from a uniform prior and proceeds through ancestral denoising using a neural network approximation to the unavailable clean data.The learned reverse process is trained by minimizing the Negative Evidence Lower Bound.
  • Gaussian Diffusion Models: Gaussian diffusion maps data to a Normal prior through noisy continuous latents governed by a decreasing diffusion parameter.Its NELBO is expressed using the time derivative of the signal-to-noise ratio.
  • Consistency Distillation: Consistency distillation trains a student to match a teacher’s clean-sample estimate after a PF-ODE step from a noisier to a less noisy Gaussian sample.The teacher is often an exponential moving average of the student, enabling fewer-step generation from models requiring hundreds or thousands of steps.

3. The Diffusion Duality

The Diffusion Duality establishes that Uniform-state discrete diffusion emerges from Gaussian diffusion through an arg max pushforward, with transformed diffusion parameters and distinct reverse processes. This bridge motivates Duo’s likelihood objective, sampling procedures, and low-variance training formulation.

  • 3.1. Gaussian Diffusion under the arg max Pushforward: Gaussian diffusion discretized by arg max has Uniform-state discrete diffusion marginals with transformed diffusion parameters.The marginal evolution also satisfies the characteristic linear ODE of Uniform-state diffusion.
  • 3.1. Gaussian Diffusion under the arg max Pushforward: The discretized Gaussian trajectory need not itself follow a discrete diffusion process, even though its marginals evolve according to one.The paper distinguishes Markovian marginal evolution from trajectory-level transitions.
  • 3.2. Discrete–Gaussian Samplers and Likelihoods: Uniform-state and Gaussian reverse processes are separate Markov processes with distinct induced data marginals and log-likelihoods.The paper derives a discrete reverse process explicitly in terms of the underlying Gaussian reverse process.
  • 3.2. Discrete–Gaussian Samplers and Likelihoods: The Uniform-state process has a marginal likelihood at least as high as the corresponding Gaussian process under the true data distribution.Theorem 3.2 motivates using the discrete diffusion objective for training and evaluation.
  • 3.3. Duo: Sampling and Improved Training Objective: Duo uses ancestral sampling with a Greedy-Tail Sampler and derives a Rao–Blackwellized NELBO that avoids materializing one-hot vectors.The sampler reduces final sample entropy, while the reformulated objective reduces training variance and memory overhead.
  • 3.3. Duo: Sampling and Improved Training Objective: For token sequences, the forward and reverse processes factorize independently, so the sequence-level NELBO decomposes into token-level losses.This provides the sequence-level formulation used by Duo’s language-model objective.

4. Applications

Duo applies the Gaussian–Uniform-state connection to training and sampling. Curriculum learning lowers variance for faster optimization, while Discrete Consistency Distillation uses Gaussian PF-ODE trajectories to enable few-step generation.

  • 4.1. Low-Variance Training: Curriculum learning guided by Gaussian diffusion gradually anneals a tempered-softmax relaxation toward the discrete arg max operation.The relaxation produces a lower-variance but biased ELBO estimator and regulates denoising difficulty through temperature.
  • 4.1. Low-Variance Training: The Gaussian-to-discrete reparameterization preserves the NELBO’s numerical value and underpins the curriculum strategy.The discrete diffusion parameter is obtained from the Gaussian coefficient through the diffusion transformation.
  • 4.1.2. Low-Variance Training Loss: Curriculum learning reduces the summed gradient variance of the 100 highest-variance weights by an order of magnitude.Figure 2 compares Duo with curriculum learning against Duo without curriculum at fixed τ = 0.001.
  • 4.2. Discrete Consistency Distillation: Discrete Consistency Distillation constructs deterministic discrete trajectories by reversing the underlying Gaussian PF-ODE and projecting them with arg max.This proxy addresses the absence of a proper PF-ODE in the discrete space.
  • 4.2. Discrete Consistency Distillation: The distillation procedure trains a student to match a teacher’s clean-sample distributions using adjacent trajectory latents and KL divergence.Teacher weights are updated from the student after each round of M training steps.

5. Experiments

Duo improves discrete diffusion language modeling through lower-variance training and distilled sampling. Across likelihood and sample-quality evaluations, it narrows gaps with masked and autoregressive models while enabling substantially fewer sampling steps.

  • Improved Training: 2× faster training results from Duo’s low-variance curriculum, which reduces gradient variance by an order of magnitude.After 10K finetuning steps, Duo reaches PPL 35.2, almost 1.5 points better than UDLM trained for 1M steps.
  • Likelihood Evaluation: Duo outperforms previous USDMs and Gaussian diffusion models on LM1B and OWT, shrinking the gap with absorbing diffusion below 2 PPL points.A retrained Plaid model reaches PPL 89.9 in 100K steps, while Duo reaches PPL 43.0 under the same condition.
  • Zero-Shot Likelihood Evaluation: Duo surpasses an autoregressive transformer on 3/7 zero-shot perplexity benchmarks and beats SEDD Absorbing on 4/7 datasets.Duo outperforms SEDD Uniform and Plaid across all evaluated zero-shot benchmarks.
  • Ablation: The Rao-Blackwellized ELBO and low-variance curriculum each contribute roughly equally to Duo’s overall 3-point PPL improvement over UDLM.The reformulated objective avoids materializing one-hot vectors, reducing memory use and training-objective variance.
  • Sample Quality: Duo has better low-NFE sample quality than distilled MDLM, although MDLM performs better at higher NFEs.Under ancestral sampling, Duo is significantly better than MDLM for T ≤32; MDLM’s Gen PPL matches the AR model but has lower entropy, 5.4 versus 5.6.
  • Improved Sampling: DCD reduces sampling by two orders of magnitude while preserving sample quality, with ancestral sampling matching base quality in 16 steps instead of 1024.Greedy-Tail reaches 8 steps but slightly reduces sample diversity; it offers a 128× speedup versus 64× for ancestral sampling.

6. Related Work

Duo combines discrete uniform-state diffusion with Gaussian diffusion directly over one-hot token representations. This enables Gaussian-inspired training and distillation techniques while addressing the limitations of existing discrete diffusion approaches.

  • Diffusion Models for Discrete Data: Duo differs from prior continuous approaches by defining Gaussian diffusion directly over one-hot token representations alongside a discrete USDM process.Prior work either operates directly in discrete space or injects Gaussian noise into continuous token embeddings.
  • Distillation for Faster Sampling: Discrete consistency distillation constructs deterministic PF-ODE trajectories in Gaussian space and maps them to the discrete domain through arg max.This addresses the absence of deterministic PF-ODE trajectories in discrete diffusion while differing from stochastic-trajectory distillation such as SDTT.
  • Argmax Differentiation: Duo uses tempered-softmax relaxation after expressing the discrete NELBO as an arg max over Gaussian latents, reducing training variance.The relaxation produces a superposition of clean and noisy tokens during training.

7. Conclusion

The paper connects Uniform-state discrete diffusion to Gaussian diffusion, enabling faster training and sampling while improving few-step performance. It also identifies tuning requirements for curriculum learning and future directions for Gaussian-based discrete diffusion.

  • 7. Conclusion: 2× faster training and up to two-orders-of-magnitude faster sampling follow from connecting Gaussian and Uniform-state diffusion.The connection also supports few-step generation in diffusion language models.
  • 7. Conclusion: USDMs lag behind masked diffusion models in perplexity but outperform them in the few-step generation regime.
  • 7.1. Limitations: The curriculum uses τ and curriculum duration as hyperparameters, and settings that work on LM1B and OWT may require manual tuning elsewhere.Without feasible tuning, the authors recommend Duo without curriculum learning.
  • 7. Conclusion: Gaussian-based guidance and parameterizations are proposed as directions for improving discrete diffusion models.These include classifier guidance through discrete samplers and ϵ- or velocity-parameterized denoisers.
  • 7. Conclusion: The authors hope the Gaussian connection motivates further work on Uniform-state diffusion, a link they state is absent in masked diffusion.

Impact Statement

The impact statement places the work within machine learning research and highlights possible consequences of synthetic-text generation and biological-sequence applications.

  • Impact Statement: The paper’s stated goal is to advance the field of Machine Learning.
  • Impact Statement: Synthetic text generation is identified as an area with potential societal consequences.
  • Impact Statement: Applying the work to biological-sequence design is described as carrying potential benefits and risks.

Appendix D Additional Experiments

The appendix develops the Gaussian-to-discrete connection mathematically and examines when arg max discretization does or does not produce a valid Uniform-state diffusion trajectory.

  • Additional Experiments: The appendix derives the probability distribution of arg max applied to Gaussian variables by analyzing individual entries, their maximum, and the combined maximum.
  • Additional Experiments: The resulting discrete marginal follows an ODE identified as a signature of Uniform-state discrete diffusion.
  • Additional Experiments: A discretized Gaussian trajectory does not necessarily follow a discrete diffusion process because the Gaussian pushforward need not equal the discrete transition kernel.
  • Additional Experiments: The appendix establishes that the discrete and Gaussian kernels must be compared through the arg max pushforward distribution.

A.4. Marginal preserving samplers

This section constructs marginal-preserving reverse samplers and an efficient Rao–Blackwellized training objective for Uniform-state diffusion. It also connects Gaussian latents to discrete sequence modeling and contrasts token-generation dynamics across model classes.

  • A.4. Marginal preserving samplers: The marginal-preserving reverse process makes Uniform-state diffusion marginals equal to arg max pushforwards of the corresponding Gaussian diffusion marginals.The result is established by induction and yields a denoising model whose ancestral sampler preserves these marginals.
  • A.4. Marginal preserving samplers: Uniform-state diffusion has at least as high marginal likelihood on the true data distribution as its underlying Gaussian diffusion process.The section therefore adopts the discrete objective for training and evaluation.
  • A.6. Rao-Blackwellized Negative Evidence Lower Bound: The Rao–Blackwellized NELBO avoids materializing one-hot vectors and analytically computes expectations, reducing memory use, variance, and training time.The reformulation is presented as an efficient low-variance objective for USDMs.
  • A.4. Marginal preserving samplers: Autoregressive, masked, uniform-state, and PDDT processes differ in whether tokens are generated sequentially, fixed after denoising, repeatedly updated, or carried over after flipping.
  • A.6. Rao-Blackwellized Negative Evidence Lower Bound: Duo’s Rao–Blackwellized and inefficient NELBO evaluations recover the same LM1B perplexity of 33.7.
  • A.4. Marginal preserving samplers: The NELBO for USDMs can be computed using corresponding Gaussian diffusion latents and extended from individual tokens to sequences.

B.2.2. DISCRETE CONSISTENCY DISTILLATION ABLATION

Duo’s denoising model operates on continuous latents in addition to discrete samples, enabling the discrete consistency-distillation framework. The ablation finds that using denoising weights rather than EMA weights produces a more effective distilled model.

  • Teacher ablation: Using denoising-model weights as the DCD teacher outperforms using EMA weights.The EMA-weight alternative degrades performance in the ablation.
  • Model design: Duo’s denoising model operates on both continuous latents and discrete samples.The model is implemented as a Transformer, with discrete inputs represented through soft vocabulary lookups.
  • Model design: Discrete inputs are embedded as convex combinations of vocabulary embeddings.This soft-lookup representation supports the denoiser’s operation on continuous latent inputs.

C.3. Low Discrepancy Sampler

Duo reduces training-estimate variance through low-discrepancy timestep sampling and restricts training to informative diffusion-time windows. The setup also specifies the tokenizer, datasets, model architecture, and optimization protocol used in the experiments.

  • Low-discrepancy sampling: Low-discrepancy timestep sampling spreads minibatch times across [0,1], reducing ELBO variance.Each sequence samples from a different partition of the unit interval rather than sampling independently.
  • Data setup: Experiments tokenize One Billion Words with BERT and OpenWebText with GPT-2, using wrapped sequence lengths of 128 and 1,024, respectively.The OpenWebText setup reserves the last 100k documents for validation.
  • Experimental setup: Uniform diffusion models use a 12-layer, 768-dimensional, 12-head modified diffusion Transformer without tied input-output word embeddings.The reported optimization uses AdamW, batch size 512, a 3e-4 learning rate, and 0.1 dropout.
  • Training window: Duo restricts Gaussian-latent training to [β,γ] and discrete training to α_t ∈ [0.05,0.95] to avoid weak learning signals.For GPT-2 and BERT tokenizers, the Gaussian window is [0.03,0.15].

D.1. Curriculum Learning Ablation

Curriculum learning lowers gradient and loss variance in Duo, producing a lower-variance training trajectory. The ablation examines this effect alongside the bias–variance trade-off induced by the temperature parameter τ.

  • Temperature ablation: The temperature parameter τ controls a bias–variance trade-off in training.The ablation varies τ on LM1B to study how closely the softmax approximates arg max.
  • Training curves: Duo’s curriculum-learning loss curve exhibits low-variance training compared with UDLM and MDLM.The curve is lower partly because it is a biased estimate of the ELBO.
  • Variance reduction: Curriculum learning drastically lowers Duo’s summed gradient variance, top-weight variance, and loss variance.Table 4 compares curriculum learning with no curriculum learning at fixed τ = 0.001.
  • Sample quality: Duo achieves lower Gen PPL than SEDD-Uniform and lower Gen PPL than masked diffusion models, with a slight entropy trade-off against the latter.Figure 9 reports sample entropy in brackets alongside Gen PPL.

D.3. Discrete Consistency Distillation: Quantitative Sample Quality Analysis

The quantitative distillation analysis shows that teacher choice affects DCD quality, while Duo consistently dominates MDLM distillation in the low-sampling-step regime. Distillation also increases Duo’s sample entropy while decreasing entropy for SDTT-distilled MDLM.

  • Teacher ablation: Using denoising weights rather than EMA weights as the DCD teacher yields a more effective distilled model.Figure 10 and Table 6 compare Gen PPL and entropy for the two teacher choices.
  • Sample quality: Duo distilled with DCD always achieves lower Gen PPL than MDLM distilled with SDTT at low sampling steps.The comparison is reported after successive distillation rounds, with exact Gen PPL and diversity values in Table 7.
  • Sample diversity: SDTT-distilled MDLM entropy decreases with distillation, whereas CDC-distilled Duo entropy increases.Lighter curves denote larger sampling-step counts, emphasizing the low-step regime.
  • Qualitative analysis: Qualitative samples are presented for Duo, DDT-distilled Duo, MDLM, and SDTT-distilled MDLM at T ∈ {8,1024}.The samples are described as non-cherry-picked and generated with the Greedy-Tail sampler.
Loading 2506.10892v3…