FLS: A New Local Search Algorithm for K-means with Smaller Search Space

1 Introduction

Clustering problems have received much attention over the past decades. The goal of clustering is to partition the given dataset into clusters such that the points within the same cluster are more similar to each other, while the points in different clusters have less similarity. Among different objective functions, k-means is one of the most widely used, which aims to minimize the sum of the squared distances between points and their closest clustering centers.

Since the k-means problem is NP-hard, heuristic and approximation algorithms are two of the most widely used methods. The most popular heuristic algorithm solving this problem is the well-known Lloyd's algorithm [Lloyd, 1982]. Though depending on the initialization of center set selection and lacking theoretical guarantee, Lloyd's algorithm performs well in practice due to its fast speed and simplicity. By using k-means++ seeding technique [Arthur and Vassilvitskii, 2007] as the initial center set, empirical performance of Lloyd's algorithm becomes more stable, since k-means++ provides O(log k)-approximate solution in expectation.

∗Corresponding author.

For approximation algorithms, local search is one of the most popular methods. The current best approximation ratio for the k-means problem based on local search is 9 + ϵ given in [Kanungo et al., 2004]. However, the algorithm in [Kanungo et al., 2004] needs to enumerate O(nk) swap pairs in each round. Even with single swap strategy, it is difficult to handle large datasets. Lattanzi and Sohler [Lattanzi and Sohler, 2019] proposed LS++ algorithm that combines k-means++ with local search to further improve the performance of local search method. The basic idea of LS++ is to iteratively sample centers for swap using D2 -sampling method. They proved that LS++ can return a constant factor approximate solution in expectation after O(k log log k) rounds of sampling and swap. The total running time of LS++ is O(ndk2 log log k). Choo et al. [Choo et al., 2020] proved that O(k) rounds of sampling and swap are enough for getting a constant approximate solution. Although LS++ provides a way to reduce the search space of local search, it is easy for LS++ to fall into local optimum while there is still a considerable gap between the local optimum and the minimum clustering cost. For LS++, 100000k log log k rounds of sampling and swap return a 509-approximate solution [Lattanzi and Sohler, 2019].

In this paper, we propose a new local search algorithm for the k-means problem, which aims to achieve good solutions with smaller search space. Meryerson et al. [Meyerson et al., 2004] proposed β-average data distribution assumption, in which each optimal cluster has size at least βn k for some constant β ∈ (0, 1]. We observe that k-means benchmarks [Fr¨anti and Sieranoja, 2018] and many real-world datasets 1 follow the β-average data distribution. Based on β-average data distribution, we develop new searching strategies to reduce the search space. Firstly, we analyze the local structure of the current solution, and use nearest neighbor search for local adjustment. By single swap strategy, O(k) swap pairs are constructed for local optimization. Secondly, random sampling is used to jump out of the local optimum of nearest neighbor search. With β-average assumption, clustering cost can be improved with high probability using uniform and random sampling. Thirdly, once the solution is trapped when nearest neighbor search and random sampling both fail, randommutation is performed to restart the searching process to find better solutions until convergence. The main contributions of this paper are summarized as follows.

1 https://archive.ics.uci.edu/ml/index.php

• We introduce a fast and practical local search algorithm, named FLS, which reduces the search space of swap pairs from O(nk) to O(k2 ) while maintaining the clustering quality using modified neighbor search techniques and random mutation. The total running time of our algorithm is O(ndk3 log 1ϵ ), where ϵ is a parameter to control approximation ratio. Smaller search space provides great potential for local search acceleration in practice.
• Under "average" cluster size assumption [Meyerson et al., 2004; Ding and Huang, 2021] that each optimal cluster has size at least βn k for some constant β ∈ (0, 1], we prove that our algorithm has approximation ratio (100 + ϵ) in expectation.
• FLS scales well on large datasets and makes improvements on time and clustering cost compared with stateof-art local search methods, especially on those "hard instances" defined in the k-means benchmarks [Fr¨anti and Sieranoja, 2018].

