Distributionally Robust K-Means Clustering

2026-04-13

AI Summary

This work addresses the limited robustness of traditional K-means clustering against outliers, distributional shifts, and small sample sizes by reframing it as a Lloyd–Max quantization problem for empirical distributions. The authors introduce, for the first time, a Wasserstein-2 ball to construct an ambiguity set over probability distributions and formulate a minimax optimization model that minimizes the worst-case expected squared distance. Leveraging duality theory and block coordinate descent, they derive a smooth, weighted soft clustering algorithm that naturally generalizes hard assignments to soft ones. The proposed method exhibits substantially enhanced robustness to noise and outliers, achieves strong performance on standard benchmarks and large-scale synthetic datasets, and enjoys theoretical guarantees of monotonic convergence and local linear convergence.

Technology Category

• Machine Learning: Clustering
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Search and Optimization: Distributed Search

Application Category

• Web Mining and Content Analysis: Robustness and generalizability of Web mining methods
• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

Abstract

K-means clustering is a workhorse of unsupervised learning, but it is notoriously brittle to outliers, distribution shifts, and limited sample sizes. Viewing k-means as Lloyd--Max quantization of the empirical distribution, we develop a distributionally robust variant that protects against such pathologies. We posit that the unknown population distribution lies within a Wasserstein-2 ball around the empirical distribution. In this setting, one seeks cluster centers that minimize the worst-case expected squared distance over this ambiguity set, leading to a minimax formulation. A tractable dual yields a soft-clustering scheme that replaces hard assignments with smoothly weighted ones. We propose an efficient block coordinate descent algorithm with provable monotonic decrease and local linear convergence. Experiments on standard benchmarks and large-scale synthetic data demonstrate substantial gains in outlier detection and robustness to noise.

Problem

• K-means clustering
• outliers
• distribution shifts
• sample size limitations
• robustness

Innovation

• Distributionally Robust Optimization
• K-Means Clustering
• Wasserstein Distance
• Soft Clustering
• Block Coordinate Descent