Title: Learning Mixtures of Gaussians Using the DDPM Objective
Venue: NeurIPS 2023
Accepted: yes
Keywords: Mixtures of Gaussians

score-based generative models

provable learning of score

Expectation-Maximization

DDPM generative model
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.

Learning Mixtures of Gaussians Using the DDPM Objective

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 gradientbased 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:

We show gradient descent with random initialization learns mixtures of two spherical Gaussians in d dimensions with 1 / poly ( d ) -separated centers.

We show gradient descent with a warm start learns mixtures of K spherical Gaussians with Ω ( √ 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 expectation-maximization (EM) algorithm and spectral methods.

1 Introduction

In recent years diffusion models [SSDK + 20, SDWMG15, SE19] have emerged as a powerful framework for generative modeling and now form the backbone of notable image generation systems like DALL · E 2 [RDN + 22], Imagen [SCS + 22], and Stable Diffusion [RBL + 22]. At the heart of this framework is a reduction from distribution learning to denoising or score estimation . That is, in order to generate new samples from a data distribution q given a collection of independent samples, it suffices to learn the score function, i.e., the gradient of the log-density of the data distribution when convolved with varying levels of noise (see Section 1.3). A popular and well-studied objective for score matching is the denoising diffusion probabilistic model (DDPM) objective due to [HJA20]. Optimizing this objective amounts to solving the following type of problem: given a noisy observation x of a sample x from q , estimate the mean of the posterior distribution over x .

˜ Are there natural data distributions under which GD provably achieves accurate score estimation?

˜ While a number of theoretical works [DBTHD21, BMR22, CLL22, DB22, LLT22, LWYL22, Pid22, WY22, CCL + 23b, CDD23, LLT23, CCL + 23a, LWCC23, BDD23] have established rigorous convergence guarantees for diffusion models under mild assumptions on the data distribution, these works assume the existence of an oracle for score estimation and leave open whether one can actually provably implement such an oracle for interesting families of data distributions. In practice, the algorithm of choice for score estimation is simply to train a student network via gradient descent (GD) to fit a set of examples ( x, x ) . We thus ask:

In this work, we consider the setting where q is given by a mixture of Gaussians . Concretely, we assume that there exist centers µ ∗ 1 , . . . , µ ∗ K ∈ R d such that

We answer the above question in the affirmative for this class of distributions:

Theorem 1 (Informal, see Theorems 7 and 13) . Gradient descent on the DDPM objective with random initialization efficiently learns the parameters of an unknown mixture of two spherical Gaussians with 1 / poly ( d ) -separated centers.

Theorem 2 (Informal, see Theorem 16) . When there is a warm start of the centers, gradient descent on the DDPM objective efficiently learns the parameters an unknown mixture of K spherical Gaussians with Ω ( √ log(min( K,d ))) -separated centers.

The DDPM objective is described in Algorithm 1. The term 'efficiently' above means that both the running time and sample complexity of our algorithm is polynomial in the dimension d , the inverse accuracy 1 / ε , and the number of components K . In the informal discussion, we often work with population gradients for simplicity, but in our proofs we show that empirical estimates of the gradient suffice (full details can be found in the Appendix).

Algorithm 1: GMMDENOISER( t, { µ (0) i } K i =1 , H )

1 Initialize the parameters for the score estimate at θ t = { µ i,t } i =1 (see Eq. (9) for how the estimate s θ depends on the parameters θ , and Eq. (8) for the definition of µ (0) i,t )

Input: Noise scale t , initialization { µ (0) i } K i =1 , number of gradient descent steps H

(0) K

2 Run gradient descent on the DDPM objective L t ( s θ t ) for H steps where

3 return θ ( H ) t = { µ ( H ) i,t } K i =1 where θ ( H ) t denotes the parameters after H steps of GD.

We refer to Section 1.3 for a formal description of the quantities used in Algorithm 1. Note that there are by now a host of different algorithms for provably learning mixtures of Gaussians (see Section 1.1). For instance, it is already known that expectation-maximization (EM) achieves the quantitative guarantees of Theorems 1 and 2 [DTZ17, XHM16, KC20, SN21], and in fact even stronger guarantees are known via the method of moments. Unlike works based on the method of moments however, our algorithm is practical. And unlike works based on EM, it is based on an approach which is empirically successful for a wide range of realistic data distributions. Furthermore, as we discuss in Section 1.2, the analysis of Algorithm 1 leverages an intriguing and, to our knowledge, novel connection from score estimation to EM, as well as to another notable approach for learning mixture models, namely spectral methods. Roughly speaking, at large noise levels, the gradient updates in Algorithm 1 are essentially performing a type of power iteration, while at small noise levels, the gradient updates are performing the 'M' step in the EM algorithm.

1.1 Related work

Theory for diffusion models. A number of works have given convergence guarantees for DDPMs and variants [DBTHD21, BMR22, CLL22, DB22, LLT22, LWYL22, Pid22, WY22, CCL + 23b, CDD23, LLT23, LWCC23, BDD23, CCL + 23a]. These results show that, given an oracle for accurate score estimation, diffusion models can learn essentially any distribution over R d (e.g. [CCL + 23b, LLT23, CLL22] show this for arbitrary compactly supported distributions). Additionally, two recent works [EAMS22, MW23] have used Eldan's stochastic localization [Eld13, Eld20], which is a reparametrization in time and space of the reverse SDE for DDPMs, to give sampling algorithms for certain distributions arising in statistical physics. As we discuss next, these works are end-to-end in that they also give provable algorithms for score estimation via approximate message passing, though the statistical task they address is not distribution learning.

Provable score estimation. There is a rich literature giving Bayes-optimal algorithms for various natural denoising problems via methods inspired by statistical physics, like approximate message passing (AMP) (e.g. [MV21, CFM21, BM11, Kab03, DMM09, DMM10]) and natural gradient descent (NGD) on the TAP free energy [CFM21, EAMS22, Cel22]. The abovementioned works [EAMS22, MW23] (see also [Cel22]) build on these techniques to give algorithms for the denoising problems that arise in their implementation of stochastic localization. These works on denoising via AMP or NGD are themselves part of a broader literature on variational inference, a suitable literature review would be beyond the scope of this work, see e.g. [BKM17, WJ + 08, MM09].

We are not aware of any provable algorithms for score estimation explicitly in the context of distribution learning . That said, it may be possible to extract a distribution learning result from [EAMS22]. While their algorithm was for sampling from the Sherrington-Kirkpatrick (SK) model given the Hamiltonian rather than training examples as input, if one is instead given training examples drawn from the SK measure, then at sufficiently high temperature one can approximately recover the Hamiltonian [AG22]. In this case, a suitable modification [EAMS22] should be able to yield an algorithm for approximately generating fresh samples from the SK model given training examples.

Learning mixtures of Gaussians. The literature on provable algorithms for learning Gaussian mixture models is vast, dating back to the pioneering work of Pearson [Pea94], and we cannot do justice to it here. We mention only works whose quantitative guarantees are closest in spirit to ours and refer to the introduction of [LL22] for a comprehensive overview of recent works in this direction. For mixtures of identity-covariance Gaussians in high dimensions, the strongest existing guarantee is a polynomial-time algorithm [LL22] for learning the centers as long as their pairwise separation slightly exceeds Ω ( √ log K ) based on a sophisticated instantiation of method of moments inspired by the quasipolynomial-time algorithms of [DKS18, HL18, KSS18]. By the lower bound in [RV17], this is essentially optimal. In contrast, our Theorem 2 only applies given one initializes in a neighborhood of the true parameters of the mixture. We also note the exponential-time spectral algorithm of [SOAJ14] and quasipolynomial-time tensor-based algorithm of [DK20], which achieve density estimation even in the regime where the centers are arbitrarily closely spaced and learning the centers is information-theoretically impossible.

A separate line of work has investigated the 'textbook' algorithm for learning Gaussian mixtures, namely the EM algorithm [BWY17, DS07, DTZ17, XHM16, YYS17, ZLS20, KC20, SN21]. Notably, for balanced mixtures of two Gaussians with the same covariance, [DTZ17] showed that finite-sample EM with random initialization converges exponentially quickly to the true centers. For mixtures of K Gaussians with identity covariance, [KC20, SN21] showed that from an initialization sufficiently close to the true centers, finite-sample EM converges exponentially quickly to the true centers as long as their pairwise separation is Ω ( √ log K ) . In particular, [SN21] establish this local convergence as long as every center estimate is initialized at distance at most ∆ / 2 away from the corresponding true center, where ∆ is the minimum separation between any pair of true centers; this radius of convergence is provably best possible for EM.

Lastly, we note that there are many works giving parameter recovery algorithms mixtures of Gaussians with general mixing weights and covariances, all of which are based on method of moments [KMV10, HP15, Kan21, BS15, MV10, LM23, BDJ + 22, DHKK20]. Unfortunately, for general mixtures of K Gaussians, these algorithms run in time at least d O ( K ) , and there is strong evidence [DKS17, BRST21] that this is unavoidable for computationally efficient algorithms.

1.2 Technical overview

We begin by describing in greater detail the algorithm we analyze in this work. For the sake of intuition, in this overview we will focus on the case of mixtures of two Gaussians ( K = 2) where the centers are well-separated and symmetric about the origin, that is, the data distribution is given by

At the end of the overview, we briefly discuss the key challenges for handling smaller separation and general K .

Loss function, architecture of the score function and student network. The algorithmic task at the heart of score estimation is that of denoising. Formally, for some noise level t > 0 , we are given

a noisy sample where X 0 is a clean sample drawn from the data distribution q , and Z t ∼ N (0 , Id) . Conditioning on X t induces some posterior distribution over the noise Z t , and our goal is to form an estimate s for the mean of this posterior which achieves small error on average over the randomness of X 0 and Z t . That is, we would like to minimize the DDPM objective , which up to rescaling is given by 1

As discussed in the introduction, the algorithm of choice for minimizing this objective in practice is gradient descent on some student network. To motivate our choice of architecture, note that when the data distribution is given by (1), the true minimizer of L t is, up to scaling,

See Appendix A for the derivation. Notably, Eq. (2) is exactly a two-layer neural network with tanh activation. As a result, we use the same architecture for our student network when running gradient descent. That is, given weights µ ∈ R d , our student network is given by s µ ( x ) /defines tanh( µ /latticetop x ) µ -x . The exact gradient updates on µ are given in Lemma C.2.

As we discuss next, depending on whether the noise level t is large or small, this update closely approximates the update in one of two well-studied algorithms for learning mixtures of Gaussians: power method and EM respectively.

Learning mixtures of two Gaussians. We first provide a brief overview of the analysis and then go into the details of the analysis. We start with mixtures of two Gaussians of the form (1) where ‖ µ ∗ ‖ is Ω (1) . In this case, we analyze the following two-stage algorithm. We first use gradient descent on the DDPM objective with large t starting from random initialization. We show that gradient descent in this 'high noise' regime resembles a type of power iteration and gives µ that has a nontrivial correlation with µ ∗ t . Starting from this µ , we then run gradient descent with small t . We show that the gradient descent in this 'small noise' regime corresponds to the EM algorithm and converges exponentially quickly to the ground truth.

Large noise level: connection to power iteration. When t is large, we show that gradient descent on the DDPM objective is closely approximated by power iteration. More precisely, in this regime, the negative gradient of L t ( s µ ) is well-approximated by

where r is a scalar that depends on µ (See Lemma 8). So the result of a single gradient update with step size η starting from µ is given by

This shows us that each gradient step can be approximated by one step of power iteration (without normalization) on the matrix (1 -η r ) Id + 2 η µ ∗ t µ ∗/latticetop t . It is know that running enough iterations of the latter from a random initialization will converge in angular distance to the top eigenvector, which in this case is given by µ ∗ t . This suggests that if we can keep the approximation error in (3) under control, then gradient descent on µ will also allow us to converge to a neighborhood of the ground truth. We implement this strategy in Lemma 10. Next, we argue that once we are in a neighborhood of the ground truth, we can run GD on the DDPM objective at low noise level to refine our estimate.

Low noise level: connection to the EM algorithm. When t is small, we show that gradient descent on the DDPM objective is closely approximated by EM. Here, our analysis uses the fact that µ ∗ is sufficiently large and requires that we initialize µ to have sufficiently large correlation with the true direction µ ∗ t . We can achieve the latter using the larget analysis in the previous section.

Provided we have this, when t is small it turns out that the negative gradient is well-approximated by

Note that the expectation is precisely the 'M'-step in the EM algorithm for learning mixtures of two Gaussians (see e.g. Eq. (2.2) of [DTZ17]). We conclude that a single gradient update with step size η starting from µ is given by mixing the old weights µ with the result of the 'M'-step in EM:

[XHM16] and [DTZ17] showed that EM converges exponentially quickly to the ground truth µ ∗ t from a warm start, and we leverage ingredients from their analysis to prove the same guarantee for gradient descent on the DDPM objective at small noise level t (see Lemma 12).

Extending to small separation. Next, suppose we instead only assume that ‖ µ ∗ ‖ is Ω (1 / poly( d )) , i.e. the two components in the mixture may have small separation. The above analysis breaks down for the following reason: while it is always possible to show that gradient descent at large noise level converges in angular distance to the ground truth, if ‖ µ ∗ ‖ is small, then we cannot translate this to convergence in Euclidean distance.

We circumvent this as follows. Extending the connection between gradient descent at large t and power iteration, we show that a similar analysis where we instead run projected gradient descent over the ball of radius ‖ µ ∗ ‖ yields a solution arbitrarily close to the ground truth, even without the EM step. 2 The projection step can be thought of as mimicking the normalization step in power iteration.

It might appear to the reader that this projected gradient-based approach is strictly superior to the two-stage algorithm described at the outset. However, in addition to obviating the need for a projection step when separation is large, our analysis for the two-stage algorithm has the advantage of giving much more favorable statistical rates. Indeed, we can show that the sample complexity of the two-stage algorithm has optimal dependence on the target error ( 1 / ε 2 ), whereas we can only show a suboptimal dependence ( 1 / ε 8 ) for the single-stage algorithm.

Extending to general K . The connection between gradient descent on the DDPM objective at small t and the EM algorithm is sufficiently robust that for general K , our analysis for K = 2 can generalize once we replace the ingredients from [XHM16] and [DTZ17] with the analogous ingredients in existing analyses for EM with K Gaussians. For the latter, it is known that if the centers of the Gaussians have separation Ω ( √ log min( K,d )) , then EM will converge from a warm start [KC20, SN21]. By carefully tracking the error in approximating the negative gradient with the 'M'-step in EM, we are able to show that gradient descent on the DDPM objective at small t achieves the same guarantee.

1.3 Preliminaries

Diffusion models. Throughout the paper, we use either q or q 0 to denote the data distribution and X or X 0 to denote the corresponding random variable on R d . The two main components in diffusion models are the forward process and the reverse process . The forward process transforms samples from the data distribution into noise, for instance via the Ornstein-Uhlenbeck (OU) process :

where ( W t ) t ≥ 0 is a standard Brownian motion in R d . We use q t to denote the law of the OU process at time t . Note that for X t ∼ q t ,

The reverse process then transforms noise into samples, thus performing generative modeling. Ideally, this could be achieved by running the following stochastic differential equation for some choice of terminal time T :

where now W t is the reversed Brownian motion. In this reverse process, the iterate X ← t is distributed acccording to q T -t for every t ∈ [0 , T ] , so that the final iterate X ← T is distributed according to the

data distribution q 0 . The function ∇ x ln q t is called the score function , and because it depends on q which is unknown, in practice one estimates it by minimizing the score matching loss

A standard calculation (see e.g. Appendix A of [CCL + 23b]) shows that this is equivalent to minimizing the DDPM objective in which one wants to predict the noise Z t from the noisy observation X t , i.e.

While we have provided background on diffusion models for context, in this work we focus specifically on the optimization problem (5).

Mixtures of Gaussians. We consider the case of learning mixtures of K equally weighted Gaussians:

where µ ∗ i denotes the mean of the i th Gaussian component. We define θ ∗ = { µ ∗ 1 , µ ∗ 2 . . . , µ ∗ K } . For the mixtures of two Gaussians, we can simplify the data distribution as

Note that distribution in Eq. (7) is equivalent to the distribution Eq. (6) with K = 2 because shifting the latter by its mean will give the former distribution, and furthermore the necessary shift can be estimated from samples. The following is immediate:

Lemma 3. If q 0 is a mixture of K Gaussians as in Eq. (6) , then for any t > 0 , q t is the mixture of K Gaussians given by

See Appendix A for a proof of this fact. We can see that the means of q t get rescaled according to the noise level t . We also define θ ∗ t = { µ ∗ 1 ,t , µ ∗ 2 ,t , . . . , µ ∗ K,t } .

Lemma 4. The score function for distribution q t , for any t > 0 , is given by

See Appendix A for the calculation.

Recall that ∇ x log q t ( x ) is the minimizer for the score-matching objective given in Eq. (4). Therefore, we parametrize our student network architecture similarly to the optimal score function. Our student architecture for mixtures of K Gaussians is

where θ t = { µ 1 ,t , µ 2 ,t , . . . , µ K,t } denotes the set of parameters at the noise scale t . For mixtures of two Gaussians, we simplify the student architecture as follows:

As θ t only depends on µ t in the case of mixtures of two Gaussians, we simplify the notation of the score function from s θ t ( x ) to s µ t ( x ) in that case. We use ˆ µ t and ˆ µ ∗ t to denote the unit vector along the direction of µ t and µ ∗ t respectively. Note that we often use µ t (or θ t ) to denote the current iterate of gradient descent on the DDPM objective and µ ′ t to denote the iterate after taking a gradient descent step from µ t .

Expectation-Maximization (EM) algorithm. The EM algorithm is composed of two steps: the E-step and the M-step. For mixtures of Gaussians, the E-step computes the expected log-likelihood based on the current mean parameters and the M-step maximizes this expectation to find a new estimate of the parameters.

Fact 5 (See e.g., [DTZ17, YYS17, KC20] for more details) . When X is the mixture of K Gaussian and { µ 1 , µ 2 , . . . , µ K } are current estimates of the means, the population EM update for all i ∈ { 1 , 2 , . . . , K } is given by

The EM update for mixtures of two Gaussians given in Eq. (7) simplifies to

An analogous version of the EM algorithm, called the gradient EM algorithm, takes a gradient step in the direction of the M-step instead of optimizing the objective in the M-step fully.

Fact 6 (See e.g., [YYS17, SN21] for more details) . For all i ∈ { 1 , 2 , . . . , K } , the gradient EMupdate for mixtures of K Gaussian is given by

where η is the learning rate.

2 Warmup: mixtures of two Gaussians with constant separation

In this section, we formally state our result for learning mixtures of two Gaussians with constant separation. This case highlights the main proof techniques, namely viewing gradient descent on the DDPM objective as power iteration and as the EM algorithm.

2.1 Result and algorithm

Theorem 7. There is an absolute constant c > 0 such that the following holds. Suppose a mixture of two Gaussians with the mean parameter µ ∗ satisfies ‖ µ ∗ ‖ > c . Then, for any ε > 0 , there is a procedure that calls Algorithm 1 at two different noise scales t and outputs ˜ µ such that ‖ ˜ µ -µ ∗ ‖ ≤ ε with high probability. Moreover, the algorithm has time and sample complexity poly( d ) / ε 2 (see Theorem C.1 for more precise quantitative bounds).

Algorithm. The algorithm has two stages. In the first stage we run gradient descent on the DDPM objective described in Algorithm 1 from a random Gaussian initialization and noise scale t 1 for a fixed number of iterations H where t 1 = Θ (log d ) ('high noise') and H = poly( d, 1 / ε ) . In the second stage, the procedure uses the output of the first step as initialization and runs Algorithm 1 at a 'low noise' scale of t 2 = Θ (1) .

2.2 Proof outline of Theorem 7

We provide a proof sketch of correctness of the above algorithm and summarize the main technical lemmas here. All proofs of the following lemmas can be found in Appendix C.

Part I: Analysis of high noise regime and connection to power iteration. We show that in the large noise regime, the negative gradient -∇ L t ( s t ) is well-approximated by 2 µ ∗ t µ ∗/latticetop t µ t -3 ‖ µ t ‖ 2 µ t . Recall that this result is the key to showing the resemblance between gradient descent and power iteration. Concretely, we show the following lemma:

Lemma 8 (See Lemma C.3 for more details) . For t = Θ (log d ) , the gradient descent update on the DDPM objective L t ( s t ) can be approximated with 2 µ ∗ t µ ∗/latticetop t µ t -3 ‖ µ t ‖ 2 µ t :

From Lemma 8, it immediately follows that µ ′ t , the result of taking a single gradient step starting from µ t , is well-approximated by the result of taking a single step of power iteration for a matrix whose leading eigenvector is µ ∗ t :

The second key element is to show that as a consequence of the above power iteration update, the gradient descent converges in angular distance to the leading eigenvector. Concretely, we show the following lemma:

Lemma 9 (Informal, see Lemma C.5 for more details) . Suppose µ ′ t is the iterate after one step of gradient descent on the DDPM objective from µ t . Denote the angle between µ t and µ ∗ t to be θ and between µ ′ t and µ ∗ t to be θ ′ . In this case, we show that

where κ 1 < 1 and κ 2 ≤ 1 / poly( d ) .

Note tan θ ′ < tan θ implies that θ ′ < θ or equivalently 〈 ˆ µ ′ t , ˆ µ ∗ t 〉 > 〈 ˆ µ t , ˆ µ ∗ t 〉 . Thus, the above lemma shows that by taking a gradient step in the DDPM objective, the angle between µ t and µ ∗ t decreases. By iterating this, we obtain the following lemma:

Lemma 10 (Informal, see Lemma C.6 for more details) . Running gradient descent from a random initialization on the DDPM objective L t ( s µ ) for t = O (log d ) gives µ t for which 〈 ˆ µ t , ˆ µ ∗ t 〉 is Ω (1) .

Note that we cannot keep running gradient descent at this high noise scale and hope to achieve µ such that ‖ µ -µ ∗ ‖ is O ( ε ) . This is because Lemma 9 can only guarantee that the angle between µ t and µ ∗ t is O ( ε ) , but this does not imply ‖ µ -µ ∗ ‖ is O ( ε ) . Instead, as described in Part II, we will proceed with a smaller noise scale.

Part II: Analysis of low noise regime and connection to EM. In the low noise regime, we run Algorithm 1 using the output from Part I as our initialization. Our analysis here shows that whenever the initialization µ t satisfies the condition of 〈 ˆ µ t , ˆ µ ∗ t 〉 being Ω (1) , ‖ µ t -µ ∗ t ‖ contracts after every gradient step. To start with, we show that the result of a population gradient step on the DDPM objective L t ( s µ ) results in the following:

where µ ′ t is the parameter after a gradient step, η is the learning rate, and function G is given by

Note we use the population gradient here only for simplicity; in the Appendix we show that empirical estimates of the gradient suffice. After some calculation, we can show that

Using Fact 5, we know that E x ∼ N ( µ ∗ t , Id) [tanh( µ /latticetop t x ) x ] is precisely the result of one step of EM starting from µ t , and it is known [DTZ17] that the EM update contracts the distance between µ t and µ ∗ t as follows:

It remains to control the second term in Eq. (10), for which we prove the following:

Lemma 11 (Informal, see Lemma C.9 for more details) . When ‖ µ ∗ ‖ = Ω (1) and the noise scale t = Θ (1) , then for every µ with 〈 ˆ µ, ˆ µ ∗ 〉 being Ω (1) , the following inequality holds:

Combining Eq. (11) and Lemma 11 with Eq. (10), we have

∥ ∥ µ ′ t -µ ∗ t ∥ ∥ ≤ (1 -η (1 -λ 1 -λ 2 )) ‖ µ t -µ ∗ t ‖ . (12) We can set parameters to ensure that λ 1 + λ 2 < 1 and therefore that ‖ µ t -µ ∗ t ‖ contracts with each gradient step. Applying Lemma 11 and Eq. (12), we obtain the following lemma summarizing the behavior of gradient descent on the DDPM objective in the low noise regime.

Lemma 12 (Informal) . For any ε > 0 and for the noise scale t = Θ (1) , starting from an initialization µ t for which 〈 ˆ µ t , ˆ µ ∗ t 〉 = Ω (1) , running gradient descent on the DDPM objective L t ( s µ ) will give us mean parameter ˜ µ such that ‖ ˜ µ -µ ∗ ‖ ≤ O ( ε ) .

Combining Lemma 10 and Lemma 12, we obtain our first main result, Theorem 7, for learning mixtures of two Gaussians with constant separation. For the full technical details, see Appendix C.

3 Extensions: small separation and more components

3.1 Mixtures of two Gaussians with small separation

In this section, we briefly sketch how the ideas from Section 2 can be extended to give our second main result, namely on learning mixtures of two Gaussians even with small separation . We defer the full technical details to Appendix D.

Theorem 13. Suppose a mixture of two Gaussians has mean parameter µ ∗ that satisfies ‖ µ ∗ ‖ = Ω ( 1 poly( d ) ) . Then, for any ε > 0 , there exists a modification of Algorithm 1 that provides an estimate µ such that ‖ µ -µ ∗ ‖ ≤ O ( ε ) with high probability. Moreover, the algorithm has time and sample complexity poly( d ) / ε 8 (see Theorem D.1 for more precise quantitative bounds).

Algorithm modification. The algorithm that we analyze runs projected gradient descent on the DDPM objective but only in the high noise scale regime where t = O (log d ) . At each step, we project the iterate µ to the ball of radius R , where R is an empirical estimate for ‖ µ ∗ ‖ obtained by drawing samples x 1 , . . . , x n from the data distribution and forming R /defines ( 1 n ∑ n i =1 ‖ x i ‖ 2 -d ) 1 / 2 .

Observe that if R is approximately equal to ‖ µ ∗ t ‖ , then the projection step in our algorithm ensures that our final estimate µ t satisfies this additional condition of ‖ µ t ‖ ≈‖ µ ∗ t ‖ . It is not hard to show that R 2 is an unbiased estimate of ‖ µ ∗ t ‖ 2 , so standard concentration shows that taking n = poly( d, 1 ε ) suffices to ensure that R is sufficiently close to ‖ µ ∗ t ‖ .

Proof sketch. Lemma 9 and Lemma 10 apply even when the components of the mixture have small separation, and they show that running gradient descent on the DDPM objective results in µ t and µ ∗ t being O (1) close in angular distance. Although our analysis can be extended to show that gradient descent can achieve O ( ε ) angular distance, this does not guarantee that ‖ µ t -µ ∗ t ‖ is O ( ε ) . If in addition to being O ( ε ) close in angular distance, we also have that ‖ µ t ‖ ≈‖ µ ∗ t ‖ , then it is easy to see that ‖ µ t -µ ∗ t ‖ is indeed O ( ε ) .

3.2 Mixtures of K Gaussians, from a warm start

In this section, we state our third main result, namely for learning mixtures of K Gaussians given by Eq. (6) from a warm start, and provide an overview of how the ideas from Section 2 can be extended to obtain this result.

/negationslash

Assumption 15. (Initialization) For each component i ∈ { 1 , 2 , . . . , K } , we have an initialization µ (0) i with the property that ‖ µ (0) i -µ ∗ i ‖ ≤ C ′ √ log(min( K,d )) for sufficiently small absolute constant C ′ > 0 .

Assumption 14. (Separation) For a mixture of K Gaussians given by Eq. (6) , for every pair of components i, j ∈ { 1 , 2 , . . . , K } with i = j , we assume that the separation between their means ‖ µ ∗ i -µ ∗ j ‖ ≥ C √ log(min( K,d )) for sufficiently large absolute constant C > 0 .

Theorem 16. Suppose a mixture of K Gaussians satisfies Assumption 14. Then, for any ε = Θ (1 / poly( d )) , running gradient descent on the DDPM objective (Algorithm 1) at low noise scale t = O (1) and with initialization satisfying Assumption 15 results in mean parameters { µ i } K i =1 such that with high probability, the mean parameters satisfy ‖ µ i -µ ∗ i ‖ ≤ O ( ε ) for each i ∈ { 1 , 2 , . . . , K } . Additionally, the runtime and sample complexity of the algorithm is poly( d, 1 / ε ) (see Theorem E.1 for more precise quantitative bounds).

We provide a brief overview of the proof here. The full proof can be found in Appendix E.

Proof sketch. For learning mixtures of two Gaussians, we have already established the connection between gradient descent on the DDPM objective and the EM algorithm. For mixtures of K Gaussians, however, in a local neighborhood around the ground truth parameters θ ∗ , we show an equivalence between gradient EM (recall gradient EM performs one-step of gradient descent on the 'M' step objective) and gradient descent on the DDPM objective. In particular, our main technical lemma (Lemma E.4) shows that for noise scale t = Θ (1) and for any µ i that satisfies ‖ µ i -µ ∗ i ‖ ≤ O ( √ log(min( K,d ))) , we have

Therefore, the iterate µ ′ i,t resulting from a single gradient step on the DDPM objective L t ( s θ t ) with learning rate η is given by

Comparing Fact 6 with Eq. (13), we see the correspondence in this regime between gradient descent on the DDPM objective to gradient EM. Using this connection and an existing local convergence guarantee from the gradient EM literature [SN21, KC20], we obtain our main theorem for mixtures of K Gaussians. Full details can be found in Appendix E.

Acknowledgments

SC would like to thank Sinho Chewi, Khashayar Gatmiry, Frederic Koehler, and Holden Lee for enlightening discussions on sampling and score estimation. KS and AK are supported by the NSF AI Institute for Foundations of Machine Learning (IFML). SC is supported by NSF Award 2103300.

References

[AG22] Ahmed El Alaoui and Jason Gaitonde. Bounds on the covariance matrix of the sherrington-kirkpatrick model. arXiv preprint arXiv:2212.02445 , 2022.

[BDD23] Joe Benton, George Deligiannidis, and Arnaud Doucet. Error bounds for flow matching methods. arXiv preprint arXiv:2305.16860 , 2023.

[BDJ + 22] Ainesh Bakshi, Ilias Diakonikolas, He Jia, Daniel M Kane, Pravesh K Kothari, and Santosh S Vempala. Robustly learning mixtures of k arbitrary gaussians. In Proceedings of the 54th Annual ACM SIGACT Symposium on Theory of Computing , pages 1234-1247, 2022.

[BKM17] David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians. Journal of the American statistical Association , 112(518):859-877, 2017.

[BM11] Mohsen Bayati and Andrea Montanari. The dynamics of message passing on dense graphs, with applications to compressed sensing. IEEE Transactions on Information Theory , 57(2):764-785, 2011.

[BMR22] Adam Block, Youssef Mroueh, and Alexander Rakhlin. Generative modeling with denoising auto-encoders and Langevin sampling. arXiv preprint 2002.00107 , 2022.

[BRST21] Joan Bruna, Oded Regev, Min Jae Song, and Yi Tang. Continuous lwe. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing , pages 694-707, 2021.

[BS15] Mikhail Belkin and Kaushik Sinha. Polynomial learning of distribution families. SIAM Journal on Computing , 44(4):889-911, 2015.

[BWY17] Sivaraman Balakrishnan, Martin J Wainwright, and Bin Yu. Statistical guarantees for the em algorithm: From population to sample-based analysis. 2017.

[CCL + 23a] Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast. arXiv preprint arXiv:2305.11798 , 2023.

[CCL + 23b] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference on Learning Representations , 2023.

[CDD23] Sitan Chen, Giannis Daras, and Alexandros G Dimakis. Restoration-degradation beyond linear diffusions: A non-asymptotic analysis for ddim-type samplers. arXiv preprint arXiv:2303.03384 , 2023.

[Cel22] Michael Celentano. Sudakov-fernique post-amp, and a new proof of the local convexity of the tap free energy. arXiv preprint arXiv:2208.09550 , 2022.

[CFM21] Michael Celentano, Zhou Fan, and Song Mei. Local convexity of the tap free energy and amp convergence for z2-synchronization. arXiv preprint arXiv:2106.11428 , 2021.

[CLL22] Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: user-friendly bounds under minimal smoothness assumptions. arXiv preprint arXiv:2211.01916 , 2022.

[DB22] Valentin De Bortoli. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research , 2022.

[DBTHD21] Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion Schrödinger bridge with applications to score-based generative modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 17695-17709. Curran Associates, Inc., 2021.

[DHKK20] Ilias Diakonikolas, Samuel B Hopkins, Daniel Kane, and Sushrut Karmalkar. Robustly learning any clusterable mixture of gaussians. arXiv preprint arXiv:2005.06417 , 2020.

[DK20] Ilias Diakonikolas and Daniel M Kane. Small covers for near-zero sets of polynomials and learning latent variable models. In 2020 IEEE 61st Annual Symposium on Foundations of Computer Science (FOCS) , pages 184-195. IEEE, 2020.

[DKS17] Ilias Diakonikolas, Daniel M Kane, and Alistair Stewart. Statistical query lower bounds for robust estimation of high-dimensional gaussians and gaussian mixtures. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) , pages 73-84. IEEE, 2017.

