MRS: Ein schneller Sampler für mittelwertregressierende Diffusion basierend auf ODE und SDE-Lösern

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

February 11, 2025

Autoren: Ao Li, Wei Fang, Hongbo Zhao, Le Lu, Ge Yang, Minfeng Xu

cs.AI

Zusammenfassung

Bei Anwendungen von Diffusionsmodellen ist die kontrollierbare Generierung von praktischer Bedeutung, aber auch herausfordernd. Aktuelle Methoden zur kontrollierbaren Generierung konzentrieren sich hauptsächlich auf die Modifikation der Score-Funktion von Diffusionsmodellen, während die Mean Reverting (MR) Diffusion direkt die Struktur der stochastischen Differentialgleichung (SDE) modifiziert, was die Integration von Bildbedingungen einfacher und natürlicher macht. Allerdings sind aktuelle, training-freie schnelle Sampler nicht direkt auf MR Diffusion anwendbar. Daher erfordert MR Diffusion Hunderte von NFEs (Anzahl der Funktionsauswertungen), um hochwertige Proben zu erhalten. In diesem Artikel schlagen wir einen neuen Algorithmus namens MRS (MR Sampler) vor, um die Anzahl der für MR Diffusion benötigten Proben-NFEs zu reduzieren. Wir lösen die SDE rückwärts in der Zeit und die mit MR Diffusion verbundene gewöhnliche Differentialgleichung des Wahrscheinlichkeitsflusses (PF-ODE) und leiten halbanalytische Lösungen her. Die Lösungen bestehen aus einer analytischen Funktion und einem Integral, das durch ein neuronales Netzwerk parametrisiert ist. Basierend auf dieser Lösung können wir hochwertige Proben in weniger Schritten generieren. Unser Ansatz erfordert kein Training und unterstützt alle gängigen Parametrisierungen, einschließlich Rauschvorhersage, Datenvorhersage und Geschwindigkeitsvorhersage. Umfangreiche Experimente zeigen, dass der MR Sampler bei zehn verschiedenen Bildwiederherstellungsaufgaben eine hohe Probenqualität bei einer Beschleunigung um das 10- bis 20-fache beibehält. Unser Algorithmus beschleunigt das Probenverfahren von MR Diffusion und macht es in der kontrollierbaren Generierung praktikabler.

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.