High-accuracy and dimension-free sampling with diffusions
This work addresses a key bottleneck in diffusion models for machine learning practitioners by enabling faster, high-accuracy sampling, though it is incremental as it builds on existing methods like low-degree approximation and collocation.
The paper tackles the problem of slow sampling in diffusion models due to high iteration complexity scaling polynomially with dimension and accuracy, and proposes a new solver that achieves polylogarithmic scaling in accuracy and dimension-free iteration complexity, depending only on the effective radius of the target distribution.
Diffusion models have shown remarkable empirical success in sampling from rich multi-modal distributions. Their inference relies on numerically solving a certain differential equation. 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 a subtle interplay between low-degree approximation and the collocation method (Lee, Song, Vempala 2018), and we prove that its iteration complexity scales \emph{polylogarithmically} in $1/\varepsilon$, yielding the first ``high-accuracy'' guarantee for a diffusion-based sampler that only uses (approximate) access to the scores of the data distribution. In addition, our bound 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.