[DKS18] Ilias Diakonikolas, Daniel M Kane, and Alistair Stewart. List-decodable robust mean estimation and learning mixtures of spherical gaussians. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing , pages 1047-1060, 2018.

[DMM09] David L Donoho, Arian Maleki, and Andrea Montanari. Message-passing algorithms for compressed sensing. Proceedings of the National Academy of Sciences , 106(45):18914-18919, 2009.

[DMM10] David L Donoho, Arian Maleki, and Andrea Montanari. Message passing algorithms for compressed sensing: I. motivation and construction. In 2010 IEEE information theory workshop on information theory (ITW 2010, Cairo) , pages 1-5. IEEE, 2010.

[DS07] Sanjoy Dasgupta and Leonard J Schulman. A probabilistic analysis of em for mixtures of separated, spherical gaussians. Journal of Machine Learning Research , 8:203-226, 2007.

[DTZ17] Constantinos Daskalakis, Christos Tzamos, and Manolis Zampetakis. Ten steps of em suffice for mixtures of two gaussians. In Conference on Learning Theory , pages 704-710. PMLR, 2017.

[EAMS22] Ahmed El Alaoui, Andrea Montanari, and Mark Sellke. Sampling from the sherringtonkirkpatrick gibbs measure via algorithmic stochastic localization. In 2022 IEEE 63rd Annual Symposium on Foundations of Computer Science (FOCS) , pages 323-334. IEEE, 2022.

