Why Masking Diffusion Works: Condition on the Jump Schedule for Improved Discrete Diffusion

TL;DR

This paper explains why masking diffusion outperforms other discrete diffusion models by leveraging the fixed jump times in discrete Markov processes and introduces schedule-conditioned diffusion (SCUD) to enhance performance across various data types.

Contribution

The paper introduces SCUD, a generalized discrete diffusion framework that incorporates known jump time distributions, improving upon masking diffusion and classical methods.

Findings

SCUD outperforms masking diffusion on multiple data types.

Incorporating jump time distributions enhances diffusion model performance.

Discrete Markov process jump timing is crucial for diffusion efficiency.

Abstract

Discrete diffusion models, like continuous diffusion models, generate high-quality samples by gradually undoing noise applied to datapoints with a Markov process. Gradual generation in theory comes with many conceptual benefits; for example, inductive biases can be incorporated into the noising Markov process, and access to improved sampling algorithms. In practice, however, the consistently best performing discrete diffusion model is, surprisingly, masking diffusion, which does not denoise gradually. Here we explain the superior performance of masking diffusion by noting that it makes use of a fundamental difference between continuous and discrete Markov processes: discrete Markov processes evolve by discontinuous jumps at a fixed rate and, unlike other discrete diffusion models, masking diffusion builds in the known distribution of jump times and only learns where to jump to. We show that we can similarly bake in the known distribution of jump times into any discrete diffusion model. The resulting models - schedule-conditioned discrete diffusion (SCUD) - generalize classical discrete diffusion and masking diffusion. By applying SCUD to models with noising processes that incorporate inductive biases on images, text, and protein data, we build models that outperform masking.