math.STSep 18, 2026

Schedule optimization for tau-leaping in masked discrete diffusion

Authors: Cecilia Secchi, Giacomo Zanella

Organizations: Bocconi University, Department of Decision Sciences, Milan, Italy. · Bocconi University, Department of Decision Sciences and BIDSA, Milan, Italy.

Abstract

Masked diffusions are popular generative models for discrete distributions. Unlike standard autoregressive sampling, they reveal several coordinates in parallel, approximating each block's joint conditional law by a product of one-coordinate conditionals. The resulting procedure, usually called tau-leaping, reduces computational cost but introduces a factorization error (εfact\varepsilon_\text{fact}), even with perfectly learned predictors. We study the resulting tradeoff between generative accuracy and computational cost, focusing on how to choose a denoising schedule to minimize εfact\varepsilon_\text{fact} for a fixed sampling budget. To do so, we establish an exact integral representation of εfact\varepsilon_\text{fact} separating the schedule from the target's dependence structure, summarized by a dependence density ρρ. This representation yields recursive stationarity equations for optimal schedules and allows us to quantify how estimation errors in ρρ affect schedule selection. As the dimension NN and sampling budget grow, we characterize the optimal schedule and quantify the cost of random block sizes relative to a deterministic planner. We highlight a fundamental dichotomy: if ρρ converges uniformly to a strictly positive continuous profile as N→∞N\to\infty, schedule optimization can only improve the leading constant of εfact\varepsilon_\text{fact}, while if ρρ degenerates, schedule optimization can improve the asymptotic order. Examples based on stationary processes and exchangeable mixtures illustrate these regimes.

Figures & tables

Explore similar work

Aug 13, 2026cs.LG

The data geometry of masking diffusion: Certified-optimal schedules via unmasking growth complexity

We study masking diffusion for discrete sampling and introduce a path-resolved measure of data geometry called the \emph{unmasking growth complexity} ({\textsf{UGC}\xspace}). Its local increments directly control Kullback--Leibler (KL) discretization error, yielding a unified analysis of Bernoulli-subset and fixed-cardinality unmasking schemes. In log-reveal-odds coordinates, this structure yields optimized single-block and multi-block schedules, and quantifies the gains from adapting computational effort to data geometry. Crucially, we show how {\textsf{UGC}\xspace} increments can be estimated from samples via KL increments along coupled reveal trajectories. This leads to \emph{certified-optimal} samplers that achieve a prescribed KL error with high probability and iteration complexity within a constant factor of the corresponding oracle procedure. Collapsing the \ugc path yields the aggregate {\textsf{UGC}\xspace} mass, which connects to classical multivariate dependence measures and complexity measures from previous analyses of discrete diffusion. In the fine-partition limit, the squared integral of the square-root {\textsf{UGC}\xspace} density determines the sharp leading-order optimal Euler discretization error. Examples exhibit substantial dimension-dependent gains over coarse schedules, including Ω~(d)\widetildeΩ(\sqrt{d}) improvements achievable with a constant number of adaptively placed blocks.
Feb 16, 2026cs.LG

Efficient Sampling with Discrete Diffusion Models: Sharp and Adaptive Guarantees

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 ττ-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 ττ-leaping algorithm achieves an iteration complexity of order O~(d/ε)\tilde O(d/\varepsilon), with dd the ambient dimension of the target distribution, eliminating linear dependence on the vocabulary size SS and improving existing bounds by a factor of dd; 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 ττ-leaping sampler whose convergence rate is governed by an intrinsic information-theoretic quantity, termed the effective total correlation, which is bounded by dlog⁡Sd \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.
May 21, 2026cs.LG

Noise Schedule Design for Diffusion Models: An Optimal Control Perspective

We develop a principled framework for analyzing and designing noise schedules in diffusion models. We show that one can recast this design problem as an optimal control problem, whose state is the Fisher information of the diffusion process which evolves according to an ODE and the control input is the noise schedule. The objective of the optimal control problem is a functional involving the Fisher information, which is shown to be an upper bound on the Kullback-Leibler sampling error. By solving this optimal control problem, we obtain sufficient conditions on noise schedules under which state-of-the-art O~(d/n)\tilde{\mathcal{O}} (d/n) sampling error is achievable, where dd is the data dimension and nn is the number of discretization steps. While existing theoretical work also prove that O~(d/n)\tilde{\mathcal{O}}(d/n) sampling error bounds are achievable, these results hold for specific noise schedules, which do not include the schedules used in practice. Under a further parametric assumption on the data distribution, we show that one can obtain closed-form expressions for the noise schedules. These noise schedules generalize standard empirical schedules such as exponential and sigmoid schedules by allowing additional parameters that can be tuned. Systematically tuning the parameters of these schedules yields new schedules that achieve superior FID scores on image generation benchmarks.