[Eld13] Ronen Eldan. Thin shell implies spectral gap up to polylog via a stochastic localization scheme. Geometric and Functional Analysis , 23(2):532-569, 2013.

[Eld20] Ronen Eldan. Taming correlations through entropy-efficient measure decompositions with applications to mean-field approximation. Probability Theory and Related Fields , 176(3-4):737-755, 2020.

[HJA20] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

[HL18] Samuel B Hopkins and Jerry Li. Mixture models, robustness, and sum of squares proofs. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing , pages 1021-1034, 2018.

[HP15] Moritz Hardt and Eric Price. Tight bounds for learning a mixture of two gaussians. In Proceedings of the forty-seventh annual ACM symposium on Theory of computing , pages 753-760, 2015.

[Kab03] Yoshiyuki Kabashima. A cdma multiuser detection algorithm on the basis of belief propagation. Journal of Physics A: Mathematical and General , 36(43):11111, 2003.

[Kan21] Daniel M Kane. Robust learning of mixtures of gaussians. In Proceedings of the 2021 ACM-SIAM Symposium on Discrete Algorithms (SODA) , pages 1246-1258. SIAM, 2021.

[KC20] Jeongyeol Kwon and Constantine Caramanis. The em algorithm gives sampleoptimality for learning mixtures of well-separated gaussians. In Conference on Learning Theory , pages 2425-2487. PMLR, 2020.

