Toward Global Convergence of Gradient EM for Over-Paramterized Gaussian Mixture Models

Authors: weihang  xu
maryam  fazel
simon s du
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 09:22:14

Abstract: We study the gradient Expectation-Maximization (EM) algorithm for Gaussian Mixture Models (GMM) in the over-parameterized setting, where a general GMM with $n>1$ components learns from data that are generated by a single ground truth Gaussian distribution. 
While results for the special case of 2-Gaussian mixtures are well-known, a general global convergence analysis for arbitrary $n$ remains unresolved and faces several new technical barriers since the convergence becomes sub-linear and non-monotonic. To address these challenges, we construct a novel likelihood-based convergence analysis framework and rigorously prove that gradient EM converges globally with a sublinear rate $O(1/\sqrt{t})$. This is the first global convergence result for Gaussian mixtures with more than $2$ components. The sublinear convergence rate is due to the algorithmic nature of learning over-parameterized GMM with gradient EM. We also identify a new emerging technical challenge for learning general over-parameterized GMM: the existence of bad local regions that can trap gradient EM for an exponential number of steps.

Main paper: 
 introduction: 
 Learning Gaussian Mixture Models (GMM) is a fundamental problem in machine learning with broad applications. In this problem, data generated from a mixture of n ≥ 2 ground truth Gaussians are observed without the label (the index of component Gaussian that data is sampled from), and the goal is to retrieve the maximum likelihood estimation of Gaussian components. The Expectation Maximization (EM) algorithm is arguably the most widely-used algorithm for this problem. Each iteration of the EM algorithm consists of two steps. In the expectation (E) step, it computes the posterior probability of unobserved mixture membership label according to the current parameterized model. In the maximization (M) step, it computes the maximizer of the Q function, which is the likelihood with respect to posterior estimation of the hidden label computed in the E step.Gradient EM, as a popular variant of EM, is often used in practice when the maximization step of EM is costly or even intractable. It replaces the M step of EM with taking one gradient step on the Q function. Learning Gaussian Mixture Models with EM/gradient EM is an important and widely-studied problem. Starting from the seminal work [Balakrishnan et al., 2014] , a flurry of work Daskalakis et al. [2017] , Xu et al. [2016] , Dwivedi et al. [2018a] , Kwon and Caramanis [2020] , Dwivedi et al. [2019] have studied the convergence guarantee for EM/gradient EM in various settings. However, these works either only prove local convergence, or consider the special case of 2-Gaussian mixtures. A general global convergence analysis of EM/gradient EM on n-Gaussian mixtures still remains unresolved. Jin et al. [2016] is a notable negative result in this regard, where the authors show that on GMM with n ≥ 3 components, randomly initialized EM will get trapped in a spurious local minimum with high probability.Over-parameterized Gaussian Mixture Models. Motivated by the negative results, a line of work considers the over-parameterized setting where the model uses more Gaussian components than the ground truth GMM, in the hope that it might help the global convergence of EM and bypass the negative result. In such over-parameterized regime, the best that people know so far is from [Dwivedi et al., 2018b] . This work proves global convergence of 2-Gaussian mixtures on one single Gaussian ground truth. The authors also show that EM has a unique sub-linear convergence rate in this over-parameterized setting (compared with the linear convergence rate in the exact-parameterized setting [Balakrishnan et al., 2014] ). This motivates the following natural open question:Can we prove global convergence of the EM/gradient EM algorithm on general n-Gaussian mixtures in the over-parameterized regime?In this paper, we take a significant step towards answering this question. Our main contributions can be summarized as follows:• We prove global convergence of the gradient EM algorithm for learning general n-component GMM on one single ground truth Gaussian distribution. This is, to the best of our knowledge, the first global convergence proof for general n-component GMM. Our convergence rate is sub-linear, reflecting an inherent nature of over-parameterized GMM (see Remark 3 for details). • We propose a new analysis framework that utilizes the likelihood function for proving convergence of gradient EM. Our new framework tackles several emerging technical barriers for global analysis of general GMM. • We also identify a new geometric property of gradient EM for learning general n-component GMM: There exists bad initialization regions that traps gradient EM for exponentially long, resulting in an inevitable exponential factor in the convergence rate of gradient EM.We consider the canonical Gaussian Mixture Models with weights π = (π 1 , . . . , π n ) ( n i=1 π i = 1), means µ = (µ ⊤ 1 , . . . , µ ⊤ n ) ⊤ and unit covariance matrices I d in d-dimensional space. Following a widely-studied setting [Balakrishnan et al., 2014 , Yan et al., 2017 , Daskalakis et al., 2017] , we set the weights π and covariances I d in student GMM as fixed, and the means µ = (µ ⊤ 1 , . . . , µ ⊤ n ) ⊤ as trainable parameters. We use GMM(µ) to denote the GMM model parameterized by µ, which can be described with probability density function (PDF) p µ : R d → R ≥0 asEQUATIONwhere ϕ(•|µ, Σ) is the PDF of N (µ, Σ),π 1 + • • • + π n = 1, π i > 0, ∀i ∈ [n].The EM algorithm is one of the most popular algorithms for retrieving the maximum likelihood estimator (MLE) on latent variable models. In general, EM and gradient EM address the following problem: given a joint distribution p µ * (x, y) of random variables x, y parameterized by µ * , observing only the distribution of x, but not the latent variable y, the goal of EM and gradient EM is to retrieve the maximum likelihood estimator μMLE ∈ arg max µ log p µ (x).The focus of this paper is the non-convex optimization analysis, so we consider using population gradient EM algorithm to learn GMM (1), where the observed variable is x ∈ R d and latent variable is the index of membership Gaussian in GMM. We follow the standard teacher-student setting where a student model GMM(µ) with n ≥ 2 Gaussian components learns from data generated from a ground truth teacher model GMM(µ * ). We consider the over-parameterized setting where the ground truth model GMM(µ * ) is a single Gaussian distribution N (µ * , I d ), namely µ * = (µ * ⊤ , . . . , µ * ⊤ ) ⊤ . We can then further assume w.l.o.g. that µ * = 0. Our problem could be seen as a strict generalization of Dwivedi et al. [2018b] , where they studied using mixture model of two Gaussians with symmetric means (they set constraint µ 2 = -µ 1 ) to learn one single Gaussian.At time step t = 0, 1, 2, . . ., given with parameters µ(t) = (µ 1 (t) ⊤ , . . . , µ n (t) ⊤ ) ⊤ , population gradient EM updates µ via the following two steps ψ i (x|µ(t)) = Pr[i|x] = π i exp -∥x-µi(t)∥ 2 2 k∈[n] π k exp -∥x-µ k (t)∥ 2 2 .(2)• M step: Define Q(•|, µ(t)) as Q(µ|µ(t)) = E x∼N (0,I d ) n i=1 -ψ i (x|µ(t)) ∥x -µ i ∥ 2 2 ,Gradient EM with step size η > 0 performs the following update:EQUATIONThe membership weight function x → ψ i (x|µ) represents the posterior probability of data point x being sampled from the i th Gaussian of GMM(µ). For ease of notation, we sometimes simply write ψ i (x|µ) as ψ i (x) when the choice of µ is obvious.Since the task of gradient EM is to find the MLE over ground truth distribution p µ * , we can define the MLE loss function for gradient EM asEQUATION)The loss L is the Kullback-Leibler (KL) divergence between the ground truth GMM and the student model GMM. Since finding MLE is equivalent to minimizing the KL divergence between model and the ground truth, the goal of gradient EM is equivalent to finding the global minimum of loss L. In other words, proving that gradient EM finds the MLE is equivalent with proving the convergence of L to 0. However, we are going to present another reason why loss function L is important, for it is also closely related to the dynamics of gradient EM.Gradient EM is gradient descent on L. We present the following important observation. The proof is deferred to appendix. Fact 1. For any µ, ∇Q(µ|µ) = ∇L(µ).Fact 1 states that the gradient of Q function that gradient EM optimizes in each iteration is identical to the gradient of loss function L. This observation is very useful since it implies that gradient EM is equivalent to gradient descent (GD) algorithm on L. This observation is not a new discovery of ours but actually a wide-spread folklore (see [Jin et al., 2016] ). However, our new contribution is to observe Fact 1 is very helpful for analyzing gradient EM, and to construct a new convergence analysis framework for gradient EM based on it.In this paper, we adopt the following notational conventions. We denote {1, 2, . . . , n} with[n]. µ = (µ ⊤ 1 , . . . , µ ⊤ n ) ⊤ ∈ R nddenotes the parameter vector of GMM obtained by concatenating Gaussian mean vectors µ 1 , . . . , µ n together. For any vector µ, µ(t) denotes its value at time step t, sometimes we omit this iteration number t when its choice is clear and simply abbreviate µ(t) as µ. We define a shorthand of expectation taken over the ground truth GMME x∼N (0,I d ) [•] as E x [•].For any vector v ̸ = 0, we use v := v/∥v∥ to denote the normalization of v. We define (with a slight abuse of notation) i max := arg max i∈[n] {∥µ i ∥} as the index of µ i with the maximum norm, and µ max := ∥µ imax ∥ = max i∈[n] {∥µ i ∥} as the maximum norm of µ i . In particular, µ max (t) = max{∥µ 1 (t)∥, . . . , ∥µ n (t)∥}. Similarly, π min := min i∈[n] π i and π max := max i∈[n] π i denotes the minimal and maximal π i , respectively. We use ∇ µi L to denote the gradient of µ i on L, and ∇L = (∇ µ1 L ⊤ , . . . , ∇ µn L) ⊤ denotes the collection of all gradients. Finally we define a potential function U : R nd → R for GMM(µ) asU (µ) = i∈[n] ∥µ i ∥ 2 .Here we provide a brief summary of the major technical barriers for our global convergence analysis and our techniques for overcoming them.New likelihood-based analysis framework. The traditional convergence analysis for EM/gradient EM in previous works Balakrishnan et al. [2014] , Yan et al. [2017] , Kwon and Caramanis [2020] proceeds by showing the distance between the model and the ground truth GMM in the parameter space contracts linearly in every iteration. This type of approach meets new challenges in the over-parameterized n-Gaussian mixture setting since the convergence is both sub-linear and nonmonotonic. To address these problems, we propose a new likelihood-based convergence analysis framework: instead of proving the convergence of parameters, our analysis proceeds by showing the likelihood loss function L converges to 0. The new analysis framework is more flexible and allows us to overcome the aforementioned technical barriers.Gradient lower bound. The first step of our global convergence analysis constructs a gradient lower bound. Using some algebraic transformation techniques, we convert the gradient projection ⟨L(µ), µ⟩ into the expected norm square of a random vector ψ(x). (See Section (4) for the full definition). Although lower bounding the expectation of ψ is very challenging, our key idea is that the gradient of ψ has very nice properties and can be easily lower bounded, allowing us to establish the gradient lower bound.Local smoothness and regularity condition. After obtaining the gradient lower bound, the missing component of the proof is a smoothness condition of the loss function L. Since proving the smoothness of L is hard in general, we define and prove a weaker notion of local smoothness, which suffices to prove our result. In addition, we design and use an auxiliary function U to show that gradient EM trajectory satisfies the locality required by our smoothness lemma.2 Related workQuestion: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: See the summary of main contributions in Section 1 and main results in Section 3.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
