Improved Sampling Schedules for Discrete Diffusion Models

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 significantly less understood than those of their continuous counterparts. In this work, we bridge this gap by analyzing the reverse process dynamics through the lens of thermodynamic entropy production. We propose the entropy production rate as a rigorous proxy for quantifying information generation, deriving as a byproduct a bound on the Wasserstein distance between intermediate states and the data distribution. Leveraging these insights, we introduce two novel sampling schedules that are uniformly spaced with respect to their corresponding physics-inspired metrics: the Entropic Discrete Schedule (EDS), which is defined by maintaining a constant rate of information gain, and the Wasserstein Discrete Schedule (WDS), which is defined by taking equal steps in terms of the Wasserstein distance. We empirically demonstrate that our proposed schedules significantly outperform state-of-the-art strategies across diverse application domains, including synthetic data, music notation, vision and language modeling, consistently achieving superior performance at a lower computational budget.

💡 Research Summary

The paper addresses a fundamental gap in the theoretical understanding of discrete diffusion models, which have become a powerful tool for generative modeling of sequence‑type data such as text, music, and categorical images. While continuous diffusion models have been extensively studied from a thermodynamic perspective—most notably through the concept of “Neural Entropy”—the reverse (denoising) dynamics of discrete diffusion have lacked a comparable rigorous framework.

Core Contributions

1. Thermodynamic Entropy Production for Discrete Diffusion – By modeling the forward corruption process as a continuous‑time Markov chain (CTMC) with time‑dependent generator →Qₜ, the authors derive the total entropy production Hₜₒₜ as the time integral of the log‑ratio between backward and forward path probabilities. They then apply the adiabatic–non‑adiabatic decomposition (Esposito & Van den Broeck, 2010) to split Hₜₒₜ into an adiabatic term Hₐd (associated with maintaining non‑equilibrium currents) and a non‑adiabatic term Hₙₐ (associated with relaxation toward the instantaneous stationary distribution πₜ). For reversible processes such as uniform diffusion, Hₐd vanishes, leaving Hₙₐ as the sole, finite measure of information generated during denoising. In contrast, absorbing (masked) diffusion exhibits divergent Hₐd but retains a finite Hₙₐ, making the latter the appropriate proxy for the “information cost” of the model.

2. Wasserstein Speed Limit – Leveraging recent work by Van Vu & Saito (2023), the authors bound the L₁‑Wasserstein distance between the initial noise distribution and the target data distribution by an integral of the product of a mobility factor M(t) (essentially the squared transition rates) and the instantaneous entropy production. This inequality formalizes a “speed limit”: for a fixed transport distance, the minimal entropy production determines the fastest possible trajectory through probability space. Because Hₐd may diverge for masked diffusion, the authors empirically replace the total entropy with the non‑adiabatic component Hₙₐ, yielding a practical surrogate bound W_θ.

3. Two Physics‑Inspired Sampling Schedules – Using the cumulative progress functions derived from the entropy and Wasserstein analyses, the paper proposes two novel discretization strategies for the reverse diffusion steps:

   • Entropic Discrete Schedule (EDS): Compute the cumulative non‑adiabatic entropy C(t)=∫₀ᵗ Hₙₐ(τ)dτ and place K sampling points so that each interval contributes an equal amount of entropy (C(tₖ)=k·C(T)/K). This yields larger time steps in high‑entropy regions (where the model is removing mostly noise) and finer steps near phase transitions where semantic structure emerges.
   • Wasserstein Discrete Schedule (WDS): Analogously, integrate the instantaneous “velocity” v(t)=M(t)·Hₙₐ(t) to obtain a cumulative Wasserstein distance W_cum(t). The schedule then enforces equal increments of this distance across K steps, ensuring that each step traverses the same amount of probability‑mass transport.

Both schedules are fully data‑driven; they require only the estimated non‑adiabatic entropy (or its mobility‑weighted version) from the trained score network, without any hand‑crafted heuristics.

4. Extensive Empirical Validation – The authors evaluate EDS and WDS on four diverse domains: (i) synthetic 2‑D Gaussian mixtures, (ii) music notation (MIDI) sequences, (iii) image generation (CIFAR‑10 / ImageNet‑like), and (iv) language modeling (WikiText‑103). Baselines include linear, cosine, DDIM‑style, fast‑solver, and recent distillation samplers. Metrics such as FID, BLEU, NLL, and token‑level accuracy consistently show that the proposed schedules achieve equal or superior quality with dramatically fewer steps (often 20–30 vs. 100). Notably, WDS excels at preserving mode diversity in image generation, while EDS provides smoother convergence in language tasks.

5. Practical Monitoring and Early Stopping – By tracking Hₙₐ and W_θ during sampling, the authors demonstrate that the entropy production curve plateaus once most structure has been recovered. This observation enables a principled early‑termination criterion, reducing average inference cost by ~30 % without measurable quality loss.

Implications and Limitations
The work bridges a theoretical divide, showing that discrete diffusion’s reverse dynamics can be understood through the same thermodynamic lens that has clarified continuous diffusion. It provides a concrete, mathematically grounded method for allocating computational budget where it matters most—during low‑entropy, high‑information phases of generation. However, the framework assumes a well‑behaved CTMC representation and relies on accurate estimation of Hₙₐ from the score network; bias in the network could affect the tightness of the Wasserstein bound. Moreover, the current formulation fixes the number of steps K a priori; future work could explore adaptive step counts based on real‑time entropy or transport estimates.

Conclusion
By defining a non‑adiabatic entropy production rate for discrete diffusion, deriving a Wasserstein‑based speed limit, and converting these physical quantities into two novel sampling schedules (EDS and WDS), the authors deliver both a deeper theoretical understanding and a practical performance boost for discrete generative models. The schedules outperform traditional heuristics across multiple modalities while offering a built‑in metric for early stopping, marking a significant step toward principled, efficient sampling in the rapidly expanding field of discrete diffusion modeling.