1.1 Other RelatedWorks

Arya et al. [Arya et al., 2004] proposed the first local search heuristic for the k-median problem and the facility location problem. They proved that local search method can achieve 3 + ϵ approximation in polynomial time for the k-median problem. Gupta and Tangwongsan [Gupta and Tangwongsan, 2008] gave a simpler analysis for the k-median problem. As for the k-means problem, Kanungo et al. [Kanungo et al., 2004] proposed the first local search algorithm with ratio (9 + ϵ). Friggstad et al. [Friggstad et al., 2019] proved that local search yields a PTAS for the k-means problem in doubling metrics. Local search strategy can be extended to different variants of clustering problems, including clustering with outliers, clustering with penalties, individual fair kclustering, capacitated clustering, etc. [Angel et al., 2015; Zhang et al., 2019; Mahabadi and Vakilian, 2020; Gupta et al., 2017].

2 Preliminaries

Throughout the paper, we use P = {p1, p2, ..., pn} and k to denote the given data set and the number of clusters, respectively. We use C ∗ to denote an optimal clustering center set for a given k-means instance (P, k). For two points p, q, let d(p, q) be their squared distance. For a point p and a center set C, we also use d(p,C) = minc∈C d(p, c) to denote the squared distance from p to the closest center of p in C. Given a center set C, we use ∆(P,C) = ∑p∈P d(p,C) to denote the clustering cost induced by C. The goal of clustering is to find a set C of k centers such that the following objective function is minimized: ∆(P, C) = ∑p∈P d(p, C).

For each point ch ∈ C, we use Ph to denote the cluster constructed by ch. For each point c∗h ∈ C ∗ , we use P ∗ h to denote the optimal cluster constructed by c∗h . For an optimal center c∗h ∈ C ∗ , we denote sc∗h = argminch∈C d(c∗h , ch) as its closest point in C. For a center ch ∈ C, we denote och = argminc∗h∈C d(ch, c∗h ) as its closest point in C ∗ .

The theoretical guarantee of our algorithm is based on Assumption 1. We note that the proposed data distribution assumption is very reasonable in practice and has been used for clustering analysis in [Ding and Huang, 2021; Meyerson et al., 2004]. Given an instance (P, k) of the kmeans problem, we assume that each optimal cluster should be large enough, i.e. the minimum optimal cluster size is larger than βn k for a given constant β between 0 and 1. Otherwise, for any optimal cluster with cluster size smaller than βn k , as pointed out in [Meyerson et al., 2004], the points in this cluster can be regarded as outliers.

Assumption 1. Given an instance (P, k) of the k-means problem, and a constant β ∈ (0, 1], we say (P, k) is β-average if the smallest optimal cluster has size at least βn k .

3 The Proposed FLS Algorithm

The general idea of our algorithm is as follows. For an instance (P, k) of the k-means problem, initial center set C can be constructed using LS++ seeding [Lattanzi and Sohler, 2019] with O(k) rounds of swaps. Then, we apply nearest neighbor search to find improvements. In each round of nearest neighbor search, we only use the nearest 10 points of each center in C to construct candidate swap pairs and there are only O(k) swap pairs. If one of the swap pairs induces an O( 1k ) reduction on current clustering cost, such swap is conducted. If nearest neighbor search fails to make improvements on clustering cost, it indicates that the solution reaches a local optimum around the center set. In this case, random sampling is used to find a set Q of O(k log 1 η ) points in P for a given parameter η. Swap pairs are constructed between Q and C. If one of the swap pairs induces an O( 1k ) cost reduction, it means that by applying sampling, we can jump out of the local optimum. Under Assumption 1, it can be proved that sampling O(k log 1 η ) points for swap pairs construction makes a swap success with probability at least 1 - η. Once nearest neighbor search and sampling methods both fail, we conduct a random mutation, where each center point is swapped with an arbitrary point in P\C with probability 1k . Thus, one center is changed in expectation in a single mutation round. Parameter R is used to control the total rounds of swaps, and parameter T is the given execution time (measured by seconds) such that FLS can execute at most T seconds. The specific FLS algorithm is given in Algorithm 4.