2-gaussian mixtures: 
 There is a vast literature studying the convergence of EM/gradient EM on 2-component GMM. The initial batch of results proves convergence within a infinitesimally small local region [Xu and Jordan, 1996, Ma et al., 2000] . Balakrishnan et al. [2014] proves for the first time convergence of EM and gradient EM within a non-infinitesimal local region. Among the later works on the same problem, Klusowski and Brinda [2016] improves the basin of convergence guarantee, Daskalakis et al. [2017] , Xu et al. [2016] proves the global convergence for 2-Gaussian mixtures. These works focused on the exact-parameterization scenario where the number of student mixtures is the same as that of the ground truth. More recently, Wu and Zhou [2019] proves global convergence of 2-component GMM without any separation condition. Their result can be viewed as a convergence result in the over-parameterized setting where the student model has two Gaussians and the ground truth is a single Gaussian. On the other hand, their setting is more restricted than ours because they require the means of two Gaussians in the student model to be symmetric around the ground truth mean. Weinberger and Bresler [2021] extends the convergence guarantee to the case of unbalanced weights. Another line of work Dwivedi et al. [2018b Dwivedi et al. [ , 2019 Dwivedi et al. [ , 2018a] ] studies the over-parameterized setting of using 2-Gaussian mixture to learn a single Gaussian and proves global convergence of EM. Our result extends this type of analysis to the general case of n-Gaussian mixtures, which requires significantly different techniques. We note that going beyond Gaussian mixture models, there are also works studying EM algorithms for other mixture models such as a mixture of linear regression Kwon et al. [2019] .Another line of results focuses on the general case of n Gaussian mixtures. Jin et al. [2016] provides a counter-example showing that EM does not converge globally for n > 2 (in the exact-parameterized case). Dasgupta and Schulman [2000] prove that a variant of EM converges to MLE in two rounds for n-GMM. Their result relies on a modification of the EM algorithm and is not comparable with ours. [Chen et al., 2023] analyzes the structure of local minima in the likelihood function of GMM. However, their result is purely geometric and does not provide any convergence guarantee.A series of paper Yan et al. [2017] , Zhao et al. [2018] , Kwon and Caramanis [2020] , Segol and Nadler follow the framework proposed by Balakrishnan et al. [2014] to prove the local convergence of EM for n-GMM. While their result applies to the more general n-Gaussian mixture ground truth setting, their framework only provides local convergence guarantee and cannot be directly applied to our setting.This paper gives an O 1/ √ t bound for fitting over-parameterized Gaussian mixture models to a single Gaussian. Recall that to learn a single Gaussian, if one's student model is also a single Gaussian, then one can obtain an exp(-Ω(t)) rate because the loss is strongly convex. This slowdown effect due to over-parameterization has been observed for Gaussian mixtures in Dwivedi et al. [2018a] , Wu and Zhou [2019] , but has also been observed in other learning problems, such as learning a two-layer neural network Xu and Du [2023] , Richert et al. [2022] and matrix sensing problems [Xiong et al., 2023 , Zhang et al., 2021 , Zhuo et al., 2021] .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: As a theoretical work, the major assumptions and limitations of our results are presented in the introduction part of Section 1.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
main results: 
 In this section, we present our main theoretical result, which consists of two parts: In Section 3.1 we present our global convergence analysis of gradient EM, in Section 3.2 we prove that an exponentially small factor in our convergence bound is inevitable and cannot be removed. All omitted proofs are deferred to the appendix.We first present our main result, which states that gradient EM converges to MLE globally. ∈ R + . Recall that µ max (0) = max{∥µ 1 (0)∥, . . . , ∥µ n (0)∥} and U (0) = i∈[n] ∥µ i (0)∥ 2 are two initialization constants. Remark 3. Without over-parameterization, for learning a single Gaussian, one can obtain a linear convergence exp(-Ω (t)). We would like to note that the sub-linear convergence rate guarantee of gradient EM stated in Theorem 2 (L(µ(t)) ≤ O(1/ √ t)) is due to the inherent nature of the algorithm. Dwivedi et al. [2018b] studied the special case of using 2 Gaussian mixtures with symmetric means to learn a single Gaussian and proved that EM has sublinear convergence rate when the weights π i are equal. Since Theorem 2 studies the more general case of n Gaussian mixtures, this type of subexponential convergence rate is the best than we can hope for. Remark 4. The convergence rate in Theorem 2 has a factor exponentially small in the initialization scale (γ ∝ exp(-16U (0))). We would like to stress that this is again due to algorithmic nature of the problem rather than the limitation of analysis. In Section 3.2, we prove that there exists bad regions with exponentially small gradients so that when initialized from such region, gradient EM gets trapped locally for exp(Ω(U (0))) number of steps. Therefore, a convergence speed guarantee exponentially small in U (0) is inevitable and cannot be improved. Remark 5. Theorem 2 is fundamentally different from convergence analysis for EM/gradient EM in previous works Yan et al. [2017] , Dwivedi et al. [2019] , Balakrishnan et al. [2014] which proved monotonic linear contraction of parameter distance ∥µ(t) -µ * ∥. But our result also implies global convergence since loss function L converging to 0 is equivalent to convergence of gradient EM to MLE. Remark 6. The convergence result in Theorem 2 is for population gradient EM, but it also implies global convergence for sample-based gradient EM as the sample size tends to infinity. For a similar reduction from population EM to sample EM, see Section 2.2 of [Xu et al., 2016] .In this section we prove that a factor exponentially small in initialization scale (exp(-Θ(U (0)))) is inevitable in the global convergence rate guarantee of gradient EM. Particularly, we show the existence of bad regions such that initialization from this region traps gradient EM for exponentially long time before final convergence. Our result is the following theorem. Theorem 7 (Existence of bad initialization region). For any n ≥ 3, define μ( 0) = (µ ⊤ 1 (0), . . . , µ ⊤ n (0)) as follows: µ 1 (0) = 12 √ de 1 , µ 2 (0) = -12 √ de 1 , µ 3 (0) = • • • = µ n (0) = 0, where e 1 is a standard unit vector. Then population gradient EM initialized with means μ(0) and equal weights π 1 = . . . = π n = 1/n will be trapped in a bad local region around μ(0) for exponentially long timeT := 1 30η e d = 1 30η exp(Θ(U (0))).More rigorously, for any0 ≤ t ≤ T, ∃i ∈ [n] such that ∥µ i (t)∥ ≥ 10 √ d.Theorem 7 states that, when initialized from some bad points µ(0), after exp(Θ(U (0))) number of time steps, gradient EM will still stay in this local region and remain 10 √ d distance away from the global minimum µ = 0. Therefore an exponentially small factor in convergence rate is inevitable. Remark 8. Theorem 7 eliminates the possibility of proving any polynomial convergence rate of gradient EM from arbitrary initialization. However, it is still possible to prove that, with some specific smart initialization schemes, gradient EM avoids the bad regions stated in Theorem 7 and enjoys a polynomial convergence rate. We leave this as an interesting open question for future analysis.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
proof overview: 
 In this section, we provide a technical overview of the proof in our main result (Theorem 2 and Theorem 7).Proving the global convergence of gradient EM for general n-Gaussian mixture is highly nontrivial. While there have been many previous works [Balakrishnan et al., 2014 , Yan et al., 2017 , Dwivedi et al., 2018b] studying either local convergence or the special case of 2-Gaussian mixtures, they all focus on showing the contraction of parametric error. Namely, their proof proceeds by showing the distance between the model parameter and the ground truth contracts, usually by a fixed linear ratio, in each iteration of the algorithm. However, this kind of approach faces various challenges for our general problem where the convergence is both sublinear and non-monotonic. Since the convergence rate is sublinear (see Remark 3), showing a linear contraction per iteration is no longer possible. Since the convergence is non-monotonic1 , we also cannot show a strictly decreasing parametric distance.To address these challenges, we propose a new convergence analysis framework for gradient EM by proving the convergence of likelihood L instead of the convergence of parameters µ. There are several benefits for considering the convergence from the perspective of MLE loss L. Firstly, it naturally addresses the problem of non-monotonic and sub-linear convergence since we only need to show L decreases as the algorithm updates. Also, since gradient EM is equivalent with running gradient descent on loss function L (see Section 1.3), we can apply techniques from the optimization theory of gradient descent to facilitate our analysis.We first briefly outline our proof of Theorem 2.Proof roadmap. Our proof of Theorem 2 consists of three steps. Firstly, we prove a gradient lower bound for L (Theorem 12). Then we prove that the MLE L is locally smooth (Theorem 13). Finally, we combine the gradient lower bound and the smoothness condition to prove the global convergence of L with mathematical induction.Step 1: Gradient lower bound.Our first step aims to show that the gradient norm of L(µ) is lower bounded by the distance of µ to the ground truth. To do this, we need a few preliminary results. Inspired by Chen et al. [2023] , we use Stein's identity [Stein, 1981] to perform an algebraic transformation of the gradient. Recalling the definition of ψ i in (2), we have the following lemma. Lemma 9. For any GMM(µ), i ∈ [n], the gradient of Q satisfies∇ µi L(µ) = ∇ µi Q(µ|µ) = E x   ψ i (x) k∈[n] ψ k (x)µ k   .The gradient expression above is equivalent with the form in (3), but is easier to manipulate. Using the transformed gradient in Lemma 9, we have the following corollary.Corollary 10. Define vector ψµ (x) := i∈[n] ψ i (x)µ i . For any GMM(µ), the projection of the gradient of ∇L(µ) onto µ satisfies⟨∇L(µ), µ⟩ = ⟨∇ µ Q(µ|µ), µ⟩ = i∈[n] ⟨∇ µi Q(µ|µ), µ i ⟩ = E x ψµ (x) 2 .Corollary 9 is important since it converts the projection of gradient ∇L(µ) onto µ to the expected norm square of a vector ψµ . Since a lower bound of the gradient projection implies a lower bound of the gradient, we only need to construct a lower bound for ⟨∇L(µ),µ⟩ = E x ψµ (x) 2 . Since ψµ (x)2 is always non-negative, we already know that the gradient projection is non-negative. Butlower bounding E x ψµ (x)2 is still highly nontrivial since the expression of ψ is complicated and hard to handle. However, our key observation is that, although ψ itself is hard to bound, its gradient has nice properties and can be handled gracefully:EQUATIONThe gradient ( 5) is nicely-behaved. One can see immediately from (5) that the matrix ∇ x ψµ (x) is positive-semi-definite, and its eigenvalues can be directly bounded. To utilize these properties, we use the following algebraic trick to convert the task of lower bounding ψ itself into the task of lower bounding its gradient.EQUATION)Recall that x = x ∥x∥ . See detailed derivation in (23). Using (5), combined with the properties of ∇ x ψµ (x), we can obtain the following lemma (Recall that U = i∈[n] ∥µ i ∥ 2 .):Lemma 11. For any GMM(µ) we haveE x ∥ ψµ (x)∥ 2 ≥ exp (-8U ) 40000d(1 + 2µ max √ d) 2   i,j∈[n] π i π j ∥µ i -µ j ∥ 2   2 .On top of Lemma 11, we can easily lower bound the gradient projection in the following lemma, finishing the first step of our proof.Lemma 12 (Gradient projection lower bound). For any GMM(µ) we have⟨∇ µ Q(µ|µ), µ⟩ = E x [∥ ψµ (x)∥ 2 ] = Ω exp (-8U ) π 2 min d(1 + µ max √ d) 2 µ 4 max .Step 2: Local smoothness.To construct a global convergence analysis for gradient-based methods, after obtaining a gradient lower bound, we still need to prove the smoothness of loss L. (Recall that global smoothness of function f means that there exists constant C such that ∥∇f (x 1 )-∇f (x 2 )∥ ≤ C∥x 1 -x 2 ∥, ∀x 1 , x 2 .) However, proving the smoothness for L in general is very challenging since the membership function ψ i cannot be bounded when µ is unbounded. To address this issue, we prove that L is locally smooth, i.e., the smoothness between two points µ and µ ′ is satisfied if both ∥µ∥ and ∥µ -µ ′ ∥ are upper bounded. Our result is the following theorem.Theorem 13 (Local smoothness of loss function). At any two points µ= (µ ⊤ 1 , . . . , µ ⊤ n ) ⊤ and µ + δ = ((µ 1 + δ 1 ) ⊤ , . . . , (µ n + δ n ) ⊤ ) ⊤ , if ∥δ i ∥ ≤ 1 max {6d, 2∥µ i ∥} , ∀i ∈ [n],then the loss function L satisfies the following smoothness property: for any i ∈ [n] we haveEQUATION)Step 3: putting everything together.Given the gradient lower bound and the smoothness condition, we still need to resolve two remaining problems. The first one is that the gradient lower bound in Lemma 12 is given in terms of µ, which we need to convert to a lower bound in terms of L(µ). For this we need the following upper bound of L.Theorem 14 (Loss function upper bound). The loss function can be upper bounded asL(µ) ≤ i∈[n] π i 2 ∥µ i ∥ 2 ≤ µ 2 max 2 .The second problem is that our local smoothness theorem requires µ to be bounded, therefore we need to show a regularity condition that for each i, µ i (t) stays in a bounded region during gradient EM updates. This is not easy to prove for each individual µ i due to the same non-monotonic issue mentioned in Section 4.1. To establish such a regularity condition, we use the potential function.U to solve this problem. We prove that U remains bounded along the gradient EM trajectory, implying each µ i remains well-behaved. With this regularity condition, combined with the previous two steps, we finish the proof of Theorem 2 via mathematical induction.Proving Theorem 7 is much simpler. The idea is natural: we found that there exists some bad regions where the gradient of L is exponentially small, characterized by the following lemma.Lemma 15 (Gradient norm upper bound). For any µ satisfying ∥µ 1 ∥, ∥µ 2 ∥ ≥ 10 √ d, ∥µ 3 ∥, . . . , ∥µ n ∥ ≤ √ d, the gradient of L at µ can be upper bounded as∥∇ µi L(µ)∥ ≤ 2(∥µ 3 ∥ + • • • + ∥µ n ∥) + 2 exp(-d)(∥µ 1 ∥ + ∥µ 2 ∥), ∀i ∈ [n].Utilizing Lemma 15, we can prove Theorem 7 by showing that initialization from these bad regions will get trapped in it for exponentially long, since the gradient norm is exponentially small. The full proof can be found in Appendix B.2.  
experiments: 
 In this section we experimentally explore the behavior of gradient EM on GMMs.Convergence rates. We choose the experimental setting of d = 5, η = 0.7. We use n = 2, 5, 10 Gaussian mixtures to learn data generated from one single ground truth Gaussian distribution N (µ * , I d ), respectively. Since a closed form expression of the population gradient is intractable, we approximate the gradient step via Monte Carlo method, with sample size 3.5 × 10 5 . The mixing weights of student GMM are randomly sampled from a standard Dirichlet distribution and set as fixed during gradient EM update. The covariances of all component Gaussians are set as the identity matrix. We recorded the convergence of likelihood function L (estimated also by Monte Carlo method on fresh samples each iteration) and parametric distance i∈[n] π i ∥µ i -µ * ∥2 along gradient EM trajectory. The results are reported in Figure 1 (left and middle panel). Both the likelihood L and the parametric distance converges sub-linearly.Weight configurations. We train 3-component GMM with 3-different weight configurations and report 4 runs each configuration in Figure 1 (right). Blue:EQUATION). More evenly distributed weights result in faster convergence. Initialization geometry. We empirically study the bad initialization point µ(0) described in Theorem 7 2 by plotting the gradient norm at µ(0) w.r.t. different dimension d in Figure 2 (left). As theoretically analyzed, the gradient norm ∥∇L(µ(0))∥ at µ(0) decreases exponentially in dimension d.Statistical rates. The statistical rate for EM/gradient EM is another interesting research problem, which we observe empirically in Figure 2 (right). We run gradient EM on 5-component GMM with equal weights. x-axis: number of training samples, y-axis: parametric error after convergence. For each sample size, we run 50 times and report the average. The statistical errors are reported in the blue line. The red line (function Θ(n -1/4 )) and green line (linear regression output fitting blue points) are references. The trajectory approximately follows the law of accuracy ∝ n -1/4 . While [Wu and Zhou, 2019] rigorously proves the asymptotic statistical rate of Õ(n -1/4 ) for the special case of 2-GMMs, our experiments imply that the same rate might also apply to the general case of multi-component GMMs. 
