Schedule optimization for tau-leaping in masked discrete diffusion
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 (), 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 for a fixed sampling budget. To do so, we establish an exact integral representation of 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 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 , schedule optimization can only improve the leading constant of , while if degenerates, schedule optimization can improve the asymptotic order. Examples based on stationary processes and exchangeable mixtures illustrate these regimes.