[KMV10] Adam Tauman Kalai, Ankur Moitra, and Gregory Valiant. Efficiently learning mixtures of two gaussians. In Proceedings of the forty-second ACM symposium on Theory of computing , pages 553-562, 2010.

[KSS18] Pravesh K Kothari, Jacob Steinhardt, and David Steurer. Robust moment estimation and improved clustering via sum of squares. In Proceedings of the 50th Annual ACM SIGACT Symposium on Theory of Computing , pages 1035-1046, 2018.

[LL22] Allen Liu and Jerry Li. Clustering mixtures with almost optimal separation in polynomial time. In Proceedings of the 54th Annual ACM SIGACT Symposium on Theory of Computing , pages 1248-1261, 2022.

[LLT22] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022.

[LLT23] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-based generative modeling for general data distributions. In International Conference on Algorithmic Learning Theory , pages 946-985. PMLR, 2023.

[LM23] Allen Liu and Ankur Moitra. Robustly learning general mixtures of gaussians. Journal of the ACM , 2023.

[LWCC23] Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv preprint arXiv:2306.09251 , 2023.

[LWYL22] Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: understanding and extending diffusion generative models. arXiv preprint arXiv:2208.14699 , 2022.

[MM09] Marc Mezard and Andrea Montanari. Information, physics, and computation . Oxford University Press, 2009.

