Learning-Augmented $k$-means Clustering

Abstract

$k$-means clustering is a well-studied problem due to its wide applicability. Unfortunately, there exist strong theoretical limits on the performance of any algorithm for the $k$-means problem on worst-case inputs. To overcome this barrier, we consider a scenario where "advice" is provided to help perform clustering. Specifically, we consider the $k$-means problem augmented with a predictor that, given any point, returns its cluster label in an approximately optimal clustering up to some, possibly adversarial, error. We present an algorithm whose performance improves along with the accuracy of the predictor, even though naïvely following the accurate predictor can still lead to a high clustering cost. Thus if the predictor is sufficiently accurate, we can retrieve a close to optimal clustering with nearly optimal runtime, breaking known computational barriers for algorithms that do not have access to such advice. We evaluate our algorithms on real datasets and show significant improvements in the quality of clustering.

ICLR 2022

References in corpus (8)

• Compressed Sensing using Generative Models
• Improving Online Algorithms via ML Predictions
• A Model for Learned Bloom Filters, and Optimizing by Sandwiching
• A Unified Approach to Robust Mean Estimation
• Learning-based Support Estimation in Sublinear Time
• Same-Cluster Querying for Overlapping Clusters
• Relaxed Oracles for Semi-Supervised Clustering
• Improved Clustering with Augmented k-means