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

A Sub-4 Approximation for Fair $k$ -Means
Sep 7, 2026Kangke Cheng, Guanlin Mo, Shihong Song +1K-Means Approximation Algorithms
University of Science and Technology of China, Hefei, China · School of Informatics, University of Edinburgh, Edinburgh, UK
An interpretable Good--Turing restart criterion for k-means++
Jul 9, 2026Renato Cordeiro de AmorimK-Means Top-K
School of Computer Science and Electronic Engineering, University of Essex, Colchester, UK.
Randomizing the Number of Centers in k-means++
Jul 28, 2026Vaclav RozhonK-Means Approximation Algorithms

Sep 7, 2026, cs.CGJ/K move · Enter open · S save

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

Kangke Cheng, Guanlin Mo, Shihong Song, Hu Ding

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

Open paper AskPDF arXiv

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.

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