Track: long paper (up to 8 pages)
Keywords: discrete diffusion, deep generative models
TL;DR: Absorbing state discrete diffusion models that enable re-masking benefit from inference-time scaling, improving sample quality and controlled generation.
Abstract: Part of the success of diffusion models stems from their ability to perform iterative
refinement, i.e., repeatedly correcting outputs during generation. However, modern
masked discrete diffusion lacks this capability: when a token is generated, it cannot
be updated again, even when it introduces an error. Here, we address this limitation
by introducing the remasking diffusion model (ReMDM) sampler, a method that
can be applied to pretrained masked diffusion models in a principled way and that
is derived from a discrete diffusion model with a custom remasking backward
process. Most interestingly, ReMDM endows discrete diffusion with a form of
inference-time scaling. By increasing the number of sampling steps, ReMDM
generates natural language outputs that approach the quality of autoregressive
models, whereas when the computation budget is limited, ReMDM better maintains
quality. ReMDM also improves sample quality of masked diffusion models for
discretized images, and in scientific domains such as molecule design, ReMDM
facilitates diffusion guidance and pushes the Pareto frontier of controllability
relative to classical masking and uniform noise diffusion.
Submission Number: 43
Loading