Remasking Discrete Diffusion Models with Inference-Time Scaling

Part of the success of diffusion models stems from their ability to perform iterativerefinement, i.e., repeatedly correcting outputs during generation. However, modernmasked discrete diffusion lacks this capability: when a token is generated, it cannotbe updated again, even when it introduces an error. Here, we address this limitationby introducing the remasking diffusion model (ReMDM) sampler, a method thatcan be applied to pretrained masked diffusion models in a principled way and thatis derived from a discrete diffusion model with a custom remasking backwardprocess. Most interestingly, ReMDM endows discrete diffusion with a form ofinference-time scaling. By increasing the number of sampling steps, ReMDMgenerates natural language outputs that approach the quality of autoregressivemodels, whereas when the computation budget is limited, ReMDM better maintainsquality. ReMDM also improves sample quality of masked diffusion models fordiscretized images, and in scientific domains such as molecule design, ReMDMfacilitates diffusion guidance and pushes the Pareto frontier of controllabilityrelative to classical masking and uniform noise diffusion.