Statistically Optimal $K$-means Clustering via Nonnegative Low-rank Semidefinite Programming
Authors: Yubo Zhuang
Xiaohui Chen
Yun Yang
Richard Y. Zhang
Venue: ICLR 2024 oral (2024)
Keywords: clustering
Burer-Monteiro
semidefinite programming
Primary area: optimization

Abstract: $K$-means clustering is a widely used machine learning method for identifying patterns in large datasets. Recently, semidefinite programming (SDP) relaxations have been proposed for solving the $K$-means optimization problem, which enjoy strong statistical optimality guarantees. However, the prohibitive cost of implementing an SDP solver renders these guarantees inaccessible to practical datasets. In contrast, nonnegative matrix factorization (NMF) is a simple clustering algorithm widely used by machine learning practitioners, but it lacks a solid statistical underpinning and theoretical guarantees. In this paper, we consider an NMF-like algorithm that solves a nonnegative low-rank restriction of the SDP-relaxed $K$-means formulation using a nonconvex Burer--Monteiro factorization approach. The resulting algorithm is as simple and scalable as state-of-the-art NMF algorithms while also enjoying the same strong statistical optimality guarantees as the SDP. In our experiments, we observe that our algorithm achieves significantly smaller mis-clustering errors compared to the existing state-of-the-art while maintaining scalability.

