Statistically Optimal $K$-means Clustering via Nonnegative Low-rank Semidefinite Programming

Regarding the separation assumption:

Thank you for pointing that out! Our derivation relies on the fact that the optimum solutions of SDP and BM coincide under the separation assumption, which acts as the sufficient condition. Interestingly, in practice, we also observe linear convergence of BM even with small separation. Hence, we anticipate a similar derivation of convergence analysis under a weak separation assumption, which will be a focus of our future research. One main difficulty, however, lies in analyzing the SDP under partial recovery, for which a sharp bound does not currently exist in the literature to the best of our knowledge. We have incorporated this discussion in Appendix C of the revised document.

Regarding the theoretical results for other distributions:

That is a good point. The results of convergence analysis can be derived similarly for sub-Gaussian distributions, except that the separation condition would become weaker. The threshold of exact recovery for Gaussian distributions is based on the high dimensional concentration bounds AND the rotation-invariance property of Gaussian. Therefore, similar threshold of exact recovery for sub-Gaussian distributions could be derived accordingly. However, we anticipate a larger separation condition for sub-Gaussian distributions as the high dimensional concentration bounds for sub-Gaussians might not as tight as the bounds for Gaussians. On the other hand, the noise level of sub-Gaussian distributions is higher than Gaussian distributions. Therefore, we need larger signal (separation) to cluster the distributions successfully.

Regarding the question ``It is unclear how Theorem 1 allows to verify the claims...":

In previous works, the SDP formulation of $K$-means was shown to have strong statistical guarantees, but it was solved using non-scalable algorithms like the primal-dual interior-point method. At a high level, our approach computes a solution for the same SDP formulation of $K$-means using a new scalable algorithm reminiscent of NMF. On one hand, our approach ''enjoys the same strong statistical optimality" because it computes the same SDP solution as prior SDP approaches. On the other hand, each iteration of our proposed algorithm is essentially identical to NMF. Theorem 1 and our experiments establish that the algorithm rapidly converges to the SDP solution. Therefore, our proposed algorithm is indeed ``simple and scalable as state-of-the-art NMF" for solving the SDP.

As mentioned in the last paragraph on page 6, the goal of our main theorem (Theorem 1) is ``to demonstrate that projected GD can efficiently solve the primal subproblem $\min_{U\in\Omega} \mathcal{L}_\beta (U,y)$", which exhibits rapid linear convergence locally. This fact, in conjunction with Prop. 1 (Existence and Quality of Primal Minimizer) and Prop. 2 (Linear Convergence of Dual Multipliers), indicates the local linear convergence of Algorithm 1. Here, our condition on initialization requires only a constant element-wise relative error bound.

Moreover, the third plot in Figure 2 also indicates the linear convergence of our algorithm, as justified by our theoretical analysis. In this setting of large separation, the global optimums of the BM and SDP formulations align. Therefore, BM is superior compared to existing clustering methods in that it can achieve exact recovery like the SDP, while maintaining linear time complexity similar to KM, NMF, and SC. This fact is summarized in Figure 1 of the paper, where our algorithm achieves significantly smaller mis-clustering errors compared to existing state-of-the-art methods given the same CPU running time. Hence, we state in the abstract, ``The resulting algorithm is just as simple and scalable as state-of-the-art NMF algorithms, while also enjoying the same strong statistical optimality guarantees as the SDP", a claim supported by both theoretical analysis and numerical results.

Regarding the question ``there should be more discussion or comparison of computational complexity...":

In Section 5, we compared the computational complexity and mis-clustering errors of our algorithm with existing clustering methods for Gaussian mixture models, as detailed in the paragraph "Performance for BM for GMM". This analysis considers small separation with sample sizes ranging from $n=400$ to $n=57,600$. The results are summarized in the first two plots of Figure 2. In the last few sentences of the discussion and comparison, we note that the mis-clustering error of SDP and BM coincides, while NMF, KM, and SC show large variance and are far from achieving exact recovery, as evident from the first plot of Figure 2. The second plot of Figure 2, as mentioned, ``indicates that SDP and SC have super-linear time complexity, while the log-scale curves of our BM approach, KM, and NMF are nearly parallel, indicating that they all achieve linear time complexity".

Furthermore, in our revised version (see Appendix B), we have added comparisons of CPU time costs and mis-clustering errors with increasing dimension $p$ or with increasing cluster number $K$ across different methods. For more details, please refer to "Reply to Reviewer v53M".

Regarding the question ``the connection between theoretical analysis in Section 4 and the num. exp. in Section 5...":

The convergence of Algorithm 1, accompanied by empirical results, is analyzed in the second experiment (Linear Convergence of BM) in Section 5. From the third plot in Figure 2, we observe that our BM approach achieves linear convergence in this setting. Additionally, the curves for $r = K$, $r = 2K$, and $r = 20K$ are closely aligned, suggesting that the choice of $r$ does not significantly impact the convergence rate under large (cluster) separation. Conversely, the theorem recommends a smaller step size $\alpha$ and a larger augmentation parameter $\beta$ for large sample sizes $n$. This guidance allows for the tuning of parameters for simple GMMs with small sample sizes, which can be adapted for applying BM to real datasets with varying sample sizes $n$.

Regarding the discussion of practical limitations:

Thank you for pointing that out! In practice, when applying our algorithm to datasets with very small separations (signal-to-noise ratio), our algorithm, like SDP and NMF, may fail to yield informative clustering results. We anticipate that this issue could be addressed with the application of different rounding procedures. Specifically, for cases with small separation where the solutions to BM, SDP, and NMF no longer represent exact membership assignments, it becomes crucial to consider and compare different rounding processes to extract informative membership assignments. We intend to pursue this as a future research goal. Due to space limitations, we have included a discussion on this topic in Appendix C of the revised document.
Regarding computational aspects, unlike SDP, our algorithm (BM) can be solved in linear time with respect to sample size $n,$ where parallel computing can be applied to further enhance the computational efficiency. On the other hand, BM formulation (eq.(6)) keeps the matrix $A$ (contains the information of dataset) from the SDP formulation (eq.(4)), which indicates that BM can deal with the same data types as those for SDP. Therefore we can adapt our BM formulation to manifold data, measure-valued data, or data with heterogeneous covariance in the GMM context.

Regarding the initialization conditions:

Thanks for your advice. The initialization criteria mentioned in Theorem 1 is derived from Theorem 3 in Appendix D.3., which provides a more rigorous form of the initialization criteria. Following your suggestion, we have added details in Appendix D.3 of the revised document to highlight the relationship between these two expressions of the initialization criteria.

Regarding the question ``Given the theoretical grounding and experimental results presented in the paper...":

As detailed in Section 5, we compared the computational complexity and mis-clustering errors of our algorithm with existing clustering methods for Gaussian mixture models in the paragraph ``Performance for BM for GMM". This comparison considers small separation with sample sizes ranging from $n=400$ to $n=57,600$. The results, summarized in the first two plots of Figure 2, reveal that the mis-clustering error of SDP and BM coincides, while NMF, KM ($K$-means++, the fastest clustering algorithm scalable to large datasets), and SC show large variance and fall short of exact recovery, as depicted in the first plot of Figure 2. The second plot in Figure 2, as mentioned, indicates that SDP and SC exhibit super-linear time complexity, while the log-scale curves of our BM approach, KM, and NMF are nearly parallel, suggesting that they all achieve linear time complexity.

Furthermore, we have added comparisons of CPU time costs and mis-clustering errors with increasing dimensions $p$ or increasing cluster numbers $K$ across different methods in Appendix B of the revised document. For more details, please refer to ``Reply to Reviewer v53M".

Overall, our proposed algorithm achieves linear time complexity with respect to both sample size $n$ and dimension $p$, (as well as super-linear complexity with respect to $K$) which is comparable to other state-of-the-art algorithms such as $K$-means++ and NMF, while maintaining the same superior statistical performance as SDP.

Regarding the question ``It is a little abstract to understand Propositions 1 and 2...":

Thank you for your advice. We have incorporated more details about Prop. 1 and Prop. 2 in the revised version. Due to space constraints, we have included additional explanations of Prop. 1 and Prop. 2 in Appendix A in the revised paper.

Regarding the time complexity with respect to number of clusters $K$:

We acknowledge that the dependence on $K$ in our theoretical analysis is very likely not sharp. Our added simulation results suggest that the dependence of our runtime on $K$ appears to be of linear order $O(K)$, which is much better than $O(K^6)$ in our theoretical guarantee. To compare the CPU time costs and mis-clustering errors when the number of clusters increases across different methods, we did the comparisons of CPU time cost across different methods when $K=5,10,20,40,50$ in the revision. Under the same setting as our second experiment (``Performance for BM for GMM") in Section 5, except that now we consider fixed sample size $n=1000$, dimension $p=50$. SC is not considered as it would fail for large $K$ in our experiments. The results are summarized in the first plot of Figure 6 in Appendix B in the revision, where we can observe that the log-scale curve of BM is nearly parallel to the log-scale curve of KM and NMF, indicating that the growth of CPU time cost for BM with respect to $K$ is reasonable (nearly $O(K)$) and would not achieve the loose upper bound $O(K^6)$ 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 as large as the order $O(n^{3.5})$. More details can be found in Appendix B in the revision.

On the theoretical side, it is common to treat the cluster number $K$ as constant. Indeed, the best known theoretical guarantee for exact recovery via SDP necessarily requires a small value of $K=O(\log(n)/\log\log(n))$. So the time complexity bound for our BM is at most $\text{polylog}(n)$. It is a widely open conjecture whether a sharp threshold exists without a statistical--computational gap when $K \gg \log(n)$. For these reasons, we did not attempt to optimize the dependence on $K$ in the theory in our paper.

Regarding ``one-shot encoding":

Thanks for pointing out. We have corrected it in the revision.

Regarding the selection of the rank $r$ in the algorithm:

The second numerical experiment in the paper demonstrates that the choice of $r$ does not significantly impact the convergence rate. Therefore, ideally, we would choose the rank $r$ as small as $r=K$. However, in practice, we found that $r=K$ results in many local minima for small separation. Consequently, we slightly increase $r$ and choose $r=2K$ for all applications, which provided overall good performance. More details regarding the choices of parameters have been added to Appendix A in the revision.

In real applications, the number of clusters $K$ might be unknown as well. In these situations, we can adopt some common methods recommended in the literature for $K$-means to choose $K$. For example, we may use the popular elbow method, which first runs fast clustering algorithms such as the $K$-means++ for a range of values of $k$ (number of clusters) and then plots the total within-cluster sum of square against $k$. After that, we may choose a smallest value of $k$ that maintains a relatively low sum of square of the distances.

Regarding the question ``In experiments, I am noticing that the misclustering error of SDP...":

That is a valid point. On one hand, Table 1 shows that BM yields slightly better mis-clustering errors compared to SDP. However, the differences are not significant since the error bars for BM and SDP overlap. On the other hand, it is noteworthy that BM imposes a more stringent constraint, $U \ge 0$, compared to SDP's constraint of $UU^T \ge 0$. This makes BM a tighter formulation relative to the original K-means formulation. Consequently, it is possible for BM to provide a better approximation to the $K$-means formulation, especially when the separation is small.