Rimodellazione dei Modelli di Diffusione Discreti con Scalatura al Momento dell'Inferenza

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

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

cs.AI

Abstract

Parte del successo dei modelli di diffusione deriva dalla loro capacità di eseguire un raffinamento iterativo, ovvero correggere ripetutamente gli output durante la generazione. Tuttavia, la moderna diffusione discreta mascherata manca di questa capacità: quando un token viene generato, non può essere aggiornato nuovamente, anche quando introduce un errore. Qui, affrontiamo questa limitazione introducendo il campionatore ReMDM (Remasking Diffusion Model), un metodo che può essere applicato a modelli di diffusione mascherati pre-addestrati in modo principe e che deriva da un modello di diffusione discreta con un processo inverso di rimascheramento personalizzato. Particolarmente interessante è il fatto che ReMDM conferisce alla diffusione discreta una forma di scalabilità computazionale al momento dell'inferenza. Aumentando il numero di passaggi di campionamento, ReMDM genera output in linguaggio naturale che si avvicinano alla qualità dei modelli autoregressivi, mentre quando il budget computazionale è limitato, ReMDM mantiene meglio la qualità. ReMDM migliora anche la qualità dei campioni dei modelli di diffusione mascherata per immagini discretizzate e, in ambiti scientifici come la progettazione di molecole, ReMDM facilita la guida della diffusione e sposta la frontiera di Pareto della controllabilità rispetto alla mascheratura classica e alla diffusione con rumore uniforme. Forniamo il codice insieme a un post sul blog nella pagina del progetto: 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.