Schedule optimization for tau-leaping in masked discrete diffusion

Cecilia Secchi, Giacomo Zanella

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 ($\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 $\varepsilon_\text{fact}$ for a fixed sampling budget. To do so, we establish an exact integral representation of $\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 $N$ 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…