3.1 Data Structure

In order to conduct nearest neighbor search of the FLS algorithm quickly, we use the selection algorithm [Blum et al., 1973] such that the 10-nearest neighbors of each center in the current center set C of size k can be found in time O(nkd). During the swap of local search, one crucial step is to recalculate the cost after swap when replacing an original center with a new one. For each point p ∈ P, the new distance from p to the modified center set C should be calculated. Thus, during the local search process, we always maintain the closest and the second closest centers to each data point p ∈ P. In this way, the cost change can be calculated more efficiently by checking whether the closest center

Algorithm 1 LS++-Seeding

Input: An instance (P, k) of the k-means problem Output: A center set C of size k

• 1: Initialize center set C using k-means++.
• 2: for i = 1, 2, ...,O(k) do
• 3: Sample a point p ∈ P according to probability d(p,C) ∆(P,C) .
• 4: if ∃ q ∈ C such that ∆(P,C\{q} ∪ {p}) < (1 - 1 100k )∆(P,C) then
• 5: Let q ∈ C be the point such that ∆(P, C\ {q}∪{p}) is minimized.
• 6: C = C\{q} ∪ {p}.
• 7: end if
• 8: end for
• 9: return C.

Algorithm 2 Search-Swap

Input: An instance (P, k) of the k-means problem, a center set C of size k

Output: A center set C1 of size k

• 1: Set C1 = C.
• 2: Search for the 10-nearest neighbors for each c ∈ C. For each center c ∈ C, let Nc denote the neighborhood of c.
• 3: Construct swap pairs set G = {(c, j) : c ∈ C, j ∈ Nc}.
• 4: if ∃ (c, j) ∈ G such that ∆(P,C\{c} ∪ {j}) < (1 - 1 100k )∆(P,C) then
• 5: Let (c, j) be the swap pair such that ∆(P\{c} ∪ {j}) is minimized.
• 6: C1 = C\{c} ∪ {j}.
• 7: end if
• 8: return C1.

is swapped out. In each iteration, it takes time O(ndk) to find the closest and the second closest centers of each data point, since there may exist some heavy clusters with size Ω(n). Thus, LS++-Seeding and Search-Swap algorithms take time O(ndk2 ) and O(ndk), respectively. Sampling-Swap takes time O(ndk3 log 1ϵ ), and random mutation takes time O(k). Since λ, η, R are all constants, the total running time of our algorithm is O(ndk3 log 1ϵ ).

3.2 Analysis

Now we will give the analysis of Algorithm 4. Let {P ∗ 1 , ..., P ∗ k } and {c∗1 , ..., c∗k } be the sets of optimal clusters and the corresponding clustering centers, respectively. We define R(j) = {p ∈ P ∗ j : d(p, c∗j ) ≤ 2 ∆(P ∗ j ,{c∗j}) |P ∗ j | } for each optimal cluster P ∗ j as the set of points close to its clustering center c∗j . We show that R(j) takes a large fraction of P ∗ j .

Lemma 1. Let P ∗ j be an arbitrary optimal cluster and T(j) = {p ∈ P ∗ j : d(p, c∗j ) ≤ α ∆(P ∗ j ,{c∗j}) |P ∗ j | } be the set of points within P ∗ j that are close to the optimal center c∗j , then it holds that |T(j)| ≥ (1- 1α )|P ∗ j |

Algorithm 3 Sampling-Swap

Input: An instance (P, k) of the k-means problem, a center set C of size k, parameters ϵ, λ and η

Output: A center set C1 of size k

• 1: Set C1 = C.
• 2: for i = 1, 2, ...,O( k 1-η log 1ϵ ) do
• 3: Randomly and uniformly sample a set S ⊆ P with size kλ log 1 η .
• 4: Construct candidate swap pairs set G = {(c, j) : c ∈ C, j ∈ S}.
• 5: if ∃ (c, j) ∈ G such that ∆(P, C\ {c} ∪ {j}) < (1 - 1 100k )∆(P,C) then
• 6: Let (c, j) be the swap pair such that ∆(P\{c} ∪ {j}) is minimized.
• 7: C1 = C\{c} ∪ {j}.
• 8: return C1.
• 9: end if
• 10: end for
• 11: return C1.

