Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations

Hugo Lavenant, Giacomo Zanella

TL;DR

This work analyzes generating samples from a discrete distribution using masked diffusion style unmasking with factorized conditional approximations. By decomposing the KL sampling error into learning and factorization components, it derives informative bounds and links the factorization error to an information profile f(i). In the random-order setting, it provides tight upper bounds and introduces a data-driven calculus-of-variations framework to obtain near-optimal non-constant schedules, with explicit results for geometric schedules and scaling limits. The results offer principled guidance for choosing schedule sizes to balance computation and accuracy, including a practical data-driven schedule formula, and extend to both diverging and bounded unmasked-variable regimes, with extensions to future enhancements such as remasking and adaptive planners.

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.

Error Bounds and Optimal Schedules for Masked Diffusions with Factorized Approximations

TL;DR

This work analyzes generating samples from a discrete distribution using masked diffusion style unmasking with factorized conditional approximations. By decomposing the KL sampling error into learning and factorization components, it derives informative bounds and links the factorization error to an information profile f(i). In the random-order setting, it provides tight upper bounds and introduces a data-driven calculus-of-variations framework to obtain near-optimal non-constant schedules, with explicit results for geometric schedules and scaling limits. The results offer principled guidance for choosing schedule sizes to balance computation and accuracy, including a practical data-driven schedule formula, and extend to both diverging and bounded unmasked-variable regimes, with extensions to future enhancements such as remasking and adaptive planners.

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.

Paper Structure

This paper contains 25 sections, 14 theorems, 61 equations, 1 algorithm.

Table of Contents

  1. Set-up, background and objectives
  2. General unmasking algorithms
  3. Factorized approximations and sources of error
  4. Computational cost and arbitrary planners
  5. Objective and Contributions
  6. Concurrent work
  7. Error decomposition
  8. Decomposition of KL error
  9. Worst-case bounds on the factorization error
  10. The random-order case
  11. Rewriting the factorization error with the information profile
  12. Resulting upper bounds
  13. Explicit computations with geometric schedules and lower bounds
  14. The case of a fixed, randomly-generated ordering
  15. Optimal schedules and scaling limits
  16. ...and 10 more sections

Key Result

Lemma 1

If $p^{\theta}(\mathbf{x}_z;\mathbf{x}_{z'})=\pi(\mathbf{x}_{z}|\mathbf{x}_{z'})$ for any $z, z'$ disjoint subsets of $\{ 1, \ldots, N \}$, then $p_{\textrm{alg}}(\mathbf{x})=\pi(\mathbf{x})$, regardless of the choice of $\nu^{\theta}$.

Theorems & Definitions (17)

  • Lemma 1
  • Proposition 2
  • Proposition 3
  • Lemma 4
  • Lemma 5
  • Lemma 6
  • Theorem 7
  • Lemma 8
  • Proposition 9
  • Proposition 10
  • ...and 7 more