MRS：基於ODE和SDE求解器的均值回歸擴散快速取樣器

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.