SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sitan Chen
sitan@seas.harvard.edu
klivans@cs.utexas.edu
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 / \mathrm{poly}(d)$ -separated centers.
2. We show gradient descent 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, the expectation-maximization (EM) algorithm and spectral methods.
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  $\widetilde{x}$  of a sample  $x$  from  $q$ , estimate the mean of the posterior distribution over  $x$ .
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,\widetilde{x})$ . We thus ask:
Are there natural data distributions under which GD provably achieves accurate score estimation?
In this work, we consider the setting where  $q$  is given by a mixture of Gaussians. Concretely, we assume that there exist centers  $\mu_1^*,\ldots ,\mu_K^*\in \mathbb{R}^d$  such that
$$
q = \frac {1}{K} \sum_ {i = 1} ^ {K} \mathcal {N} \left(\mu_ {i} ^ {*}, \mathrm {I d}\right).
$$
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  $\Omega(\sqrt{\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 / \varepsilon$ , 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,\{\mu_i^{(0)}\}_{i = 1}^K,H)$
Input: Noise scale  $t$ , initialization  $\{\mu_i^{(0)}\}_{i = 1}^K$ , number of gradient descent steps  $H$
1. Initialize the parameters for the score estimate at  $\theta_t^{(0)} = \{\mu_{i,t}^{(0)}\}_{i=1}^K$  (see Eq. (9) for how the estimate  $s_\theta$  depends on the parameters  $\theta$ , and Eq. (8) for the definition of  $\mu_{i,t}^{(0)}$ )
2 Run gradient descent on the DDPM objective  $L_{t}(s_{\theta_{t}})$  for  $H$  steps where
$$
L _ {t} (s _ {\theta_ {t}}) = \mathbb {E} \Big [ \bigg \| s _ {\theta_ {t}} (X _ {t}) + \frac {Z _ {t}}{\sqrt {1 - \exp (- 2 t)}} \bigg \| ^ {2} \Big ],
$$
3 return  $\theta_t^{(H)} = \{\mu_{i,t}^{(H)}\}_{i=1}^K$  where  $\theta_t^{(H)}$  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.
Theory for diffusion models. A number of works have given convergence guarantees for DDPMs and variants [DBTHD21, BMR22, CLL22, DB22, LLT22, LWYL22, Pid22, WY22,  $\mathrm{CCL}^{+}23\mathrm{b}$ , CDD23, LLT23, LWCC23, BDD23,  $\mathrm{CCL}^{+}23\mathrm{a}$ ]. These results show that, given an oracle for accurate score estimation, diffusion models can learn essentially any distribution over  $\mathbb{R}^d$  (e.g.  $[\mathrm{CCL}^{+}23\mathrm{b}$ , 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  $\Omega (\sqrt{\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  $\Omega (\sqrt{\log K})$ . In particular, [SN21] establish this local convergence as long as every center estimate is initialized at distance at most  $\Delta /2$  away from the corresponding true center, where  $\Delta$  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,  $\mathrm{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
$$
q = \frac {1}{2} \mathcal {N} \left(\mu^ {*}, \operatorname {I d}\right) + \frac {1}{2} \mathcal {N} \left(- \mu^ {*}, \operatorname {I d}\right). \tag {1}
$$
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
$$
X _ {t} = \exp (- t) X _ {0} + \sqrt {1 - \exp (- 2 t)} Z _ {t},
$$
where  $X_0$  is a clean sample drawn from the data distribution  $q$ , and  $Z_{t} \sim \mathcal{N}(0, \mathrm{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
$$
L _ {t} (s) = \mathbb {E} _ {X _ {0}, Z _ {t}} \| s (X _ {t}) - Z _ {t} \| ^ {2}.
$$
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,
$$
\tanh  \left(\left\langle \mu_ {t} ^ {*}, x \right\rangle\right) \mu_ {t} ^ {*} - x, \quad \text {w h e r e} \mu_ {t} ^ {*} \triangleq \mu^ {*} \exp (- t). \tag {2}
$$
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  $\mu \in \mathbb{R}^d$ , our student network is given by  $s_{\mu}(x) \triangleq \tanh(\mu^\top x)\mu - x$ . The exact gradient updates on  $\mu$  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  $\| \mu^{*}\|$  is  $\Omega (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  $\mu$  that has a nontrivial correlation with  $\mu_t^*$ . Starting from this  $\mu$ , 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_{\mu})$  is well-approximated by
$$
- \nabla_ {\mu} L _ {t} (s _ {\mu}) \approx (2 \mu_ {t} ^ {*} \mu_ {t} ^ {* \top} - r \mathrm {I d}) \mu ,
$$
where  $r$  is a scalar that depends on  $\mu$  (See Lemma 8). So the result of a single gradient update with step size  $\eta$  starting from  $\mu$  is given by
$$
\mu^ {\prime} \triangleq \mu - \eta \nabla_ {\mu} L _ {t} \left(s _ {\mu}\right) \approx \left(\left(1 - \eta r\right) \operatorname {I d} + 2 \eta \mu_ {t} ^ {*} \mu_ {t} ^ {* \top}\right) \mu . \tag {3}
$$
This shows us that each gradient step can be approximated by one step of power iteration (without normalization) on the matrix  $(1 - \eta r)\mathrm{Id} + 2\eta \mu_t^*\mu_t^{*\top}$ . It is known 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  $\mu_t^*$ . This suggests that if we can keep the approximation error in (3) under control, then gradient descent on  $\mu$  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  $\mu^{*}$  is sufficiently large and requires that we initialize  $\mu$  to have sufficiently large correlation with the true direction  $\mu_{t}^{*}$ . We can achieve the latter using the large-  $t$  analysis in the previous section.
Provided we have this, when  $t$  is small it turns out that the negative gradient is well-approximated by
$$
- \nabla_ {\mu} L _ {t} (s _ {\mu}) \approx \mathbb {E} _ {X \sim \mathcal {N} (\mu_ {t} ^ {*}, \mathrm {I d})} [ \tanh  (\langle \mu , X \rangle) X ] - \mu .
$$
<sup>1</sup>The real DDPM objective is slightly different, see (5). The latter is what we actually consider in this paper, but this distinction is unimportant for the intuition in this overview.
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  $\eta$  starting from  $\mu$  is given by mixing the old weights  $\mu$  with the result of the "M"-step in EM:
$$
\mu^ {\prime} \triangleq \mu - \eta \nabla_ {\mu} L _ {t} (s _ {\mu}) \approx (1 - \eta) \mu + \eta \underbrace {\mathbb {E} _ {X \sim \mathcal {N} (\mu_ {t} ^ {*} , \mathrm {I d})} [ \tanh  (\langle \mu , X \rangle) X ]} _ {\text {" M" s t e p i n t h e E M a l g o r i t h m}}.
$$
[XHM16] and [DTZ17] showed that EM converges exponentially quickly to the ground truth  $\mu_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  $\| \mu^{*}\|$  is  $\Omega (1 / \mathrm{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  $\| \mu^{*}\|$  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  $\| \mu^{*}\|$  yields a solution arbitrarily close to the ground truth, even without the EM step. 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 / \varepsilon^2)$ , whereas we can only show a suboptimal dependence  $(1 / \varepsilon^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  $\Omega(\sqrt{\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.
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  $\mathbb{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:
$$
\mathrm {d} X _ {t} = - X _ {t} \mathrm {d} t + \sqrt {2} \mathrm {d} W _ {t} \quad \text {w i t h} \quad X _ {0} \sim q _ {0},
$$
where  $(W_{t})_{t\geq 0}$  is a standard Brownian motion in  $\mathbb{R}^d$ . We use  $q_{t}$  to denote the law of the OU process at time  $t$ . Note that for  $X_{t}\sim q_{t}$ ,
$$
X _ {t} = \exp (- t) X _ {0} + \sqrt {1 - \exp (- 2 t)} Z _ {t} \quad \text {w i t h} X _ {0} \sim q _ {0}, Z _ {t} \sim \mathcal {N} (0, \operatorname {I d}).
$$
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$ :
$$
\mathrm {d} X _ {t} ^ {\leftarrow} = \left\{X _ {t} ^ {\leftarrow} + 2 \nabla_ {x} \ln q _ {T - t} \left(X _ {t} ^ {\leftarrow}\right) \right\} \mathrm {d} t + \sqrt {2} \mathrm {d} W _ {t} \quad \text {w i t h} \quad X _ {0} ^ {\leftarrow} \sim q _ {T},
$$
where now  $W_{t}$  is the reversed Brownian motion. In this reverse process, the iterate  $X_{t}^{\leftarrow}$  is distributed according to  $q_{T - t}$  for every  $t\in [0,T]$ , so that the final iterate  $X_{T}^{\leftarrow}$  is distributed according to the
2Note that although  $\mu^{*}$  is unknown, we can estimate its norm from samples.
data distribution  $q_{0}$ . The function  $\nabla_{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
$$
\min  _ {s _ {t}} \mathbb {E} _ {X _ {t} \sim q _ {t}} [ \| s _ {t} (X _ {t}) - \nabla_ {x} \ln q _ {t} (X _ {t}) \| ^ {2} ]. \tag {4}
$$
A standard calculation (see e.g. Appendix A of  $\mathrm{[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.
$$
\min  _ {s _ {t}} L _ {t} \left(s _ {t}\right) = \mathbb {E} _ {X _ {0}, Z _ {t}} \left[ \left\| s _ {t} \left(X _ {t}\right) + \frac {Z _ {t}}{\sqrt {1 - \exp (- 2 t)}} \right\| ^ {2} \right]. \tag {5}
$$
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:
$$
q = q _ {0} = \frac {1}{K} \sum_ {i = 1} ^ {K} \mathcal {N} \left(\mu_ {i} ^ {*}, \mathrm {I d}\right), \tag {6}
$$
where  $\mu_{i}^{*}$  denotes the mean of the  $i^{\mathrm{th}}$  Gaussian component. We define  $\theta^{*} = \{\mu_{1}^{*},\mu_{2}^{*}\dots ,\mu_{K}^{*}\}$ . For the mixtures of two Gaussians, we can simplify the data distribution as
$$
q = q _ {0} = \frac {1}{2} \mathcal {N} \left(\mu^ {*}, \operatorname {I d}\right) + \frac {1}{2} \mathcal {N} \left(- \mu^ {*}, \operatorname {I d}\right). \tag {7}
$$
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
$$
q _ {t} = \frac {1}{K} \sum_ {i = 1} ^ {K} \mathcal {N} \left(\mu_ {i, t} ^ {*}, \operatorname {I d}\right) \quad \text {w h e r e} \quad \mu_ {i, t} ^ {*} \triangleq \mu_ {i} ^ {*} \exp (- t). \tag {8}
$$
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  $\theta_{t}^{*} = \{\mu_{1,t}^{*},\mu_{2,t}^{*},\ldots ,\mu_{K,t}^{*}\}$ .
Lemma 4. The score function for distribution  $q_{t}$ , for any  $t > 0$ , is given by