conclusion: 
 This paper gives the first global convergence of gradient EM for over-parameterized Gaussian mixture models when the ground truth is a single Gaussian, and rate is sublinear which is exponentially slower than the rate in the exact-parameterization case. One fundamental open problem is to study when one can obtain global convergence of EM or gradient EM for Gaussian mixture models when the ground truth has multiple components. The likelihood-based convergence framework proposed in this paper might be an helpful tool towards solving this general problem. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [No] Justification: Our experiment only shows the phenomenon on small-scale synthetic data, so we did not record the computation resource used.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: We conform, in every respect, with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: This is a theoretical work. It societal impact lies within its potential pratical applications.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: This is a theoretical work and has no such risks. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [NA] Justification: This is a theoretical work and does not use existing assets. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: Proof of Fact 1. It is well known that (see Section 1 of Wu and Zhou [2019] )Q(µ ′ |µ) = E x∼p µ * [log(p µ ′ (x)) -D KL (p µ (•|x)||p µ ′ (•|x)) -H(p µ (•|x))] ,where p µ (•|x) denotes the distribution of hidden variable y (in our case of GMM the index of Gaussian component) conditioned on x, and H denotes information entropy.Since µ ′ = µ is a global minimum of D KL (p µ (•|x)||p µ ′ (•|x)), we have ∇D KL (p µ (•|x)||p µ (•|x)) = 0. Also ∇H(p µ (•|x)) = 0 since H(p µ (•|x)) is a constant. Therefore ∇Q(µ|µ) = E x∼p µ * [∇ log(p µ (x))] = ∇L(µ).The proof of Lemma 9 uses ideas from Theorem 1 of Chen et al. [2023] and relies on Stein's identity, which is given by the following lemma. Lemma 16 (Stein [1981] ). For x ∼ N (µ, σ 2 I d ) and differentiable function g : R d → R we haveE[g(x)(x -µ)] = σ 2 E[∇ x g(x)],if the two expectations in the above identity exist. Now we are ready to prove Lemma 9. Lemma 9. For any GMM(µ), i ∈ [n], the gradient of Q satisfies∇ µi L(µ) = ∇ µi Q(µ|µ) = E x   ψ i (x) k∈[n] ψ k (x)µ k   .Proof. Applying Stein's identity (Lemma 16), for each i ∈ [n] we have∇ µi Q(µ|µ) = E x∼N (0,I d ) [ψ i (x)(µ i -x)] = E x∼N (0,I d ) [ψ i (x)] µ i -E x∼N (0,I d ) [ψ i (x)x] = E x∼N (0,I d ) [ψ i (x)] µ i -E x∼N (0,I d ) [∇ x ψ i (x)].Recall thatψ i (x) = Pr[i|x] = π i exp -∥x-µi∥ 2 2 k∈[n] π k exp -∥x-µ k ∥ 2 2 .The gradient ∇ x ψ i (x) could be calculated asEQUATIONnote that we used k∈[n] ψ i (x) = 1.Then we have∇ µi Q(µ|µ) = E x [ψ i (x)] µ i -E x [∇ x ψ i (x)] = E x [ψ i (x)] µ i -E x   ψ i (x)   µ i - k∈[n] ψ k (x)µ k     = E x   ψ i (x) k∈[n] ψ k (x)µ k   .Proof of Corollary 10.⟨∇ µ Q(µ|µ), µ⟩ = i∈[n] ⟨∇ µi Q(µ|µ), µ i ⟩ = i∈[n] E x   ψ i (x) k∈[n] ψ k (x)µ k   , µ i = i∈[n] k∈[n] E x ⟨ψ i (x)ψ k (x)µ k , µ i ⟩ = E x    i∈[n] ψ i (x)µ i 2    = E x ψµ (x) 2 .Lemma 17. For any constant c satisfying 0 < c ≤ 1 3d , we haveE x∼N (0,I d ) [exp (c∥x∥)] ≤ 1 + 5 √ dc.Proof. Note that E x∼N (0,I d ) [exp (c∥x∥)] = M ∥x∥ (c) is the moment-generating function of ∥x∥.To upper bound the value of a moment generating function at c, we use Lagrange's Mean Value Theorem:EQUATIONwhere ξ ∈ [0, c]. Note that M ∥x∥ (0) = 1, So the remaining task is to bound M ′ ∥x∥ (ξ). We bound this expectation using truncation method as:EQUATIONwhereV d = π d/2 Γ(d/2+1) is the volume of d-dimensional unit sphere. Since ∥x∥ ≥ 1 ⇒ c∥x∥ -∥x∥ 2 2 ≤ 1 3d ∥x∥ -∥x∥ 2 2 ≤ -∥(1-1/(2d))x∥ 2 2 , we have ∥x∥≥1 ∥x∥(2π) -d/2 exp c∥x∥ - ∥x∥ 2 2 dx ≤ ∥x∥≥1 ∥x∥(2π) -d/2 exp - ∥ 2d-1 2d x∥ 2 2 dx = ∥y∥≥ 2d-1 2d 2d 2d -1 ∥y∥(2π) -d/2 exp - ∥y∥ 2 2 2d 2d -1 d dy ≤ 2d 2d -1 d+1 E y∼N (0,I d ) [∥y∥] = 2d 2d -1 d+1 √ 2Γ d+1 2 Γ d 2 ≤ 4 √ d,where we used 2d 2d-1 d+1 ≤ 4 and the log convexity of Gamma function at the last line. Plugging this back to (10), we getEQUATIONPlugging ( 11) into (9), we obtain the final boundE x [exp (2∥δ i ∥(∥x∥ + ∥µ i ∥)) -1] = M ∥x∥ (c) = M ∥x∥ (0) + M ′ ∥x∥ (ξ)c ≤ 1 + 5 √ dc.Lemma 18. Recall that U = i∈[n] ∥µ i ∥ 2 . For any fixed x ∈ R d , x ̸ = 0 and any µ we have1 t=-1 ψ i (tx|µ)ψ j (tx|µ)dt ≥ 1 2µ max ∥x∥ π i π j exp (-4U ) (1 -exp (-4µ max ∥x∥)) .Proof.EQUATIONB Proofs for Section 3 and 4B.1 Proofs for global convergence analysis Theorem 13. At any two points µ= (µ ⊤ 1 , . . . , µ ⊤ n ) ⊤ and µ + δ = ((µ 1 + δ 1 ) ⊤ , . . . , (µ n + δ n ) ⊤ ) ⊤ , if ∥δ i ∥ ≤ 1 max {6d, 2∥µ i ∥} , ∀i ∈ [n],then the loss function L satisfies the following smoothness property: for any i ∈ [n] we haveEQUATIONProof. Note that exp (-∥δ i ∥(∥x∥+ ∥µ i ∥)) exp - ∥δ i ∥ 2 2 ≤ exp -∥x-(µi+δi)∥ 2 2 exp -∥x-µi∥ 2 2 = exp ⟨x -µ i , δ i ⟩ - ∥δ i ∥ 2 2 ≤ exp (∥δ i ∥(∥x∥ + ∥µ i ∥)) exp - ∥δ i ∥ 2 2 .Therefore ψ i (x|µ + δ) can be bounded asEQUATIONSimilarly, we haveEQUATIONRecall that by Lemma 9 we haveEQUATIONwhere the last inequality is because ψ i , ψ k ≤ 1 and applying ( 15) and ( 16).The remaining task is to bound E x [exp (2∥δ i ∥(∥x∥+ ∥µ i ∥)) -1]. Since 2∥δ i ∥ ≤ 1 3d, we can use Lemma 17 to bound it asEQUATION)where we used exp(1 + x) ≤ 1 + 2x, ∀x ∈ [0, 1] at the last line. Plugging this back to (17), we get∥∇ µi+δi L(µ + δ) -∇ µi L(µ)∥ ≤ k∈[n] ∥δ k ∥ + k∈[n] E x [exp (2∥δ i ∥(∥x∥ + ∥µ i ∥)) -1] ∥µ k ∥ ≤ k∈[n] ∥δ k ∥ + k∈[n](30EQUATION)Theorem 14. The loss function can be upper bounded asL(µ) ≤ i∈[n] π i 2 ∥µ i ∥ 2 ≤ µ 2 max 2 .Proof. Since the logarithm function is concave, by Jensen's inequality we haveL(µ) = D KL (p µ * ||p µ ) = -E x log p µ (x) p µ * (x) = -E x   log   i π i exp -∥x-µi∥ 2 2 exp -∥x∥ 2 2     ≤ -E x   i π i log   exp -∥x-µi∥ 2 2 exp -∥x∥ 2 2     = - i π i E x ⟨x, µ i ⟩ - ∥µ i ∥ 2 2 = i∈[n] π i 2 ∥µ i ∥ 2 ≤ µ 2 max 2 .Lemma 12. For any GMM(µ) we have⟨∇ µ Q(µ|µ), µ⟩ = E x [∥ ψµ (x)∥ 2 ] ≥ Ω exp (-8U ) π 2 min d(1 + µ max √ d) 2 µ 4 max .Case 1. There exists k ∈[n] such that ∥µ k -µ imax ∥ ≥ µmax 2 .Then by Lemma 19 and Lemma 11 we haveE x ∥ ψµ (x)∥ 2 ≥ exp (-8U ) 40000d(1 + 2µ max √ d) 2   i,j∈[n] π i π j ∥µ i -µ j ∥ 2   2 ≥ exp (-8U ) 40000d(1 + 2µ max √ d) 2 π min 8 µ 2 max 2 = exp (-8U ) π 2 min 2560000d(1 + 2µ max √ d) 2 µ 4 max . Case2. For ∀k ∈ [n], ∥µ imax -µ k ∥ < µmax 2 .Then by Lemma 20 we haveE x ∥ ψµ (x)∥ 2 ≥ 1 4 µ 2 max ≥ Ω(exp(-8µ 2 max )µ 4 max ) ≥ Ω(exp(-8U )µ 4 max ) ≥ Ω exp(-8U )π 2 min d(1+µmax √ d) 2 µ 4 max , (since e -x x ≤ 1, ∀x). Lemma 19. For any GMM(µ), if there exists k ∈ [n] such that ∥µ k -µ imax ∥ ≥ µmax 2 , then we have i,j∈[n] π i π j ∥µ i -µ j ∥ 2 ≥ π min 8 µ 2 max .Proof. By Cauchy-Schwarz inequality, we have ∥a∥ 2 + ∥b∥ 2 ≥ 1 2 ∥a -b∥ 2 , so for ∀i ∈ [n] we havej∈[n] π j ∥µ i -µ j ∥ 2 ≥ π imax ∥µ i -µ imax ∥ 2 + π k ∥µ i -µ k ∥ 2 ≥ π min 2 ∥(µ i -µ imax ) -(µ i -µ k )∥ 2 = π min 2 ∥µ k -µ imax ∥ 2 . Therefore i,j∈[n] π i π j ∥µ i -µ j ∥ 2 = i∈[n] π i j∈[n] π j ∥µ i -µ j ∥ 2 ≥ i∈[n] π i π min 2 ∥µ k -µ imax ∥ 2 ≥ π min 8 µ 2 max ,where the last inequality is because ∥µ k -µ imax ∥ ≥ µmax 2 and i π i = 1.Lemma 20. For any GMM(µ), if for ∀k ∈ [n] we have ∥µ imax -µ k ∥ < µmax 2 , thenE x ∥ ψµ (x)∥ 2 ≥ 1 4 µ 2 max .Proof. For any k ∈ [n], by Cauchy-Schwarz inequality we haveEQUATIONwhere the last inequality is becauseEQUATIONwhere we used k∈[n] ψ k (x) = 1 at the last inequality.Lemma 11. For any GMM(µ) we haveE x ∥ ψµ (x)∥ 2 ≥ exp (-8U ) 40000d(1 + 2µ max √ d) 2   i,j∈[n] π i π j ∥µ i -µ j ∥ 2   2 .Proof. The key idea is to consider the gradient of ψµ , which can be calculated asEQUATIONwhere we used (8) in the second identity.By Cauchy-Schwarz inequality, we have ∥a∥ 2 + ∥b∥ 2 ≥ 1 2 ∥a -b∥ 2 , which impliesEQUATION)where we used ∂ ∂t ψµ (tx) = ∇ ψµ (tx)x at the second to last identity. Careful readers might notice that the term1 t=-1 ∥x∥ • x ⊤ ∇ ψµ (tx)xdt 2is not well-defined when x = 0, but we can still calculate its expectation over the whole probability space since the integration is only singular on a zero-measure set.For each x ̸ = 0, by ( 22) we havex ⊤ ∇ ψµ (tx)x = 1 2 i,j∈[n] ψ i (tx)ψ j (tx)⟨µ i -µ j , x⟩ 2 . So E x ∥ ψµ (x)∥ 2 ≥ 1 16 E x      1 t=-1 ∥x∥ i,j∈[n] ψ i (tx)ψ j (tx)⟨µ i -µ j , x⟩ 2 dt   2    = 1 16 E x      ∥x∥ i,j∈[n] ⟨µ i -µ j , x⟩ 2 1 t=-1 ψ i (tx)ψ j (tx)dt   2    ≥ 1 16 E x      ∥x∥ i,j∈[n] ⟨µ i -µ j , x⟩ 2 1 2µ max ∥x∥ π i π j exp (-4U ) (1 -exp (-4µ max ∥x∥))   2    = exp (-8U ) 64 E x      i,j∈[n] π i π j ⟨µ i -µ j , x⟩ 2 1 -exp (-4µ max ∥x∥) µ max   2    ≥ exp (-8U ) 64   i,j∈[n] π i π j E x ⟨µ i -µ j , x⟩ 2 1 -exp (-4µ max ∥x∥) µ max   2 (24)where we used Lemma 18 at the fourth line and Cauchy-Schwarz inequality at the last line.The last step is to lower bound E x ⟨µ i -µ j , x⟩ 2 (1 -exp (-4µ max ∥x∥)) /µ max . Since x is sampled from N (0, I d ), which is spherically symmetric, we know that the two random variables {x, ∥x∥} are independent. ThereforeE x ⟨µ i -µ j , x⟩ 2 1 -exp (-4µ max ∥x∥) µ max = E x ⟨µ i -µ j , x⟩ 2 E x 1 -exp (-4µ max ∥x∥) µ max .(25) For the first term in (25), we have E x ⟨µ i -µ j , x⟩ 2 = ∥µ i -µ j ∥ 2 /d since x is spherically symmetrically distributed. By norm-concentration inequality of Gaussian [Dasgupta and Schulman, 2000] we know that Pr ∥x∥ ≥√ d 2≥ 1/50, ∀d. The second term in ( 25) can be therefore lower bounded asEQUATION)Plugging ( 26) into ( 25), we getEQUATION)Now we can plug ( 27) into ( 24) and getE x ∥ ψµ (x)∥ 2 ≥ exp (-8U ) 64   i,j∈[n] π i π j E x ⟨µ i -µ j , x⟩ 2 1 -exp (-4µ max ∥x∥) µ max   2 ≥ exp (-8U ) 64   i,j∈[n] π i π j 1 -exp -2µ max √ d 50dµ max ∥µ i -µ j ∥ 2   2 ≥ exp (-8U ) 64   i,j∈[n] π i π j 1 - 1 1+2µmax √ d 50dµ max ∥µ i -µ j ∥ 2   2 = exp (-8U ) 40000d(1 + 2µ max √ d) 2   i,j∈[n] π i π j ∥µ i -µ j ∥ 2   2 (28)where we used the inequality ∀t ≥ 0, e -t ≤ 1 1+t at the second to last line. Theorem 2. Consider training a student n-component GMM initialized from µ( 0) = (µ 1 (0) ⊤ , . . . , µ n (0) ⊤ ) ⊤ to learn a single-component ground truth GMM N (0, I d ) with popula- tion gradient EM algorithm. If the step size satisfies η ≤ O exp(-8U (0))π 2 min n 2 d 2 ( 1 µmax(0) +µmax(0)) 2 , then gradient EM converges globally with rate L(µ(t)) ≤ 1 √ γt , where γ = Ω η exp(-16U (0))π 4 min n 2 d 2 (1+µmax(0) √ dn) 4 ∈ R + . Recall that µ max (0) = max{∥µ 1 (0)∥, . . . , ∥µ n (0)∥} and U (0) = i∈[n] ∥µ i (0)∥ 2 are two initialization constants.Proof. We use mathematical induction to prove Theorem 2, by proving the following two conditions inductively:EQUATIONNote that (30) directly implies the theorem, so now we just need to prove ( 29) and ( 30) together.The induction base for t = 0 is trivial. Now suppose the conditions hold for time step t, consider t + 1. By induction hypothesis (29) we have ∥µ i (t)∥ ≤ µ max (t) ≤ √ nµ max (0), ∀t.Proof of (30). Since ∇ µ Q(µ|µ) = ∇ µ L(µ), we can apply classical analysis of gradient descent [Nesterov et al., 2018] asEQUATIONNote that the gradient norm can be upper bounded asEQUATION∥} . So we can apply Theorem 13 and get∥∇ µi L(µ(t)) -∇ µi L(µ(t) -sη∇ µi L(µ(t)))∥ ≤ nµ max (t)(30 √ d + 4µ max (t))∥sη∇ µi L(µ(t))∥ + k∈[n] ∥sη∇ µ k L(µ(t))∥.Therefore for ∀s ∈ [0, 1],EQUATION)Plugging ( 32) into (31EQUATION)By Lemma 12 we can lower bound the gradient norm asEQUATIONCombining ( 34) and ( 33), we haveEQUATIONNote that the above inequality implies L(µ(t + 1)) ≤ L(µ(t)), thereforeEQUATION≥Ω η exp (-16U (0)) π 4 min n 2 d 2 (1 + µ max (0) √ dn) 4 = γ.On the other hand, by induction hypothesis we have1 L 2 (µ(t)) ≥ γt + 1 L 2 (µ(0)), combined with the above inequality, we have 0)) , which finishes the proof of (30).1 L 2 (µ(t+1)) ≥ 1 L 2 (µ(t)) + γ ≥ γ(t + 1) + 1 L 2 (µ(Proof of (29). The dynamics of potential function U can be calculated asEQUATIONBy induction hypothesis, the first term I 1 can be bounded by Lemma 12 asEQUATIONThe second term I 2 is a perturbation term that can be upper bounded by Lemma 9 asEQUATIONwhere we use triangle inequality twice at the second and third line, and Cauchy-Schwarz inequality twice at the fourth and fifth line.Putting (38), ( 37) and ( 36) together, we getU (µ(t + 1)) ≤ U (µ(t)) -ηΩ exp (-8U (0)) π 2 min n 2 d(1 + µ max (0) √ nd) 2 U 2 (µ(t)) + η 2 U (µ(t)).Consider two cases: 29) holds in both cases, our proof is done.a). If U (0) 2 ≤ U (µ(t)) ≤ U (0), then U (µ(t + 1)) ≤ U (µ(t)) -ηU (µ(t)) Ω exp (-8U (0)) π 2 min n 2 d(1 + µ max (0) √ nd) 2 U (µ(t)) -η ≤ U (µ(t)) -ηU (µ(t)) Ω exp (-8U (0)) π 2 min n 2 d(1 + µ max (0) √ nd) 2 n 2 µ 2 max (0) -η ≤ U (µ(t)) ≤ nµ 2 max (0), note that we used η ≤ O exp(-8U (0))π 2 min n 2 d(1+µmax(0) √ nd) 2 n 2 µ 2 max (0). b). If U (µ(t)) < 1 2 U (0), then U (µ(t + 1)) ≤ (1 + η 2 )U (µ(t)) ≤ 2U (µ(t)) ≤ U (0). Since (B.2 Proofs for Section 3.2 Lemma 15. For any µ satisfying ∥µ 1 ∥, ∥µ 2 ∥ ≥ 10 √ d, ∥µ 3 ∥, . . . , ∥µ n ∥ ≤ √ d, the gradient of L at µ can be upper bounded as∥∇ µi L(µ)∥ ≤ 2(∥µ 3 ∥ + • • • + ∥µ n ∥) + 2 exp(-d)(∥µ 1 ∥ + ∥µ 2 ∥), ∀i ∈ [n].Proof. Recall that the gradient has the form ∇ µi L(µ) = E x ψ i (x) k∈[n] ψ k (x)µ k , hence its norm can be upper bounded asEQUATIONFor any ∥x∥ ≤ 2 √ d and i > 2, we have exp(-∥x-µ i ∥ 2 /2) ≥ exp(-(∥x∥ + ∥µ i ∥) 2 /2) ≥ exp(-9d/2), while for i ∈ {1, 2}, exp(-∥x -µ i ∥ 2 /2) ≤ exp(-(∥µ i ∥ -∥x∥) 2 /2) ≤ exp(-(10 √ d -2 √ d) 2 /2) = exp(-32d). Since ψ i (x) ∝ exp(-∥x -µ i ∥ 2 /2) we have ∥x∥ ≤ 2 √ d ⇒ ψ i (x) ≤ exp(-∥x -µ i ∥ 2 /2) exp(-∥x -µ 1 ∥ 2 /2) ≤ exp(-32d) exp(-9d/2) ≤ exp(-25d), ∀i ∈ {1, 2}.Therefore the first term in ( 36) can be bounded asE x k∈[n] ψ k (x)∥µ k ∥ ∥x∥ ≤ 2 √ d ≤ (∥µ 3 ∥ + • • • + ∥µ n ∥) + exp(-25d)(∥µ 1 ∥ + ∥µ 2 ∥).On the other hand, by tail bound of the norm of Gaussian vectors (see Lemma 8 of [Yan et al., 2017] ) we have Pr ∥x∥ > 2√ d ≤ exp(-d).Putting everything together, (39) can be further bounded as∥∇ µi L(µ)∥ ≤ (∥µ 3 ∥ + • • • + ∥µ n ∥) + exp(-25d)(∥µ 1 ∥ + ∥µ 2 ∥) + exp(-d) i∈[n] ∥µ i ∥ ≤ 2(∥µ 3 ∥ + • • • + ∥µ n ∥) + 2 exp(-d)(∥µ 1 ∥ + ∥µ 2 ∥).Theorem 7. For any n ≥ 3, define μ(0) = (µ ⊤ 1 (0), . . . , µ ⊤ n (0)) as follows:µ 1 (0) = 12 √ de 1 , µ 2 (0) = -12 √ de 1 , µ 3 (0) = • • • = µ n (0) = 0, where e 1 is a standard unit vector. Then population gradient EM initialized with means μ(0) and equal weights π 1 = . . . = π n = 1/n will be trapped in a bad local region around μ(0) for exponentially long time T = 1 30η e d = 1 30η exp(Θ(U (0))). More rigorously, for any0 ≤ t ≤ T, ∃i ∈ [n] such that ∥µ i (t)∥ ≥ 10 √ d,Proof. We prove the following statement inductively: ∀ 0 ≤ t ≤ T :EQUATION(40) states that during the gradient EM update, µ 1 will keep stationary at 0. while the symmetry between µ 2 , . . . , µ n will be preserved.The induction base is trivial. Now suppose ( 41), (40) holds for 0, 1, . . . , t, we prove the case for t + 1.Proof of (40). Due to the induction hypothesis, one can see from direct calculation that ∀x, we have ψ i (x|µ(t)) = ψ i (-x|µ(t)) for i = 3, . . . , n, and ψ 1 (x|µ(t)) = ψ 2 (-x|µ(t)).Consequently for ∀i > 2 we have∇ µi L(µ(t)) = E x   ψ i (x|µ(t)) k∈[n] ψ k (x|µ(t))µ k (t)   = E x [ψ i (x)(ψ 1 (x)µ 1 (t) + ψ 2 (x)µ 2 (t))] = 1 2 E x [ψ i (x)(ψ 1 (x)µ 1 (t) + ψ 2 (x)µ 2 (t)) + ψ i (-x)(ψ 1 (-x)µ 1 (t) + ψ 2 (-x)µ 2 (t))] = 1 2 E x [ψ i (x)(ψ 1 (x)(µ 1 (t) + µ 2 (t)) + ψ 2 (x)(µ 2 (t) + µ 1 (t)))] = 0 ⇒ µ 1 it + 1) = µ i (t) = 0.Similarly, for µ 1 , µ 2 we have∇ µ1 L(µ(t)) = E x   ψ 1 (x|µ(t)) k∈[n] ψ k (x|µ(t))µ k (t)   = E x [ψ 1 (x)(ψ 1 (x)µ 1 + ψ 2 (x)µ 2 )] = E x [ψ 2 (-x)(ψ 2 (-x)µ 1 + ψ 1 (-x)µ 2 )] = -E x [ψ 2 (-x)(ψ 2 (-x)µ 2 + ψ 1 (-x)µ 1 )] = -∇ µ2 L(µ(t)).This combined with the induction hypothesis implies µ 2 (t + 1) = -µ 1 (t + 1), (40) is proved.Proof of (41).By induction hypothesis, we have ∀i, ∥µ i (t)-µ i (0)∥ ≤ ηt•(60√ de -d ) ≤ ηT •(60 √ de -d ) ≤ 2 √ d. So ∀i ∈ {1, 2}, ∥µ i (t)∥ ≤ ∥µ i (0)∥ + 2 √ d < 15 √ d.Then by Lemma 15, ∀i ∈ [n] we have∥∇ µi L(µ(t))∥ ≤ 2(∥µ 3 ∥+• • •+∥µ n ∥)+2 exp(-d)(∥µ 1 ∥+∥µ 2 ∥) ≤ 4 exp(-d)•15 √ d = 60 √ de -d ,note that here we used µ3 (t) = • • • = µ n (t) = 0.Therefore by the induction hypothesis we have∥µ i (t + 1) -µ i (0)∥ ≤ ηt • (60 √ de -d ) + η∥∇ µi L(µ(t))∥ ≤ η(t + 1) • (60 √ de -d ), (41) is proven. By (41), ∀0 ≤ t ≤ T , for i = 1, 2 we have ∥µ i (t)∥ ≥ ∥µ i (0)∥ -∥µ i (t) -µ i (0)∥ ≥ 12 √ d - ηT (60 √ de -d ) ≥ 12 √ d -2 √ d = 10 √ d.Our proof is done.NeurIPS Paper Checklist38th Conference on Neural Information Processing Systems (NeurIPS 2024).To see this, consider n =2, µ1 = 0, µ2 = (1, 0, . . . , 0) ⊤ , then the norm of µ1 strictly increases after one iteration.To prevent numerical underflow issues, we change the constant 12 in µ(0) to 2. 


