Learning Gaussian mixtures with diffusion models

Gaussian mixtures are a "simple" family of multimodal distributions, for which the problem of distribution learning has been extensively studied and specialized algorithms have been developed. Diffusion models are a general framework for generative modeling with widespread empirical success in learning complex multimodal distributions, but scarce theoretical guarantees. This begs the question: can diffusion models provably learn Gaussian mixtures?

We show that indeed they can--an algorithm based on diffusion models learns a mixture of k isotropic Gaussians in R^n with quasi-polynomial time- and sample-complexity (exp(poly log((n+k)/epsilon))), under a mild minimum weight assumption. This gives a completely different, analytic proof of a result previously known only using a specialized algebraic approach. We introduce higher-order Gaussian noise sensitivity to show that the score functions of a Gaussian mixture can be learned sequentially via piecewise polynomial regression. Our results extend beyond discrete mixtures to a non-parametric family--Gaussian convolutions of distributions supported on sets with small covering number, where no sub-exponential algorithms were previously known.

Joint work with Khashayar Gatmiry and Jonathan Kelner (MIT). https://arxiv.org/abs/2404.18869