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

Wenyan Luo; Dustin G. Mixon

stat.MLDec 5, 2025

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

Authors: Wenyan Luo, Dustin G. Mixon

Organizations: Department of Mathematics, The Ohio State University, Columbus, Ohio, USA · Translational Data Analytics Institute, The Ohio State University, Columbus, Ohio, USA

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 several theoretical guarantees and a variety of numerical experiments. On the theory front, we first 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. We also propose initialization schemes that achieve one-step recovery of planted clusters. To conclude, we present numerical experiments that corroborate our theoretical results.

Figures & tables

Algorithm	Separation	No. Clusters	Reference
thresholding	$\Delta>4$	any	folklore
SDP	$\Delta>2\sqrt{2}\cdot(1+1/\sqrt{d})$	any	Theorem 3 in [ 4 ]
	$\Delta>2+k^{2}/d$	any	Theorem 9 in [ 20 ]
	$\Delta>2+\Omega(\sqrt{k/d})$	any	Corollary 2 in [ 21 ]
LP	$\Delta>1+\sqrt{3}$	$k=2$	Theorem 14 in [ 15 ]
BalLOT	$\Delta>2\sqrt{2}$	$k=2$	Corollary 11

Table 1 : Balanced

k

-means clustering guarantees under the balanced stochastic ball model.

Panel of Figure 1 — Figure 1 : The two solid circles with unit radius are centered at $\bm{\mu}^{\natural}_{1}$ and $\bm{\mu}^{\natural}_{2}$ , and these centers have distance $\Delta$ . The two dashed circles denote the initialization neighborhood of radius $\delta$ that achieves one-step recovery. By the Pythagorean theorem, the sufficient condition $\Delta\cos\theta>2$ from Lemma 15 is equivalent to $\delta<((\frac{\Delta}{2})^{2}-1)^{1/2}$ .

Panel of Figure 2 — Figure 2 : Exact recovery under the balanced stochastic ball model with $n=100$ and $k=2$ . For each value of the separation parameter $\Delta$ , we run $500$ trials and plot the fraction in which the planted clustering is exactly recovered. The oracle method runs one exact BalLOT assignment step initialized at the planted centers. (left) $d=2$ . (right) $d=10$ . See Experiment 16 for details.

Figure 3 : Median runtime in Experiment 17 . For each $n\in\{2^{2},2^{3},\ldots,2^{27}\}$ , we draw $n$ data points in $\mathbb{R}^{2}$ from a balanced stochastic ball model with two clusters. BalLOT, E-BalLOT, and Lloyd’s algorithm exhibit near-linear empirical runtimes in this experiment, while the other methods are super-linear.

Panel of Figure 4 — Figure 4 : Estimating centers in a balanced Gaussian mixture model when the planted means have i.i.d. $N(0,4)$ coordinates. The first boxplot records the distance from the $k$ -means++ initialization to the planted means. The remaining boxplots record the distance after running the indicated algorithms, with oracle BalLOT initialized at the planted means. (left) $d=2$ . (right) $d=10$ . See Experiment 18 .

Panel of Figure 5 — Figure 5 : Probability of exact one-step recovery in Experiments 19 and 20 . White denotes probability $1$ and black denotes probability $0$ . (left) $k=2$ . (right) $k=3$ . The sufficient thresholds from Theorem 7 are overlaid.

Panel of Figure 6 — Figure 6 : First-step and final misclustering rates for BalLOT under a balanced stochastic ball model with $d=25$ , $n=2000$ , and $k=2$ . (left) First-step misclustering rates on a logarithmic scale, together with the asymptotic upper bound from Theorem 8 (b). (right) First-step and final misclustering rates. The dashed curve is an upper bound, not an exact limit. See Experiment 21 .

Panel of Figure 7 — Figure 7 : The phase transition in Experiment 22 . (left) Heat map of $\Delta\cos\theta$ , with white denoting $\Delta\cos\theta\geq 2$ . (right) Heat map of the average first-step misclustering rate, with white denoting one-step recovery. The red dashed curve marks $\Delta\cos\theta=2$ .

