Schedule optimization improves masked discrete diffusion sampling accuracy
Schedule optimization for tau-leaping in masked discrete diffusion
Machine Learning
Summary
Sampling in certain machine learning models involves revealing data parts step-by-step, but doing this at once for multiple parts can create errors. The authors studied these errors and how they change depending on the schedule that decides how many parts are revealed at each step. They developed ways to estimate these errors and find the best schedules to reduce them. Their findings show that carefully planned schedules can improve the accuracy of the sampling process, especially when the pattern of dependencies among data parts changes in certain ways.
What this means in practice
- •For machine learning engineers: Optimize sampling schedules in masked discrete diffusion models to reduce approximation errors for more reliable predictions.
- •For data scientists: Improve efficiency and accuracy of probabilistic model sampling by tailoring schedules according to data dependency profiles.
A theory result. No direct application yet.
Authors
Cecilia Secchi, Giacomo Zanella
Abstract
Masked discrete diffusion models are commonly accelerated using the so-called tau-leaping discretization method, which reveals several coordinates in parallel at each sampling step. The sampler replaces the joint conditional law of each revealed block by a product distribution, incurring a factorization error $\varepsilon_\text{fact}$ present even with perfectly learned predictors. We analyze the standard sampler on $N$ coordinates with $K$ sampling steps, whose random block sizes depend on a denoising schedule. Our analysis uses an exact integral representation of $\varepsilon_\text{fact}$ in terms of a distribution-dependent dependence density $ρ$, which records how conditional dependence evolves as the revealed fraction of coordinates grows. We develop estimators for this profile and quantify how estimation errors affect schedule selection. We derive recursive stationarity equations for the finite-$K$ optimization problem and, under a monotonicity condition, characterize its unique optimizer. In the joint limit $N,K\to\infty$, we obtain an explicit characterization of the optimal limiting smooth schedule and quantify the cost of random block sizes relative to a deterministic planner. When $ρ_N$ converges uniformly to a strictly positive continuous profile, optimizing over fixed smooth schedules can improve the leading constant but not the $N/K$ scaling of $\varepsilon_\text{fact}$. By contrast, if $ρ_N$ degenerates, suitable schedules can improve the asymptotic order relative to the uniform schedule. Examples based on stationary processes and exchangeable mixtures illustrate these two regimes.