Published as a conference paper at ICLR 2024
STATISTICALLY OPTIMAL K-MEANS CLUSTERING VIA NONNEGATIVE LOW-RANK SEMIDEFINITE PROGRAM-
MING
Yubo Zhuang*, Xiaohui Cheni*, Yun Yang*, Richard Zhang
University of Illinois at Urbana-Champaan Champaign* University of Southern California (yubo2, yy84, ryz]@illinois edu, xiachuiclusc.edi
ABSTRACT
K-means clustering is widely used machine learning method for identifying pat- terns in large datasets. Recently, semidefinite programming (SDP) relaxations have been proposed for solving the K-means optimization problem, which enjoy strong statistical optimality guarantees. However, the prohibitive cost of implementing an SDP solver renders these guarantees inaccessible to practical datasets. In contrast, nonnegative matrix factorization (NMF) is simple clustering algorithm widely used by machine learning practitioners, but it lacks solid statistical underpinning and theoretical guarantees. In this paper, we consider an NMFF-lie algorithm that solves nonnegative low-rank restriction of the SDP-relaxed K --means formulation using nonconvex Burer-Monteiro factorization approach. The resulting algorithm is as simple and scalable as sfate-of-the-art NMF algorithms while also enjoying the same strong statistical optimality guarantees the SDP. In experiments, we observe that our algorithm achieves significantly smaller mis-clustering errors compared to the existing ofte-of-theear while maintaining scalability.
INTRODUCTION
Clustering remains a common unsupervised learning technique, for which the basic objective is to assign similar data points the same group. Given data in the Euclidean space, a widely used clustering method is K -means clustering, which quantifies "similarity" in terms of distances between the given data and learned clustering centers, known as centroids (MacQueen, 1967). In order to divide the data points Xy...., X,, E R into K groups, K K-means clustering aims to minimize the following cost function:
n
min
IX 1x KE[K]
min
(1)
8 1 ERP
where 3g is the centroid of the k-th cluster for l [ f1, ..... K) It is well-known that exactly solving problem (1) is NP-hard in the worst-case (Dasgupta 2007 Aloise et al. (009). so computa- tionally tractable approximation algorithms and relaxed formulations have been extensively studied in the literature. Notable examples include Lloyd's algorithm (Lloyd 1982), spectral clustering (von)
Luxburg 2007 Ng al. 200T). nonnegative matrix factorization (NMF) (He al. 2011] Kuang et al. Wang and Zhang 2012), and semidefinite programming (SDP) (Peng and Wei 2007
Mixon et 2017 Royer [01 Fei and Chen 2018 Giraud and Verzelen 2018).
Among those popular relaxations, the SDP approach enjoys the strongest statistical guarantees under the standard Gaussian mixture model in that it achieves an information-theorett sharp threshold for exact recovery of the true cluster partition (Chen and Yang 2021). Unfortunately. the SDP and its strong statistical guarantees remain completely inaccessible to real-world datasets, owing to the prohibitively high costs of solving the resulting SDP relaxation. Given data points, the SDP is matrix optimization problem, over a dense n X n membership matrix Z, that is constrained to be both positive semidefinite Z 2 0 well as elementwise nonnegative Z 2 0. Even ignoring the constraints, basic but fundamental difficulty is the need to store and optimize over the n² individual elements of the matrix. Even small dataset with R 1000, such the banknote authentication

---

Published as a conference paper at ICLR 2024
Computational complexity & Statistical performance trade-off 1000
CPU time [seconds]
10
NLR KM
SC
2
SDP
10:3
10 Mis-clustering error
10a
Figure 1: Log-scale trade-off plot of CPU time versus mis-clustering error, under an increasing number of data points n. Here, NLR corresponds to our proposed non-negative low-rank factorization method, KM corresponds to K (Arthur and Vassilviskki| (007)) SC corresponds to spectral clustering (Ng et al.) 2001), SDP (Peng and Weil 2007) uses the SDPNAL+ solver (Yang al., 2015), and NMF corresponds non-negative factorization (Ding et al. 2005). We follow the same setting as the first experiment in Section where the theoretically optimal mis- clustrusg error decays to zero as n increases to infinity.
dataset (Dua and Graff 2017 . translates into an SDP with n² z 10 optimization variables, which is right the very limit of state-of-the-ee SDP solvers like SDPNAL+ (Yang et al. 2015).
On the other hand, NMF remains one of the simplest and practically useful approaches to clustering due to its scalability (He et al. [201T: Kuang et al. 2015). When the clustering problem at hand exhibits an appropriate oon-dimensional structure, NMF gains significant computational savings by imposing elementwise nonnegativity over an n x r low-rank factor matrix U O, in order to imply positive semidefiniteness Z ž 0 and elementwise nonnegativity Z 2 0 over the n x n membership matrix Z UUT While highly scalable, there remains unfortunately very little statistical underpinning behind NMF-based algorithms.
Our contributions. In this paper, we propose an efficient, large-scale, NMF-like algorithm for the K- means clustering problem, that meanwhile enjoys the same sharp exact recovery guarantees provided by SDP relaxations. We are motivated by the fact that the three classical approaches to K- K-means clustering, namely spectral clustering, NMF, and SDP, can all be interpreted as techniques for solving slightly different relaxations of the same underlying K-means formulated as a mixed integer program; exposition in Section This gives hope to break the existing computational and statistical bottleneck by investigating the intersection of these three classical approaches.
At its core, our proposed algorithm is a primal-dual gradient descent- ascent algorithm to optimize over a nonnegative factor matrix, inside an augmented Lagrangian method (ALM) solution of the SDP. The resulting iterations closely resemble the projected gradient descent algorithms widely used for NMF and spectral clustering in the existing literature; in fact, we show that the latter can be recovered from our algorithm by relaxing suitable constraints. We prove that the new algorithm enjoys local linear convergence within primal-dual neighborhood of the SDP solution, which is unique whenever the centroids satisfy a well-separation condition from (Chen and Yang 2021). In practice, we observe that the algorithm converges globally at a linear rate, As shown in Figure I our algorithm achieves substantially smaller mis-clustering errors compared to the existing state-of -f-the-rrt.
The main novelty of our algorithm is the use of projected gradient descent to solve the difficult primal update inside the ALM. Indeed, this primal update had been the main critical challenge faced by prior work; while a similar nonnegative low-rank ALM was previously proposed Kulis et al. (2007). ) their inability to solve the primal update to high accuracy resulted in substantial slow-down to the overall algorithm, which behaved more like inexact ALM. In contrast, our projected gradient descent method can solve the primal update at a rapid linear rate to machine precision (see Theorem I], so overall algorithm is able to enjoy the rapid primal-dual linear convergence that is predicted and justified by classical theory for exact ALMs. As shown in Figure4iin in Appendix B our algorithm is the first ALM that can solve the SDP relaxation to arbitrarily high accuracy.
Organization. The rest of the paper is organized as follows. In Section 2 we present some background on several equivalent and relaxed formulations of the K-means clustering problem. In
2

---

Published as a conference paper at ICLR 2024
introduce nonnegative low-rank SDP for solving the Burer- Monteiroo formulation of -maans problem. In Section 4 we establish the linear convergence guarantee for our proposed primal-dual gradient descent-ascent algorithm in the exact recovery regime. Numerical experiments are reported in Section 5 Proof details are deferred to the Supplementary Material.
Section 3
2 BACKGROUND ON K-MEANS AND RELATED COMPUTATIONAL METHODS
Due to the parallelogram law in Euclidean space, the centroid- -bseed formulation of the K -means clustering problem in ① has equivalent partitioo-based formulation (cf. [Zhuagg al. (2022) as
max GI.....x
(X X,)=
G
Gx mm]
GR
L.EGG
where the Euclidean inner product (X,, Xj) X X represents the similarity between the vectors X; and X,, the clusters (Gx) form a partition of the data index [n]. |G denotes the cardinality of Gk and Li denotes disjoint union. The objective function in (2) is the log-likelihood function of the
cluster labels (modulo constant) by profiling out the centroids in ① as nuisance parameters under the standard Gaussian mixture model with common isotropic covariance matrix across all the K components (Zhuang et al. (023)) Using the one-hot encoding of the partition Gi,..., Gk. we can uniquely associate it (up to label permutation) with a binary assignment matrix H (hax)
{0, 1) 1)a* such that hus 1 if E Gk and has = 0 otherwise. Then the K K-means clustering problem in (2) can be written as a mixed integer program (MIP) as follows:
min (A, (A, H E (0,1)"XK, HI, - 1m
33)
where = is the n X n negative Gram matrix of the data Xpxx = (X = diag(|G,|-1 ...., G is a normalization matrix and the constraint H1, = 1, with 1, being the n--dimensionl vector of all ones reflects the row sum constraint of assignment, i.e., each row of H. contains exactly nonzero entry with value one.
It is well-known that the problem in (3) is NP-hard in the worst-case (Dasgupta 2007 Aloise et al., 2009), so various relaxations for the K-means constraint are formulated in the literature to tractably approach problem clustering (von Luxburg. 2007 Ng et al. 200T) which only maintains the weakened orthonormal
with the same objective function. Popular methods include: (i) spectral
constraint H" H Is for H = /B (ti) nonnegative matrix factorization (NMF) (He et al.
2011 Kuang et al. 2015) which only enforces the elementwise nomnegativity constraint on H 2 0: (iii) semidefinite programming (SDP) (Peng and Wei, 2007 Royer 2017 Giraud and Verzelen 2018) which reparameterizes the assignment matrix as the positive semidefinite (psd) membership matrix
Z : HBH ž 0 and additionally preserves the constraints tr(Z) K, Z1, 1, and Z 2
namely we solve
(A. : Lz( = K, E1, 1, Z 0
min
ZES
where St stands for the convex cone of the n X n real symmetric psd matrices.
Among those popular relaxations, the SDP approach enjoys the strongest statistical guarantees under the Gaussian mixture model (I4). It is shown JyChhn and Yang (2021) that the above SDP achieves an information- -hhoorettc limit for exact recovery of the true cluster partition. Precisely, the sharp threshold on the centroid-separation for phase transition is given by
e² = =40
1
Kp log
log
(5)
in the sense that for any a > 0 and K Olog(()/ log log(n)), if the minimal sentrinnn Omin == min [|e, - /ull22 2 (1 + O) then with high probability solution of (4) is unique and it perfectly recovers the cluster partition structure; while if θ < (1 - a then the maximum likelihood estimator (and thus any other estimator) fails to exactly recover the cluster partition. In simpler words, the SDP relaxed K K-means optimally solves the clustering problem with zero mis-clustering error as soon is possible to do so, i.e., there is relaxation gap.
3

---

Published as a conference paper at ICLR 2024
It is important point out that mii + oo is necessarily needed to achieve exact recovery with high probability. When the centroid -serartttno diverges, spectral clustering is shown to achieve the minimax mis-clustering error rate exp(-(1 + o(1) min// /8) under the condition p o(n2mmin) (or p o min! with extra assumptions on singular values of the population matrix E[X) (Löfffer at. 202T). This result implies that, in the low-dimensional setting, spectral clustering asymptotically
e
attains the sharp threshold (5) as lim inf log > 1. However, it remains unclear how spectral
cluster performs in theory under the hig-dimensioonl regime when nO? = Opp) because running
K K-means on the eigen-embedded data would also depend on the structure of the population singular
values (Han et al., 2023). Thus, practice, spectral clustering is less appealing for high dimensional data, in view of more robust alternatives such as the SDP relaxation in (4)
Despite the K-means clustering problem solved via SDP enjoys statistical optimality, it optimizes over an X dense psd matrix to make it completely inaccessible to practical datasets with even moderate size of a few thousand data points. By contrast, the NMF approach using a typical workstation hardware can handle much larger scale datasets as in the documents and images clustering applications (Cichocki and Phan, 2009 Kim et al., 2014: Xu et al., 2003). Nonetheless, little is known in the literature about the provable statistical guarantee for general NMF approach. For the clustering problem, our goal is to break the computational and statistical bottleneck by simultaneously leveraging the implicit psd structure of the membership matrix and nonnegativity constraint. In Section 3 we propose a novel NMF-like algorithm that achieves statistical optimality by solving a nonnegative low-rank restricted SDP formulation.
3 OUR PROPOSAL: K-MEANS VIA NONNEGATIVE LOW-RANK SDP
The non-negative low-rank (NLR) factorization is a nonconvex approach for solving SDP when its solution is expected be low rank (Burer and Monteiro 2003). Standard SDP problem with only equality constraints an n X r matrix with r < n. In our SDP relaxed K-means formulation (4). the true membership matrix Z" is a block diagonal matrix containing K blocks, where each block has size n& n& with all entries equal to ny and ny = =Ga2 is the size of the true cluster G Thus we can exploit this low-rank structure and recast the SDP in (4) as nonconvex optimization problem over U E Ruxr;
the psd matrix variable Z E St. can be factorized as Z UU where U is
{(A.UU)) 1111= K, U = In, U
min VER
(6)
where we replaced the constraint UUT: 2 O with the stronger constraint U 2 0 that is easier to enforce (Kulis al., (007) as the NMF setting. Note that the NLR reformulation (6) can be viewed
as restriction of the rank-constrained SDP formulation since U > 0 implies 2 = -U 2 0. Even though the NLR method turns a convex SDP into a nonconvex problem, it leverages the low-rank structure to achieve substantial computational savings. In addition, the solution Ú to (6) has a natural statistical interpretation: rows of Ú can be interpreted as latent positions of n data points when projected into R'. and thus can be used to conduct other downstream data analyses in reduced dimensions. Under this latent embedding perspective, formulation (6) can be viewed as a constrained PCA respecting the clustering structure through constraint || K.
The NLR formulation (6) also has a close connection to a class of NaF-d clustering algorithms, which instead solve the following optimization problem
{I.A + UUT 1| U2:
min UER
The NMF formulation (7) advocated in a number of papers (Ding et al., 2005 Kuang et al. 2015
Zhu et al., 2018) for data clustering in the NMF literature due to its remarkable computational scalability, and is equivalent to a simpler version of (6) by dropping the two equality constraints. Our empirical results in Section 5 Figure show that keeping these two equality constraints is beneficial: the resulting method significantly improves the classification accuracy while maintaining comparable computational scalability. More importantly, we have the theoretical certification that the resulting method from (6) achieves the information-thecnm thheoretccc limit for exact recovery of the true cluster labels (see Section
We will now derive a simple primal-dual gradient descent-ascent algorithm to solve the NLR formu- lation in 66. To begin, note that can first turn the nonsmooth inequality constraint U > together
4

---

Published as a conference paper at ICLR 2024
Algorithm 1 Primal-dual algorithm for solving NLR formulation (6) of K-means clustering
Input. Dissimilarity matrix A -X" X, clustering parameter K 2 1, rank parameter r 2 K augmentation parameter 0, step size a > 0. Output. A second-order locally optimal point U. Algorithm. Initialize y 0. Do the following:
1. (Primal descent; projected gradient descent steps) Recursively run the following until convergence: U U: for 1 E N,
U++1= IIl VuL:(U', .v) with ll( (VV V Raxr
where VvLa(U" v) 2 ld Iny" y and y - y 3(-UUU1 1, 1.,) Upon convergence at iteration count ln, set Unew = -
2. (Dual ascent; augmented Lagrangian step) Update dual variable via
Unew + B(U U 1, - 1,p.
3. (Stopping criterion) If max [ Ul Uee UT 1n 1a |1 falls below some tolerance threshold, then return Unew- Otherwise, set t Unew and +- Iuew, and repeat Steps and 2.
with the trace constraint to the subset
Ω := (U E ER | K, U 02.
(8)
The projection operator to Q can be easily computed as
VK. VK-(V++
(V) = arg min |IU Vllr
(9)
UEn
II(V)+Il
where (V)+ = max (V.0) denotes elementwise positive part of V. Then the NLR can be converted the equality-constraine problem over Ω:
min VER
(A, 10 11
(10)
Using the standard augmented Lagrangian method (Nocedal and Wright 2006 Chapter 17), that the above 6 is also equivalent to
(L- Id, A,UUT) 3 ( LIUUTI 1n L In).
min Uen
(11)
where B >> is penalty parameter and optimal dual variable for the trace constraint |0|} = tr(2)= = K. The augmented Lagrangian for problem is
(0, A) for some proper choice of À motivated from the
3 Ls(U,)) = (L ld, + A,U + (v. U , - 1 - 1, - -1n
(12)
Therefore, we consider the augmented Lagrangian iterations
Unew = arg min CU, Unew = + 3(U U new In 1,n).
(13)
UEl
Now, we consider minimizing CU,U y) over U. with a fixed 3 and Here, we observe that (U, a quartic polynomial over U, and therefore its gradient can be easily derived as Ls (2.A + 2L ld + 1T + TU where == y + 3UU 1,.. Combining this with the insight that it is easy to project onto Ω we can perform the following projected gradient descent iterations Unew = [U - a V Ca.
until (U, y) is minimized. The complete algorithm is summarized in Algorithm I whose per iteration time and space complexity are both Onr)). We point out that the NMF formulation 7 is simpler version of our NLR formulation (6). projected gradient descent algorithm for solving the former corresponds to the primal decent step Algorithm with y = 0, 3 = 0 and a simpler IT Tg (V) (V)+- Upon obtaining the optimal solution U, rounding procedure
projection operator
will be applied. More details about the rounding procedure and the choice of tuning parameters are mentioned in Appendix

---

Published as a conference paper at ICLR 2024
4 THEORETICAL ANALYSIS
In this section, we establish the local linear convergence rate of the NLR algorithm (Algorithm D for solving the K K-means clustering. We shall work with the standard Gaussian mixture model (GMM) that assumes the data XI,...,,,, X, are generated from the following mechanism: if i E G then
X; 1 + Es,
(14)
where Gi....,Gy G is true (unknown) partition of [n]] we wish to recover, p1..... MK R" are the cluster centers and & ~ (00,o2pp) are i.idd. Gaussian noises. Let nı = |G&] denote the size of G3 and by convention nin 0. Since in the exact recovery regime (Assumption A below), there is no relaxation gap any global optimum U* of the NLR problem (4) at = K corresponds to the unique global optimum 2* of the SDP problem (6) through the relation Z = U'U*" it is sufficient to focus our analysis on the r K case. We note that the K case corresponds to an exact parameterization of the matrix rank, which is standard in the analysis of algorithms based on the
NLR formulation (Ge al., 2017: Chi et al. (019).
Algorithm Dis formulated as an exact augmented Lagrangian method, in which the primal subproblem mingen (p(U,)) y) in Step 1 is solved to sufficiently high accuracy (i.e., machine precision) as to be viewed an exact solution for Step 2 This contrasts with inexact methods, in which the primal subproblem over U is only solved to coarse accuracy (i.e.. 1-2 digits), usually owing to slow convergence in Step 1. We will soon prove in this section that projected gradient descent enjoys rapid linear convergence within local neighborhood of the primal-dual solution, and it is therefore very reasonable to run Step until it reaches the numerical floor.
Under exact minimization in Step 1, it is a standard result that the dual multipliers converge at a linear rate in Step under second-order sufficiency conditions. This, in turn, implies the linear convergence of the primal minimizers in Step 1 towards the true solution. The wording for the following is taken directly from Proposition and 2 of Bertsekas (1976), though a more modern proof without constraint qualification assumptions is given in Fernández and Solodov (2012). (See also Bertsekas (2014) Proposition 2.7)) More details and explanations can be found in Appendix[A]
Proposition 1 (Existence and quality of primal minimizer). Let U* denote a local minimizing point that satisfies the second-order sufficieni conditions for an isolated local minimum with respect to multipliers y". Then, there exists scalar 3 2 such that, for every B > 8 the augmented Lagrangian C3(U, y) has a unique minimizing point U B) within an open ball centered at U*.
Furthermore, there exists some sealar M > such that (U,(y, (3-3)-U - U- S (M/B)lly - y²1l-
Proposition 2 (Linear convergence of dual multipliers). Under the same assumptions Proposi- tion I define the sequence
yk+1 = yk + 3(ULUT - 11,) where Uv = Ulk, 3) = arg min Co (U,W)) yk).
U
Then, there exists a radius R > such that, if 11y0 y01l < R, then yk + y converges linearly
| y'll
M
lip k-
1yyk 101l
3
Propositions [ aadd are directly applicable Algorithm 1 [whnn r K, because the global minimum U is made isolated by the nonnegativity constraint U 2 0. In fact, it is possible to generalize both results to all r 2 K, even when the global minimum U* is no longer isolated, although this would require further specializing Propositions D and 2 to our specific problem. The key insight is that, in the r > K case, the closest global minimum U to the current iterate U continues to satisfy all the properties satisfied by the isolated global minimum when = K. This is essentially the same idea as (Ge et al. 2017, Definition 6) and (Chi et al. 2019 Lemma 4). Nevertheless, we leave the extension to the K case as future work and focus on the K case in this paper.
In view of Propositions Iepposstionsaand]]] [ all of the difficulty that remains is to show that projected gradient descent is able to solve the primal subproblem minven Ls(U,y) y) efficiently. Below is our main theorem for establishing the local linear rate of convergence of the proposed NLR solution for K-means clustering under GMM. In fact, this local linear convergence result holds for all r > K.
Note that with finite precision arithmetic, any "exact solution" is exact only up to numerical precision. Our of "exact" in this context is consistent with the classical literature, e.g., Bertsekas (176).

---

Published as a conference paper at ICLR 2024
Assumption A (Exact recovery regime). For notation simplicity, consider the equal cluster size case, where the minimal separation min := m sj ||4 - pxl12 satisfies Omi 2 (1 + a)
for some á > 0, where θ is the information-theorrticc optimal threshold given in (
In Appendix [ we provide the theorem (Theorem 9 in Appendix C.4)
for the general case where the cluster sizes are unbalanced (Assumption 1 in Appendix C.33p p. Analogous to the minimal
separation, we define the maximal separation max == maxi <j jikks ||æ - |4x 1xxl22. For any block partition sizes m (m1. m2,........................ ,mk) satisfying my 2 1, L, m& = we let Gm denote the
set of all m-block diagonal matrices with nonnegative entries (see Appendix [c.1 for a precise definition). For any within block weights ag = (Ax.,1, 41.2..... m. we also let U denote the associated optimal solution, defined as (17p in Appendix C.1 Let a = [ min { } à = max {ax,e),
(ox J. and O be the orthogonal transformation group on R" We will consider a fixed dual variable ÿ that is close to y² -Ltt
"Le U? denote the initialization of the algorithm.
of and U denote the unique stationary poin of the projected gradient descent updating formula under
this dual value y. i.e., U satisfies U = II (ū - (B.
Theorem 1 (Local convergence of projected gradient descent). Suppose y" is the optimum dual for the SDP problem (see Assumption 2 in Appendix C.3) and Assumption A holds. Assume O logn), 0max S C00mm for some C > 0, and there exists some block partition sizes m and an associated optimal solution U satisfying ā S ca, c > Moreover, assume that the runing parameters (L, B) in augmented Lagrangian (12) satisfy 5 = O igin/ /K*), ), L = O min! K) and step size a of the projected gradient descent satisfies a = O n0 win )- If fl y" S 8 for
some 8 >> 0, and the initialization discrepancy A = 1 U satisfies
|| -- = O(K/vnr) and || || f| = Or 0.5 K 5.5 min K
2 2
min/ log 1))
then it holds with probability at least - t2 that for any t 2 I O(K3),
U 9m and inf QEO,.UQEGm
| iQli r S
inf
10-- - tQQ1
QEO.;0099..
for 1 OK- Here, C1 and 02 are some constants.
Proof sketches. The proof strategy is to divide the convergence of the projected gradient descent for solving the primal subproblem into two phases. In phase one, we show that after at most iterations, the iterates will become block diagonal, that is, fall into set 9m provided that the initialization is close to certain optimum point with some block diagonal form in 9m. We then show that once the iterate becomes m-block diagonal, the algorithm enters phase two, where the iterate remains m-block diagonal. Moreover, the projected gradient descent attains a linear convergence rate, since the objective function C(., 8) is restricted strongly convex at Ú within Gm- More precisely, there exists some 3 > 0, such that for any U E 9m. we have
(V) 1) A) 2 33 || for A = UU- - -Ú,
where Qu = argmin geo.. I UQ || p. More details can be found in Appendix C
Implications of the theorem. According to Theorem if we choose L = (n0 min/ K),
O /) and a O(K²n0² ), then phase one will last at most / 7= 0 iterations; and after entering phase two, the contraction rate of the projected gradient descent is y = 1 1 O(K-)) These choices make the iteration complexity of the projected gradient descent for solving the primal subproblem of order ((K*) (modulo logarithmic factors), and the overall time complexity of the primal-dual algorithm becomes of order nr) In addition, given that the stationary point Ü satisfies ||,|| = O(VK/n) and Ü | = O(vK), the initialization condition only demands a constant relative error with respect to n, making it reasonable requirement. For example, under mild conditions, rapid K K-means algorithms like Lloyd's algorithm (Lloyd 1982 Lu and Zhou 2016) can be employed to construct an initialization that meets this condition with high probability,
NUMERICAL EXPERIMENTS
We present numerical results to assess the effectiveness of the proposed NLR method. We first conduct two simulation experiments using GMM to evaluate the convergence and compare the performance
Note that its existence and uniqueness is guaranteed in the proof of Theorem i when ÿ close to y"
7

---

Published as a conference paper at ICLR 2024
Statistica al performance
Computational complexity
Convergence of NLR
Reaffe oforence otweal
N6-di20hn..o0 onot
00
Ene boctro1
337
00
10"
107
50
104
107
32
107
Figure 2: Log scale plots with error bars of mis-clustering error (leftmost) and time cost (in the middle) as sample size n increases and the convergence of NLR over iterations (rightmost). The plots are partial for SC (SDP) due to their huge space (time) complexity when the sample size is large.
of NLR with other methods. Then, we perform comparison using two real datasets. One of the competing methods in our comparison is clustering based on solving the NMF formulation (7) Specifically, we employ the projected gradient descent algorithm, which is a simplified version of Algorithm I discussed in Section 3 to implement this method. We adopt random initialization and set the same r for both NMF and NLR to ensure a fair comparison. Finally, we conduct further experiments on three datasets in UCL.
Performance for NLR for GMM. Our goal is to compare the statistical performance and time complexity of NLR, NMF, the SDPNAL+ solver (Yang et al. 2015) for solving SDP, spectral clustering (SC), and +-means++ (KM) under GMM. In our setting, we choose cluster numbers K 4 and place the centers of Gaussian distributions the vertices of simplex such that
2 min = - with = 0.64, where + is the sharp threshold defined in Equation (5) The sample 57, 600, the dimension is p = 20, and we set the rank parameter of the NL. to be r = 2K. The results are summarized in the first two plots of Figure 2 with each case repeated 50 times. From the first plot, we observe that the mis-clustering error of SDP and NLR coincides; the eгrоr remains stable and decreases as the sample size n increases. However, NMF, KM, and SC exhibit large variances and are far from achieving exact recovery. The second plot indicates that SDP has super-linear time complexity, while the log-scale curves of our NLR approach, KM, and NMF are nearly parallel, indicating that they all achieve linear time complexity. In particular, SDP becomes time-consuming when is larger than 2, 000. Further comparisons of statistical performance and time complexity with increasing dimension p or varying numbers of clusters K can be found in Appendix B
size ranges from n 400
Linear convergence of NLR. To analyze the convergence of Algorithm I we maintain the same setting as described earlier. However, we now fix π 1, 000 and consider minimum distance
between centers O min = y 5 with 7 = 1.44. We explore three cases: r = K, r = 2K, and 20K, and set the number of iterations for Step 1 (primal update) to be 100. The convergence of our NLR approach is shown in the third plot of Figure [ with each result based on 30 replicates. We measure the relative difference of U compared to optimal solution U* using the Frobenius
Statistical performance for CyTOF
Statistical performance for CIFAR-10
05
BK
7 BK
0.4
100
0.35
0.4
Wis-clsterring erfor
0.3
Mis-cliterrigg
0.3
0.25
0.7
0.1
5C
NMF
KM
NLR
SDP
SC
NMF
KM
NLR
SDP
Figure 3: Boxplots of mis-clustering error (with means) for CyTOF dataset (on the left) and CIFAR-10 10 (on the right) among five different methods. The plots are partial for SC (SDP) due their huge space (time) complexity when the sample size is large.
8

---

Published as a conference paper at ICLR 2024
Table I: Mis-clustering error (SD) for clustering three datasets in UCI: Msplice, Heart and DNA. We
randomly sample n = 1, 000 (n = 300 for Heart) many data points for 10 replicates. DNA, (DNAg) stands for the perturbation with t-distribution (skewed normal distribution) random noise.
NMF
KM
NLR
SC
SDP
Msplice
0.372 (0.073)
0.283 (0.112)
0.161 (0.034)
0.288 (0.024)
0.194 (0.091)
Heart
0.470 (0.024)
0.348 (0.057)
0.230 (0.014)
0,472 (0.011)
0.248 (0,007)
DNA
0.449 (0.063)
0.294 (0.082)
0.188 (0.020)
0.291 (0.014)
0.196 (0.022)
DNA1
0.416 (0.048)
0.337 (0.084)
0.243 (0.055)
0.326 (0.039)
0.263 (0.044)
DNA2
0.435 (0.061)
0.268 (0.035)
0.235 (0.031)
0.317 (0.025)
0.252 (0.040)
From the third plot, we observe that our NLR approach achieves linear convergence in this setting. The curves overlap closely, indicating that the choice of r does not significantly affect the convergence rate. Additional experiments for the comparison of the convergence rate between our algorithm (based on projected gradient descent) and the conventional algorithm (which lifts all constraints into the Lagrangian function) can be found in Appendix B
norm -00""~000""">>>>1000000000
CyTOF dataset. This mass cytometry (CyTOF) dataset consists of protein expression levels for N = 265, 627 cells with p 32 protein markers. The dataset contains 14 clusters or gated cell populations. In our experiment, we select the labeled data for individual H1, From these labeled data, we uniformly sample n data points from K unbalanced clusters, specifically clusters 2, 8, and 9, which together contain a total of 46, 258 samples. The results are presented in Figure 3 where we conduct all methods for 100 replicates for each value of n. We consider two cases: case corresponds to n = 1800 and case 2 corresponds to 46, 258. From the plot, we observe that NLR and SDP remain stable, while KM exhibits significant variance, and NMF produces many outliers. case 1, the mis-clustering error of NLR is comparable to that of SDP, indicating that the performance of NLR can be as good as SDP. However, when n is on the order of 1, 000, the time complexity of SDP becomes dramatically high, which is at least (nn..)).
CIFAR-10 dataset. We conduct a comparison of all methods on the CIFAR- 10 dataset, which comprises 60,000 colored images of size 32 32 3, divided into 10 clusters. Firstly. we apply the Inception v3 model with default settings and perform PCA to reduce the dimensionality to = 50. The test set consists of 10,000 samples with 10 equal-size clusters. In experiment, focus on four clusters ("bird", "deer", "ship", and "truck") with K = 4. We consider two cases: case involves random sub-sampling with n = 1, 800 to compare NLR with SDP, and case 2 uses 4, 000. The results are displayed in the second plot of Figure based on 100 replicates. Similar to the previous experiment, we observe that NLR and SDP exhibit similar behavior and achieve superior and more consistent performance compared to KM, SC, and NMF. NMF displays significant variance, while KM, SC produce many outliers.
UCI datasets. To empirically illustrate the advantages and robustness of our method against the GMM assumption, we conduct further experiments on three datasets in UCI: Msplice, Heart and DNA. The results of mis-clustering errors (with standard deviation) are summarized in Table [ where we randomly sample n = I, 000 (n = 300 for Heart dataset) many data points for a total 10 replicates. From the table we can observe the superior behavior of NLR (our algorithm) and SDP over the rest competitors. DNA, (DNAg) stands for the results of DNA dataset after perturbation with t-distribution (skewed normal distribution) random noise. In both experiments, each data point I; is perturbed with x, + 0.2c,. In the first experiment (DNA1), the injected noise e; is i.i.d. random vector with i...d. entries following t-distribution with 5 degree of freedom, which has much heavier tail than the normal. In the other experiment (DNAg), the injected noise follows a skewed normal distribution with variance and skewness 0.2. From the table we can observe that, compared to DNA, the performances of all methods for DNA, and DNA2 are impacted; but NLR performs better around all cases, which is comparable with the original SDP as expected. Moreover, from the QQ-plots of Figure 5 in Appendix B. the GMM assumption is also clearly violated for the Heart dataset; however, we can still observe the superior behavior of NLR (our algorithm) and SDP over the rest competitors.
9

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank the reviewers for their valuable comments. X.C. acknowledges support from NSF CAREER Award DMS-2347760. Y.Y was supported by NSF DMS-2210717 R.Y.Z was supported by NSF CAREER Award ECCS-2047462 and the C3.ai Digital Transformation Institute.
REFERENCES
Daniel Aloise, Amit Deshpande, Pierre Hansen, and Preyas Popat. Np-hardness of euclidean sum-of- squares clustering. Machine learning, 75(2):245-248, 2009.
David Arthur and Sergei Vassilvitskii. K-means++ the advantages of careful seeding. Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithms, pages 1027-1035, 2007.
Dimitri P Bertsekas. Multiplier methods: A survey. Automatica, 12(2):133-144 145, 1976.
Dimitri P Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014.
Samuel Burer and Renato D. C. Monteiro. A nonlinear programming algorithm for solv- ing semidefinite programs via low-rank factorization. Mathematical Programming, 95(2): 329-357, 2003.
doi: 10. 0011077/110007-00--00522 URL fottps://don 000.00.0r//0.11101 1007/ 8 810107-002-0352-2
Xiaohui Chen and Yun Yang. Cutoff for exact recovery of gaussian mixture models. IEEE Transac- fions Information Theory, 67(6))4223 4238, 2021.
Yuejie Chi, Yue M Lu, and Yuxin Chen. Nonconvex optimization meets low-rank matrix factorization: An overview. IEEE Transactions on Signal Processing, 67(20):5239-5260 2019.
Andrzej Cichocki and Anh-Huy Phan. Fast Local Algorithms for Large Scale Nonnegative Matrix and Tensor Factorizations. IEICE Transactions Fundamentals of Electronics Communications and Computer Sciences, 92(3):708-722 January 2009. doi: 0tt7://ransfunBB22A....
Sanjoy Dasgupta. The hardness of k-means clustering. Technical Report CS2007-0890, University of California, San Diego, 2007.
Chris Ding, Xiaofeng He, and Horst Simon. On the equivalence nonnegative matrix factorization and spectral clustering. In Proceedings of the 2005 SIAM international conference data mining, pages 606-610. SIAM, 2005.
Dheeru Dua and Casey Graff. UCI machine learning repository, 2017. URL http:://acchive ics uci edu/ml
Yingjie Fei and Yudong Chen. Hidden integrality of sdp relaxations for sub- gaussian mixture models. In Sébastien Bubeck, Vianney Perchet, and Philippe Rigollet, editors, Proceedings of the 31st Conference On Learning Theory, volume 75 of Proceedings of Machine Learning Research, pages 1931-1965. PMLR, 06-09 Jul 2018. URL httpp:://proceeeeeeeeee wgg.rroeeeddng........... ml www.ppessssl...
Feil8a. html
Damián Fernández and Mikhail V Solodov. Local convergence of exact and inexact augmented lagrangian methods under the second-order sufficient optimality condition. SIAM Journal on Optimization, 22(2):384 407, 2012.
Rong Ge, Chi Jin, and Yi Zheng. No spurious local minima in nonconvex low rank problems: A unified geometric analysis. In International Conference Machine Learning, pages 1233-1242. PMLR, 2017.
Christophe Giraud and Nicolas Verzelen. Partial recovery bounds for clustering with the relaxed kmeans. Math. Stat. Learn., (3/4):317-374, 2018.
10

---

Published as a conference paper at ICLR 2024
Xiao Han, Xin Tong, and Yingying Fan. Eigen selection in spectral clustering: A theory-guided practice. Journal of the American Statistical Association, 1 118(541):109 121, 2023. doi: (0.10800016211599222111111101 1917418
2021
Zhaoshui He, Shengli Xie, Rafal Zdunek, Guoxu Zhou, and Andrzej Cichocki. Symmetric non- negative matrix factorization: Algorithms and applications to probabilistic clustering. IEEE Transactions on Neural Networks, 22(12)2117-2131, 2011. doi: 10.1109/TNN.201 2.12247.
Jingu Kim, Yunlong He, and Haesun Park. Algorithms for nonnegative matrix and tensor fac- torizations: a unified view based on block coordinate descent framework. Journal of Global Optimization, 58(2):285 319, 2014. doi: 10. (00007//1099--01--0000
hRtps:/s::/dt
0xc/10.1007/s10888 013 0035-4
Da Kuang, Sangwoon Yun, and Haesun Park. Symnmf: nonnegative low-rank approximation of similarity matrix for graph clustering. Journal of Global Optimization, 62:545-555 574, 2015.
Brian Kulis, Arun C. Surendran, and John C. Platt. Fast low-rank semidefinite programming for embedding and clustering. In Marina Meila and Xiaotong Shen, editors, Proceedings of the Eleventh International Conference on Artificial Intelligence and Statistics, volume 2 of Proceedings of Machine Learning Research, pages 235-242, San Juan, Puerto Rico, 21-24 Mar 2007. PMLR. URL htt.htttss///............... http:///coceeeeegg.. ing. ml www.srrssssee.
/
html
Stuart Lloyd. Least squares quantization in pem. IEEE Transactions on Information Theory, 28: 129-137, 1982.
Matthias Löffler, Anderson Zhang, and Harrison Zhou. Optimality of spectral clustering in the Gaussian mixture model. The Annals of Statistics, 49(5):2506 2530, 2021. doi: 10. 20-A0S2044, URL /ttttpa://do..rggee 121g://00002222222000000020
Yu Lu and Harrison Zhou. Statistical and computational guarantees of lloyd's algorithm and its variants. arXiv preprint 2.000099, 2016.
J.B. MacQueen. Some methods for classification and analysis multivariate observations. Proc. Fifth Berkeley Sympos. Math. Statist. and Probability, pages 281-297, 1967.
Dustin G Mixon, Soledad Villar, and Rachel Ward. Clustering subgaussian mixtures by semidefinite programming. Information and Inference: A Journal of the IMA, 414):389-415, 03 2017. ISSN 2049-8764.
doi: 10. (0.10933/maiia/ax000... URL
org
1ax001
Andrew Y. Ng, Michael I. Jordan, and Yair Weiss. On spectral clustering: Analysis and algorithm. In Advances Neural Information Processing Systems, pages 849-856. MIT Press, 2001.
Jorge Nocedal and Stephen J. Wright. Numerical optimization. Springer series in operations research and financial engineering, Springer. New York, NY, 2. ed. edition, 2006. ISBN 978-0- 387-30303-1. URL http: gbv de DD-.e //C AACT- 10164 100
Jiming Peng and Yu Wei. Approximating K-means-type clustering via semidefinite programming. SIAM J. OPTIM, 186)),18 2007.
Martin Royer. Adaptive clustering through semidefinite programming. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 1795-1803. Curran Associates, Inc., 2017.
Ulrike von Luxburg. A tutorial on spectral clustering. Statistics and Computing, 17(4):395-416, 2007.
Yu- Yu-Xiong Wang and Yu-Jin Zhang. Nonnegative matrix factorization: A comprehensive review.
IEEE Transactions knowledge and data engineering, 25(6):1336-135; 2012.

---

Published as a conference paper at ICLR 2024
Wei Xu, Xin Liu, and Yihong Gong. Document clustering based on non-negative matrix factor- ization. In Proceedings of the 26th Anmual International ACM SIGIR Conference on Research and Development in Informaion Retrieval, SIGIR '03, page 267-273, New York, NY, USA, 2003.
Association for Computing Machinery. ISBN 1581136444 136463. doi: 10. 00.1145/86043558664488 URL
1145...111800000 https:://dl..or// wtt../ooo/....11//10100 8
Liuqin Yang, Defeng Sun, and Kim- Chn-Chhan Toh. SDPNAL+: majorized semismooth newton-
eg augmented lagrangian method for semidefinite programming with nonnegative constraints. Mathematical Programming Computation, 7(3):331-366, 2015. doi: 10.1007/112332-015-000---
htt.tttpps///oooo. wwwp:/....1..... 1007
Zhihui Zhu, Xiao Li, Kai Liu, and Qiuwei Li. Dropping symmetry for fast symmetric nonnegative matrix factorization. Advances Neural Information Processing Systems, 31, 2018.
Yubo Zhuang, Xiaohui Chen, and Yun Yang. Wasserstein K-means for clustering probability distributions. In Advances Neural Information Processing Systems, 2022.
Yubo Zhuang, Xiaohui Chen, and Yun Yang. Likelihood adjusted semidefinite programs for clustering heterogeneous data. In International Conference on Machine Learning (ICML), 2023.
12

---

Published as a conference paper at ICLR 2024
SUPPLEMENTARY MATERIAL
This supplementary material provides additional details of the algorithm and propositions, further experimental results, the proof of the main theoretical result (Theorem T presented in the paper and a concluding discussion.
A ADDITIONAL DETAILS OF THE ALGORITHM AND PROPOSITIONS
Choices of tuning parameters and rounding procedure for Algorithm | practice, we start with small step size a = 10-6 and increase a until a = 10-3 Similarly, we start with small augmentation term 3 = and increase 3 until 3 = 10 The second experiment shows that the choice of r does not significantly affect the convergence rate. In practice, many local minima for small separations. Therefore, we slightly increase r and choose r = 2K for all applications, which turns out to provide desirable statistical performance that is comparable to or slightly better than SDP.
found that K will result in
After getting the second-order locally optimal point U in Algorithm ! (NLR), rounding procedure is applied, where we extract the first K cigenvectors of U as columns of new matrix V and then use K-means to cluster the rows of V to get the final assignments. The same rounding procedure is applied to the outputs from both SDP and NMF.
Additional explanations to the propositions. Proposition| I lindicates that if U is the local minimum of problem (6) that satisfies constraint qualification assumptions (Assumption (S) in Bertsskkss (176)). then we can get unique minimizing point around U of augmented Lagrangian (U, y) for any provided that we have large augmented coefficient 3. A more modern proof without constraint qualification assumptions is given in in/eernández and Solodov (2012). Proposition [2 indicates that the dual variable ys in Algorithm I will converge to y" at exponential rate locally provided that Us solves ming L (U, yss). Therefore, Algorithm| can achieve local exponential rate provided that we can solve the local min of (z(U,)) y) for y = with an exponential rate, which is proved by Theorem
B ADDITIONAL EXPERIMENTAL RESULTS
Comparison of difffrent ways to solve NLR. We conduct an experiment under the same settings as described in Section 5] "Linear convergence of NI.R" except that here we consider a smaller scale with n = 300 and p 4. We aimed to compare the convergence rate towards the optimum per dual update between our algorithm (based on projected gradient descent) and the conventional algorithm (which lifts all constraints into the Lagrangian function). As seen in Figure 4 our algorithm successfully converges to the SDP solution, whereas the conventional algorithm's 's convergence halts at a specific step. Furthermore, the average time required to compute the minimum of the augmented Lagrangian (as denoted by (12) in the main paper) is 0.002 seconds for our algorithm. In contrast, the conventional algorithm takes an average of 8 seconds, even when employing the state-of-the-art limited memory BFGS method.
Comparison of computational complexity w.r.t. K. Here we want to compare the time complexity of NLR, NMF, the SDPNAL+ solver (Yang al., (0155) for solving SDP, spectral clustering (SC), and K-means++ (KM) under Gaussian mixture models (GMM) when the number of cluster K is moderate in practice. Under the same setting our second experiment ("Performance for NLR for GMM") in Section except that now we consider fixed sample size n 1000, dimension and different number of clusters K 5, 10, 20, 40, 50. The results are summarized in the first plot of Figure where we can observe that the log-scale curve of NLR is nearly parallel to the log-scale curve of KM and NMF, indicating that the growth of CPU time cost for NLR with respect to K is reasonable (nearly O(K)) and would not achieve the loose upper bound O(K*) derived from the analysis. The curve of computational time cost for SDP is relatively stable for different K since the dominant term of computational complexity for SDP is the sample size n, which is of as large as order ()23.)).
50.
Performance of NLR under different p. Our goal is to compare the statistical performance and time complexity of NLR. NMF, the SDPNAL+ solver (Yang et al. 2013)) for solving SDP, and K K-means++ (KM) with increasing dimension p. Here consider the same setting as second experiment
13

---

Published as a conference paper at ICLR 2024
Convergence of NLR (Regular ADMM)
Convergence of NLR (Projected GD)
10
20
Relstve diferrence to optmmal point
ReleSiee difference optimal point
10
104
10
10-
10
10
100
200
300
400
500
100
200 300
400
500
Dual update
Figure 4: Log-scale plots of the convergence of NLR algorithms over iterations per dual update. Left: the NLR algorithm where the Lagrangian function contains all constraints and the minimum of the augmented Lagrangian function solved based on limited memory BFGS. Right: Our algorithm where the minimum of the augmented Lagrangian function is solved based on projected GD.
Plot of Sample veraua Standard Normal
Plat Standard Normal
Sample Data Standard Neenal
Quenties of trpur Samole
0.6
Quartoo oftree Surple
c.8
Currdaa of Irout Sarrle
0.4
02
0.4
Standard Nomaal Quardils
Standied
Standard Normal Guantiess
90 Plat Sample Data versus Standard Normal
Plot
Standard Normal
Plot Standard
Quunder or of Sample
Dunnkic of Irout Bamble
Quarsaua of Wroou Sampke
04
0.
222
04-
62
Standard Normel Quarttioo
Standard
Figure 5: QQ-plots for Heart dataset. The first (second) row corresponds to three randomly selected covariates for the first (second) cluster in the Heart dataset.
"Performance for NLR for GMM" in Section 5 except that now we consider fixed sample size 2500. The dimension ranges from 125 to 1000. The results are summarized in Table 2 and the second plot of Figure 2 with each case repeated the mis-clustering errors for both SDP and NLR coincide and stay optimal when the dimension increases, while K-means++ has large variance and NMF would fail when the dimension as large as 1000. The second plot in Figure shows that the log-scale curve for NLR is nearly parallel to the log-scale curves for both KM and NMF. This indicates the same order of CPU time cost with respect to the dimension p for NLR, KM and NMF, which are nearly of order O(p). Similar to the case when K changes, the curve of computational time cost for SDP is relatively stable for different p since the dominant term of computational complexity for SDP is the sample size
times. From Table 2 we observe that
14

---

Published as a conference paper at ICLR 2024
CPU time cost when K changes (log-scale)
CPU time cost when changes (log-scale)
10
UPU time [seconds]
CPU time [seconds]
10
10
-- NLR
KM
10a
NMF
10
10a 200
10a
400
600
1000
K
Figure Log-scale plots with eπror bars computational cost with increasing number of clusters K (left) and computational cost with increasing dimension p (right). NLR corresponds our proposed method, KM corresponds to K K-means++, SDP uses the SDPNAL* solver, and NMF corresponds to non-negative factorization.
Table 2: Mis-clustering error (SD) vs dimension for different methods.
P
125
250
500
1000
SDP
0.0018 (0.0008)
0.0024 (0.0010)
0.0037 (0.0005)
0.0024 (0.0009)
NLR
0,0018 (0.0008)
0.0024 (0.0010)
0.0037 (0.0005)
0.0024 (0.0009)
KM
0,0754 (0.1647)
0.0721 (0.1556)
0.0634 (0.1332)
0.1244 (0. 1681)
NMF
0.0728 (0.1594)
0.0026 (0.0010)
0.0037 (0.0007)
0.7496 (0)
C PROOF OF THEOREM 1
In this section, we prove the theorem for the general case with explicit constants, which extends Theorem [ from the equal cluster size case to the unequal cluster size case (Theorem (9). The section is organized as follows. First, we will introduce the notations that will be used in the proof (Appendix C.T). Then, characterize the optimum feasible solutions of the NLR algorithm (Theorem 2 in Appendix C.2). Next, we present the convergence results at the optimum dual y y" (Appendix C.3) After that, we extend the proof of convergence at the optimum dual to
y² (Appendix C.4
The proofs of all technical lemmas are placed at the end of the section
(Appendix C.5)
C.I NOTATIONS
We shall work with the standard Gaussian mixture model (GMM) that assumes the data Xi,..., ., X, are generated from the following mechanism: if E Gg. then
X; = + E
(15)
where GI,..., GK is true (unknown) partition of [r] we wish to recover, fl, PK E RP are the cluster centers and E, (0,, Ip) are i..d. Gaussian noises. Let ny = |G&a denote the size of Gt. Let Z U (U*)T where
Viu
0
0
0
U²=
y n2
m2
B
(16)
0
1nk
V
15

---

Published as a conference paper at ICLR 2024
Define the minimum of cluster size max [mx} and m min my+-r
[ (n the maximum of cluster size Define the minimal separation O Q
2
minis || - |ul|2,, and the maximal separation max == maxissj Ilu; - Us 1= For any block partition sizes (m1, m2, ... mx) satisfying m& 1, Σx my we define the set Fm as consisting of all (orthogonal transformation) matrices Q such that U = U has the form
Ta mp a1,m In
02 1 ng: 22 Ing
(17)
1oxy.. K.KK 1nk-
mL Vk [K], E. my =. We will call any matrix U whose sparsity/nonzero pattern coincides with (T7) (up to column permutations) as an m-block diagonal matrix. Let 9m denote the of all m-block diagonal matrices with nonnegative entries. Note that Q exists for any U of the form since UUT = (U*)) We denote the set of all such Q as Fm. Now we denote ag = (AK,1, 0k .... m).,)) and U as (77. Define a ming, { == max For any matrix A E Raxr, we use As(a) to denote the matrix of the same size as the restriction of A onto the support Sa) of U that is, [As(a)]is = Aij if [Ua * 0 and [As(a]]is = otherwise; and let As(aj=
for some E1 Ci
As(a). Let O, be the orthogonal transformation group on IR..
Let Ug be the projection map on to the set Ω2 i.c..
(u)
EQ,, |10 ull S I ulI, Ω.
Let ((2) to the convex hull of Q. Recall that Q == {U Rax | = K, U 2 0}. Then C(S2) (U E Rxxr, || - K, U 2 0}. We define
(W), W aVf(U), a
In particular, we define I (U) to be the positive part of U. Let U be the initialization. For / E N define U recursively by
U Ila(U' VuCc(UU,
Denote Q, Q' Fm as
Q argmin IU - Ur Q argmin |10 Unlll V.
C.2 CHARACTERIZATION OF GLOBAL OPTIMA OF NLR PROBLEM
Theorem 2 (Feasible solutions). For any U E gaxr such that UUT Z and U 0, U must take the following block form (up to column permutations), due to the nonnegative constraint U 2 0:
1 mm 1m
0
0
02, 1 1ngs 02mm 1
(18)
0K.1 1 nx OK mn 1 Ing
for some block partition sizes m (mi, m2, mk) satisfying my 2 1, En my = r and within block weights ag: (ak,1, 0s,2; ... m, satisfying txet 2 0, Ee az = a,
for all E [K].
Proof. Recall
Via 1
nn
0
0
0
0
Vuz
m2
--=
0
:
0
0
Then for any Q (4,, ( qr] Or, the orthogonal transformation group on R", have
1
U==U"Q
OK
1 nK
16

---

Published as a conference paper at ICLR 2024
U is feasible hence U must have the form (18) (up to column permutations).
20 qi O, Vi. Note that (1). 1, (gt, qz) Vi + then U
In the rest of the proof, first show the convergence of projected gradient descent at y" as Theorem below, and then extend the theorem to a general dual variable ý in neighborhood of y in Appendix C.4
C.3 PROOF OF LINEAR CONVERGENCE WITH y y
We list below the assumption on the minimal separation for the general unequal cluster size case. This is generalized version of Assumption A, where cluster sizes are equal:
A Assumption 1 (Assumptions for Theorem II.1 in (Chen and Yang 2021)) If there exist constants
C E (0, 1), C3 > such that
(1 (1-
Din
C7 D2
4(1+
log n2 2
C3
m 2
C
m
(1 1)2
K
C3
40 (Ccss
11-00 (1-C))
+ min 2 0
1+ V 11 1 C²C3 m log + DyRn
1+
log
(1- (1- C1)²
with
(1- 11)
VVllog logn
R, (1+ C3) logn
n
2ngni where n == minke(x] {mx) and m = min m
Here D1, D2, D3 are universal constants.
This assumption requires the minimum squared separation O be of order O(log n.). Recall that we use ñ == m[X [ and 6 to denote the maximum squared separation.
Assumption 1 will be reduced Assumption A in the main paper when the cluster sizes equal. (1+200s)
(1 a). Moreover, This is due the fact that for any 0, we can find C₁ and C3 s.t. (1-C,)²
m = nn/K K in this case.
Assumption 2 (The optimum dual y Choose (,, B) to be the dual variables of the SDP defined by equations (20), (21), (22) and (25) in (Chen and Yang 2021). In particular, p C m min'
B 120, 0, G. = 0, vk E [K] and y" equals a in (Chen and Yang. 2021).
Theorem (Convergence of projected gradient descent at y²), Suppose Assumption 1 & 2 hold. Consider the projected gradient descent updating formula at the dual y":
U² = II -oo---- .y")), t E
If there exists some block partition sizes m and an associated optimal solution U with some within block weights a (0k, ak.2.... ..... m)) satisfyyng S ca, O, such that the initialization discrepancy A = T T satisfies
Cian
Ca
| D
Sloo 11o0 5 p/02 min + m
and
| <
Cza(L + nB + 0
maxx) )
71 -min
an min
(1 7)nn m
O v /K log p + log" n)) (1 7)
r 3 [ + n² + n max
where m minket
nutni n min(n Omax maxigjtksK ||14 - ykk 12,, and (4xx)
2
C min 1-1 we we take B = 1orn2a3
L= G, (m 8rn2a3) ) in 12p and choose step
size a (L+n))) then it holds with probability at least 1 n Co that for any t 2 I,
U' E 9m and inf QEfm
|1++1 -U*QUI ST inf qemm
where 72 =1
Ca and Cz,..., Cs are universal constants. (1155)))[[+
17

---

Published as a conference paper at ICLR 2024
In scenarios where the dimension p is moderate (p OVv logn))))))))))) and the cluster sizes equal, the parameters and initialization criteria can be simplified to the expression presented Theorem [ O( /n log n)) and
nt n K. The proof of Theorem consists of four parts (Appendix C.3.1 C.3.2 C.3.3 and C.3.4 First, we will prove some key propositions which are derived from the theoretical analysis of K-means SDP (Chen and Yang (02T) in Appendix [.3.
once substitute
and some important properties of the NLR formulation of the SDP in Appendix C.3.2 Then we will analyze the behavior of convergence in two phases. For Phase (.3.3), we will prove a constant speed convergence (Theorem 4 of the algorithm until it reaches the same block form as the optimum point, where we call it Phase 2. We then show that the algorithm converges exponentially quickly to the optimum solution in Phase 2 (Theorem [] which combined with Phase result leads to the overall convergence Projected-GD (Theorem 3 Proofs of all intermediate lemmas are placed in Appendix C.5
C.3.1 PROPERTIES FOR SDP
First, we list our main assumptions and propositions from (Chen and Yang 202I) below. We will assume Assumption and Assumption 2 to hold throughout the of the proof. Recall that Assumption 1 (general version of Assumption A for the case of unbalanced cluster sizes) requires the minimal separation min be larger than certain threshold (with order ((olog))) The minimal separation can be interpreted as the signal-to-noise ratio of the clustering problem. Under such assumption, the solution of SDP can achieve exact recovery under Gaussian mixture models (GMM) (Chen and Yang, 2021).
The following propositions are derived from the original SDP problem (Chen and Yang 2021), which characterize the bounds for eigenvalues of the matrices related to the dual construction of the original SDP (Proposition 3 4 and Corollary IT) that will be used mainly describe the smoothness of
augmented Lagrangian function (ALF) (Lemma and Hessian of the ALF (Lemma (5) Proposition
5 will be used for characterizing the behavior of the gradient of ALF Phase I at some optimum point (Lemma and Lemma 3
Proposition 3. For any v E R", define S(v) = v² Av, I'K == span : k E [K]]+ which is the
orthogonal complement of the linear subspace of R² spanned by the vectors l. 1ox Then Lemma VI.1 (Chen and Yang 2027) implies
P
max DET
[S(v] 2 (Vū+ vp+ 22)2
Se Vt > 0.
|+|==
Proof. Refer to the argument after Lemma IV.1 in (Chen and Yang 2021 (right below equation (26) in (Chen and Yang 202T)) The exact formula is showed right above Lemma IV.2 in (Chen and Yang
2021). Proposition 4 (Theorem II.1 in (Chen and Yang 2021)). Define
Wn == Ald + (3"""""1 + In(y")²) - B
If Asssmptionss I and 2 hold, then with high probability the following equation holds:
[|W loop I X C2(n + mmplog n), W,, > O, W,1c, 0, Vk E [K] for some constant C2 O. Furthermore, SDP achieves exact recovery with high probability least 1 D Co for some constant D4-
Proof. W, > O, W,IG. 0, vk E [K] and SDP achieves exact recovery have been shown from the proof of Theorem II.1 in (Chen and Yang 202T). From the argument after Lemma IV.1 (Chn and Yang 2021) we have the upper bound of W.,v S A||o|22 - T(v). Then from the bound of [T(v)
(Lemma IV.2 in (Chen and Yang] 2021)) we have
[WWlllop S A Cz(n + /mplog
for some constant C2 > 0 with high probability. Proposition 5. Recall the minimal separation Quan min ||pg Mallz and the maximal separation mmax := maxi [|1 142122- If Assumption and 2 hold, then we have Vjj E G1, K & I, (k, E [K]?
D2,())= == [22 1n + 111,",)),,,,, G. 1nlls E [Cin0 min: T0 with probability at least 1 12n In particular, Bij 2/C1 (0 Vi,j E [n].
18

---

Published as a conference paper at ICLR 2024
Proof. From equation (19) (21) (26) in (Chen and Yang 2021). we know that
Dka(i) = [B G2 Andl 2 /2/2g |4 mell 2 C1 C1/2n0 min'
In particular, C in Chen and Yang 2021 is defined be [Ba, Inyli and Bij E jeG, Dk, Vi E GR. J E Gj (refer the paragraph right below (22) in CCen and Yang 20211). On the other hand, from (21) in (Chen and Yang 2021) we know
(k.1)
Dxa(j) (i)
Da,(i)) = n1 + Thk X +
1kk (IIX, x,,|
||x X,11)),
2nd
2
where Xx= = mi E X; is the empirical mean of points in cluster k. From the proof of Lemma
IV.1 in (Chen and Yang 2021) recall X, μM + E and denote 0 = MA we can write
| X,l12
||
Xyl1² 110 + Ej Ex 2 -lle 2 s 110 +1 Exl1?
Hence
Da,(i) (1
x,バ - X,113) 2 00 E1 Ex
Then can get the high dimensional upper bound n& 011 for (j. by bounding the Gaussian
vectors E, using the same way as the proof of Lemma IV.1 in (Chen and Yang 2021). Thus,
DRA() < nx S n max!
and B.. DE,,(iDDLA()) EjeG, DR.(I) S 2/Cimallu mell S 2/CC/0 max
Corollary 1. Let R", S to be the set of non-zero positions for w. If S for some & [K], then
-211[u112
Inl"" ww)) 2C2(n /mplogn))]lulf
Furthermore, for general w E TR."
-2 < (2A + y y01 + (nn) ,w"""""w""")""" 2(C2 for some constants C₁, C2.
/ mplog + 1/C ||xl
Proof. From Proposition we have
Wn Ald 1 5011% + 1,()")²) B+
Hence
1 (000 ( 1. + (1())))) + A, ww 2 (B - AId, wwx").
Recall B 2 0, BG... = 0, vk E [K], then we have S S G = (B, www = 0. Thus,
(2A + (v'13, 1 + (,,))) ww")
22
On the other hand,
G + 1, A, www (W, + B AId, ww")
< 1| l| A ||
< C2(n+ + / log n) 1
Since then l|Wnllop S 2C22(n + /mplogn) from Proposition Furthermore, if g GR, Vk [K],
(' + +1x())²) + A, wwn") = (W,, + B - Ald, ww")
< Cz(n + Vmplog n) w| 1/2 max Bi ||ol
By Proposition
Bjj < 2/C E where universal constant C1 (0, 1), mmax is the maximum of separation. Thus,
[ + IT( A, ww)) < (C2(n- + mmplogn) + 1|
19

---

Published as a conference paper at ICLR 2024
C.3.2 THE NLR FORMULATION AND SOME PROPERTIES
Before diving into the proofs of Phases and 2, let recall the Augmented Lagrangian Function (ALF) to which we would apply Projected Gradient Descent (PGD), and present two lemmas for characterizing the smoothness ALF and the properties of the projection operator.
Consider the y" defined Assumption 2. Recall the ALF is given by
8 f(U) ((U,)) (L Id, A,U (u. UU" 1n) J
+
1n1li. (19)
where L E (0, A) 0. Then we can get the gradient of as
VS(U) = (2A + 27. Id + 1 + Lay")U + "I. 1, 1.)² + (UU" I
direct implication from Proposition is {)))) [U]s stands for the matrix that keeps positive entries of U*Q and set the non-positive ones zero. This is due to the fact that W,lc, = 0, vk E [K] and Wn, = Ald + {1 + 1. (v))) B A. Recall that Q, Q' E Fm are defined to be
Q]s VQ Fm: where
Q argmin
- Q' argmin
1 U Qll At
Now will consider the smoothness of the ALF. Here we consider two cases of U: the first is for general U in the feasible set ft; the second case is for the U that has the same block form as some optimum point U U' Q. As we will see, bound of the Lipschitz constant for the latter one will be smaller. Recall that {U E Ruxr =|U K, U 2 0}.
Lemma (Smoothness condition). If - U I for some Q E Fm then
(VS(V) f Q)1l - RullU U² Q"r,,
where R1 2n + V moglogn 2L + 12m3 if U 9mi Rı 202(n + vmp log
2/C max + 2L + 12n3, for some constants C₁ C2 and general U E Q where O is the maximum squared separation.
The next lemma shows that locally, the projection of PGD is equivalent to the projection to a convex set, and therefore local contraction map.
Lemma 2 (Local projection to the ball). Denote e :- 2 K
Here R
2Cz(n + Vmplog n) 21 + 2/C2 max 12n3. Then
UI - UU S =
< (A--))))----
Furthermore, if >
[1 - aV fV()) K, un(UU - -Vf(()) = (UVa(U) for some constants C1, C2.
C.3.3 LINEAR CONVERGENCE NEAR OPTIMAL POINT (PHASE 1)
Here will prove the convergence of Ut to the block form in Phase 2 (Theorem 4 which will be implied by Lemma 3 and Lemma 4 The first lemma shows that after at most iterations, the iterate will become the block form and fall into set Gm. which then enters Phase 2. The second lemma controls the distance to the optimum within Phase 1, which will help us set the appropriate neighborhood initialization to ensure the convergence of Phase 1.
Lemma 3 (One-step shrinkage towards block form). Denote A = - -aa+ where ya. is defined Choose (L, 3) such that:
as 17
2(1 - L)
C min
<B5
n
3
Suppose ||||| S <e: where is defined in Lemma 2 Assume
0 min
3
(As S
| & min 61
7 2
(20)
32X
20

---

Published as a conference paper at ICLR 2024
where 61
Cran
Caan min
Cian
min 64K max V2X nog Cs(p + log n)) V
min
min
963räm
96,ra²n
Then we have
[1] - aVV[(U)) & (0]. - Cann min vi E Gi 1) < <Ts sk, #
Furthermore,
CianO
Mer S max[[U]ir
min 0), vi sk - <T sk, 1
8
where = (W). W= =V- a f(U), for some constants C1, Cs. Consequently, after at most =
And iterations, U will enter Phase 2, i.e.U' E Im for every step t, U meets the above conditions (20).
The next lemma calculates the upper bound of the Inflation of ||U U'Ql after Projected-GD. This will be used when we have the neighborhood assumptions for the initialization of Phase 2. Then trace the assumptions back to the initialization of Phase each step of Phase 1.
know the inflation rate
Lemma (Inflation of distance to U for Phase 1). IU U*QlI S cc. then
|V U F S nl|U U***
where n= + aR1, R1 = 2C2(n (n + mmplog + 2/2C/02 O max + 2T. + 12n3, some constants Cr.C2.
Now, we will establish our condition on the initialization, given our knowledge of the inflation of || for each step in Phase 1. This is summarized in the following theorem.
Theorem 4 (Convergence of Phase 1). Recall that X x p+ 4 m min!
ming mn-n; Define
201
U be the initialization of our NLR method and define A == U - Let 1/(4a)); 1 + aR2, where R₁ 2C + / /mp log n)) + 2L + 2/C, e + 12n,3 Suppose is integer and â S ca, for some constant > 0. If we choose (L, 3) in 19) such that:
2(A L)
C Cian min
S 3 -
12rn?
Then with probability at least - Can Co, we have U E 9m: 110 - Uall| S 6, S provided that
Clan min
1
3 Eos él, T
| 0 ( (a)///oo 11100 <
|| S o.
min
32X
Here 6 =
C aann 2
Cran02 min 96,3rä
C C
min
(4 maxx /2K log Cs(p + log n))'² 'V
96,3ra
( L)VK
6o ((1+vK)R, + L)VK + /nr an max
Here, Ci, C2, Ca, C4,, Cs are some constants.
Proof. This can be implied from Lemma and Lemma 4 by assuming the neighborhood requirement
in Lemma multiplied by the factor 1/r/ for the neighborhood of the initialization ||a|| F. Note that if we know the total step of Phase which is I, then we only need assume A S
f min { 61, ( if we want || min {Fes 61, [/ Vt. < T.
21

---

Published as a conference paper at ICLR 2024
C.3.4 LINEAR CONVERGENCE NEAR OPTIMAL POINT (PHASE 2)
Here our goal to prove the linear convergence of Phase 2 when U 9m (Theorem The main idea is to show the local strong convexity at optimal points (Lemma 5 followed by a standard argument of linear convergence (Lemma (). Moreover, we also need to ensure that the iterations continue in block form during Phase 2 (Lemma
Lemma 5 (Local strong convexity). When U E Gm, if we define A = UQ² - U then
(V)))))) 2 B|
where B [ (VT vp 12 -2A L) 3n} > O. provided that L ((v++ /p+ / 1o A), 2(2 Lm/n. Furthermore, for every U satisfying -U-UU|l 6s.
we have
(VL((A)))))) 2 //2
where Eg
36
After establishing the local strong convexity, we can prove the linear (exponential) convergence of the algorithm through a standard analysis for the projected gradient descent (PGD) as below.
Lemma 6 (Local exponential convergence). When U E 9m. if we define A == UQ" - U then there exists
(0, 1),6 O, s.l.,
||V U'Q S - - Q'Q
where
Q argmin |IU @l
provided that IIU U*Q
F e, where min(fe; <.) Recall that e. defined in Lemma 2
C, defined in Lemma 5 particular, if we we choose the step size S 3/(2R(), where
R₂ (n + mmp log + 2L + 12n3, for some constant Cz. Then the contraction factor would
be 2 (1- (1-a8/2)
The convergence for Phase 2 has not yet been proved since the above lemma only demonstrates the convergence PGD once U² has reached Phase 2. However, the subsequent update U++1 might not remain in Phase 2. The following lemma ensures that Ut stays in Phase 2, provided that the initial point U in Phase 2 is located within some neighborhood an optimum point.
Lemma 7 (Iterations remain staying Phase 2). Suppose Lemma holds and U Gm- Let
ē : min {e1, [. Then we have
Uº E Gm, 14 2 1,
where
U++ = ll - - 1
as long as - - Un*") < min - 72//2E, c), where recall that min]fese...
Now, we can ensure the overall exponential convergence in Phase 2 by combining the above lemmas. Theorem 5 (Linear convergence in Phase 2). Suppose U E 9m- Let the step size: a < (12 +2 3 min - (VTi /D /2 n))lg)))]. 2(A L) + 3n). Co, we have
B
Then with probability at least 1 -Cn-Ca
| U+Q ||5 S - UUQQIIr Ak,
given ||000 - U0Q < S,, | - Umt+
S En, where 72 = (1 - -a3/2),
= min 9
A-L)VR
60 1- min
3
(1, a
368m' 2[V KR + 2(A L)vK + n max
for some constants C2 Ca, Ca-
22

---

Published as a conference paper at ICLR 2024
Proof. This can be implied from Lemma 6 and Lemma Lemma @ indicates linear convergence if the former step is in Phase 2. Lemma 7 guarantees that every step of Projected-GD will stay in Phase 2. Therefore we only need to combine the conditions for both lemmas.
Combining Phase 1 and Phase (Theorem and Theorem 5) and lifting the assumptions of initial- ization for Phase 2 to the initialization of Phase 1 (by multiplying the factor 1/1/',, obtain the following theorem. This theorem represents a stronger version of Theorem 3 presented at the begin- ning of Appendix C.3 This theorem is stronger in the sense that we can get the exact relationships between parameters. However, the presentation of the theorem would sacrifice the conciseness of the statement. Hence we present the simpler version as Theorem at the beginning of Appendix C.3
Theorem 6 (Local linear rate of convergence). Recall A'
U U Suppose ca, C. 0. Let I = (Aar))= " = 1 + aR1, Rí = 2C2(n + /mp log v.) 22 + 2/C, max + 12n33. Choose (I,8) such that:
2(A L)
SBS C Cian min
L> (V Vp+ logn),
n
12rn
and the step size a < 322C2la+mmpoos +2
B = min (Vn + vp
/2logn) 2( L) + Bn). Then with probability at least Can Cs, we have that for any
to I.
Ut E gm, and || U³ U+Q 7 UU*
provided that
|| 0 ((a) 1la S
Cian0 min A
| S n
1
min 6,
1 1-7
7
(1,
1 1-
32X
2
2 4
where 72 (1 19--ak -a3/2)) E min
/V/K
36pn' 2VKR L) V +n
61
Cran
C CianO min
Cian
min
64K(0 v 2K log
Cp(p + log n)) V
96,3ran
min 96Bra²n
Here, C1, Cz. C3, C2, C3 are some constants.
C.4 PROOF OF LINEAR CONVERGENCE WITH y Rs y²
The following theorem extends the results of Theorem 3 from dual variable y yy to any value in neighborhood of the optimum dual y", which completes the primal-dual local convergence analysis by demonstrating that the primal subproblem mingen (U y) can be efficiently solved around the optimum.
Theorem I (Convergence of projected gradient descent). Suppose y" is the optimum dual for the SDP problem (Assumption 2 in Appendix C.3) and Assumption A holds. Assume On log n), + max I CO for some C O, and there exists some block partition sizes m and associated optimal solution U with some within block weights a = (4x.1 Ak. ak.m,) satisfying S co, c >>> and a min minse[mw))(as.. and ā max max {ok,e)- We denote the initialization of the algorithm as UO; let U denote the unique (whose existence is also guaranteed) stationary point of the projected gradient descent updating formula under dual value ý close to y". ie. U satisfies
t = The (ū CV
If 1 v²1l S for some 8 > 0, and the initialization discrepancy A 1 Ua. satisfies
| (()) las < <O(K/vm and
|| S 0( 0.5 K 5.5 min (1, K -25.0 2 /lgg/llg))) 1))
then it holds with probability at least - Gnn 02 that, for any 1 12 I O(K3),
U² E 9m and inf QEO,
|1+ - ÚQ]I= Sy inf QEO,;UQES...
1 ÜQ]IF.
23

---

Published as a conference paper at ICLR 2024
where y = 1 1 providdd that 8 O (imn///))) L 0 in (12 and step size a is chosen such that a OOK (("" Here, 01 and C2 are some constants.
Proof sketches. The proof slightly modifies that of Theorem 3 where divide the convergence of the projected gradient descent for solving the primal subproblem into two phases. In phase one, we demonstrate that, for the same reason in the previous proof, the iterate will become block diagonal 9m after at most I iterations. In phase two, due the strong convexity of C( ÿ) around any Ú within Gm. the projected gradient descent will achieve a linear convergence rate. Here, the strong convexity is implied by the strong convexity of L( y²) U within Gm and the continuity of the Hessian (with respect to U) of the Augmented Lagrangian function L(U, y) with respect to U and Precise statements regarding the two phases are provided below.
By employing the same argument used for Phase I and Phase 2 regarding convergence at the optimum dual y (as per Theorem and 5), we derive the following analogs for y.
Theorem 7 (Convergence of Phase 1). Recall that X = p+ S m min: m = minkyt ny mn-mi
2n mL Define
U be the initialization of our NLR method and define A U U Let I 1/(4a)), n= 1 + aR1, where R1 = 2C2(n + /mp log n + 2L + 2/C max + 12m3. Suppose is an integer and â S ca. 0. we choose (L, B) in 19 such that:
2(A L)
C Jaan 2
B
min
2
12rn2a3
then with probability at least 1 Cnn"" Cs (the high probability argument comes from Assumption 1),
have U Gmi || -U I <c, Vt S I, provided that
Cian min
3 min eor f javej.
S S
| S m
111" - 5 s 6,
32A
for some small 8 > 0, where El
C Cran min
c Caan
Cian min
min
Cs(p + log" n))" VV
min
64K(0 max v2K log
963rän
96
(2 L)VK
6a = ((L+VK)L + +(A L)VK /Dr
an max
Here, C1 C2 C2 CA, C3 are some constants.
Theorem 8 (Linear convergence of Phase 2). Suppose U E Ims, then there exists unique Ü, such that it is the stationary point of the projected gradiens descent updating the formula for any dual y close y", that is,
Furthermore, we define
Q' argmingeo,,uqee 10- U'LUR| Vt V20 0.
Iflly* ell s 8, for some 8 > 0; the step size a 2(2C2(n+ 21+ 12m3332 3 min
3
(Vn + vD+ /2log 21, 2(A L) + 3n), then with probability least 1 C4n Cs, we have
| iQ +1 P S >/U -UQ Vk.
given - U=Q <,, Un S cu. where 72 (1-aB/2)
(A- --LK
1-7
3 4
c= min
=0 2
min (1,
@V
366n' 2 Rg + 2(A L)VK + n max
for some constants C2 Ca, Ca-
24

---

Published as a conference paper at ICLR 2024
Proof sketches. Lemma indicates that (V) 2 Bl for some 3 0. where = -V, VU, V 9m. This is due to the fact that A U - = -VV V - U+Q) Im Hence under the condition of the local projection lemma (Lemma 2 we have
(U-V-U-Q) +U*Q (U-
|| aVI(U)) Hl - AV)())
I |IU aVf(U)) (U (V)())
< IU - - -V)(VVUV--
V/(MI
s IU VII -1/2-1B/21|U VII
aB/2[|UU
= (1 - a33/2)1UU--- - -
Therefore the map + -U f(U) a (strict) contraction map. Hence there exists unique stationary point within 9m by the contraction mapping theorem. Rest of the proof is the same as the proof of linear convergence in Phase 2 at optimum dual y" (Theorem
we combine Phase and Phase together (Theorem 7 and Theorem S and lift the assumptions of initialization for Phase 2 to the initialization of Phase 1 (by multiplying the factor 1/'7' ) we will get the following theorem, which is a stronger version of Theorem [ in the sense that we do not require moderate p O log n1) and the constraint the maximal separation m COO What more, we can get the exact relationships between parameters from this theorem. However, the presentation of the theorem would sacrifice the conciseness of the statement. Hence present the simpler version as Theorem D in the main paper.
Theorem 9 (Local convergence of projected gradient descent). Recall A := - Suppose Assumption and 2 hold, and assume || y''l < 8 for some > O. Let I (4a))-1, 7 1 + ali, Ri = 2n + /mpplog Ti + 2L + 2/C 02 max + 12n8. Choose (L, such that:
2(A L)
Cran0
L> (Vn+ vD+ 2logg)
<B 12rn?a3
n
and the step size
B min (VT + vp
3(2Cala+ + 122
(olon)3) 2((A L) + 3n)- If Co O. then with probability at least 1 Can Cs,
we have that for any 1 2 1
U E Im and | - ūQt+ 1| s
provided that
Cian min
1
1 - 7
1 1-
( 1lx0 5
A
min
E,
E1,
32X
2 4
where 72 = (1 a.//2), = min {
(A-L)V
363n 2[
61 :=
C e min
0 min
C C min
min 64K(0 V 2K log + Cs(p + log2 n))'
96,3aä
96,Bra2n
Here, C1, C2 Ca. C4, Cs some constants.
C.5 PROOF OF LEMMAS
Lemma (Smoothness condition).
If | - U @l [ I 1, for some Q E Fm, then
[VS(V) V Q)1 S R.IlU U Q(Q
where R1 2n + v Vmp log n) + 2L + 12n if U Gmi Rı 202(n /mp log
2/C max + 2L + 12n for some constants C₁ C2 and general U E Ω where @ 2 is the maximum squared separation.
25

---

Published as a conference paper at ICLR 2024
Proof. By definition we have Vf(U) - Vf(U-()=
(2 222p:4/44|HHH|400000000000000000 Note that
| I UU - U Ul0p + @l@p
< || U' Q|Q -
<1+1 = 2.
Hence
2 ||UU*- ,0 11,
- -U +U- [Q S U
s 2u
From Corollary D we have for U E S2
[V QU(
< (I](2A + y1 + 1,."" )| +1,)))++)11 -
-U 1)1 - 2311(UU -UQQU
-UQ 11p 1U 1n 1l
- (2C2(n+ mog n) 2/C/ - 11n U @(@
x (202(n+ m/ 0 max + 2L + +1 - U*ÂIIF.
For U E 9m:
[Vf(U) vf(U*())l=
(II(2A + -1 1,,y") 2U - -
11/lopp 1I,
$ (2C2(n+ (2Cg(n+Vmplogn mmplog n) + U U' + -2n-)|U-U'QQ U'Qlle
- (2C2(n V m) + +2L - -
Lemma
(Local contraction to the ball). Denote
(A-L)K
Here R
2 L) K
2C + v log n) 2L + 2/C1 max 12m8. Then
| - U" Q|1 S = < 1)//A/ <0.
Furthermore, if a > 0,
[I > K, (U aV)U))V((V)) Ie(cy (V)VV(()).
for some constants C1, C2.
Proof. Let [U]s to be the matrix that keeps positive entries of T and set the nonpositive ones to zero i.e.,
vf(U) [V)())s + {U)))ss-
Note that
(IVI(U)))) = 2(1 D)]U"* 2(A L,VK, (VS(U))))U')) = -2(A T.K.
From Proposition we have
[V S 0 max
Then by Lemma
we have
(VI(U),)) VT = (VS(U) -/
[V/(U) -
IV/(U) - -//000////AVVV+11700000000011000
S [VKR, +2( L)VK K + -U
26

---

Published as a conference paper at ICLR 2024
Finally we have
(V((U,,)) < - L)K <0,
given 6. S -((KR:-2(A--))
+22 Here R = +2 max + 12233.
Note the fact
[ = ( (I+ (U) = 11 Bu)(K)) (T)(
where (U) stands for the projection of U on to the space of matrices with positive entries, gur-1 (yr-1(K)) (VK) stands for the sphere with radius VK, B" (VK) stands for the ball with radius vK. Thus we only need to show that
|1 -LL(U-VV(U)IE 2 K.
Define Ü.j a(VS(U))) i.g if - (V/ 0. And Üij U.. otherwise. Then have
U- 0 -U) (V,((U--V)(
Moreover, if [fU-aVf(U)). - 0, then
(aVS(U)), 2 U.s 20.
Hence
(U, (aV)(,)) <0,
which implies that
(U - aVU--Vf(UU))E - > |1 K.
Lemma 3 One-step shrinkage towards block form). Denote A = U U where U is defined Choose (L, (3 such that:
17
2(1 - L)
CjanO2 min
< 3
Y1
12
Suppose || S Ec, where e, is defined in Lemma 2
Assume
CianO min
|||| F S min 3 EI,
||
-
(21)
32X
where ==
C Cian min
Clanb 963rän
Cian
min
(OK(omaa /2K log
Cs(p + log" n)) V
min
963ra²n
min
Then we have
CianO min Vi -k < sk, l
U1) -llr-)) & 10
Furthermore,
C
[ I max([U](,r
min 0), Vi G1, - <T sk, 1
8
where V = (W), W =V- (VU),, for some constants C1, C5. Consequently, after at most 4aA iterations, U will enter Phase 2, ie. E 9m if for every step t, Ut meets the above conditions (2I).
Proof. Without loss generality and for notation simplicity, we assume my = mc = s r/ 0g.1 ak ak. Vk. By definition have
Vs(V) (2.4+ 27. Id 1 T) 3 LLU 1, 1, 1n)²² (UU 1, 1 1,
27

---

Published as a conference paper at ICLR 2024
Now suppose A U - U A [dl....d,] Then vi G1, 1 1 1.
[ 2L (d1); + e (2A + y1n 1n,
)ny")),,, 1a, teg (2A + y1 + +1, )ny dd,.
From Proposition 5 we know
o (2A+ y1, +
= a1,1 =,,
2 1aa,,1
The goal of the rest of the proof is to show that e." (2A + y1 + 1,y ",,,))),1111 )a1,1la, is the dominant term in (V [U)),, the rest of the terms can be absorbed given that U is located within some neighborhood of the optimum point U From calculation we have
[V((U)) = 2L (di); e (2A+ y + 11,11, + t (2A + 12,,,1 + ",n" )d1 -
+ ", 1,, (UU - 1,) la +11
-2. (.) (2A 1 +,,""
Be, I
- 1,.²² (UU" 1, -I,) 01 lo, TT,
where
r se; [1, (UU 1,,) + (UU 1, 1, Further orthogonally decompose d, Σn 2hi 1G, + wr for 2k, E RR Then we have
wy E TT.
c2A y2
1, )d1 +d = Σ FR. (1,(() + ot (2A + y1 1,,y"
2 21,111,111) 22 21,1 + e (2A + y1 + 1ny yw,,
since T4,1 0, Wk 11, = Recall X , Xn), X, E, + Me, Vi G.. Then
from high dimensional bound of Gaussian distributions we have
1ef (2.1 + "1 + 1,,""" 12
Σ &
2 T m AG..o. W..
n&
2 Σ [(m ((p) + (e: - Ex))"
Σ
ey(w))
JEGz
S 20 max V 2K log n Cs(p + log² m))l
with probability 2 1 - C1/n, for some constants C.. C5 > 0. On the other hand,
- 1,) G [1,, [1,17(( (AC0:)/"))++=UU))))) ((U (A + (I,
Be: [1, (UU - 1m)"T
+A(U
= (TT T + +(((aM)))aG, + +n (U AT ((Ua.))))1.)
Ban Σ
I (u d-(u)) me m (u, (1
89
= =Bai
2g
11 + ain
g
T s(g 1)+1
al
xg
+ n2
Σ
a11 T n2
d, + n2
Σ
ag (d+))))
1=
8=1
1=
2 Bag / -||||
sl (d+):))
(m + an
1)+1
2 2 (3ann + ann((di), - an.arø
28

---

Published as a conference paper at ICLR 2024
Similarly we can show that |ral S Brā(3n' + + Hence
[VS()) 2 21. (di): + (11,1 + 21, 21,1) - 2//
2K((0 max V 2K log + Cb + log 1
801,1 (3ram / | a1.1n((11)) ānjnyaā 71
2 (2L+ nn) Cian min
provided that A
31, 2 -3/401,
C Ciame min 32Í
||
c,2: || S1.
Here El 1=
C Canne min
Crant min
CianO min
min
64 / log + Cs(p + log" n)) V
963rän
96Bra²n
is also sufficient to assume
s ce |||| I min c 4 aavn
The BL pair need to satisfy:
2(1 L)
C min
< 3
2
12r
Thus,
CjanOa
(1h,, - alVf(U)),, S [1), a
min
vi G., 1.
8
From Lemma sk, 1 + k.
we know that ||]]| -IU aV(() VR. Hence Vi E Gi, s(k 1) < r S
VK W H1, (U))) -(1 S max [[U),
C
MR.
min 0}.
8
Note that the initialization / needs to satisfy | Staje oloo <
Gran
from the above argument,
32X
where A 0 U Then takes most total steps for U to converge the block form Gm- where
Cian min )/(a
Cian Min
min
1/(4a)).
32A
8
Lemma
(Inflation of distance to U for Phase 1). If |UU - UUQ/IF S Ec. then
V ( F S 1/lU-UU - (8.* I
where 1) = 1 + aRi, Ri = 2C2(n + v Vmplog n) + +2/C,02 6 max + 2L + 12n3, for some constants C1,
Proof. Suppose
- S cc, then from Lemma 2 we have
IV - ur
[TT[[U - -VVs(V) (U -V](U***)]
||Tecta] (U - -aV/(U)) - (Uat -V/
< II(U - -VV(UU) - -V
-
29

---

Published as a conference paper at ICLR 2024
where R1 = (n + Vmplogni + 2//C10 max + +L + 12n3.
Lemma (Local strong convexity).
U E 9m. define A = UQ" - - then
(V)( 2 B||||
where min (Vn vp- + /2 log -2(A L) + }n} > 0, provided that
((vn+ + vp + v2logn) 8 2(A L) Furthermore, YU s.t. ||U - U*| < Es. we have
(v2 A 2 B
where E 36
Proof. E 9m. Q Fm =
U W + W2, with
[d1, 1 1npl dl
0
0
d2, 1 12g
d2
W1
dx, lnk
dx. 1
W1,1
W1,s
0
W2,1
W2,s
W2
W K.1 WK.a.
for some da R. Wh.t E RR (Wa,) Img) O, E [K]. i [s]. Note that
Th mii... ā1, In.
0
22,1 lng
U'Q =
0
σ 1 Inx........................ ................................ k. 1 nk-
for some 0z E IR, k E [K] i E [s]. Recall
Q argmin IU QU*
then there exist Ca > 0, E [K], s.t.
dr Cpas. Wk [K] i [s).
Hence WiQ"
Ki}, where
1
0
Co V/z
0
W
Cg VOK
And W²Q" [W,] [W3,, where
W1
0
0
W2
W
WK
for some W& E R² E [K]. And every column of W/? belongs to the space I'k == span(1c,
k E [K]]_, which is the orthogonal complement the linear subspace of IR" spanned by the vectors 1d 1Gx- i.e.,
UQ - 0² [W] W U W?)
30

---

Published as a conference paper at ICLR 2024
Denote A1 W1 + W - U 4 W A [A1|42]. First we will show that
|AI( 1n + U 7 1,
||
2
A
1/.-
Denote
d1
0
Uj
0
TT2
n2
A I
0
d2
:
0
0
1ax
0
dx
where de E R Then we have
dj
d2
I,
dx.
where
1
de Vnndd +
In IT ds.
Tikk
Note that
|| Amin
(vraldme
1n
1
||
nu nxlldx
ni
Hence
|
=
Σ |dxl1
2 min(nx)
Σ d2
k
By calculating the Hessian at U we get
(V)) = (2. Id, + A) + y1 1,y" AAT) + ,, U-A
= (2 Id, + A) + y + 1 1, A24?)
+ (2(L Id, + +A) + + 1,y" A,A)) + (1A)) 1, 1.
2((L Id, + +A), + (2(L. Id, +A)+ +
+ (1 1,
> 2[L - (Vi+ vp+ 22 n 211 1 + R2l1 + 3n A
min (VT + +p+ 2gg 21, Rs + 3n
= A|?
with probability 2 1- - 1/n, where (A, A 2 -
121 by Proposition 3
R₂ = - L) by Proposition] I
In particular, by choosing L E (1 / /2lgg/ A). B >
RR2/2, we have O. Recall the Hessian of f at U and U
(V*)(U))))) A) (2(L. Id, + A) +)+y1 v + 1, AAT) 1,
(()()AA))) (2(L. Id, + +) + y1 + (
+
+ 11 - (U*)(U*)))) + (UUT - I,,")(")),,,,, AAT).
Then we have
-
+
-000=~(0')000 -U"')'0""1111,1 AATI
< -U -
-88mn|U-UU
31

---

Published as a conference paper at ICLR 2024
Hence
2 B/21AA8
provided that ||U U* ||
36,m
Lemma (Local exponential convergence).
If U E 9m. define A = =UQ² - U*, Then By E (0,1), > 0, s.t,
IV Q'l S r Q'
where
Q argmin |U - U* QI(,,
provided that IIU U+ F < e, where = min C.). Recall that is defined in Lemma 2
is defined in Lemma 5
In particular, if we choose the step size S 3/(2R)),, where + 2L + 12n,3 for some constant C2. Then the contraction factor would
R2 = 2n + /mplogg n
be 2 = (1 - aa/2).
Proof. By Lemma 2 we have
IU - VV)) (U
where C(s) stands for the convex hall of f2. It is known that
[[]c( IT(())l s ||u uul:
for any convex set C. Note VQ Fm, U is stationary point for Na. Then we have
IVV - Q1 |lTx(U - -VV L 0000:000---007000111110000000000
S IIU - -U==
-IUQQ - - - -/VT)
since V((Ā) = V VỘ E Opxr, E Raxr. Note that |UQQ - - IU Q' [VF((Q - V/(U")]] S R - U*|IF, for some constant R2, so we only
need to analyze the last term. And by MVT we have
-U7> /, 1 -UF-
Notice ||QQ -
< ce, from Lemma we have Vr E (0,1), {(" + (UQ)"UU)))) U*))
2 //2
Hence
-/0")," - -US
- + - -d
2 3/2 - P
Then we have
IV - uQQ = - - +V
22 -
< IIQ - + - U 20/ - -
s (1 (11 11= - -
= -
32

---

Published as a conference paper at ICLR 2024
by choosing E 3/(2R)), 72 (1 aB/22).
Lemma
7 [Iterations remain staying in Phase 2). Suppose Lemma 6 holds and U E 9m- Denote }a we have
E:= min
c1, 3
U E Im, 11 2 1,
where
U++1 = 1l -VV v22 1,
given |0 - U@* S min (1->)/26, c), recall e = min((er $2)
Proof. From the convergence of Phase 1 (Theorem we know U 9m if U 9m and
||U - -U**|| S E. Thus from definition we know U E 9m. Recall that we define Q E Fm as
Q' = argmin ||U - UUQQ
Define co == (1 - 7/2 1/. Note that U 9m. - U min((1 - //22 2,, then from the previous lemma (Lemma 6 we have
| U+Q I - U U*Q Q-I
Note U U²QQ is the projection of T on Fm that located sphere, which implies that
|U*000 U || 21|00 U
given | - U/=+1| 1I < 1/2 11p. Thus,
|
U 1| S || U*Q"Il + Q1 - U=*+ S 00 -U"Q || + 2 -Ua
$ 7 7 Un*1|F 21100 U
1 760 + 2na
8E.
Then we will finish the proof by induction. Suppose U E 9ms -1 going show that U 9m. By assumption we only need show U [ || S From Lemma n we have
Un* [ S ετ AI S t, we
- 11'
S -110"-1- ss Al L L.
Hence
1 - Ua 1|0 S - U -0"' -
Q U
$ 2 Σ
10 - U Q + 2 U
1=0
2 Σ
7 1-1 1 11 u|
1=0
<
D DISCUSSION
One limitation of our derivation in Theorem T is the separation assumption. Theorem is based on the fact that the optimum solutions of SDP and NLR coincide, where the separation assumption serves as the sufficient condition. In practice, we can observe linear convergence of NLR for small separation as well. Therefore, we anticipate similar derivation of convergence analysis for weak separation assumption, which will be our future research. On the other hand, in practice, if we apply our algorithm to the datasets where the separations (signal-to-noise ratio) are very small, our algorithm would fail to provide informative clustering results, and similar issues occur for both SDP and NMF. We expect that this problem can be solved different rounding procedures are applied. In other words, for the small separation cases where the solutions to NLR, SDP and NMF no longer represent exact membership assignments, it would be important to consider and compare different rounding processes extract informative membership assignments.
33


Reviews:
[{"review_id": "3YYnxmi2kl", "reviewer": "Reviewer_v53M", "summary": "This paper introduces a new algorithm for $k$-means problem in the Gaussian mixture model setting. The new algorithm overcomes some of the key limitations of prior algorithms for the same problem. Concretely, the SDP-based algorithm is not practical in settings where the datasets are large, and the NMF-based algorithm, despite its scalability, does not have theoretical guarantees. The new algorithm is inspired by these two methods and it is designed a way that it enjoys desirable properties of both the SDP and NMF-based methods. The authors show convergence guarantees in the exact recovery regime as well as the general setting for the algorithm and provide extensive numerical experiments that compare the new algorithm with prior approaches and demonstrate its performance.", "questions": "1. What is the criterion to select the rank $r$ in the algorithm? Is it arbitrary or is there are heuristic for this?\n\n2. In experiments, I am noticing that the misclustering error of SDP is higher than this algorithm in general. Is it supposed to be like this? My understanding was SDP should have comparable or better accuracy than BM.  \n\n3. Can the authors compare the dependency of time complexity on $K$  for this algorithm and prior methods? Perhaps it is not clear for the NMF-based method but for the other methods discussed in the paper, it maybe possible. It is helpful to understand in what regimes this algorithm can be applied instead of others. I believe the dependency of Lloyd's initialized with $k$-means++ on $K$ is not as severe as this algorithm.\n\n4. Nowadays it is normal to see datasets with high dimensions. Numerical experiments in the paper use rather small values for dimension $p$(eg: $20$). Are there experiments done with higher dimensions, perhaps in the range $p=100, 500, 1000$? This will be helpful to determine how this algorithm performs in high dimensions compared to others.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "**Originality**\n\nThe clever use of projected gradient descent to solve the primal of augmented Lagrangian efficiently shows the originality of the algorithm. I also appreciate the way authors cast the relaxed SDP formulation as a non-convex optimization problem so that a method like Burer-Monteiro can be used to find the low-rank solution. Derivations of the problem and the proof techniques are non-trivial.\n\n**Quality**\n\nSolid theoretical results on problem formulation and the convergence of the algorithm. Numerical experiments are on point and demonstrate the theoretical guarantees.\n\n**Significance**\n\n$k$-means problem in GMM setting is an important problem for the community. A scalable solution for this problem that enjoys good theoretical guarantees is significant.\n\n**Clarity**\n\nThe paper is easy to follow. The contributions are clearly stated and the content is well organized.", "weaknesses": "**Weaknesses**\n\nThe time complexity of solving the primal-dual algorithm is $O(K^6nr)$ and this becomes prohibitively large when $K$  is large. The experiments show small $K$ values(eg: $4$). Even for $K=10$, the time for convergence can grow very quickly. In some applications such as document deduplication, and entity resolution, the value of $K$ can be significantly larger than what is used in the experiments.  \n\n**Typos**\n\n1. On page 3, in the paragraph after equation $2$, \"one-shot encoding\" $\\rightarrow$ \"one-hot encoding\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oVuDaZLmJ5", "reviewer": "Reviewer_cacp", "summary": "This paper introduces a new approach to solving the K-means clustering problem, addressing the limitations of existing methods. The authors propose an efficient nonnegative matrix factorization-like algorithm, incorporating semidefinite programming (SDP) relaxations within an augmented Lagrangian framework. The algorithm optimizes a nonnegative factor matrix using primal-dual gradient descent ascent, ensuring rapid convergence and precise solutions to the challenging primal update problem. The method demonstrates strong statistical optimality guarantees comparable to SDP while being scalable and simple, similar to state-of-the-art NMF algorithms. Experimental results confirm significantly reduced mis-clustering errors compared to existing methods, marking a significant advancement in large-scale K-means clustering.", "questions": "1. Given the theoretical grounding and experimental results presented in the paper, how does the proposed algorithm compare to other state-of-the-art techniques in terms of computational efficiency and scalability, especially when dealing with large-scale datasets? \n\n2. It is a little abstract to understand Propositions 1 & 2. The authors should improve their presentation here.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "(1) $\\textbf{New algorithm design}$: The paper introduces a novel algorithm for solving the K-means clustering problem, leveraging a combination of nonnegative matrix factorization (NMF) techniques and semidefinite programming (SDP) relaxations. The proposed algorithm addresses the challenges faced by prior methods and offers a unique solution approach by integrating concepts from different areas of machine learning.\n\n(2) $\\textbf{Theoretical grounding and guarantees}$: The authors provide a strong theoretical foundation for their algorithm, demonstrating local linear convergence within a primal-dual neighborhood of the SDP solution. The paper also offers rigorous proofs, such as the ability to solve the primal update problem at a rapid linear rate to machine precision. These theoretical insights establish the reliability and efficiency of the proposed method.\n\n(3) $\\textbf{Empirical validation}$: The paper supports its claims with empirical evidence, showcasing the effectiveness of the proposed algorithm through extensive experiments. The results demonstrate substantial improvements in terms of mis-clustering errors when compared to existing state-of-the-art methods. This empirical validation strengthens the credibility of the proposed approach and highlights its practical utility in real-world applications.", "weaknesses": "$\\textbf{Insufficient discussion of practical limitations}$: The paper might not thoroughly address the practical limitations or challenges that users might face when applying the proposed algorithm in real-world scenarios. Understanding the algorithm's limitations in terms of computational resources, scalability, or specific data types is crucial for potential users and researchers.\n\n$\\textbf{Initialization condition}$: The authors base their proof of Theorem 1 on the assumption that the initialization meets a specific condition. While this assumption is discussed in the paper, it would significantly enhance the rigor and credibility of their work if the authors were to provide a rigorous proof for this initialization criterion.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "nVdPDsRN19", "reviewer": "Reviewer_hAG7", "summary": "The paper presents an iterative method to (approximately) solve the k-means clustering problem. This method is obtained by applying existing methods (\"Burer–Monteiro factorization approach\") to a particular reformulation of (a relaxed version) of k-means. \nThe authors claim that the resulting algorithm blends favorable computational and statistical guarantees of different existing methods for solving k-means.", "questions": "see above.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 3, "strengths": "Computationally efficient methods for solving k-means clustering, which is a core task of data analysis, are welcome.", "weaknesses": "* It is unclear how Theorem 1 allows to verify the claims about computational and statistical superiority of Algorithm 1 compared to existing clustering methods. In particular, how does Theorem 1 and the numerical experiments imply the claims \"..simple and scalable as state-of-the- art NMF algorithms, while also enjoying the same strong statistical optimality..\" in the abstract?\n\n* there should be more discussion or comparison of computational complexity and clustering error incurred by Algorithm 1 compared to existing clustering methods for the Gaussian mixture model Eq. (14). \n\n* the connection between theoretical analysis in Section 4 and the num. exp. in Section 5 could be made more explicit. For example, there are not many references to the theoretical results in the current Section 5. How do the numerical results confirm Theorem 1? How did the theoretical analysis guide the design choices (datasets, hyperparams of Algorithm 1) of the numerical experiments. \n\n* the use of Algorithm 1 needs more discussion: How to choose beta, alpha and r in practice? How does Algorithm 1 deliver a cluster assignment that approximately solves k-means?\n\n* use of language can be improved, e.g.,\n--  \"..is the sharp threshold defined..\" what is a \"sharp\" threshold ?; \n-- \"..can be converted to the an equality-constrained..\"\n-- what is a \"manifold-like\" subset ? \n-- what is a \"is a nonconvex approach\" ? \n-- \"...by simultaneously leverage the implicit psd structure\" \n-- \".. that achieves the statistical optimality \" \n-- \"..which reparameterizes the assignment matrix .. as the psd membership matrix..\"\n--  what is a \".. one-shot encoding \"?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BShDfuhdTy", "reviewer": "Reviewer_3o8f", "summary": "This paper proposes an NMF-like algorithm for the problem of k-means clustering. The benefit of an NMF algorithm is its simplicity and scalability, but at the same time achieve the same statistical optimality as proven for the SDP. \nThis is a clean, strong, and interesting contribution.", "questions": "I do not have important questions, but I wonder whether one can prove similar results for other distributions in (14)? I guess things will work out for sub-Gaussian distributions, but can you say something about the assumptions that are needed in your analysis in this sense?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The problem is of course classical in machine learning and very important. \nThe authors propose a simple and strong algorithm for this problem, and prove statistical guarantees. \nThe paper is well-written and the proof seem clean to me.", "weaknesses": "No evident weakness except the separation assumption, but I acknowledge that overcoming this assumption is difficult mathematically, and even with this assumption the derivations and ideas are not trivial.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)