Improved Sampling Schedules for Discrete Diffusion Models

TL;DR

This paper introduces physics-inspired sampling schedules for discrete diffusion models, grounded in thermodynamic entropy principles, leading to improved generative performance across multiple domains.

Contribution

It proposes two novel sampling schedules based on entropy production and Wasserstein distance, enhancing efficiency and quality in discrete diffusion model generation.

Findings

Outperforms existing sampling strategies in diverse tasks

Achieves higher quality with lower computational cost

Demonstrates effectiveness across synthetic, music, vision, and language data

Abstract

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.