Work

Schedule optimization for tau-leaping in masked discrete diffusion

Year: 2026
Type: preprint
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 ($\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 $ρ$ converges uniformly to a strictly positive continuous profile as $N\to\infty$, schedule optimization can only improve the leading constant of $\varepsilon_\text{fact}$, while if $ρ$ degenerates, schedule optimization can improve the asymptotic order. Examples based on stationary processes and exchangeable mixtures illustrate these regimes.
Institution: Bocconi University
Open Access status: green
OpenAlex ID: W7213937302
Cites:
Cited by:
Related to: