MRS: Un Muestreador Rápido para Difusión de Reversión a la Media Basado en Resolvedores de EDO y EDE

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

February 11, 2025

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

cs.AI

Resumen

En las aplicaciones de los modelos de difusión, la generación controlable tiene una importancia práctica significativa, pero también presenta desafíos. Los métodos actuales para la generación controlable se centran principalmente en modificar la función de puntuación de los modelos de difusión, mientras que la Difusión de Reversión a la Media (MR Diffusion) modifica directamente la estructura de la ecuación diferencial estocástica (SDE), lo que hace que la incorporación de condiciones de imagen sea más simple y natural. Sin embargo, los muestreadores rápidos actuales que no requieren entrenamiento no son directamente aplicables a MR Diffusion. Por lo tanto, MR Diffusion requiere cientos de NFEs (número de evaluaciones de función) para obtener muestras de alta calidad. En este artículo, proponemos un nuevo algoritmo llamado MRS (MR Sampler) para reducir los NFEs de muestreo de MR Diffusion. Resolvemos la SDE en tiempo inverso y la ecuación diferencial ordinaria de flujo de probabilidad (PF-ODE) asociada con MR Diffusion, y derivamos soluciones semi-analíticas. Las soluciones consisten en una función analítica y un parámetro integral parametrizado por una red neuronal. Basándonos en esta solución, podemos generar muestras de alta calidad en menos pasos. Nuestro enfoque no requiere entrenamiento y es compatible con todas las parametrizaciones principales, incluyendo la predicción de ruido, la predicción de datos y la predicción de velocidad. Experimentos extensos demuestran que MR Sampler mantiene una alta calidad de muestreo con una aceleración de 10 a 20 veces en diez tareas diferentes de restauración de imágenes. Nuestro algoritmo acelera el procedimiento de muestreo de MR Diffusion, haciéndolo más práctico en la generación controlable.

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.