BalLOT: Balanced $k$-means clustering with optimal transport

TL;DR

BalLOT introduces an optimal transport-based method for balanced $k$-means clustering, providing fast solutions with theoretical guarantees for cluster recovery and effective initialization schemes.

Contribution

The paper presents a novel optimal transport approach called BalLOT for balanced $k$-means clustering, with proven theoretical guarantees and efficient initialization strategies.

Findings

BalLOT produces integral couplings at each step for generic data.

Theoretical guarantees for exact and partial cluster recovery under the stochastic ball model.

Initialization schemes enable one-step recovery of planted clusters.

Abstract

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.