Decision: As stated in the abstract, the paper's focus is the gradient expectation-maximization algorithm for Gaussian mixture models in the overparametrized regime: the authors consider the convergence properties of the problem when a single Gaussian is fitted by a mixture, showing that the considered algorithm converges globally with sublinear rate. The paper collected a positive consensus due to its clarity and novelty. Despite the apparent simplicity of the set-up, the extension of results already available in the literature to the case of generic Gaussian mixture is interesting, as it requires the introduction and application of nontrivial technical steps for the study of the non-convex optimization dynamics of the algorithm. I therefore recommend the acceptance of the manuscript.

Reviews: review: paper_summary: paper_summary: The paper considers fitting a single Gaussian with multiple-component Gaussian mixture models (GMM) through the Gradient EM algorithm. While the two balanced over-specified Gaussian setting has been widely studied in the previous work, generalizing it to multiple-component GMM requires significant algebraic efforts. The entirety of the paper is to show the $1/\sqrt{t}$ convergence rate of the population EM algorithm. In particular, the paper characterizes the explicit convergence rate of $1/\sqrt{T}$ with constants exponential in the number of components, the phenomenon that coincides with the exponential lower bound for the parameter estimation of general GMMs with no separation.
strength_weakness: strength_weakness: -	While the lower bound result (Theorem 7) is a nice addition to the literature, I believe that the gap between this lower bound and the upper bound is large, since the upper bound is exponentially slow in the number of components. 

