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

要旨

拡散モデルの応用において、制御可能な生成は実用的な意義を持つ一方で、挑戦的な課題でもあります。現在の制御可能な生成手法は主に拡散モデルのスコア関数を修正することに焦点を当てていますが、Mean Reverting (MR) Diffusionは確率微分方程式（SDE）の構造を直接修正し、画像条件の組み込みをよりシンプルかつ自然に行うことができます。しかし、現在のトレーニング不要の高速サンプラーはMR Diffusionに直接適用できません。そのため、MR Diffusionは高品質なサンプルを得るために数百回の関数評価（NFEs）を必要とします。本論文では、MR DiffusionのサンプリングNFEsを削減するために、MRS（MR Sampler）という新しいアルゴリズムを提案します。MR Diffusionに関連する逆時間SDEと確率流常微分方程式（PF-ODE）を解き、半解析的な解を導出します。この解は、解析関数とニューラルネットワークによってパラメータ化された積分から構成されます。この解に基づいて、より少ないステップで高品質なサンプルを生成することができます。我々のアプローチはトレーニングを必要とせず、ノイズ予測、データ予測、速度予測を含むすべての主流のパラメータ化をサポートします。大規模な実験により、MR Samplerが10から20倍の高速化を維持しながら、10種類の異なる画像復元タスクにおいて高いサンプリング品質を維持することが示されました。我々のアルゴリズムはMR Diffusionのサンプリング手順を加速し、制御可能な生成においてより実用的なものとします。

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.