Panel of Figure 8 — Figure 8 : $m$ -coupon-collecting initialization in Experiment 23 . (left) One data realization and its sampled points for the initialization. (right) For each trial, we record the smallest $r$ for which all of the sampled points lie within distance $r$ of planted centers. The point $(0.5,0.061)$ indicates that $93.9\%$ of the trials satisfy $r\leq 1/2$ , the sufficient condition from Theorem 7 .

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 9 — Figure 9 : Illustration of the balanced clustering geometry.

Explore similar work

Sep 7, 2026cs.CG

A Sub-4 Approximation for Fair $k$ -Means

Fairness in clustering has attracted sustained research interest, motivated by the need to ensure equitable representation of protected groups in machine learning applications. We study fair

k

-means clustering in Euclidean space, where the proportion of each protected group in every cluster must lie within specified lower and upper bounds. These constraints make it challenging to determine both cluster centers and point assignments. We propose an approximation algorithm that combines a linear programming relaxation with geometric transformations of the input to construct candidate center sets. Given a

ρ

-approximate algorithm for weighted

k

-means and any

ε>0

, our algorithm returns a fractional solution whose cost is at most

1+(3-1/Γ)ρ+O(ε)

times the optimal integral fair cost, where

Γ\approx6.357

is an upper bound on the integrality gap of the standard Euclidean

k

-means LP. With a PTAS as the subroutine, the approximation ratio becomes

3.8427+O(ε)

, improving the previous factor of

5+O(ε)

to below

4

. The solution satisfies all fairness constraints exactly and can be rounded to an integral assignment with a bounded additive violation of fairness and no increase in cost. The same approximation guarantee extends to the

k

-sparse Wasserstein barycenter problem.

Kangke Cheng, Guanlin Mo, Shihong Song +1

University of Science and Technology of China, Hefei, China · School of Informatics, University of Edinburgh, Edinburgh, UK

K-Means Approximation Algorithms Algorithmic Fairness Wasserstein Distance

Ask PDF

Jul 9, 2026cs.LG

An interpretable Good--Turing restart criterion for k-means++

The k-means++ algorithm is commonly restarted multiple times to avoid poor local optima, yet the number of restarts is almost always chosen arbitrarily and applied uniformly regardless of data set difficulty. This undermines any comparison relying on such a choice and wastes computation on easy data sets while potentially under-serving hard ones. We introduce GTRC, a restart criterion combining a Good-Turing estimate, a proven unconditional bound, and a confidence-based bound on the probability that a further restart would improve on the current result, stopping once this probability falls below a user-specified tolerance

\varepsilon

. Across 36 data sets, GTRC reached clustering quality competitive with well-chosen fixed restart counts, while the number of restarts used varied considerably and appropriately with data set difficulty, governed by an interpretable, data-dependent signal rather than a fixed rule. GTRC offers a principled and reportable alternative to fixing the number of

k

-means++ restarts in advance. Software:https://github.com/RCdeAmorim/Good-Turing-Restart-Criterion.

Renato Cordeiro de Amorim

School of Computer Science and Electronic Engineering, University of Essex, Colchester, UK.

K-Means Top-K

Ask PDF

Jul 28, 2026cs.DS

Randomizing the Number of Centers in k-means++

The

k

-means++ algorithm is a standard and widely used seeding method for

k

-means clustering, but for a fixed number

k

of centers its worst-case expected approximation ratio is

Θ(\log k)

. We consider the same algorithm when an adversary first fixes the dataset and some

K

; the number of centers

k

is then chosen uniformly from

\{K,\ldots,2K-1\}

. We prove that

k

-means++ is an

O(1)

-approximation with constant probability in this budget-smoothed setup.

Vaclav Rozhon

K-Means Approximation Algorithms Center Token Budget Allocation

Ask PDF

Abstract

Explore similar work

A Sub-4 Approximation for Fair kkk-Means

An interpretable Good--Turing restart criterion for k-means++

Randomizing the Number of Centers in k-means++

A Sub-4 Approximation for Fair $k$ -Means