# Learning-Augmented k-means Clustering

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

## Key facts

- Cited by 26 papers in the map (data through 14 Oct 2026).
- Builds on 31 of its 74 references that are in the map.
- First posted on arXiv 27 Oct 2021.

## Builds on (most cited first, 30 of 31)

- [Learning Combinatorial Optimization Algorithms over Graphs](https://mlsanity.com/paper/1704.01665/) (2017)
- [Compressed Sensing using Generative Models](https://mlsanity.com/paper/1703.03208/) (2017)
- [Competitive caching with machine learned advice](https://mlsanity.com/paper/1802.05399/) (2018)
- [The Case for Learned Index Structures](https://mlsanity.com/paper/1712.01208/) (2017)
- [Algorithms with Predictions](https://mlsanity.com/paper/2006.09123/) (2020)
- [Spreading vectors for similarity search](https://mlsanity.com/paper/1806.03198/) (2018)
- [A Model for Learned Bloom Filters, and Optimizing by Sandwiching](https://mlsanity.com/paper/1901.00902/) (2019)
- [Faster Matchings via Learned Duals](https://mlsanity.com/paper/2107.09770/) (2021)
- [Better Guarantees for k-Means and Euclidean k-Median by Primal-Dual Algorithms](https://mlsanity.com/paper/1612.07925/) (2016)
- [Performance of Johnson-Lindenstrauss Transform for k-Means and k-Medians Clustering](https://mlsanity.com/paper/1811.03195/) (2018)
- [Approximate Nearest Neighbor Search in High Dimensions](https://mlsanity.com/paper/1806.09823/) (2018)
- [A Deep Learning Approach to Structured Signal Recovery](https://mlsanity.com/paper/1508.04065/) (2015)
- [Learning to Hash for Indexing Big Data - A Survey](https://mlsanity.com/paper/1509.05472/) (2015)
- [Geometric Median in Nearly Linear Time](https://mlsanity.com/paper/1606.05225/) (2016)
- [Training Gaussian Mixture Models at Scale via Coresets](https://mlsanity.com/paper/1703.08110/) (2017)
- [Clustering with Noisy Queries](https://mlsanity.com/paper/1706.07510/) (2017)
- [Clustering with Same-Cluster Queries](https://mlsanity.com/paper/1606.02404/) (2016)
- [Improved and Simplified Inapproximability for k-means](https://mlsanity.com/paper/1509.00916/) (2015)
- [Learning-based Support Estimation in Sublinear Time](https://mlsanity.com/paper/2106.08396/) (2021)
- [Local algorithms for interactive clustering](https://mlsanity.com/paper/1312.6724/) (2013)
- [Dimensionality Reduction for Wasserstein Barycenter](https://mlsanity.com/paper/2110.08991/) (2021)
- [A Unified Approach to Robust Mean Estimation](https://mlsanity.com/paper/1907.00927/) (2019)
- [Same-Cluster Querying for Overlapping Clusters](https://mlsanity.com/paper/1910.12490/) (2019)
- [Interactive Bayesian Hierarchical Clustering](https://mlsanity.com/paper/1602.03258/) (2016)
- [Scalable Learning-Based Sampling Optimization for Compressive Dynamic MRI](https://mlsanity.com/paper/1902.00386/) (2019)
- [Query K-means Clustering and the Double Dixie Cup Problem](https://mlsanity.com/paper/1806.05938/) (2018)
- [Approximate Clustering with Same-Cluster Queries](https://mlsanity.com/paper/1704.01862/) (2017)
- [k-center Clustering under Perturbation Resilience](https://mlsanity.com/paper/1505.03924/) (2015)
- [Relaxed Oracles for Semi-Supervised Clustering](https://mlsanity.com/paper/1711.07433/) (2017)
- [Semi-Supervised Algorithms for Approximately Optimal and Accurate Clustering](https://mlsanity.com/paper/1803.00926/) (2018)

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

- [Sorting with Predictions](https://mlsanity.com/paper/2311.00749/) (2023)
- [Algorithms for Caching and MTS with reduced number of predictions](https://mlsanity.com/paper/2404.06280/) (2024)
- [Learning-Augmented Algorithms for Online Linear and Semidefinite Programming](https://mlsanity.com/paper/2209.10614/) (2022)
- [Learning-augmented Maximum Independent Set](https://mlsanity.com/paper/2407.11364/) (2024)
- [Paging with Succinct Predictions](https://mlsanity.com/paper/2210.02775/) (2022)
- [Improved Learning-augmented Algorithms for k-means and k-medians Clustering](https://mlsanity.com/paper/2210.17028/) (2022)
- [Improved Frequency Estimation Algorithms with and without Predictions](https://mlsanity.com/paper/2312.07535/) (2023)
- [Online Classification with Predictions](https://mlsanity.com/paper/2405.14066/) (2024)
- [Learning-Augmented Search Data Structures](https://mlsanity.com/paper/2402.10457/) (2024)
- [Improved Approximations for Hard Graph Problems using Predictions](https://mlsanity.com/paper/2505.23967/) (2025)
- [On the Power of Learning-Augmented Search Trees](https://mlsanity.com/paper/2211.09251/) (2022)
- [Rethinking Warm-Starts with Predictions: Learning Predictions Close to Sets of Optimal Solutions for Faster L-/L^natural-Convex Function Minimization](https://mlsanity.com/paper/2302.00928/) (2023)
- [Learning-Augmented Moment Estimation on Time-Decay Models](https://mlsanity.com/paper/2603.02488/) (2026)
- [On Approximability of ℓ22 Min-Sum Clustering](https://mlsanity.com/paper/2412.03332/) (2024)
- [Learning-Augmented Algorithms for Boolean Satisfiability](https://mlsanity.com/paper/2505.06146/) (2025)
- [Optimizing Health Coverage in Ethiopia: A Learning-augmented Approach and Persistent Proportionality Under an Online Budget](https://mlsanity.com/paper/2509.00135/) (2025)
- [Learning-Augmented Streaming Algorithms for Correlation Clustering](https://mlsanity.com/paper/2510.10705/) (2025)
- [Learning-Augmented Ski Rental with Discrete Distributions: A Bayesian Approach](https://mlsanity.com/paper/2512.07313/) (2025)
- [Learning-augmented smooth integer programs with PAC-learnable oracles](https://mlsanity.com/paper/2602.02505/) (2026)
- [A Simple Learning-Augmented Algorithm for Online Packing with Concave Objectives](https://mlsanity.com/paper/2406.03574/) (2024)
- [Learning Augmented Graph k-Clustering](https://mlsanity.com/paper/2506.13533/) (2025)
- [Transductive and Learning-Augmented Online Regression](https://mlsanity.com/paper/2510.03917/) (2025)
- [Metric k-clustering using only Weak Comparison Oracles](https://mlsanity.com/paper/2601.19333/) (2026)
- [Better Learning-Augmented Spanning Tree Algorithms via Metric Forest Completion](https://mlsanity.com/paper/2602.24232/) (2026)
- [Learning-Augmented Algorithms for k-median via Online Learning](https://mlsanity.com/paper/2603.18157/) (2026)
- [Robust and Consistent Ski Rental with Distributional Advice](https://mlsanity.com/paper/2603.29233/) (2026)

## Abstract (arXiv)

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"ively 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.