MRS : Un échantillonneur rapide pour les diffusions à retour à la moyenne basé sur les solveurs d'EDO et d'EDS

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

Résumé

Dans les applications des modèles de diffusion, la génération contrôlée revêt une importance pratique, mais reste un défi. Les méthodes actuelles pour la génération contrôlée se concentrent principalement sur la modification de la fonction de score des modèles de diffusion, tandis que la diffusion à réversion vers la moyenne (Mean Reverting, MR) modifie directement la structure de l'équation différentielle stochastique (EDS), rendant l'incorporation de conditions d'image plus simple et plus naturelle. Cependant, les échantillonneurs rapides actuels, sans nécessité d'entraînement, ne sont pas directement applicables à la diffusion MR. Ainsi, la diffusion MR nécessite des centaines d'évaluations de fonctions (NFEs) pour obtenir des échantillons de haute qualité. Dans cet article, nous proposons un nouvel algorithme nommé MRS (MR Sampler) pour réduire les NFEs d'échantillonnage de la diffusion MR. Nous résolvons l'EDS en temps inverse et l'équation différentielle ordinaire de flux de probabilité (PF-ODE) associées à la diffusion MR, et dérivons des solutions semi-analytiques. Ces solutions consistent en une fonction analytique et une intégrale paramétrée par un réseau de neurones. Sur la base de cette solution, nous pouvons générer des échantillons de haute qualité en moins d'étapes. Notre approche ne nécessite pas d'entraînement et prend en charge toutes les paramétrisations principales, y compris la prédiction de bruit, la prédiction de données et la prédiction de vitesse. Des expériences approfondies démontrent que le MR Sampler maintient une qualité d'échantillonnage élevée avec une accélération de 10 à 20 fois sur dix tâches différentes de restauration d'images. Notre algorithme accélère la procédure d'échantillonnage de la diffusion MR, la rendant plus pratique pour la génération contrôlée.

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.