Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations
This work addresses the computational bottleneck in generative models for discrete data, offering a principled approach to optimize schedule sizes, though it is incremental in refining existing MDM methods.
The paper tackles the trade-off between computational cost and accuracy in Masked Diffusion Models (MDMs) for discrete data, providing error bounds independent of data dimensionality and identifying optimal schedules for token generation to improve efficiency.
Recently proposed generative models for discrete data, such as Masked Diffusion Models (MDMs), exploit conditional independence approximations to reduce the computational cost of popular Auto-Regressive Models (ARMs), at the price of some bias in the sampling distribution. We study the resulting computation-vs-accuracy trade-off, providing general error bounds (in relative entropy) that depend only on the average number of tokens generated per iteration and are independent of the data dimensionality (i.e. sequence length), thus supporting the empirical success of MDMs. We then investigate the gain obtained by using non-constant schedule sizes (i.e. varying the number of unmasked tokens during the generation process) and identify the optimal schedule as a function of a so-called information profile of the data distribution, thus allowing for a principled optimization of schedule sizes. We define methods directly as sampling algorithms and do not use classical derivations as time-reversed diffusion processes, leading us to simple and transparent proofs.