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: The paper studies k-means problem in a learning-augmented setting. Recall that in k-means, a point set $P \subset \mathbb{R}^d$ is given and the goal is to find a set $C$ of $k$ centers, such that $\mathrm{cost}(P,C) = \sum_{p \in P} \min_{c \in C} \|p-c\|_2^2$ is minimized. In addition to the input point set, a predicted solution, which is an approximately optimal clustering of $P$, is also provided in the proposed learning-augmented setting. The algorithm can access this solution by querying the predictor the cluster that a point $x$ belongs to.   The main result is an algorithm that can leverage the predictions. For $\alpha \in (10 \log n / \sqrt{n}, 1/7)$, given access to a predictor with label error rate $\lambda \leq \alpha$, and let $\gamma \geq 1$ be a sufficiently large constant, if the predicted solution satisfies all clusters has at least $\gamma k \log k / \alpha$ points, then the algorithm outputs a $(1+20\alpha)$-approximate solution with probability at least $3/4$, using $O(nd \log n + \mathrm{poly}(k, \log n))$ total time.  Furthermore, the algorithm only uses $\tilde{O}(k / \alpha)$ queries to the predictor for outputting the centers. The authors claim that for any $\delta \in (0,1]$, any algorithm makes $O(k^(1-\delta)/(\alpha \log n))$ queries to the predictor with label rate $\alpha$ cannot output a $(1+C \alpha)$-approximate solution for k-means problem in $2^{O(n^(1-\delta))}$ time, assuming ETH.  The authors also experimented with their algorithm on three datasets, along with three different predictors. The experiment result shows that their algorithm significantly improves the performance when using with predictor or k-means++, and has a stable accuracy against corrupted predictors.   Technically, since it is guaranteed that the predicted solution is a $(1+\alpha)$ approximately optimal solution, with some of the labels corrupted, Algorithm 1 aims to reconstruct centers of each cluster $X$ in $(1+\alpha)$-approximately optimal solution, i.e. find point $c$ such that $\mathrm{cost}(X, c) \leq (1+O(\alpha)) \mathrm{cost}(X, C_X)$. Notice that $d$ dimensions are independent of each other in k-means objective. Hence Algorithm 1 can estimate each coordinate of $c$ independently, which is implemented by Algorithm 2.  Then Algorithm 2 (randomly) divides the input (1D) data points into two parts, and one part is used for the computation of interval $I$, an estimation of the value range of the uncorrupted data points. Then the points outside $I$ are filtered out and the mean of the remaining coordinates are returned as the final estimation.