Reenmascaramiento de Modelos de Difusión Discreta con Escalado en Tiempo de Inferencia

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

Autores: Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, Volodymyr Kuleshov

cs.AI

Resumen

Parte del éxito de los modelos de difusión se debe a su capacidad para realizar refinamiento iterativo, es decir, corregir repetidamente las salidas durante la generación. Sin embargo, la difusión discreta enmascarada moderna carece de esta capacidad: cuando se genera un token, no puede ser actualizado nuevamente, incluso si introduce un error. Aquí, abordamos esta limitación presentando el muestreador del modelo de difusión con reenmascaramiento (ReMDM), un método que puede aplicarse a modelos de difusión enmascarados preentrenados de manera principlada y que se deriva de un modelo de difusión discreta con un proceso inverso de reenmascaramiento personalizado. Lo más interesante es que ReMDM dota a la difusión discreta de una forma de escalabilidad computacional en tiempo de inferencia. Al aumentar el número de pasos de muestreo, ReMDM genera salidas de lenguaje natural que se acercan a la calidad de los modelos autoregresivos, mientras que cuando el presupuesto computacional es limitado, ReMDM mantiene mejor la calidad. ReMDM también mejora la calidad de las muestras de los modelos de difusión enmascarada para imágenes discretizadas, y en dominios científicos como el diseño de moléculas, ReMDM facilita la guía de difusión y empuja la frontera de Pareto de la controlabilidad en comparación con el enmascaramiento clásico y la difusión de ruido uniforme. Proporcionamos el código junto con una publicación en el blog en la página del proyecto: https://remdm.github.io.

English

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.