Improved Sampling Schedules for Discrete Diffusion Models
Authors: Alberto Foresti, Mustapha Bounoua, Giulio Franzese, Luca Ambrogioni, Pietro Michiardi
Organizations: Department of Data Science, EURECOM, 450 Route des Chappes, 06410 Biot, France · Donders Institute for Brain, Cognition and Behaviour, Radboud University, Postbus 9102, 6500 HC Nijmegen, Netherlands
Uniform Discrete Diffusion Models have emerged as a powerful paradigm for generative modeling on sequence data; however, the information-theoretic principles governing their reverse processes remain underexplored. In this work, we bridge this gap by analyzing the reverse process dynamics through the lens of thermodynamic entropy production and dynamical activity. We use a variational expression of the discrete Wasserstein distance to characterize the dynamics of the generation process. Leveraging these insights, we first introduce the Entropic Discrete Schedule (EDS), which aligns with entropy production, and later construct the Wasserstein Discrete Schedule (WDS), which combines entropy production and dynamical activity to act as a geometric cost. Our experimental results on language modeling show that WDS provide the same results as uniform while using up to half of the computations on the MBPP benchmark. The code is available at https://github.com/AlbertoForesti/impscheds.
Figures & tables
Figure 1: Binomial distribution benchmark. We report the mean performance across 5 seeds, the bars indicate the standard error of the mean.
Figure 2: Sumi reasoning benchmark.
Figure 3: Performance comparisons on piano note generation across different metrics.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Samples ( N )
Timesteps ( Nt )
Binomial Distributions
1024
1024
Monophonic Music
64
1024
Sumi-7B (MBPP)
64
512
Appendix
Table 1: Parameters used for estimating neural entropy across different datasets.
Dataset
JYS (s)
Entropy production (s)
JYS / Entropy production
Binomial ( n=255 , linear)
13.72
12.70
1.08 ×
Binomial ( n=255 , loglinear)
12.37
10.91
1.13 ×
Binomial ( n=31 , linear)
12.24
10.51
1.16 ×
Binomial ( n=31 , loglinear)
12.27
11.07
1.11 ×
Piano
40.65
21.66
1.88 ×
Appendix
Table 2: Schedule computation time across different datasets
Task
NFE
Uniform
WDS
Performance Improvement (WDS / Uniform)
GSM8K
4
0.1327
0.1660
1.25 ×
8
0.2456
0.2472
1.01 ×
16
0.2987
0.2949
0.99 ×
32
0.3017
0.3124
1.04 ×
MBPP
16
0.0900
0.1420
1.58 ×
32
0.1800
0.2160
1.20 ×
Appendix
Table 3: Numerical results for WDS on reasoning tasks
Figure 4: Entropy production rate (top) and Wasserstein rate (bottom) for Binomial, Monophonic Music, and Sumi-7B.
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 : https://oussamazekri.fr/gdds.
Oussama Zekri, Théo Uscidda, Nicolas Boullé +1
CREST, ENSAE, Institut Polytechnique de Paris, Paris, France · Department of Mathematics, Imperial College London, London, SW7 2AZ, UK
Discrete diffusion models are a powerful class of generative models with strong performance across many domains. For efficiency, however, discrete diffusion typically parameterizes the generative (reverse) process with factorized distributions, which makes it difficult for the model to learn the target process in a small number of steps and necessitates a long, computationally expensive sampling procedure. To reduce the gap between the target and model distributions and enable few-step generation, we propose Forward-Learned Discrete Diffusion (FLDD), which introduces discrete diffusion with a learnable forward (noising) process. Rather than fixing a Markovian forward chain, we adopt a non-Markovian formulation with learnable marginal and posterior distributions. This allows the generative process to remain factorized while matching the target defined by the noising process. We train all parameters end-to-end under the standard variational objective. Experiments on various benchmarks show that, for a given number of sampling steps, our approach produces a higher quality samples than conventional discrete diffusion models using the same reverse parameterization.
Discrete diffusion models and flow matching have emerged as powerful frameworks for generative modeling over discrete state spaces, yet efficient few-step generation remains a fundamental challenge. In this work, we introduce the Discrete Average Generator, a principled extension of MeanFlow to Continuous-Time Markov Chains (CTMCs). Analogously to how MeanFlow defines an average velocity field over a time interval in continuous spaces, we define an average generator as the normalized increment of the transition kernel over a time interval. We show that this average generator satisfies a self-consistency identity, which provides the foundation for our training objective. We further develop training strategies that align with the standard training paradigm of diffusion language models while keeping the resulting objective tractable. When projected onto per-coordinate marginals, the self-consistency identity admits a closed-form expression, enabling efficient training and inference. In Potts model simulations, our objective reduces the total variation distance of the K-step sampler by up to 67%. On OpenWebText, our method achieves the lowest generative perplexity among the evaluated methods for 8 to 64 sampling steps while enabling a 16× acceleration, and achieves comparable performance to existing methods on ImageNet.
Yidong Ouyang, Zhengyan Wan, Themis Haris +8
University of California, Los Angeles · Google · Work done as a research intern at Google