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

2023-09-12

arXiv.org

AI Summary

Diffusion model sampling faces a fundamental trade-off between speed and sample quality: ODE solvers are efficient but suboptimal in performance, whereas SDE solvers achieve superior quality at high computational cost. To address this, we propose the Extended Reverse-time Stochastic Differential Equation (ER-SDE) framework—a unified theoretical model that reveals the intrinsic cause of performance disparity between ODE and SDE samplers as one-step prediction error. We derive the first exact and approximate analytical solutions for VP and VE SDEs within this framework and prove the equivalence of classical ODE and SDE solvers under ER-SDE. Leveraging semilinear SDE theory and reverse-time modeling, we design efficient numerical solvers—ER-SDE-Solvers—that inherit both deterministic efficiency and stochastic expressiveness. Evaluated on ImageNet 128×128, our method achieves a state-of-the-art FID of 8.33 in only 20 function evaluations, bridging the gap between speed and generation quality.

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 imes128$ 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}

Problem

• Balances speed and quality in diffusion models.
• Unifies ODE and SDE approaches in sampling.
• Develops efficient, high-quality ER-SDE-Solvers.

Innovation

• Extended Reverse-Time SDE formulation
• Unifies ODE and SDE sampling methods
• Develops efficient ER-SDE-Solvers