Algorithm 4 FLS

Input: An instance (P, k) of k-means, parameters ϵ, λ, η, R, T .

Output: A center set C1 of size k

• 1: Initialize C by calling algorithm LS++-Seeding and set C1 = C.
• 2: Set r = 0.
• 3: while r < R or execution time does not exceed T do
• 4: Set r = r +1.
• 5: if Search-Swap induces clustering cost reduction then 6: C = Search-Swap(P, k,C).
• 7: else if Sampling-Swap induces clustering cost reduction then
• 8: C = Sampling-Swap(P, k,C, ϵ, λ, η)
• 9: else
• 10: Let C1 be the center set with minimum clustering cost between C and C1.
• 11: for each c ∈ C do
• 12: For each c ∈ C, with probability 1k , replace c with an arbitrary point in P\C.
• 13: end for
• 14: end if
• 15: end while
• 16: return C1.

Proof.

For nearest neighbor search process in Algorithm 2, it seeks to locally improve the solution round by round. Such heuristic is fast and efficient in practice since the searching process only guarantees local optimum. In this searching step, the algorithm focuses more on finding stable local solutions instead of picking points that make clustering cost reduced sharply. Once nearest neighbor search gets stuck, clustering structure needs to be adjusted greatly, where we move to the random sampling part.

Assume that (P, k) is an instance of the k-means problem. Let C denote the current center set and C ∗ denote a set of optimal centers. Follow the work of Lanttanzi and Sohler [Lattanzi and Sohler, 2019], we give the definition of good cluster with respect to C ∗ as follows.

Definition 1. A cluster P ∗ h is called good, if there exists a pair of points (cj, c∗h ) such that c∗h ∈ C ∗ and cj ∈ C and ∆(P ∗ h , C)-Γ(P, C, cj , c∗h )-9∆(P ∗ h , {c∗h }) > 1 100k ∆(P, C), where Γ(P, C, cj , c∗h ) = ∆(P \P ∗ h , C\ {cj }) - ∆(P \P ∗ h , C) represents the reassignment cost after swapping cj out.

The above definition estimates the cost change of replacing cj with a point close to c∗h . Then, we want to show that with good probability, we can sample a point q close to c∗h for swap to reduce the current cost by a factor of O( 1k ). This prevents the algorithm from getting stuck too early.

For each center c∗h ∈ C ∗ , let sc∗h denote its closest center in C. For simplicity, we say that center c∗h is captured by sc∗h . For a center cj ∈ C, let Z(cj) be the set of centers in C ∗ captured by cj. If |Z(cj)| = 0, we say that cj is a lonely center. Let c∗h be an optimal center such that sc∗h = cj. If |Z(cj)| = 1, then (cj, c∗h ) forms a matched swap pair. If |Z(cj)| > 1, for each c∗h ∈ Z(cj), we find a lonely center cq ∈ C to make (cq, c∗h ) a matched swap pair. During the swap pair construction, each lonely center must be used at most twice. Hence, each optimal center c∗h ∈ C ∗ is matched with a center cj ∈ C such that each cj is used at most twice. For a point p ∈ P, let op and sp denote its closest center in C ∗ and C, respectively. The following lemma gives an upper bound of reassignment cost induced by a matched swap pair.

Lemma 2. Let (cj, c∗h ) be a matched pair. Then Γ(P, C, cj , c∗h ) ≤ (4 + 2λ )∆(Pj , C ∗ ) + 2λ∆(Pj , C).

Proof.

Note that the second and fourth steps follow from the triangle inequality, the third step follows from the fact that sop is the nearest point to op in C, and the fifth step follows from Cauchy Inequality. □

