MRS: Um Amostrador Rápido para Difusão de Reversão à Média baseado em ODE e Solucionadores SDE

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

Resumo

Nas aplicações de modelos de difusão, a geração controlável é de significado prático, porém desafiadora. Os métodos atuais para geração controlável focam principalmente na modificação da função de pontuação dos modelos de difusão, enquanto a Difusão de Média Reversão (MR) modifica diretamente a estrutura da equação diferencial estocástica (SDE), tornando a incorporação de condições de imagem mais simples e natural. No entanto, os atuais amostradores rápidos sem treinamento não são diretamente aplicáveis à Difusão de MR. Portanto, a Difusão de MR requer centenas de NFEs (número de avaliações de função) para obter amostras de alta qualidade. Neste artigo, propomos um novo algoritmo chamado MRS (Amostrador de MR) para reduzir as NFEs de amostragem da Difusão de MR. Resolvemos a SDE de tempo reverso e a equação diferencial ordinária de fluxo de probabilidade (PF-ODE) associada à Difusão de MR, e derivamos soluções semi-analíticas. As soluções consistem em uma função analítica e um parâmetro integralizado por uma rede neural. Com base nesta solução, podemos gerar amostras de alta qualidade em menos passos. Nossa abordagem não requer treinamento e suporta todas as parametrizações principais, incluindo previsão de ruído, previsão de dados e previsão de velocidade. Experimentos extensos demonstram que o Amostrador de MR mantém alta qualidade de amostragem com uma aceleração de 10 a 20 vezes em dez tarefas diferentes de restauração de imagem. Nosso algoritmo acelera o procedimento de amostragem da Difusão de MR, tornando-a mais prática na geração controlável.

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.