MRS: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)가 필요합니다. 본 논문에서는 MR 확산의 샘플링 NFEs를 줄이기 위해 MRS(MR 샘플러)라는 새로운 알고리즘을 제안합니다. 우리는 MR 확산과 관련된 역시간 SDE와 확률 흐름 상미분 방정식(PF-ODE)을 해결하고, 반해석적 해를 도출합니다. 이 해는 분석적 함수와 신경망으로 매개변수화된 적분으로 구성됩니다. 이 해를 기반으로 더 적은 단계로 고품질 샘플을 생성할 수 있습니다. 우리의 접근 방식은 훈련이 필요하지 않으며, 잡음 예측, 데이터 예측, 속도 예측을 포함한 모든 주류 매개변수화를 지원합니다. 광범위한 실험을 통해 MR 샘플러가 10가지 다른 이미지 복원 작업에서 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.