Learning general Gaussian mixtures with efficient score matching

TL;DR

This paper introduces a novel algorithm for learning Gaussian mixture models without separation assumptions, using diffusion models and score matching, achieving efficient sample complexity and runtime.

Contribution

It presents the first diffusion model-based approach with theoretical guarantees for learning Gaussian mixtures in polynomial time.

Findings

Achieves sample complexity polynomial in dimension and inverse accuracy

Runs in sample-polynomial time with total variation guarantees

First to apply diffusion models to unsupervised learning with theoretical bounds

Abstract

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 $d^{\mathrm{poly}(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.