[MV10] Ankur Moitra and Gregory Valiant. Settling the polynomial learnability of mixtures of gaussians. In 2010 IEEE 51st Annual Symposium on Foundations of Computer Science , pages 93-102. IEEE, 2010.

[MV21] Andrea Montanari and Ramji Venkataramanan. Estimation of low-rank matrices via approximate message passing. The Annals of Statistics , 49(1), 2021.

[MW23] Andrea Montanari and Yuchen Wu. Posterior sampling from the spiked models via diffusion processes. arXiv preprint arXiv:2304.11449 , 2023.

[Pea94] Karl Pearson. Contributions to the mathematical theory of evolution. Philosophical Transactions of the Royal Society of London. A , 185:71-110, 1894.

[Pid22] Jakiw Pidstrigach. Score-based generative models detect manifolds. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 35852-35865. Curran Associates, Inc., 2022.

[RBL + 22] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10684-10695, 2022.

[RDN + 22] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 , 2022.

[RV17] Oded Regev and Aravindan Vijayaraghavan. On learning mixtures of well-separated gaussians. In 2017 IEEE 58th Annual Symposium on Foundations of Computer Science (FOCS) , pages 85-96. IEEE, 2017.

[SCS + 22] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 , 2022.

[SDWMG15] Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

[SE19] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