-	One important result from two specified GMM is the $n^{-1/4}$ (n is the number of samples here) statistical rate after convergence. I would like to see $n^{-1/2k}$ style results in general k-component GMM settings. At least, the authors should have discussed this aspect of previous work and contrasted the implications to k-GMM settings. 

-	The experiment would have been nicer if the final statistical rates were compared.
questions: questions: -	Maybe authors can elaborate on how their results can imply learning k-GMM with small separations?

-	In Theorem 7, there is no restriction on the step size $\eta$. I believe that the lower bound should also be able to tell that $\eta$ cannot be set too large.

-	Why only on the gradient EM? Can the analysis in the paper imply some convergence rates of the standard EM algorithm as well? I think it would make the paper much stronger if it could show that the same results hold for standard EM.
limitations: limitations: N/A
paper_summary: paper_summary: The paper focuses on the setting of a Gaussian Mixture Model with several summands and an input vector produced by one Gaussian distribution, where it employs the Expectation-Maximization rule to infer the model's parameters. Since the problem of having arbitrary number of summands has been unsolved, the paper provides an innovative scheme which includes the computation of the likelihood function and shows that the EM algorithm converges with sublinear complexity. 

The authors also show that there exist neighborhoods of slow convergence rates.
strength_weakness: strength_weakness: - The experimental evaluation is used as a proof of concept and thus is limited. The authors could have (potentially) experimented with several datasets, with varying weights in the GMM, and try to benchmark their algorithm to compare the emergent convergence rates.
questions: questions: NA.
limitations: limitations: NA.
paper_summary: paper_summary: This paper talks about the gradient-EM algorithm for over-parameterized GMM. The paper mostly shows the GLOBAL convergence and its rate when using this model to learn a single Gaussian.
strength_weakness: strength_weakness: 1. The over-parametrized model may have severe overfitting problem. 
2. The based distribution is quite easy: a single normal, with known variance. In the paper, the covariance is fixed as the identity, which simplifies the problem in a deep way. Actually for symmetric 2-GMM, there are already faster algorithms to learn both mean and cov. 
3. I feel confused about the consistency and convergence in the paper. In Line 96, the convergence of KL divergence also contains the convergence of MLE, ie consistency. The convergence to the MLE is another loss function. Also in Remark 6, the convergence when sample size to infinity seems more easily ensured by WLLN.
questions: questions: Besides the weakness above, I also have following questions:
4. If you only learn the single normal, how is the algorithm compared with Dwivedi et al. 2019 or just 2-GMM? Is it necessary to use more? Is it overfitting so the performance seems better?
5. I don’t get why the paper introduces Fact 1. It seems obvious. 
6. The mean is convergent to 0 (true) instead of the MLE.
limitations: limitations: Besides above, 
7. the citation format is not uniform.
paper_summary: paper_summary: The paper studies the convergence of EM for learning mixtures of Gaussians.  Specifically, they consider a simplified setting where the Gaussians are in $d$-dimensions and all have covariance $I_d$.  They consider an overparameterized version of the problem where they parametrize the mixture they are trying to learn by a mixture of $n$ Gaussians with means $\mu_1, \dots , \mu_n$ and the ground truth distribution generating the data just consists of a single Gaussian $N(\mu^* , I_d)$.  The paper analyzes the dynamics of gradient EM for this problem.  The main result of the paper is proving that for this overparametrized variant, gradient EM converges to the true distribution at a rate of $1/\sqrt{t}$ with additional constants depending exponentially on the distance between the initialized means and the true mean, which they show is necessary.

