Scalable Second-order Riemannian Optimization for $K$-means Clustering

ICLR 2026 Conference Submission15150 Authors

Keywords: K-means clustering, manifold optimization, Newton's method, nonconvex

TL;DR: Smooth Riemannian formulation for $K$-means clustering; rapid convergence to second-order critical points with linear per-iteration cost.

Abstract: 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.

Primary Area: optimization

Submission Number: 15150