Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations

TL;DR

This paper provides theoretical error bounds for Masked Diffusion Models, analyzes the impact of schedule size variations, and introduces optimal scheduling strategies based on data information profiles, enhancing efficiency and understanding.

Contribution

It offers the first general error bounds for MDMs independent of sequence length and proposes optimal schedule strategies based on data information profiles.

Findings

Error bounds depend only on tokens per iteration, not sequence length.

Optimal schedule sizes can be derived from data information profiles.

Non-constant schedules improve sampling efficiency.

Abstract

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.