MRS: Быстрый сэмплер для возвращающихся к среднему диффузий на основе решателей ОДУ и СДУ

MRS: A Fast Sampler for Mean Reverting Diffusion based on ODE and SDE Solvers

February 11, 2025

Авторы: Ao Li, Wei Fang, Hongbo Zhao, Le Lu, Ge Yang, Minfeng Xu

cs.AI

Аннотация

В приложениях диффузионных моделей управляемая генерация имеет практическое значение, но также представляет собой сложную задачу. Современные методы управляемой генерации в основном сосредоточены на модификации функции оценки диффузионных моделей, тогда как метод Mean Reverting (MR) Diffusion напрямую изменяет структуру стохастического дифференциального уравнения (SDE), что делает включение условий на изображение более простым и естественным. Однако текущие быстрые сэмплеры, не требующие обучения, не применимы напрямую к MR Diffusion. В результате MR Diffusion требует сотен NFEs (количество вычислений функции) для получения высококачественных сэмплов. В данной статье мы предлагаем новый алгоритм под названием MRS (MR Sampler) для сокращения количества NFEs при сэмплировании в MR Diffusion. Мы решаем обратное по времени SDE и обыкновенное дифференциальное уравнение потока вероятности (PF-ODE), связанные с MR Diffusion, и выводим полуаналитические решения. Эти решения состоят из аналитической функции и интеграла, параметризованного нейронной сетью. На основе этого решения мы можем генерировать высококачественные сэмплы за меньшее количество шагов. Наш подход не требует обучения и поддерживает все основные параметризации, включая предсказание шума, предсказание данных и предсказание скорости. Многочисленные эксперименты демонстрируют, что MR Sampler сохраняет высокое качество сэмплирования с ускорением в 10–20 раз для десяти различных задач восстановления изображений. Наш алгоритм ускоряет процесс сэмплирования в MR Diffusion, делая его более практичным для управляемой генерации.

English

In applications of diffusion models, controllable generation is of practical significance, but is also challenging. Current methods for controllable generation primarily focus on modifying the score function of diffusion models, while Mean Reverting (MR) Diffusion directly modifies the structure of the stochastic differential equation (SDE), making the incorporation of image conditions simpler and more natural. However, current training-free fast samplers are not directly applicable to MR Diffusion. And thus MR Diffusion requires hundreds of NFEs (number of function evaluations) to obtain high-quality samples. In this paper, we propose a new algorithm named MRS (MR Sampler) to reduce the sampling NFEs of MR Diffusion. We solve the reverse-time SDE and the probability flow ordinary differential equation (PF-ODE) associated with MR Diffusion, and derive semi-analytical solutions. The solutions consist of an analytical function and an integral parameterized by a neural network. Based on this solution, we can generate high-quality samples in fewer steps. Our approach does not require training and supports all mainstream parameterizations, including noise prediction, data prediction and velocity prediction. Extensive experiments demonstrate that MR Sampler maintains high sampling quality with a speedup of 10 to 20 times across ten different image restoration tasks. Our algorithm accelerates the sampling procedure of MR Diffusion, making it more practical in controllable generation.