In the following, we analyze that if current clustering cost ∆(P,C) is large enough, there must exist at least one good cluster.

Lemma 3. If ∆(P,C) > 100OPT, then there must exist at least one good cluster induced by some matched pair (cj , c∗h ) Proof. Assume that there is no good cluster. For each c∗h ∈ C ∗ , denote its matched swap pair as (cmh , c∗h ). By the defini-

tion of good cluster and Lemma 2, it holds that

Thus, by summing up all matched pairs, we have

According to Lemma 3, if the local optimum after nearest neighbor search has clustering cost larger than 100OPT, there exists at least one good cluster P ∗ h induced by amatched pair (cj, c∗h ). In the following, we argue that by random sampling O( kβ log 1 η ) points from P for swap pairs construction, with probability at least 1 - η, the clustering cost can be reduced by a factor of O( 1k ).

Lemma 4. With probability at least 1 - η, by sampling O( kβ log 1 η ) points randomly from P , the sampled set contains at least one point fromR(h) for an optimal center c∗h .

Proof. For an optimal cluster P ∗ h , since |P ∗ h | ≥ βn k , by Lemma 1 and the definition of R(h), we know that |R(h)| ≥ βn 2k . Let ζ = |R(h)| |P| . If we randomly sample a set V of points from P, the probability that V contains at least one point from R(h) is at least 1 - (1 - ζ)|R(h)| . If the probability that at least one point from R(h) is sampled is at least 1-η, then 1- (1-ζ)|V | ≥ 1-η. Thus, V has size at least log 1η log 1 1-ζ ≤ 1 ζ log 1 η . Since ζ = |R(h)| |P| ≥ β 2k , if |V | ≥ 2k β log 1 η , V must contain at least one point fromR(h) with probability at least 1- η. □

By Lemma 4, we know that by sampling O( kβ log 1 η ) points from P, there exists at least one point q close to the optimal center c∗h in the sampled set, where (cj, c∗h ) induces a good cluster. By swapping q with cj , we will show that the current clustering cost will be reduced by a factor of O( 1k ). Observe that

where the second inequality follows from the fact that q ∈ R(h) is within distance 2∆(P ∗ h ,{c∗h }) |P ∗ h | to c∗h and triangle inequality, and the last step follows from the definition of good cluster. Hence, following the results in [Rozhoˇn, 2020], by repeating the sampling process O(k log 1ϵ ) times, clustering cost can be reduced to (100OPT + ϵ) in expectation.

4 Experiments

In this section, we compare our algorithm with other local search based methods for the k-means problem.

Datasets. We use the k-means benchmarks [Fr¨anti and Sieranoja, 2018] and other datasets from UCI Machine Learning Repository 2 to conduct our experiments. The datasets used in our experiments are summarized in Table 1. For any instance of the k-means problem, the β value in Assumption 1 for an instance is defined as the ratio between the minimum optimal cluster size and the average cluster size nk using ground truth provided by benchmarks. It is easy to see that as long as the β value of given instance is between 0 and 1, then the instance satisfies Assumption 1. If there is no provided ground truth, we use linear programming solver Gurobi to calculate an optimal solution. Based on those "hard instances" defined in [Fr¨anti and Sieranoja, 2018], we test whether different local search algorithms can find good solutions. Fr¨anti and Sieranoja [Fr¨anti and Sieranoja, 2018] tested different k-means datasets, and defined "hard instances" as those with high separation, large number of clusters or unbalanced distribution. "Hard instances" include the instances in datasets A2, A3, Dim32, Birch2 and Unbalance. We also compare the scalability of different local search algorithms on large datasets such as Covertype and SUSY, which contains 581012 and 5000000 points, respectively.

2 https://archive.ics.uci.edu/ml/index.php

Algorithms and Parameters. In our experiment, we consider three algorithms: our algorithm described in Algorithm 4, LS++ proposed in [Lattanzi and Sohler, 2019] and LS proposed in [Kanungo et al., 2004], which are summarized as follows.

