Fast Convergence for High-Order ODE Solvers in Diffusion Probabilistic Models

TL;DR

This paper provides a rigorous convergence analysis of high-order deterministic samplers based on probability flow ODEs for diffusion models, highlighting how approximation errors and score function regularity affect sampling accuracy.

Contribution

It introduces and analyzes $p$-th order Runge-Kutta schemes for diffusion models, establishing bounds on sampling errors considering score function approximation and solver step size.

Findings

Total variation distance bound depends on score approximation error and step size.

Numerical experiments confirm bounded derivatives of learned score functions in practice.

High-order solvers improve sampling efficiency in diffusion probabilistic models.

Abstract

Diffusion probabilistic models generate samples by learning to reverse a noise-injection process that transforms data into noise. A key development is the reformulation of the reverse sampling process as a deterministic probability flow ordinary differential equation (ODE), which allows for efficient sampling using high-order numerical solvers. Unlike traditional time integrator analysis, the accuracy of this sampling procedure depends not only on numerical integration errors but also on the approximation quality and regularity of the learned score function, as well as their interaction. In this work, we present a rigorous convergence analysis of deterministic samplers derived from probability flow ODEs for general forward processes with arbitrary variance schedules. Specifically, we develop and analyze $p$-th order (exponential) Runge-Kutta schemes, under the practical assumption that the first and second derivatives of the learned score function are bounded. We prove that the total variation distance between the generated and target distributions can be bounded as \begin{align*} O\bigl(d^{\frac{7}{4}}\varepsilon_{\text{score}}^{\frac{1}{2}} +d(dH_{\max})^p\bigr), \end{align*} where $\varepsilon^2_{\text{score}}$ denotes the $L^2$ error in the score function approximation, $d$ is the data dimension, and $H_{\max}$ represents the maximum solver step size. Numerical experiments on benchmark datasets further confirm that the derivatives of the learned score function are bounded in practice.