MRS：基于常微分方程与随机微分方程求解器的均值回归扩散快速采样器

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

February 11, 2025

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

cs.AI

摘要

在扩散模型的应用中，可控生成不仅具有实际意义，同时也面临挑战。当前的可控生成方法主要集中于修改扩散模型的评分函数，而均值回归（MR）扩散则直接调整了随机微分方程（SDE）的结构，使得图像条件的融入更为简便自然。然而，现有的无需训练快速采样器并不直接适用于MR扩散，因此MR扩散需要数百次函数评估（NFEs）才能获得高质量样本。本文提出了一种名为MRS（MR采样器）的新算法，旨在减少MR扩散的采样NFEs。我们求解了与MR扩散相关的反向时间SDE及概率流常微分方程（PF-ODE），并推导出半解析解。这些解由一个解析函数和一个由神经网络参数化的积分组成。基于此解，我们能够在更少的步骤中生成高质量样本。我们的方法无需训练，并支持所有主流参数化方式，包括噪声预测、数据预测和速度预测。大量实验表明，MR采样器在十种不同的图像修复任务中，均能保持高采样质量，同时实现10至20倍的加速。本算法显著提升了MR扩散的采样效率，使其在可控生成领域更具实用性。

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.