Neumaskierung diskreter Diffusionsmodelle mit Skalierung zur Inferenzzeit

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

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

cs.AI

Zusammenfassung

Ein Teil des Erfolgs von Diffusionsmodellen rührt von ihrer Fähigkeit her, iterative Verfeinerung durchzuführen, d.h. Ausgaben während der Generierung wiederholt zu korrigieren. Moderne maskierte diskrete Diffusion verfügt jedoch nicht über diese Fähigkeit: Sobald ein Token generiert wurde, kann es nicht mehr aktualisiert werden, selbst wenn es einen Fehler einführt. Hier adressieren wir diese Einschränkung durch die Einführung des Remasking-Diffusionsmodell (ReMDM)-Samplers, einer Methode, die auf prinzipielle Weise auf vortrainierte maskierte Diffusionsmodelle angewendet werden kann und die aus einem diskreten Diffusionsmodell mit einem benutzerdefinierten Remasking-Rückwärtsprozess abgeleitet ist. Besonders interessant ist, dass ReMDM der diskreten Diffusion eine Form von Inferenzzeit-Rechenleistungsskalierung verleiht. Durch die Erhöhung der Anzahl der Sampling-Schritte generiert ReMDM natürliche Sprachausgaben, die sich der Qualität autoregressiver Modelle annähern, während bei begrenztem Rechenbudget die Qualität besser erhalten bleibt. ReMDM verbessert auch die Probenqualität maskierter Diffusionsmodelle für diskretisierte Bilder, und in wissenschaftlichen Bereichen wie dem Moleküldesign erleichtert ReMDM die Diffusionsführung und verschiebt die Pareto-Grenze der Steuerbarkeit im Vergleich zu klassischer Maskierung und gleichmäßiger Rauschdiffusion. Wir stellen den Code zusammen mit einem Blogbeitrag auf der Projektseite bereit: 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.