SAMPLER FOR A MASKED DIFFUSION MODEL

Abstract

Masked diffusion models (MDMs), a variant of discrete diffusion formulations, generally employ a gradual unmasking process that can generate tokens in any order. These MDMs are useful for generating discrete data, such as text, images, and other sequential data. However, MDM sampling, which is performed in continuous time, traditionally requires each sampling step to perform a forward traverse of the network, even though a single sampling step may not result in any changes to any token in the sequence. The present disclosure provides a first-hitting sampler for an MDM that can more efficiently make predictions for unmasking tokens in an input sequence for at least one or more sampling steps.

Images