[SN21] Nimrod Segol and Boaz Nadler. Improved convergence guarantees for learning gaussian mixture models by em and gradient em. Electronic journal of statistics , 15(2):4510-4544, 2021.

[SOAJ14] Ananda Theertha Suresh, Alon Orlitsky, Jayadev Acharya, and Ashkan Jafarpour. Near-optimal-sample estimators for spherical gaussian mixtures. Advances in Neural Information Processing Systems , 27, 2014.

[SSDK + 20] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

[Ver] Roman Vershynin. High-Dimensional Probability: An Introduction with Applications in Data Science . Number 47 in Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press.

[VW04] Santosh Vempala and Grant Wang. A spectral algorithm for learning mixture models. Journal of Computer and System Sciences , 68(4):841-860, 2004. Special Issue on FOCS 2002.

[WJ + 08] Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends® in Machine Learning , 1(1-2):1305, 2008.

[WY22] Andre Wibisono and Kaylee Y. Yang. Convergence in KL divergence of the inexact Langevin algorithm with application to score-based generative models. arXiv preprint 2211.01512 , 2022.

[XHM16] Ji Xu, Daniel J Hsu, and Arian Maleki. Global analysis of expectation maximization for mixtures of two gaussians. Advances in Neural Information Processing Systems , 29, 2016.

