Learning Mixtures of Gaussians Using the DDPM Objective

Kulin Shah, Sitan Chen, Adam Klivans

Keywords: Mixtures of Gaussians, score-based generative models, provable learning of score, Expectation-Maximization, DDPM generative model

TL;DR: We prove that gradient descent based on the Denoising Diffusion Probabilistic Model (DDPM) objective can efficiently learn the accurate score function of Mixtures of Gaussians.

Abstract:

Recent works have shown that diffusion models can learn essentially any distribution provided one can perform score estimation.
Yet it remains poorly understood under what settings score estimation is possible, let alone when practical gradient-based algorithms for this task can provably succeed.

In this work, we give the first provably efficient results for one of the most fundamental distribution families, Gaussian mixture models.
We prove that GD on the denoising diffusion probabilistic model (DDPM) objective can efficiently recover the ground truth parameters of the mixture model in the following two settings:
1. We show GD with random initialization learns mixtures of two spherical Gaussians in $d$ dimensions with $1/\text{poly}(d)$-separated centers.
2. We show GD with a warm start learns mixtures of $K$ spherical Gaussians with $\Omega(\sqrt{\log(\min(K,d))})$-separated centers.

A key ingredient in our proofs is a new connection between score-based methods and two other approaches to distribution learning, EM and spectral methods.

Submission Number: 13472