Scalable Second-order Riemannian Optimization for $K$-means Clustering
This addresses the challenge of balancing constraint feasibility and objective optimality in K-means clustering for researchers and practitioners in machine learning, though it is incremental as it builds on existing optimization frameworks.
The paper tackles the K-means clustering problem by reformulating it as a smooth unconstrained optimization over a submanifold and solving it with a second-order cubic-regularized Riemannian Newton algorithm, achieving significantly faster convergence than state-of-the-art first-order methods while maintaining similar statistical accuracy.
Clustering is a hard discrete optimization problem. Nonconvex approaches such as low-rank semidefinite programming (SDP) have recently demonstrated promising statistical and local algorithmic guarantees for cluster recovery. Due to the combinatorial structure of the $K$-means clustering problem, current relaxation algorithms struggle to balance their constraint feasibility and objective optimality, presenting tremendous challenges in computing the second-order critical points with rigorous guarantees. In this paper, we provide a new formulation of the $K$-means problem as a smooth unconstrained optimization over a submanifold and characterize its Riemannian structures to allow it to be solved using a second-order cubic-regularized Riemannian Newton algorithm. By factorizing the $K$-means manifold into a product manifold, we show how each Newton subproblem can be solved in linear time. Our numerical experiments show that the proposed method converges significantly faster than the state-of-the-art first-order nonnegative low-rank factorization method, while achieving similarly optimal statistical accuracy.