[YYS17] Bowei Yan, Mingzhang Yin, and Purnamrita Sarkar. Convergence analysis of gradient em for multi-component gaussian mixture. arXiv preprint arXiv:1705.08530 , 2017.

[ZLS20] Ruofei Zhao, Yuanzhi Li, and Yuekai Sun. Statistical convergence of the em algorithm on gaussian mixture models. 2020.

Peer reviews:
comments: SUMMARY:
The authors propose to leverage the DDPM objective to learn Mixtures of Gaussians and prove that gradient descent on the DDPM objective can efficiently recover the ground truth parameters of the mixture model under certain assumptions.

STRENGTHS:
Several interesting insights are revealed, such as those associated with Eq. (5) and those related to large/low noise levels.

WEAKNESSES:
(1) The novelties of the proposed method over existing works are not clearly stated. For example, in Line 52, it's stated that ``EM achieves the quantitative guarantees of Theorems 1 and 2 Daskalakis et al. (2017); Xu et al. (2016);...'' So what are the advantages of using the DDPM objective?

(2) It's not clear how practical the proposed method is for general Mixtures of Gaussians with $K\ge 3$ components, where a global optimum is generally impossible for maximizing the loglikelihood [1]. Note that a contradiction between the presented paper and [1] might emerge here, because minimizing the DDPM objective is equivalent to maximizing a hierarchically constructed lower bound of the loglikelihood [2]. 

(3) No empirical experiments are given.

[1] Chi Jin, Yuchen Zhang, Sivaraman Balakrishnan, Martin J Wainwright, and Michael I Jordan. Local maxima in the likelihood of Gaussian mixture models: Structural results and algorithmic consequences. Advances in neural information processing systems, 29, 2016.

[2] Calvin Luo. Understanding Diffusion Models: A Unified Perspective. Arxiv, 2022.

QUESTIONS:
Please address the questions in the above ``Weaknesses.''

Minor.

(1) "poweri" in Line 156 is a typo.

