Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

2026-02-16

AI Summary

This work addresses the lack of theoretical guarantees for sampling efficiency in discrete diffusion models on high-dimensional structured data. By leveraging continuous-time Markov chains, it establishes the first tight convergence bounds for the τ-leaping sampler under both uniform and masking noise processes, revealing an unavoidable linear dependence on the dimensionality \(d\). To overcome this limitation, the paper proposes an adaptive sampler that exploits low-dimensional structures, achieving an iteration complexity of \(\widetilde{O}(d/\varepsilon)\) under uniform diffusion—eliminating dependence on vocabulary size \(S\)—and attaining sublinear adaptive convergence for structured data under masking diffusion. Key technical contributions include refined KL divergence analysis, an effective total correlation measure, and careful control of the score entropy loss.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Structured Learning
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Diffusion models over discrete spaces have recently shown striking empirical success, yet their theoretical foundations remain incomplete. In this paper, we study the sampling efficiency of score-based discrete diffusion models under a continuous-time Markov chain (CTMC) formulation, with a focus on $\tau$-leaping-based samplers. We establish sharp convergence guarantees for attaining $\varepsilon$ accuracy in Kullback-Leibler (KL) divergence for both uniform and masking noising processes. For uniform discrete diffusion, we show that the $\tau$-leaping algorithm achieves an iteration complexity of order $\tilde O(d/\varepsilon)$, with $d$ the ambient dimension of the target distribution, eliminating linear dependence on the vocabulary size $S$ and improving existing bounds by a factor of $d$; moreover, we establish a matching algorithmic lower bound showing that linear dependence on the ambient dimension is unavoidable in general. For masking discrete diffusion, we introduce a modified $\tau$-leaping sampler whose convergence rate is governed by an intrinsic information-theoretic quantity, termed the effective total correlation, which is bounded by $d \log S$ but can be sublinear or even constant for structured data. As a consequence, the sampler provably adapts to low-dimensional structure without prior knowledge or algorithmic modification, yielding sublinear convergence rates for various practical examples (such as hidden Markov models, image data, and random graphs). Our analysis requires no boundedness or smoothness assumptions on the score estimator beyond control of the score entropy loss.

Problem

• discrete diffusion models
• sampling efficiency
• convergence guarantees
• τ-leaping
• KL divergence

Innovation

• discrete diffusion models
• τ-leaping sampler
• convergence guarantees
• effective total correlation
• sampling efficiency