Масштабирование моделей дискретного диффузионного процесса на этапе вывода с повторным маскированием

Remasking Discrete Diffusion Models with Inference-Time Scaling

March 1, 2025

Авторы: Guanghan Wang, Yair Schiff, Subham Sekhar Sahoo, Volodymyr Kuleshov

cs.AI

Аннотация

Часть успеха диффузионных моделей связана с их способностью выполнять итеративное уточнение, то есть многократно корректировать выходные данные в процессе генерации. Однако современные маскированные дискретные диффузионные модели лишены этой возможности: когда токен сгенерирован, его нельзя обновить снова, даже если он содержит ошибку. В данной работе мы устраняем это ограничение, представляя метод сэмплера ReMDM (remasking diffusion model), который может быть применен к предобученным маскированным диффузионным моделям в принципиальном виде и который выводится из дискретной диффузионной модели с пользовательским процессом обратного маскирования. Наиболее интересно то, что ReMDM наделяет дискретные диффузионные модели возможностью масштабирования вычислений на этапе вывода. Увеличивая количество шагов сэмплирования, ReMDM генерирует выходные данные на естественном языке, приближающиеся по качеству к авторегрессивным моделям, тогда как при ограниченном бюджете вычислений ReMDM лучше сохраняет качество. ReMDM также улучшает качество сэмплов маскированных диффузионных моделей для дискретизированных изображений, а в научных областях, таких как дизайн молекул, ReMDM облегчает управление диффузией и расширяет Парето-границу управляемости по сравнению с классическим маскированием и диффузией с равномерным шумом. Мы предоставляем код вместе с блог-постом на странице проекта: 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.