Learning general Gaussian mixtures with efficient score matching

We study the problem of learning mixtures of k Gaussians in d dimensions. We make no separation assumptions on the underlying mixture components: we only require that the covariance matrices have bounded condition number and that the means and covariances lie in a ball of bounded radius. We give an algorithm that draws dpoly(k/varepsilon) samples from the target mixture, runs in sample-polynomial time, and constructs a sampler whose output distribution is varepsilon-far from the unknown mixture in total variation. Prior works for this problem either (i) required exponential runtime in the dimension d, (ii) placed strong assumptions on the instance (e.g., spherical covariances or clusterability), or (iii) had doubly exponential dependence on the number of components k. Our approach departs from commonly used techniques for this problem like the method of moments. Instead, we leverage a recently developed reduction, based on diffusion models, from distribution learning to a supervised learning task called score matching. We give an algorithm for the latter by proving a structural result showing that the score function of a Gaussian mixture can be approximated by a piecewise-polynomial function, and there is an efficient algorithm for finding it. To our knowledge, this is the first example of diffusion models achieving a state-of-the-art theoretical guarantee for an unsupervised learning task.