# Wasserstein K-means for clustering probability distributions

arXiv: https://arxiv.org/abs/2209.06975 · HTML: https://mlsanity.com/paper/2209.06975/

## Key facts

- Cited by 16 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 47 references that are in the map.
- First posted on arXiv 14 Sep 2022.

## Builds on (most cited first, 12 of 12)

- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Statistical and Computational Guarantees of Lloyd's Algorithm and its Variants](https://mlsanity.com/paper/1612.02099/) (2016)
- [Wasserstein Regression](https://mlsanity.com/paper/2006.09660/) (2020)
- [Debiased Sinkhorn barycenters](https://mlsanity.com/paper/2006.02575/) (2020)
- [On Robust Optimal Transport: Computational Complexity and Barycenter Computation](https://mlsanity.com/paper/2102.06857/) (2021)
- [Partial recovery bounds for clustering with the relaxed Kmeans](https://mlsanity.com/paper/1807.07547/) (2018)
- [Cutoff for exact recovery of Gaussian mixture models](https://mlsanity.com/paper/2001.01194/) (2020)
- [Principal Geodesic Analysis for Probability Measures under the Optimal Transport Metric](https://mlsanity.com/paper/1506.07944/) (2015)
- [Uncoupled isotonic regression via minimum Wasserstein deconvolution](https://mlsanity.com/paper/1806.10648/) (2018)
- [Hybrid Wasserstein Distance and Fast Distribution Clustering](https://mlsanity.com/paper/1812.11026/) (2018)
- [Improved Complexity Bounds in Wasserstein Barycenter Problem](https://mlsanity.com/paper/2010.04677/) (2020)
- [Sketch-and-Lift: Scalable Subsampled Semidefinite Program for K-means Clustering](https://mlsanity.com/paper/2201.08226/) (2022)

## Built on by (most cited first, 16 of 16)

- [Towards Cross-Tokenizer Distillation: the Universal Logit Distillation Loss for LLMs](https://mlsanity.com/paper/2402.12030/) (2024)
- [Statistically Optimal K-means Clustering via Nonnegative Low-rank Semidefinite Programming](https://mlsanity.com/paper/2305.18436/) (2023)
- [A New Robust Partial p-Wasserstein-Based Metric for Comparing Distributions](https://mlsanity.com/paper/2405.03664/) (2024)
- [Diffeomorphic Measure Matching with Kernels for Generative Modeling](https://mlsanity.com/paper/2402.08077/) (2024)
- [Geometric Sparse Coding in Wasserstein Space](https://mlsanity.com/paper/2210.12135/) (2022)
- [Likelihood Adjusted Semidefinite Programs for Clustering Heterogeneous Data](https://mlsanity.com/paper/2209.15097/) (2022)
- [Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization](https://mlsanity.com/paper/2501.14635/) (2025)
- [Sobolev Gradient Ascent for Optimal Transport: Barycenter Optimization and Convergence Analysis](https://mlsanity.com/paper/2505.13660/) (2025)
- [Scalable Second-order Riemannian Optimization for K-means Clustering](https://mlsanity.com/paper/2509.21675/) (2025)
- [Analysis of Utterance Embeddings and Clustering Methods Related to Intent Induction for Task-Oriented Dialogue](https://mlsanity.com/paper/2212.02021/) (2022)
- [The geometry of financial institutions -- Wasserstein clustering of financial data](https://mlsanity.com/paper/2305.03565/) (2023)
- [An Agglomerative Clustering of Simulation Output Distributions Using Regularized Wasserstein Distance](https://mlsanity.com/paper/2407.12100/) (2024)
- [Random Normed k-Means: A Paradigm-Shift in Clustering within Probabilistic Metric Spaces](https://mlsanity.com/paper/2504.03928/) (2025)
- [A Particle-Flow Algorithm for Free-Support Wasserstein Barycenters](https://mlsanity.com/paper/2509.11435/) (2025)
- [A Distribution Testing Approach to Clustering Distributions](https://mlsanity.com/paper/2512.08376/) (2025)
- [K-means with learned metrics](https://mlsanity.com/paper/2603.14601/) (2026)

## Abstract (arXiv)

Clustering is an important exploratory data analysis technique to group objects based on their similarity. The widely used K-means clustering method relies on some notion of distance to partition data into a fewer number of groups. In the Euclidean space, centroid-based and distance-based formulations of the K-means are equivalent. In modern machine learning applications, data often arise as probability distributions and a natural generalization to handle measure-valued data is to use the optimal transport metric. Due to non-negative Alexandrov curvature of the Wasserstein space, barycenters suffer from regularity and non-robustness issues. The peculiar behaviors of Wasserstein barycenters may make the centroid-based formulation fail to represent the within-cluster data points, while the more direct distance-based K-means approach and its semidefinite program (SDP) relaxation are capable of recovering the true cluster labels. In the special case of clustering Gaussian distributions, we show that the SDP relaxed Wasserstein K-means can achieve exact recovery given the clusters are well-separated under the 2-Wasserstein metric. Our simulation and real data examples also demonstrate that distance-based K-means can achieve better classification performance over the standard centroid-based K-means for clustering probability distributions and images.