High-accuracy and dimension-free sampling with diffusions
Khashayar Gatmiry, Sitan Chen, Adil Salim
TL;DR
This work advances high-accuracy diffusion-based sampling by proving that a collocation-based solver, built on low-degree time-derivative bounds of the reverse-process score, achieves polylogarithmic dependence on $1/\varepsilon$ and dimension-free performance up to an effective radius. By partitioning the probability-flow ODE into short windows and applying Picard iterations with a polynomial time basis (via collocation), the authors obtain a sampler that outputs TV-accurate samples for distributions of the form $q=q_{\sf pre} \ast \mathcal{N}(0,\sigma^2 I)$ with radius $R$ in $\tilde{O}((R/\sigma)^2 \operatorname{polylog}(1/\varepsilon))$ rounds, assuming sub-exponential score error and Lipschitz estimates. A coupling-based robustness analysis ensures the low-degree approximation remains valid along the algorithm, and a Langevin-based corrector upgrades the Wasserstein guarantee to total variation. The results hold for Gaussian-mixture-like targets and highlight a dimension-free pathway to efficient, high-accuracy diffusion sampling with practical implications for score-based models and related stochastic samplers.
Abstract
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.