Her-maskeren van Discrete Diffusiemodellen met Schaling tijdens Inferentie

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

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

cs.AI

Samenvatting

Een deel van het succes van diffusiemodellen komt voort uit hun vermogen om iteratieve verfijning uit te voeren, d.w.z. het herhaaldelijk corrigeren van uitvoer tijdens het generatieproces. Moderne gemaskeerde discrete diffusie mist echter deze mogelijkheid: zodra een token is gegenereerd, kan deze niet meer worden bijgewerkt, zelfs niet wanneer deze een fout introduceert. Hier pakken we deze beperking aan door de remasking diffusiemodel (ReMDM) sampler te introduceren, een methode die op een principiële manier kan worden toegepast op vooraf getrainde gemaskeerde diffusiemodellen en die is afgeleid van een discreet diffusiemodel met een aangepast remasking achterwaarts proces. Het meest interessante is dat ReMDM discrete diffusie voorziet van een vorm van schaalbaarheid van rekentijd tijdens inferentie. Door het aantal samplingstappen te verhogen, genereert ReMDM natuurlijke taaluitvoer die de kwaliteit van autoregressieve modellen benadert, terwijl wanneer het rekenbudget beperkt is, ReMDM de kwaliteit beter behoudt. ReMDM verbetert ook de steekproefkwaliteit van gemaskeerde diffusiemodellen voor gediscretiseerde afbeeldingen, en in wetenschappelijke domeinen zoals molecuulontwerp vergemakkelijkt ReMDM diffusiebegeleiding en verlegt het de Pareto-grens van beheersbaarheid ten opzichte van klassieke masking en uniforme ruis-diffusie. We bieden de code samen met een blogpost op de projectpagina: 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.