Learning Mixtures of Gaussians Using Diffusion Models

2024/04/29 by Gatmiry, Khashayar, Kelner, Jonathan, Lee, Holden · 9 citations
#Data Structures and Algorithms (cs.DS) #FOS: Computer and information sciences #FOS: Mathematics #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Probability (math.PR) #Statistics Theory (math.ST)

paper · doi:10.48550/arxiv.2404.18869

Abstract

We give a new algorithm for learning mixtures of k Gaussians (with identity covariance in ℝⁿ) to TV error ε, with quasi-polynomial (O(n^{poly log((n+k)/(ε))})) time and sample complexity, under a minimum weight assumption. Our results extend to continuous mixtures of Gaussians where the mixing distribution is supported on a union of k balls of constant radius. In particular, this applies to the case of Gaussian convolutions of distributions on low-dimensional manifolds, or more generally sets with small covering number, for which no sub-exponential algorithm was previously known. Unlike previous approaches, most of which are algebraic in nature, our approach is analytic and relies on the framework of diffusion models. Diffusion models are a modern paradigm for generative modeling, which typically rely on learning the score function (gradient log-pdf) along a process transforming a pure noise distribution, in our case a Gaussian, to the data distribution. Despite their dazzling performance in tasks such as image generation, there are few end-to-end theoretical guarantees that they can efficiently learn nontrivial families of distributions; we give some of the first such guarantees. We proceed by deriving higher-order Gaussian noise sensitivity bounds for the score functions for a Gaussian mixture to show that that they can be inductively learned using piecewise polynomial regression (up to poly-logarithmic degree), and combine this with known convergence results for diffusion models.

Cited by

• Generalization of Diffusion Models Arises with a Balanced Representation Space
• Convergence Dynamics of Over-Parameterized Score Matching for a Single Gaussian
• Dimension-free error estimate for diffusion model and optimal scheduling
• Denoising growth complexity: Data geometry and certified schedules for diffusion sampling
• DiffEM: Learning from Corrupted Data with Diffusion Models via Expectation Maximization
• On the Edge of Memorization in Diffusion Models
• LLM-Prior: A Framework for Knowledge-Driven Prior Elicitation and Aggregation
• Implicit Regularisation in Diffusion Models: An Algorithm-Dependent Generalisation Analysis
• Breaking AR's Sampling Bottleneck: Provable Acceleration via Diffusion Language Models

Related

• Learning general Gaussian mixtures with efficient score matching
• Learning Mixtures of Gaussians Using the DDPM Objective
• A Nearly Optimal and Agnostic Algorithm for Properly Learning a Mixture of k Gaussians, for any Constant k
• Settling the Polynomial Learnability of Mixtures of Gaussians
• Polynomial Learning of Distribution Families