ICML2024

Minimax Optimality of Score-based Diffusion Models: Beyond the Density Lower Bound Assumptions

Kaihong Zhang, Heqi Yin, Feng Liang, Jingbo Liu

40 citations

Abstract

We study the asymptotic error of score-based diffusion model sampling in large-sample scenarios from a non-parametric statistics perspective. We show that a kernel-based score estimator achieves an optimal mean square error of O~(n1td+22(td21))\widetilde{O}\left(n^{-1} t^{-\frac{d+2}{2}}(t^{\frac{d}{2}} \vee 1)\right) for the score function of p0N(0,tId)p_0*\mathcal{N}(0,t\boldsymbol{I}_d), where nn and dd represent the sample size and the dimension, tt is bounded above and below by polynomials of nn, and p0p_0 is an arbitrary sub-Gaussian distribution. As a consequence, this yields an O~(n1/2td4)\widetilde{O}\left(n^{-1/2} t^{-\frac{d}{4}}\right) upper bound for the total variation error of the distribution of the sample generated by the diffusion model under a mere sub-Gaussian assumption. If in addition, p0p_0 belongs to the nonparametric family of the β\beta-Sobolev space with β2\beta\le 2, by adopting an early stopping strategy, we obtain that the diffusion model is nearly (up to log factors) minimax optimal. This removes the crucial lower bound assumption on p0p_0 in previous proofs of the minimax optimality of the diffusion model for nonparametric families.