There has been a long line of work on understanding the convergence of EM or gradient EM for learning mixtures of Gaussians.  Without overparametrization, provable convergence is known for mixtures of two Gaussians and it is also known that convergence fails in general for mixtures of three or more components.  For overparamterized settings, a previous work [Dwivedi et. al. 2018] shows that if we parametrize a mixture of two Gaussians and try to learn a ground truth distribution consisting of a single Gaussian, then EM converges at a $1/\sqrt{t}$ rate (as long as the mixing weights are set to be different).  This is in contrast to when we parametrize with only a single Gaussian and EM converges exponentially fast.  The results of the current paper can be seen as generalizing the results of [Dwivedi et. al. 2018] to more than two components.  The paper empirically validates their theoretical results with experiments on simple synthetic datasets.
strength_weakness: strength_weakness: The results of the paper only work when the ground truth is "trivial" i.e. a single Gaussian.

The results are qualitatively similar to previous work on overparametrized mixtures of two Gaussians.  The contributions of the paper are mostly technical and it is a bit difficult to find a nice conceptual takeaway \--- the previous work for two components already showed that overparametrization can lead to drastically slower convergence.  It would be much more exciting and novel, say, if we could prove something when the ground truth were not just a single Gaussian.
questions: questions: .
limitations: limitations: Yes

