Learning-Augmented $k$-means Clustering

1 INTRODUCTION . Clustering is a fundamental task in data analysis that is typically one of the first methods used to understand the structure of large datasets . The most common formulation of clustering is the kmeans problem where given a set P ⊂ Rd of n points , the goal is to find a set of centers C ⊂ Rd of k points to minimize the objective cost ( P , C ) = ∑ p∈P minc∈C ‖p− c‖22 . ( 1 ) Despite decades of work , there exist strong theoretical limitations about the performance of any algorithm for the k-means problem . Finding the optimal set C is NP-hard even for the case of k = 2 ( Dasgupta , 2008 ) and even finding an approximate solution with objective value that is within a factor 1.07 of the optimal solution is NP-hard ( Cohen-Addad & S. , 2019 ; Lee et al. , 2017 ) . Furthermore , the best-known practical polynomial time algorithms can only provably achieve a large constant factor approximation to the optimal clustering , e.g. , the 50-approximation in Song & Rajasekaran ( 2010 ) , or use techniques such as linear programming that do not scale , e.g. , the 6.357- approximation in Ahmadian et al . ( 2020 ) . A natural approach to overcome these computational barriers is to leverage the fact that in many applications , the input is often not arbitrary and contains auxiliary information that can be used to construct a good clustering , e.g. , in many applications , the input can be similar to past instances . Thus , it is reasonable to create a ( possibly erroneous ) predictor by using auxiliary information or through clusterings of similar datasets , which can inform the proper label of an item in our current dataset . Indeed , inspired by the developments in machine learning , many recent papers have studied algorithms augmented with predictions ( Mitzenmacher & Vassilvitskii , 2020 ) . Such algorithms utilize a predictor that , when invoked , provides an ( imperfect ) prediction for future inputs . The predictions are then used by the algorithm to improve performance ( see references in Section 1.3 ) . Hence , we consider the problem of k-means clustering given additional access to a predictor that outputs advice for which points should be clustered together , by outputting a label for each point . The goal is to find k centers that minimize objective ( 1 ) and assign each point to one of these centers . The question is then whether one can utilize such predictions to boost the accuracy and runtime of clustering of new datasets . Our results demonstrate the answer in the affirmative . Formal learning-augmented problem definition . Given a set P ⊆ Rd of n points , the goal is to find a set of k points C ( called centers ) to minimize objective ( 1 ) . In the learning-augmented setting , we assume we have access to a predictor Π that provides information about the label of each point consistent with a ( 1+α ) -approximately optimal clustering C. We say that a predictor has label error rate λ ≤ α if for each label i ∈ [ k ] : = { 1 , . . . , k } , Π errs on at most a λ ≤ α fraction of all points in cluster i in C , and Π errs on at most a λ ≤ α fraction of all points given label i by Π . In other words , Π has at least ( 1− λ ) precision and recall for each label . Our predictor model subsumes both random and adversarial errors by the predictor . For example if the cluster sizes are somewhat well-balanced , then a special case of our model is when Π ( p ) outputs the correct label of point p ∈ P with some probability 1 − λ and otherwise outputs a random label in [ k ] with probability λ . The example where the predictor outputs an adversarial label instead of a random label with probability λ also falls under our model . For more detail , see Theorems 2.1 and 3.4 . We also adjust our algorithm to have better performance when the errors are random rather than adversarial in the supplementary material . 1.1 MOTIVATION FOR OUR WORK . We first motivate studying k-means clustering under the learning-augmented algorithms framework . Overcoming theoretical barriers . As stated above , no polynomial time algorithm can achieve better than a constant factor approximation to the optimal clustering . In addition , the best provable approximation guarantees by polynomial time algorithms have a large constant factor ( for example the 50 approximation in Song & Rajasekaran ( 2010 ) ) , or use methods which do not scale ( such as the linear programming based algorithm in Ahmadian et al . ( 2020 ) which gives a 6.357-approximation ) . Therefore , it is of interest to study whether a natural assumption can overcome these complexity barriers . In our work , we show that knowing the true labels up to some possibly adversarial noise can give us arbitrarily good clusterings , depending on the noise level , which breaks these computational barriers . Furthermore , we present an algorithm that runs in nearly linear time , rather than just polynomial time . Lastly , we introduce tools from the robust statistics literature to study k-means clustering rather than the distance-based sampling procedure that is commonly analyzed ( this is the basis of kmeans++ ) . This new toolkit and connection could have further applications in other learning-augmented clustering problems . Practical considerations . In practice , good predictors can be learned for datasets with auxiliary information . For a concrete example , we can take any dataset that has a train/test split and use a clustering on the training dataset to help us cluster the testing portion of the dataset . Therefore , datasets do not have to be specifically curated to fit our modelling assumption , which is a requirement in other modelling formulations that leverage extra information such as the SSAC model discussed in Section 1.3 . A predictor can also be created from the natural class of datasets that vary over time , such as Census data or spectral clustering for temporal graphs ( graphs slowly varying over time ) . For this class of datasets , a clustering from an earlier time step can function as a predictor for later time steps . Lastly , we can simply use the labels given by another clustering algorithm ( such as kmeans++ ) or heuristic as a predictor . Therefore , predictors are readily and easily available for a wide class of natural datasets . Following the predictor alone is insufficient . Given a predictor that outputs noisy labels , it is conceivable that its output alone can give us a good clustering relative to optimal . However , this is not the case and naı̈vely using the label provided by the predictor for each point can result in an arbitrarily bad solution , even when the predictor errs with low probability . For example , consider a cluster of n2 points at the origin and a cluster of n 2 points at x = 1 . Then for k = 2 , choosing centers at the origin and at x = 1 induces a k-means clustering cost of zero . However , even for a predictor that errs with probability 1n , some point will be mislabeled with constant probability , which results in a positive k-means clustering cost , and so does not provide a relative error approximation . Thus , using the provided labels by the predictor can induce an arbitrarily bad clustering , even as the label error rate of the predictor tends to zero . This subtlety makes the model rich and interesting , and requires us to create non-trivial clustering algorithms . Predictors with adversarial errors . Since the predictor is separate from the clustering algorithm , interference with the output of the predictor following the clustering algorithm ’ s query can be a source of non-random noise . Thus any scenario in which communication is performed over a noisy channel ( for example , if the predictor is hosted at one server and the algorithm is hosted at another server ) is susceptible to such errors . Another source of adversarial failure by the predictor is when the predictor is trained on a dataset that can be generated by an adversary , such as in the context of adversarial machine learning . Moreover , our algorithms have better guarantees when the predictor does not fail adversarially , e.g. , see the supplementary material ) . 1.2 OUR RESULTS . In this paper we study “ learning-augmented ” methods for efficient k-means clustering . Our contributions are both theoretical and empirical . On the theoretical side , we introduce an algorithm that provably solves the k-means problem almost optimally , given access to a predictor that outputs a label for each point p ∈ P according to a ( 1 + α ) -approximately optimal clustering , up to some noise . Specifically , suppose we have access to a predictor Π with label error rate λ upper bounded by a parameter α . Then , Algorithm 1 outputs a set of centers C̃ in Õ ( knd ) time1 , such that cost ( P , C̃ ) ≤ ( 1 +O ( α ) ) · cost ( P , Copt ) , where Copt is an optimal set of centers . We improve the runtime in Section 3 by introducing Algorithm 3 , which has the same error guarantees , but uses Õ ( nd ) runtime , which is nearly optimal since one needs at least nd time to read the points for dense inputs ( Theorem 3.4 , and Remark A.14 ) . To output labels for all points , Algorithm 3 requires n queries to the predictor . However , if the goal is to just output centers for each cluster , then we only require Õ ( k/α ) queries . This is essentially optimal ; we show in Theorem 3.5 that any polynomial time algorithm must perform approximately Ω̃ ( k/α ) queries to output a 1+α-approximate solution assuming the Exponential Time Hypothesis , a well known complexity-theoretic assumption ( Impagliazzo & Paturi , 2001 ) . Note that one could ignore the oracle entirely , but then one is limited by the constant factor hardness for polynomial time algorithms , which we bypass with a small number of queries . Surprisingly , we do not require assumptions that the input is well-separated or approximation-stable ( Braverman et al. , 2011 ; Balcan et al. , 2013 ) , which are assumed in other works . Finally in the supplementary material , we also give a learning-augmented algorithm for the related problem of k-median clustering , which has less algebraic structure than that of k-means clustering . We also consider a deletion predictor , which either outputs a correct label or a failure symbol ⊥ and give a ( 1 + α ) -approximation algorithm even when the “ deletion rate ” is 1− 1/poly ( k ) . On the empirical side , we evaluate our algorithms on real and synthetic datasets . We experimentally show that good predictors can be learned for all of our varied datasets , which can aid in clustering . We also show our methodology is more robust than other heuristics such as random sampling . 1.3 RELATED WORK . Learning-augmented algorithms . Our paper adds to the growing body of work on learningaugmented algorithms . In this framework , additional “ advice ” from a possibly erroneous predictor is used to improve performance of classical algorithms . For example , a common predictor is a “ heaviness ” predictor that outputs how “ important ” a given input point is . It has been shown that such predictors can be learned using modern machine learning techniques or other methods on training datasets and can be successfully applied to similar testing datasets . This methodology has found applications in improving data structures ( Kraska et al. , 2018 ; Mitzenmacher , 2018 ) , streaming algorithms ( Hsu et al. , 2019 ; Jiang et al. , 2020 ) , online algorithms ( Lykouris & Vassilvtiskii , 2018 ; Purohit et al. , 2018 ) , graph algorithms ( Dai et al. , 2017 ) , and many other domains ( Mousavi et al. , 2015 ; Wang et al. , 2016 ; Bora et al. , 2017 ; Sablayrolles et al. , 2019 ; Dong et al. , 2020 ; Sanchez et al. , 2020 ; Eden et al. , 2021 ) . See Mitzenmacher & Vassilvitskii ( 2020 ) for an overview and applications . Clustering with additional information . There have been numerous works that study clustering in a semi-supervised setting where extra information is given . Basu et al . ( 2004 ) gave an active learning framework of clustering with “ must-link ” / “ can not -link ” constraints , where an algorithm is allowed to interact with a predictor that determines if two points must or can not belong to the same cluster . Their objective function is different than that of k-means and they do not give theoretical bounds on the quality of their solution . Balcan & Blum ( 2008 ) and Awasthi et al . ( 2017 ) studied an interactive framework for clustering , where a predictor interactively provides feedback about whether or not to split a current cluster or merge two clusters . Vikram & Dasgupta ( 2016 ) also worked with an interactive oracle but for the Bayesian hierarchical clustering problem . These works differ from ours in their assumptions since their predictors must answer different questions about partitions of the input points . In contrast , Howe ( 2017 ) used logistic regression to aid k-means clustering but do not give any theoretical guarantees . 1The notation Õ hides logarithmic factors . The framework closest in spirit to ours is the semi-supervised active clustering framework ( SSAC ) introduced by Ashtiani et al . ( 2016 ) and further studied by Kim & Ghosh ( 2017 ) ; Mazumdar & Saha ( 2017 ) ; Gamlath et al . ( 2018 ) ; Ailon et al . ( 2018 ) ; Chien et al . ( 2018 ) ; Huleihel et al . ( 2019 ) . The goal of this framework is also to produce a ( 1 + α ) -approximate clustering while minimizing the number of queries to a predictor that instead answers queries of the form “ same-cluster ( u , v ) ” , which returns 1 if points u , v ∈ P are in the same cluster in a particular optimal clustering and 0 otherwise . Our work differs from the SSAC framework in terms of both runtime guarantees , techniques used , and model assumptions , as detailed below . We briefly compare to the most relevant works in the SSAC framework , which are Ailon et al . ( 2018 ) and Mazumdar & Saha ( 2017 ) . First , the runtime of Ailon et al . ( 2018 ) is O ( ndk9/α4 ) even for a perfectly accurate predictor , while the algorithm of Mazumdar & Saha ( 2017 ) uses O ( nk2 ) queries and runtime Õ ( ndk2 ) . By comparison , we use significantly fewer queries , with near linear runtime Õ ( nd ) even for an erroneous predictor . Moreover , a predictor of Mazumdar & Saha ( 2017 ) independently fails each query with probability p so that repeating with pairs containing the same point can determine the correct label of a point whereas our oracle will always repeatedly fail with the same query , so that repeated queries do not help . The SSAC framework uses the predictor to perform importance sampling to obtain a sufficient number of points from each cluster whereas we use techniques from robust mean estimation , dimensionality reduction , and approximate nearest neighbor data structures . Moreover , it is unclear how the SSAC predictor can be implemented in practice to handle adversarial corruptions . One may consider simulating the SSAC predictor using information from individual points by simply checking if the labels of the two input points are the same . However , if a particular input is mislabeled , then all of the pairs containing this input can also be reported incorrectly , which violates their independent noise assumption . Finally , the noisy predictor algorithm in Ailon et al . ( 2018 ) invokes a step of recovering a hidden clique in a stochastic block model , making it prohibitively costly to implement . Lastly , in the SSAC framework , datasets need to be specifically created to fit into their model since one requires pairwise information . In contrast , our predictor requires information about individual points , which can be learned from either a training dataset , from past similar datasets , or from another approximate or heuristic clustering and is able to handle adversarial corruptions . Thus , we obtain significantly faster algorithms while using an arguably more realistic predictor . Approximation stability . Another approach to overcome the NP-hardness of approximation for k-means clustering is the assumption that the underlying dataset follows certain distributional properties . Introduced by Balcan et al . ( 2013 ) , the notion of ( c , α ) -approximate stability ( Agarwal et al. , 2015 ; Awasthi et al. , 2019 ; Balcan et al. , 2020 ) requires that every c-approximation is α-close to the optimal solution in terms of the fraction of incorrectly clustered points . In contrast , we allow inputs so that an arbitrarily small fraction of incorrectly clustered points can induce arbitrarily bad approximations , as previously discussed , e.g. , in Section 1.1 . 2 LEARNING-AUGMENTED k-MEANS ALGORITHM Preliminaries . We use [ n ] to denote the set { 1 , . . . , n } . Given the set of cluster centers C , we can partition the input points P into k clusters { C1 , . . . , Ck } according to the closest center to each point . If a point is grouped in Ci in the clustering , we refer to its label as i . Note that labels can be arbitrarily permuted as long as the labeling across the points of each cluster is consistent . It is well-known that in k-means clustering , the i-th center is given by the coordinate-wise mean of the points in Ci . Given x ∈ Rd and a set C ⊂ Rd , we define d ( x , C ) = minc∈C ‖x − c‖2 . Note that there may be many approximately optimal clusterings but we consider a fixed one for our analysis .

Summary: This work aims to find a solution of the k-means clustering problem based on (prior) predicted labels. Here the predicted labels could be obtained from some clustering algorithms or some supervised models, with additional noise. A polynomial time procedure (Algorithm 1) is proposed as follows: for each fitted component from the predicted labels, a robust mean is computed (in a coordinate wise way). These robust means are output as the final clustering solution. The running time is $O(knd\log n)$. The paper rigorously establishes a theoretical guarantee on the approximation ratio of the solution assuming the predictor label has a bounded label error. To further improve the running time, the authors utilize dimension reduction technique to cluster $O(k/\alpha)$ points in dimension $O(\log n)$ and then obtain the label for original data points with approximate nearest neighbor data structure. This modified approach (Algorithm 3) has running time $O(nd\log n+\text{poly}(k,\log n))$ and attains a solution with a similar theoretical guarantee. To empirically evaluate the proposed method, the authors perform experiments on synthetic data and a few real datasets. The experiments demonstrate that the proposed method with $k$-means++ initialization achieves better performance than $k$-means++; moreover, the performance is competitive and robust even when the predictor labels are corrupted.