Remasking Discrete Diffusion Models with Inference-Time Scaling
Authors: name: Guanghan Wang; name: Yair Schiff; name: Subham Sekhar Sahoo; name: Volodymyr Kuleshov
Published: 2025-03-01 02:37:51

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 compute 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. We provide the code along with a blog post on the
project page: https://remdm.github.io.