Elucidating the solution space of extended reverse-time SDE for diffusion models

TL;DR

This paper introduces the Extended Reverse-Time SDE framework for diffusion models, unifying ODE and SDE sampling methods, and proposes ER-SDE-Solvers that achieve state-of-the-art quality and efficiency in image generation.

Contribution

It formulates a unified ER SDE approach, provides theoretical insights into sampling speed and quality, and develops ER-SDE-Solvers that outperform existing methods.

Findings

ER-SDE-Solvers achieve state-of-the-art FID scores on ImageNet 128x128.

Theoretical analysis explains the rapid sampling of ODEs and high quality of SDEs.

VP and VE SDEs are shown to have comparable performance.

Abstract

Sampling from Diffusion Models can alternatively be seen as solving differential equations, where there is a challenge in balancing speed and image visual quality. ODE-based samplers offer rapid sampling time but reach a performance limit, whereas SDE-based samplers achieve superior quality, albeit with longer iterations. In this work, we formulate the sampling process as an Extended Reverse-Time SDE (ER SDE), unifying prior explorations into ODEs and SDEs. Theoretically, leveraging the semi-linear structure of ER SDE solutions, we offer exact solutions and approximate solutions for VP SDE and VE SDE, respectively. Based on the approximate solution space of the ER SDE, referred to as one-step prediction errors, we yield mathematical insights elucidating the rapid sampling capability of ODE solvers and the high-quality sampling ability of SDE solvers. Additionally, we unveil that VP SDE solvers stand on par with their VE SDE counterparts. Based on these findings, leveraging the dual advantages of ODE solvers and SDE solvers, we devise efficient high-quality samplers, namely ER-SDE-Solvers. Experimental results demonstrate that ER-SDE-Solvers achieve state-of-the-art performance across all stochastic samplers while maintaining efficiency of deterministic samplers. Specifically, on the ImageNet $128\times128$ dataset, ER-SDE-Solvers obtain 8.33 FID in only 20 function evaluations. Code is available at \href{https://github.com/QinpengCui/ER-SDE-Solver}{https://github.com/QinpengCui/ER-SDE-Solver}