[Paper Review] Learning Mixtures of Gaussians Using the DDPM Objective

Kulin Shah, Sitan Chen+1 more|arXiv (Cornell University)|Jul 3, 2023

Gaussian Processes and Bayesian InferenceComputer Science3 citations

TL;DR

This paper provides the first provably efficient gradient descent algorithms for learning mixtures of spherical Gaussians using the DDPM objective. It shows that with random initialization, two-component mixtures with $1/\text{poly}(d)$-separated centers can be recovered, and with a warm start, $K$-component mixtures with $\Omega(\sqrt{\log(\min(K,d))})$-separated centers are learnable, establishing a novel link between score-based diffusion and classical methods like EM and spectral learning.

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 along these lines for one of the most fundamental distribution families, Gaussian mixture models. We prove that gradient descent 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 gradient descent with random initialization learns mixtures of two spherical Gaussians in $d$ dimensions with $1/ ext{poly}(d)$-separated centers. 2) We show gradient descent with a warm start learns mixtures of $K$ spherical Gaussians with $Ω(\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, the EM algorithm and spectral methods.

Motivation & Objective

• To establish provable convergence guarantees for gradient descent on the DDPM objective in learning Gaussian mixture models.
• To close the gap between theoretical score estimation and practical training in diffusion models for structured distributions.
• To demonstrate that gradient descent can efficiently recover mixture parameters under mild separation conditions.
• To connect score-based diffusion learning to classical distribution learning methods like EM and spectral methods.

Proposed method

• Uses gradient descent on the DDPM objective to estimate the score function of a Gaussian mixture model.
• Analyzes the population gradient descent dynamics to show convergence to true mixture parameters.
• Establishes a theoretical connection between the DDPM objective and the EM algorithm via score estimation.
• Leverages spectral methods to analyze the low-noise regime and derive separation bounds.
• Proves that empirical gradients suffice for convergence, ensuring practical applicability.
• Introduces a novel analysis framework linking score-based generative modeling to classical statistical learning techniques.

Experimental results

Research questions

• RQ1 Can gradient descent on the DDPM objective provably recover the parameters of a mixture of two spherical Gaussians with polynomially separated centers?
• RQ2 Does the DDPM objective enable provable recovery of $K$-component spherical Gaussian mixtures when initialized near the true parameters?
• RQ3 How does the DDPM-based score estimation method relate to classical approaches like EM and spectral methods?
• RQ4 What is the minimal separation between mixture components required for efficient recovery via gradient descent on the DDPM objective?
• RQ5 Can empirical gradient descent on the DDPM objective achieve the same convergence guarantees as population-level analysis?

Key findings

• Gradient descent with random initialization efficiently learns two-component spherical Gaussian mixtures with $1/\text{poly}(d)$-separated centers.
• With a warm start, gradient descent on the DDPM objective learns $K$-component spherical Gaussian mixtures with $\Omega(\sqrt{\log(\min(K,d))})$-separated centers.
• The convergence is provably efficient, with both sample complexity and running time polynomial in $d$, $1/\varepsilon$, and $K$.
• A new theoretical connection is established between score-based diffusion and the EM algorithm in the context of Gaussian mixture learning.
• The analysis shows that empirical gradients suffice for convergence, validating the practicality of the approach.
• The work provides the first provable success of gradient-based score estimation for a fundamental family of distributions—Gaussian mixtures.