High-accuracy and dimension-free sampling with diffusions

Khashayar Gatmiry, Sitan Chen, Adil Salim

Keywords: diffusion models, high-accuracy sampling, picard iteration, low-degree polynomials, probability flow ODE

Abstract: Diffusion models have shown remarkable empirical success in sampling from rich multi-modal distributions. Their inference relies on solving a certain differential equation initialized at pure noise. However, this differential equation cannot be solved in closed form, and its resolution via discretization typically requires many small iterations to produce \emph{high-quality} samples. More precisely, prior works have shown that the iteration complexity of discretization methods for diffusion models scales polynomially in the ambient dimension and the inverse accuracy $1/\varepsilon$. In this work, we propose a new solver for diffusion models relying on the collocation method~\cite{lee2018algorithmic}, and we prove that its iteration complexity scales \emph{logarithmically} in $1/\varepsilon$, and does not depend explicitly on the ambient dimension. More precisely, the dimension affects the complexity of our solver through the \emph{effective radius} of the support of the target distribution only. Our solver constitutes the first "high-accuracy" diffusion-based sampler that only uses approximate access to the scores of the data distribution.

Primary Area: generative models

Submission Number: 22826