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.
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), 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 for a fixed sampling budget. To do so, we establish an exact integral representation of ε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→∞, schedule optimization can only improve the leading constant of εfact, while if ρ degenerates, schedule optimization can improve the asymptotic order. Examples based on stationary processes and exchangeable mixtures illustrate these regimes.
Figures & tables
Figure 1 : (a) Dependence density ρN for the Markov chain at N=64,128,256,512 , with the limiting profile g . (b) ρN for the Beta (1,1) -Bernoulli exchangeable model at N=64,128,256,512 . Panels (a)-(b) display u∈[0,0.15] . (c) εfact -optimal versus linear schedule at N=256 , K=9 . (d) Improvement ratio εfactlin/εfactopt versus N with K=⌈log2N⌉+1 .
Figure 2 : (a) Estimated dependence densities ρˉN(u) for mixture-of-products targets with N∈{8,16,32,64,128,256} . The vertical axis is truncated at 70 . (b) Linear and numerically optimized schedules for N=128 and K=7 . (c) Representation-based and direct Monte Carlo estimates of εfact for both schedules with K=log2N . The curves use the estimated profiles, and the boxplots summarize 100 independent Monte Carlo estimates. Both axes are logarithmic. (d) Estimated improvement ratio RN,K versus N , with K=log2N , on logarithmic axes.
ECE & CSL University of Illinois Urbana-Champaign · Computer Science Department Carnegie Mellon University · ECE, CSL & NCSA University of Illinois Urbana-Champaign