The probability flow ODE is provably fast

TL;DR

This paper proves that the probability flow ODE for score-based generative models converges in polynomial time, offering theoretical guarantees and improved dimension dependence over previous stochastic methods.

Contribution

It provides the first polynomial-time convergence proof for the probability flow ODE with a corrector step, using novel techniques for deterministic dynamics analysis.

Findings

Polynomial-time convergence guarantees for probability flow ODE.

Better dimension dependence ($O(\sqrt{d})$) compared to prior $O(d)$ results.

Potential advantages of the ODE framework over SDE-based models.

Abstract

We provide the first polynomial-time convergence guarantees for the probability flow ODE implementation (together with a corrector step) of score-based generative modeling. Our analysis is carried out in the wake of recent results obtaining such guarantees for the SDE-based implementation (i.e., denoising diffusion probabilistic modeling or DDPM), but requires the development of novel techniques for studying deterministic dynamics without contractivity. Through the use of a specially chosen corrector step based on the underdamped Langevin diffusion, we obtain better dimension dependence than prior works on DDPM ($O(\sqrt{d})$ vs. $O(d)$, assuming smoothness of the data distribution), highlighting potential advantages of the ODE framework.