• LS. This is the local search method proposed by Kanungo et al. [Kanungo et al., 2004], which searches swap pairs on whole given dataset and performs swaps iteratively until convergence.
• LS++. This is the algorithm proposed by Lattanzi and Sohler [Lattanzi and Sohler, 2019], which uses D2 - sampling in each round to sample a point and construct a candidate swap pair set of size k.
• FLS. This is our algorithm described in Algorithm 4. Parameters λ, η are set to be 0.5 and 0.5, respectively. For the sampling process in Algorithm 3, we only repeat the sampling for 3 times.

Experimental Setup. We run all the above algorithms on each dataset 5 times and give the average results. All the algorithms use the same initial center set obtained by 25 rounds of LS++. During the local search process, we also study the parallelization of search and swap. In each round of search and swap, the constructed swap pairs can be parallelized for calculation using multiple cores and the one with the maximum cost reduction is taken as the final swap pair. LS requires high computing resources since it constructs O(nk) swap pairs in each round while the upper bound of parallelization for LS++ is k cores. Our algorithm constructs O(k2 ) swap pairs which can make full use of computing resources in practice. In experiments, we only use k cores for parallelization for the sake of fairness since LS++ only contains k swap pairs in each round. For hardware, we use machines with 72 Intel Xeon Gold 6230 CPUs.

Methodology. In our experiments, we run each algorithm for 5 times with fixed execution time and give the gaps between costs of solutions and corresponding optimal clustering costs. Gap is defined as the clustering cost relative to the optimal clustering cost as gap = ∑5i=1 15 Si-OPT OPT , where Si is the clustering cost of solutions returned by different algorithms in the i-th experiment. For large datasets, since optimal solutions cannot be found by solver, we use the solution with the minimum cost among LS, LS++ and FLS as the ground truth solution, and compute the differences between the ground truth solution and other solutions. To compare the running time, we give the time needed to reach different gaps for each algorithm in Table 3.

4.1 Results

Table 2 compares the clustering costs of LS, LS++ and FLS at several different time points. On small datasets, our algorithm finds optimal solutions on two of the datasets, while LS++ fails to find any optimal solution. LS++ performs well at the beginning. However, for larger execution time, LS++ is easy to get stuck in local optimum. It is more and more difficult for LS++ to find improvements as the gap between clustering cost and optimal clustering cost narrows. As for LS, clustering costs hardly change since LS constructs O(nk) swap pairs for enumeration in each round, and it is impossible for LS to find good solutions in fixed execution time. Solutions returned by our algorithm are very close to the optimal clustering costs on small datasets, and there is only a gap ranging from 0% to 0.88%. On "hard instances", our algorithm improves the performance of LS++ from 0.22% to 5.73% on clustering costs. On large datasets, our algorithm scales well and achieves the minimum final clustering costs, which improves the performance of LS++ from 1.03% to 3.82% on clustering costs.

Table 3 summarizes the time spent to reach a fixed clustering cost. FLS finds four datasets with gap 0% while LS++ fails to find any of them. When the quality of solution is high, for example the gap to optimal solutions ranges from 0% to 2%, FLS can find high-quality solutions quickly. On the contrary, LS++ fails to find high-quality solutions on most datasets. It takes LS++ nearly forty and eight times, compared to ours, to find solutions with gap 1% on datasets Yeast and Dim32, respectively. From Table 3, we can get that our algorithm FLS is faster to find near-optimal solutions (solutions with small gaps to optimal solutions) than LS++.

5 Conclusion

In this paper, we propose a fast local search algorithm for kmeans. By developing new searching techniques and random mutation, the search space of local search in each round is reduced greatly while the quality of solution is maintained. Under "average" assumption of data distribution in practice, we prove that our algorithm returns a constant factor approximate solution in expectation. Finally, we show experimentally that our algorithm achieves better performance on both small and large datasets compared with other local search methods.

Acknowledgments

This work was supported by National Natural Science Foundation of China (Grant Nos. 62172446, 61872450), 111 Project (Grant No. B18059).