Comments: title: Thank you for the response
Thank You and Response to the Question on Population Setting
Rebuttal
comment: I thank the authors for the clarification and for addressing my concerns. The additional experimental result on the statistical rate also looks interesting. I have adjusted my evaluation score accordingly.
Thank you for reading our paper and feedback! Please let us know if you have any further questions.
Thank you for recognizing our contribution and for raising the score!

> About our choice of the population data model.

We study the population setting to focus on the non-convex optimization dynamics of gradient EM algorithm. Indeed, using the population model is a standard approach in previous literature of EM analysis [1, 2], and generally in non-convex optimization [3,4].  As discussed in Remark 6, it also implies (asymptotically) the optimization convergence for reality, i.e., sample-based EM.

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

[2]. Sivaraman Balakrishnan, Martin J. Wainwright, and Bin Yu. Statistical guarantees for the em algorithm: From population to sample-based analysis. In Annals of Statistics, 2014.

[3]. Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International Conference on Machine Learning, 2017.

[4]. Mo Zhou, Rong Ge. A local convergence theory for mildly over-parameterized two-layer neural network. In Conference on Learning Theory, 2021.
Thank you for the addressing my concerns! Part of my questions and concerns are clarified by the authors. Although the ground truth is standard 1-d Gaussian, the work is valuable. Therefore, I would like to lift the score to 6.
But I still have questions about the MLE and the true (0). The author responds that the algorithm is implemented on the population data. No mentioning the possibility and reality, if the population data available, the model chosen seems inappropriate.
Thank you for the response and addressing my concerns/questions.  My overall assessment remains the same.
Thanks for reading our response and providing feedback! We look forward to any further comments and updates.
I thank the authors for their answers. They provided experiments as an attachment to their rebuttal answer. I will further study the responses to the other reviewers and update my review.
Dear authors, dear reviewers,

the discussion period has begun as the authors have provided their rebuttals.
I encourage the reviewers to read all the reviews and the corresponding rebuttals: the current period might be an opportunity for further clarification on the paper results and in general to engage in an open and constructive exchange.

Many thanks for your work.
The AC