MRS: Een Snelle Sampler voor Mean Reverting Diffusie gebaseerd op ODE- en SDE-Oplossers

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

February 11, 2025

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

cs.AI

Samenvatting

In toepassingen van diffusiemodellen is controleerbare generatie van praktisch belang, maar ook uitdagend. Huidige methoden voor controleerbare generatie richten zich voornamelijk op het aanpassen van de scorefunctie van diffusiemodellen, terwijl Mean Reverting (MR) Diffusion direct de structuur van de stochastische differentiaalvergelijking (SDE) wijzigt, waardoor de integratie van beeldcondities eenvoudiger en natuurlijker wordt. Echter, huidige trainingsvrije snelle samplers zijn niet direct toepasbaar op MR Diffusion. Hierdoor vereist MR Diffusion honderden NFEs (aantal functie-evaluaties) om hoogwaardige samples te verkrijgen. In dit artikel stellen we een nieuw algoritme voor, genaamd MRS (MR Sampler), om de sampling NFEs van MR Diffusion te verminderen. We lossen de omgekeerde-tijd SDE en de waarschijnlijkheidsstroom gewone differentiaalvergelijking (PF-ODE) op die geassocieerd zijn met MR Diffusion, en leiden semi-analytische oplossingen af. De oplossingen bestaan uit een analytische functie en een integraal geparameteriseerd door een neuraal netwerk. Op basis van deze oplossing kunnen we hoogwaardige samples genereren in minder stappen. Onze aanpak vereist geen training en ondersteunt alle gangbare parameterisaties, inclusief ruisvoorspelling, datavoorspelling en snelheidsvoorspelling. Uitgebreide experimenten tonen aan dat MR Sampler een hoge samplingkwaliteit behoudt met een versnelling van 10 tot 20 keer over tien verschillende beeldhersteltaken. Ons algoritme versnelt de samplingprocedure van MR Diffusion, waardoor het praktischer wordt in controleerbare generatie.

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.