Reaplicando Máscaras em Modelos de Difusão Discreta com Escalonamento no Tempo de Inferência

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

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

cs.AI

Resumo

Parte do sucesso dos modelos de difusão decorre de sua capacidade de realizar refinamento iterativo, ou seja, corrigir repetidamente as saídas durante a geração. No entanto, a difusão discreta mascarada moderna carece dessa capacidade: quando um token é gerado, ele não pode ser atualizado novamente, mesmo que introduza um erro. Aqui, abordamos essa limitação ao introduzir o amostrador ReMDM (Remasking Diffusion Model), um método que pode ser aplicado a modelos de difusão mascarada pré-treinados de forma fundamentada e que é derivado de um modelo de difusão discreta com um processo de retropropagação de remascaramento personalizado. Mais interessante ainda, o ReMDM confere à difusão discreta uma forma de escalonamento de computação no tempo de inferência. Ao aumentar o número de etapas de amostragem, o ReMDM gera saídas de linguagem natural que se aproximam da qualidade dos modelos autoregressivos, enquanto, quando o orçamento de computação é limitado, o ReMDM mantém melhor a qualidade. O ReMDM também melhora a qualidade das amostras de modelos de difusão mascarada para imagens discretizadas e, em domínios científicos, como o design de moléculas, o ReMDM facilita a orientação por difusão e empurra a fronteira de Pareto da controlabilidade em relação ao mascaramento clássico e à difusão de ruído uniforme. Disponibilizamos o código junto com um post no blog na página do projeto: 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.