BalLOT: Balanced $k$-means clustering with optimal transport
This work addresses balanced clustering, a fundamental issue in data analysis, with incremental improvements in method and theory.
The paper tackles the problem of balanced k-means clustering by introducing BalLOT, an optimal transport-based alternating minimization method, which is shown to be fast and effective through numerical experiments and theoretical guarantees, including integral couplings for generic data and recovery of planted clusters under the stochastic ball model.
We consider the fundamental problem of balanced $k$-means clustering. In particular, we introduce an optimal transport approach to alternating minimization called BalLOT, and we show that it delivers a fast and effective solution to this problem. We establish this with a variety of numerical experiments before proving several theoretical guarantees. First, we prove that for generic data, BalLOT produces integral couplings at each step. Next, we perform a landscape analysis to provide theoretical guarantees for both exact and partial recoveries of planted clusters under the stochastic ball model. Finally, we propose initialization schemes that achieve one-step recovery of planted clusters.