(2) "Gaussian"  in Line 214 is a typo.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 1
DATE: 2023-06-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows that the diffusion model can be used to learn mixtures of Gaussians. In particular, they show that GD on the DDPM objective can efficiently recover the ground truth parameters for both mixtures of two spherical Gaussians and mixtures of $K$ spherical Gaussians under different assumptions. The key ingredient in their analysis is the observation that there is an inherent connection between score-based models and EM, spectral methods.\n\nSTRENGTHS:\nThis is a very interesting paper and can be served as a starting point to further explore the power of DDPM from the optimization perspective. Firstly, mixtures of Gaussians are one of the most important and fundamental distribution families of interest, and showing DDPM works for mixtures of Gaussians can be very strong support for its solidness. On the other hand, different from prior work, this work does not require the existence of an oracle for score estimation. Instead, they show that empirical estimation suffices. Lastly, the connection between DDPM and the classic algorithms in this field, i.e., EM and spectral methods can be of independent interest. This might be used to show the effectiveness of DDPM.\n\nWEAKNESSES:\nMy main concern is that the model and algorithm are a little bit artificial.\n\n- First, the model used for mixtures of two Gaussians is $s_{\\theta_t}(x)=\\tanh \\left(\\mu_t^{\\top} x\\right) \\mu_t-x$. Here $-x$ can be understood as the residual connection and $\\tanh(\\cdot)$ is also a reasonable activation. However, using the same weight $\\mu_t$ for both layers is a little bit restrictive. A far more general model is $s_{\\theta_t}(x)=\\tanh \\left(u_t^{\\top} x\\right) v_t-x$. My conjecture is that it is not hard to extend the current result to this \"asymmetric\" version. The reason is that there might exist some kind of balanceness between $u_t$ and $v_t$ so that they are roughly the same. \n\n- Second, the authors consider two noise regimes: first they use a large noise and run several iterations of GD, then they switch to the small noise regime. This mimics the practical noise regime. However, in practice,  people decrease the noise level in a continuous manner, like exponentially decaying. Is it possible to adapt the current analysis to a more realistic noise-decaying regime?\n\nSome technical issues: \n- When claiming high noise regime, only the upper bound is provided, see line 252. I think there should be also a lower bound for $t_1$. \n\nLast question: The result in this paper only show that DDPM is as good as EM and spectral methods. A far more important question is why empirically DDPM outperforms these methods. Can DDPM outperforms the existing methods for the mixtures of Gaussians in some aspects like computational time?\n\nQUESTIONS:\nPlease see the weakness part.\n\nLIMITATIONS:\nPlease see the weakness part.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new approach to learning Gaussian mixture models using the denoising diffusion probabilistic model (DDPM) objective. The authors provide a 2-part algorithm that allows one to reconstruct the parameters of a mixture of Gaussians. The first part of the algorithm uses gradient descent with \"high nose\" scale of $t$ and random initialization to learn parameter, correlated to the true ones, while the second step refines these parameters, using those obtained in the previous step as a \"warm start\" and using the small $t$ mode.\n\nSTRENGTHS:\nThe paper gives the relation between gradient descent with loss from denoising diffusion probabilistic model and a) power iteration in case of large $t$ and b) Expectation-Maximization algorithm in case of not large $t$.\n\nSome number of theorems and lemmas are proved.\n\nWEAKNESSES:\n1. The presentation of the results is weak. The article is poorly structured, with no \"conclusions\" or \"discussion\" section. A lot of fuzzy language. Such wording may be acceptable in informal descriptions of results, but even in the section with formal descriptions phrases like \"with high probability\" are used. The main text of the article in the main pdf file differs from the version in the file with the Appendix.\n\n2. The matrix $(1-\\eta r)I +2\\eta \\mu^*_t {\\mu^*_t}^T$ has eigenvalues close to one. The largest of them (for which the eigenvector in power iteration (6) is sought) differs only by an additional term proportional to $e^{-2t}\\eta$. This term is very small, since by definition $t$ in this case is large, and the learning rate $\\eta$ is usually chosen very small ($\\sim10^{-2}-10^{-4}$). In such a situation power iteration can converge to an arbitrary vector. Moreover, as expression (6) is an approximation, this convergence is even more uncertain.\n\n3. There are no numerical experiments, even for the simplest cases. However, the real convergence of power iterations with the matrix (7) is not obvious.\n\n 4. There is no separate  comparison with other, including state-of-the-art, approaches to recovering Gaussian mixture parameters. Although paper has rather large literature review section.\n\n\nMinor:\n\nline 156 typo: \"poweri iteration\"\n\n\"EM\" abbreviation is deciphered only on page 6, but widely used starting from the abstract.\n\nQUESTIONS:\n1. Have you performed numerical experiments to check the convergence of the method and compare it with other approaches?\n\n2. In reality, for some model numerical parameter values, how many iterations are needed for power iteration described in (6)--(7) to converge?\n\n3. line 246: (Sec. 2.1 with Formal statements) What is \"high probability\"?\n\n\nMinor\n\nWhat authors mean by \"well-separated\" (line 131)\n\nWhat authors mean by _student_ network (line 142)\n\nLIMITATIONS:\n1. Only rough asymptotic estimates are given. Real constants, which can be in practical experiments are not given, however, they can be large. \n2. The proof for the case of large $t$ (the relation of DDPM to power iteration) is based on an approximate, not exact, formula (6) and thus it is not conclusive.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nProofs are given for showing that the true mean parameters of Gaussian mixture models (GMMs) with identity covariance matrices can be recovered when using gradient descent to optimize the DDPM objective. It is argued that it is not well understood if score-based models can provably estimate the parameters of the data distributions, which is the objective of this submission. The findings and results are connected to the EM algorithm.\n\nSTRENGTHS:\nThe paper attempts contribute with fundamental understanding of diffusion models, focusing on an important class data distributions, GMMs. Diffusion models are clearly of high interest to the NeurIPS community and theoretical results of the type provided in the submission could be significant, or of interest to researchers in the field. \n\nThe literature review is thorough and it appears that the topic of the paper is original, in the sense that the theoretical results do not exist in the literature. It is good that the results are connected to other well-studied algorithms, like the EM algorithm.\n\nThe clarity of the purpose of the submission benefits from the bullet list in the abstract.\n\nWEAKNESSES:\nThe structure of the submission is unfortunately not very strong, which makes the arguments in the submission tedious to follow. I am missing sufficient background to the power method and the EM algorithm, a background of the former is not provided, and there is no conclusion section or discussion. Such sections would help the reader to understand the contributions of the submission in terms of the lingo and terminology introduced in the paper. Also, it would give the paper a more concise ending, as it now ends with a proof sketch. \n\nMy score of the paper is mainly based on the structure and quality (presentation) of the paper, which I believe is below acceptable for a NeurIPS paper. However, my confidence level is relatively low as theoretical results for diffusion/score-based models is slightly outside my area of expertise. That is, am not sure how significant the results in the submission are--hopefully other reviewers can aid to evaluate this part. Nonetheless, as it stands, my opinion is that the inadequate presentation of the work outweighs (my understanding of) the significance of the results in the presented work.\n\nTo help the authors in the revision of the paper, I below enumerate some suggestions on how to improve the structure and some of my concerns. I hope they are helpful.\n\n* GD is frequently used to denote gradient descent. But it is not defined as an acronym in the abstract nor in the paper. I inferred the definition of GD only after reading quite a bit of the paper. This confused me, as I did not know what GD referred to. I suggest introducing the acronym in the abstract and write gradient descent (GD) in line 35 where GD is first used.\n\n* In line 35, what is the meaning of a \"natural\" data distribution? I could find an explanation of it and I could not derive what it ought to specify. Is \"natural\" an important distinction?\n\n* Sentences like \"we use this as our student network architecture when running gradient descent.\" (line 146) gives me expectation that some experiments have been carried out where these architectures have been used. As there are no experiments in the paper, this confused me. A potential reformulation could be \"one could use\" as you have in fact not used it. Similar confusion arouse at line 162 \"we run projected gradient descent\".\n\n* Related to the previous bullet, why not run experiments to validate the results on synthetic data? \n\n* Is the heading in line 164 supposed to be a paragraph or a section? Now it is an empty paragraph.\n\n* I am not familiar with the power method and no background or explanation of it is given in the paper. I think this needs to be included to make the paper more stand alone.\n\n* Projected gradient descent is introduced after the background section without proper explanation. Is projected gradient descent typically used to train diffusion models?\n\n* The EM algorithm is central to the findings in the paper. As such, I think it should be more carefully explained. Especially, including a Background section earlier than Sec 1.3 would strengthen the paper and make it easier to follow the arguments made in the earlier sections. For instance, \"We show that the gradient descent in this \"small noise\" regime corresponds to the EM algorithm\" in line 158 would become easier to understand.\n\n* In the setting where $K=2$, is there an additional, consistent assumption that the means of the Gaussian are centered around the origin? I.e. the means are $-\\mu^*$ and $\\mu^*$? I read the section above Eq. (2) as this being an example of a GMM \"for the sake of intuition\". But then this parameter setting is specified again in Eq. (13). Does this mean that the results in the paper are constrained to Gaussians centered around the origin? If so, this should be specified in the relating theorems? I.e. Thm 1 and 2 and their formal counterparts.\n\n**Minor issues** \n* line 46: the flow if the sentence could be improved by using punctuation after \"$d$, the dimension.\"\n* The acronym EM is most often used without definition. Although this is more clear what it refers to, \"expectation-maximization\" is spelled out in other parts of the paper. I suggest introducing the definition and sticking to it throughout the paper.\n\nQUESTIONS:\n* In Sec. 1.2, it would be nice to see a discussion or explanation of why the covariance matrices Id and not $\\sigma^2$I_d. That is, the variance elements would not necessarily equal to one, but the Gaussians would still be spherical. Would your results still hold for this setting? Otherwise, shouldn't the setting itemized in the abstract specify that your results apply to unit-spherical Gaussians?\n\nLIMITATIONS:\nI had some questions on the assumptions of the parameters of the Gaussians in the GMMs. Otherwise I believe the limitations were addressed.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}