Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

The probability flow ODE is provably fast

Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, Adil Salim

TL;DR

This work provides the first polynomial-time convergence guarantees for the probability flow ODE in score-based generative models when using predictor steps interleaved with corrector steps. By introducing two schemes, DPOM (overdamped) and DPUM (underdamped), it demonstrates that the underdamped corrector yields a substantially better dimension dependence, achieving Õ(L^2√d/ε) iterations compared to Õ(L^3 d/ε^2) for the overdamped variant. The analysis hinges on a Wasserstein-to-TV regularization framework and a sharpened score perturbation lemma, enabling control of discretization and score-estimation errors under mild assumptions. These results strengthen the theoretical foundations of the probability flow ODE, illustrating its practical potential relative to SDE-based DDPMs, especially in high dimensions. The paper also provides preliminary numerical evidence and outlines open questions about the necessity of correctors and extensions to non-smooth or manifold-supported data distributions.

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.

The probability flow ODE is provably fast

TL;DR

This work provides the first polynomial-time convergence guarantees for the probability flow ODE in score-based generative models when using predictor steps interleaved with corrector steps. By introducing two schemes, DPOM (overdamped) and DPUM (underdamped), it demonstrates that the underdamped corrector yields a substantially better dimension dependence, achieving Õ(L^2√d/ε) iterations compared to Õ(L^3 d/ε^2) for the overdamped variant. The analysis hinges on a Wasserstein-to-TV regularization framework and a sharpened score perturbation lemma, enabling control of discretization and score-estimation errors under mild assumptions. These results strengthen the theoretical foundations of the probability flow ODE, illustrating its practical potential relative to SDE-based DDPMs, especially in high dimensions. The paper also provides preliminary numerical evidence and outlines open questions about the necessity of correctors and extensions to non-smooth or manifold-supported data distributions.

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 ( vs. , assuming smoothness of the data distribution), highlighting potential advantages of the ODE framework.
Paper Structure (29 sections, 19 theorems, 88 equations, 2 figures, 2 algorithms)

This paper contains 29 sections, 19 theorems, 88 equations, 2 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Our contributions
  3. Related works
  4. Preliminaries
  5. Score-based generative modeling
  6. Probability flow ODE (predictor steps)
  7. Corrector steps
  8. Overdamped Langevin.
  9. Underdamped Langevin.
  10. Diffusions as corrector steps.
  11. Results
  12. Assumptions
  13. Algorithms
  14. Convergence guarantees
  15. Proof overview
  16. ...and 14 more sections

Key Result

Theorem 1

Assume that the score function along the forward process is $L$-Lipschitz, and that the data distribution has finite second moment. Assume that we have access to $\widetilde{O}(\varepsilon/\sqrt L)$$L^2$-accurate score estimates. Then, the probability flow ODE implementation of the reversed Ornstein

Figures (2)

  • Figure 1: A realization of DPUM for a mixture of Gaussians.
  • Figure 2: A realization of DPUM for another mixture of Gaussians.

Theorems & Definitions (23)

  • Theorem 1: Informal
  • Remark 1
  • Theorem 2: DPOM
  • Theorem 3: DPUM
  • Lemma 1: Score perturbation
  • Lemma 2: Reparameterization
  • Lemma 3: Score perturbation for OU
  • Corollary 1
  • Lemma 4
  • Remark 2
  • ...and 13 more