Fast Local Search Algorithms for Clustering with Adaptive Sampling and Bandit Strategies

We thank the reviewer for the positive rating and thoughtful comments. Below, we address the concerns.

Question 1: How reasonable are the sub-Gaussian assumptions for Theorems 1 and 2?

Response: We thank the reviewer for this important question. For the $k$-means objective, it is based on squared distances. This makes our assumption-free weighted sampling strategy (originally designed for $k$-median) inapplicable, as it can lead to significant approximation errors in estimating swap gains. To address this issue, we formulate the best swap pair selection as a bandit problem. In the bandit setting, the sub-Gaussian assumption on arm gains is a standard and widely adopted condition for establishing theoretical guarantees (e.g., [1], [2] and [6], [30], [31] in our ref list). Thus, the sub-Gaussian assumption used in Theorems 1 and 2 is both reasonable and well supported by prior work.

Moreover, previous assumptions for bandit-based methods are much stronger than ours. In [6], [30], and [31] (see our ref list), it is assumed that each arm’s gain follows a fixed distribution with mean $\gamma$ and variance 1, or that the variance-to-maximum-gain ratio is bounded by a constant. In contrast, we only require the variance to be bounded by a function of the average clustering cost, making our assumption much weaker than previous ones.

To empirically verify the sub-Gaussian assumption, we use the fact that any variable with range $2\sigma$ is $\sigma$-sub-Gaussian. Thus, we assess sub-Gaussianity by computing the ratio between the gain range and the average clustering cost $\delta^2(P, C)/n$, where the gain of a point $p$ under swap pair $(u, v)$ and a center set $C$ is defined as $\delta^2(p, C \setminus u \cup v) - \delta^2(p, C)$. We run our algorithm 5 times with 400 local search steps and calculate the ratios. In Tables 1 and 2, we give the mean, variance, and maximum of these ratios. The results show that the ratio remains below 50 on most datasets except for dataset SYN and HCV_L. This indicates that the sub-Gaussian assumption is both reasonable and empirically justified in practice.

Table 1: Ratio of gain range to average cost on large datasets

Table 2: Ratio of gain range to average cost on small datasets

[1] Baharav T Z, Kang R, Sullivan C, et al. Adaptive sampling for efficient softmax approximation[C]//Proceedings of the 38th International Conference on Neural Information Processing Systems. 2024: 117580-117613.

[2] Tiwari M, Kang R, Lee J Y, et al. Faster maximum inner product search in high dimensions[C]//Proceedings of the 41st International Conference on Machine Learning. 2024: 48344-48361.

Question 2: Does the algorithm scale to high-dimensional datasets? The experiments seems to vary n, k but not d.

Response: We thank the reviewer for this insightful question. Our algorithm consists of two main components: (1) sampling via the Metropolis-Hastings process, and (2) bandit-based swap pair selection. In the Metropolis-Hastings step, samples are drawn uniformly, and only their distances to the current centers are computed, requiring time linear in the dimension $d$. Similarly, the bandit-based selection evaluates swap gains on a small sample of size $\tilde{O}(k^2)$, where each step also involves distance computations with runtime linear in $d$. Therefore, the overall runtime scales linearly with $d$, making the algorithm suitable for high-dimensional datasets.

To empirically verify this, we conducted additional experiments on several high-dimensional datasets from diverse domains. These include: (1) the Twin Gas Sensor Arrays (TGas) dataset with 40,000 points and 2,000 dimensions; (2) the CIFAR-10 image dataset represented as 10,000 flattened 3,072-dimensional vectors; and (3) the Electricity Load Diagrams (ELD) dataset with 370 instances and 140,256 dimensions. These datasets were selected to cover a wide range of $d$ values and application domains.

Tables 3-5 give the comparison results on these datasets. It can be seen that our algorithm consistently scales well in all cases, achieving an average 10x speedups over baseline local search methods. Moreover, it maintains comparable or even better clustering quality despite the increase in dimensionality. These findings demonstrate that our method is robust not only to data size $n$ and number of clusters $k$, but also to the dimensionality $d$.

We will include these additional experimental results and corresponding runtime analysis in the appendix of the revised manuscript.

Table 3: Comparison results on dataset TGas

Table 3: Comparison results on dataset CIFAR10

Table 3: Comparison results on dataset ELD

We thank the reviewer for the thoughtful comments and constructive feedbacks. Below, we address the concerns.

Weakness 1: The Sub-Gaussian Prior assumption seems not mild and difficult to verify in practice. Generally, the assumptions of an algorithm are expected to depend on the dataset and be easy to verify. However, this assumption concerns the “swap pairs” generated during the algorithm's execution, making it difficult to determine whether a given dataset satisfies the assumption in practice.

Response: We thank the reviewer for raising this important concern. Unlike the $k$-median problem with $\ell_1$-norm objective, our weighted sampling strategy (which requires no assumptions) cannot be applied to $k$-means due to approximation loss for swap gain estimation under squared distances. To overcome this issue, we model the swap pair selection process as a bandit problem. In this setting, the sub-Gaussianity assumption is both necessary for theoretical guarantees and commonly adopted in prior work (e.g., [6], [30], [31] in our ref list).

We refer to our sub-Gaussian Prior assumption as mild because it is much weaker than those used in prior work. Existing bandit-based clustering methods (e.g., [6], [30], [31]) typically assume that the gain of each swap pair follows a fixed distribution with mean $\gamma$ and variance $1$, or that the ratio between the variance and the maximum gain is bounded by a constant. In contrast, we only require that the variance is bounded by a function of the average cost, making it a milder assumption than previous ones.

To verify our assumptions empirically, we conduct the following experiments. We use the fact that any variable with range $2\sigma$ is $\sigma$-sub-Gaussian. Thus, the ratio between the gain range and the average clustering cost can be used to measure whether the sub-Gaussian assumption is satisfied. For each swap pair $(u,v)$ and center set $C$, we define the gain for a point $p$ as $\delta^2(p, C \setminus u \cup v) - \delta^2(p, C)$. Then the ratio can be calculated as the gain divide by $\delta^2(P, C)/n$. For each dataset, we run our algorithm 5 times with 400 local search steps and record the gain-range-to-average-cost ratio for each swap pair. Tables 1-2 give the mean, variance, and maximum of these ratios. The results show that the maximum ratio is typically bounded by a constant (less than 50) across most datasets, supporting the validity of the sub-Gaussian assumption. Even on datasets like SYN, where the assumption may not strictly hold, our algorithm still achieves over 20× speedup with competitive clustering quality.

Table 1: Ratio of gain range to the average cost on large datasets

Table 2: Ratio of gain range to the average cost on small datasets

Weakness 2: As mentioned in the introduction of this paper, there exist coreset-based algorithms achieving running time of $\tilde{O}(nd+k^3d)$, which is better than the new algorithm theoretically. This weakens the contribution.

Response: We thank the reviewer for raising this concern. While it is true that some coreset-based algorithms for $k$-means achieve a theoretical runtime of $\tilde{O}(nd + k^3d)$ under no distribution assumptions, this may overlook the practical overhead involved in coreset construction—particularly for large-scale datasets. Experimental results indicate that there is a sharp increase on the coreset construction time as data sizes grow (see Tables 5 and 6 in our appendix). Moreover, coreset-based methods often produce worse clustering quality than algorithms that directly optimize the original objective.

In contrast, our method bypasses the data compression step required by coreset construction and achieves comparable time complexity with much better empirical performance. As demonstrated in our experiments (see our Appendix), it delivers up to 80× speedup over coreset-based methods on large-scale datasets such as SIFT (100M points), while achieving better clustering quality. Therefore, even in the presence of coreset methods with similar time complexity, our approach contributes a more practical and scalable solution with theoretical guarantees.

Weakness 3 and Question 1: Could the authors provide a comparison between their new algorithm and a faster coreset-based approach over one million data points? For instance, a possible method is to apply uniform sampling to obtain a weak coreset and run -Means++ on the coreset. If additional experimental results could demonstrate that the proposed algorithm still outperforms such coreset-based methods, I think it will increase the contribution.

Response: We thank the reviewer for the constructive feedbacks. We agree with the reviewer that an unweighted weak coreset can be constructed efficiently via uniform sampling when the dataset exhibits the balance property. While this approach offers excellent runtime performance, it often results in worse clustering qualities. As it targets a different trade-off between accuracy and efficiency, we did not include it in our original comparison.

To address this concern, we conducted additional experiments on large-scale datasets such as USC1990 (1M), HIGGS (11M), and SIFT (100M), where we apply uniform sampling (with sizes ranging from 1,000 to 1,000,000) to construct the coreset followed by $k$-means++ to obtain the clustering solutions. The results (Tables 3-5) show that even with 1 million samples, our method consistently achieves much better clustering quality than uniform coresets while maintaining comparable runtime. On average, our method reduces the clustering cost by 5% compared to the uniform coreset baseline.

We will include these results in the revised version to provide a more complete comparison.

Table 3: Comparisons with weak coreset on USC1990

Table 4: Comparisons with weak coreset on HIGGS

Table 4: Comparisons with weak coreset on SIFT

We thank the reviewer for the positive rating and thoughtful comments. Below, we address the concerns.

Question 1: Your work is important when k is large, but in that case why is it important to have exactly k clusters, and not employ bicriteria approximation, which allows $O(k)$ clusters, or even uniform facility location where the number of clusters is controlled indirectly via a penalty term?

Response: We thank the reviewer for raising this insightful question. In the standard $k$-means objective, enforcing exactly $k$ centers is crucial for fair comparisons against prior methods. By fixing $k$, it allows for more meaningful comparisons of theoretical results and clearer assessments of algorithmic effectiveness.

We agree with the reviewer that allowing more than $k$ centers to be opened can reduce the clustering cost and improve approximation ratios. Bicriteria approximation is therefore a practical and effective alternative in real-world scenarios. Extending our framework to obtain bicriteria solutions and achieve better practical performance is an interesting direction for further studies.

As for the uniform facility location problem, to the best of our knowledge, existing local search algorithms do not achieve linear running time. This is mainly because the opening costs may affect the sampling process and complicate the design of sampling distributions. Developing local search methods with linear runtime for facility location remains an interesting open problem for future work.

Question 2: In Lemma 1, can you explain how this bound on the expectation of the tree distance $\delta_T$, but not on the expectation of its square, is useful for k-means, which corresponds to squared distances? In addition, why do you use [12] and not the well-known tree embedding that corresponds to a randomly shifted quadtree, and has stretch bound $O(d\log\Delta)$ instead of $O(d)$ in Lemma 1?

Response: We appreciate the reviewer’s insightful question.

For the $k$-means problem, our rejection sampling relies on a tree embedding to approximate the squared-distance-based sampling distribution. A key challenge is that directly bounding the expectation of squared tree distances through tree metric leads to large distortion. Prior work (e.g., [1]) shows that the expected squared distance distortion under tree embeddings can be as large as $\Omega(n)$, which adds an $O(n)$ factor loss in sampling probabilities and runtime.

To address this, we instead bound the expected tree distance under the $\ell_1$ objective, as stated in Lemma 1. Then, by carefully relating $\ell_1$ objective to squared distances through the Cauchy–Schwarz Inequality, we reduce the distortion in sampling probabilities from $O(n)$ to $\tilde{O}(\sqrt{n})$.

Regarding the quadtree embedding, the method in [12] is indeed based on the well-known quadtree embedding technique for $\ell_1$ distances, which has a distortion of $O(d \log \Delta)$. We apologize for the typo in Lemma 1, where the distortion should be $\tilde{O}(d)$ instead of $O(d)$ to reflect the hidden logarithmic dependence.

To ensure that the $\log \Delta$ factor does not affect the overall runtime, we include a preprocessing step (lines 198–199 in our paper) that rescales the dataset to reduce the aspect ratio $\Delta$ to $\mathrm{poly}(n, d)$. This step runs in $\tilde{O}(nd)$ time and guarantees that $\log \Delta = \tilde{O}(1)$. As a result, the distortion becomes $\tilde{O}(d)$, which only contributes an $\tilde{O}(1)$ overhead and does not change the overall time complexity of our algorithm. Thus, the quadtree embedding remains compatible with our rejection sampling framework.

We will clarify this in the revised version.

[1] Cohen-Addad et al. Streaming euclidean k-median and k-means with o(log n) space. Proc. FOCS 2023, pp. 883-908.

Question 3：In Lemma 4, what is the intuition for the term $O(\sqrt{n})$, I mean where it is coming from? In addition, this lemma shows that each point $x$ is sampled almost with the right probability, up to factor $O(\sqrt{n})$; can we use such samples and rejection sampling to obtain one sample with a precisely correct probability? Doesn't this require computing many distances in $R^d$ (not in the tree) which would take too much time?

Response: We thank the reviewer for the thoughtful question.

The $O(\sqrt{n})$ factor loss arises from approximating squared distances using $\ell_1$ objective in the sampling probability. For our rejection sampling process, to avoid the $O(n)$ factor loss caused by squared distance embedding, the $\ell_1$ distances are used. However, we still need to relate the $\ell_1$-based sampling distribution to the one suitable for squared distances. This introduces an $O(\sqrt{n})$ factor loss through Cauchy–Schwarz inequality.

For the sampling process, the samples drawn from rejection sampling cannot recover the precise sampling probabilities. However, our goal is to ensure that we can sample a “good” point for local search swaps with constant probability. For these points, it is observed that the clustering cost of them account for a certain fraction of the total clustering cost (as shown in Lemma B.3). Since our sampling probabilities are within an $O(1/\sqrt{n})$ factor, by repeating the rejection sampling process for $\tilde{O}(\sqrt{n})$ rounds, it suffices to select a good point with constant probability.

As for computation cost, each round of rejection sampling only requires computing the true squared distance from a single candidate point to its nearest center in $\mathbb{R}^d$. Thus, the total cost over $\tilde{O}(\sqrt{n})$ rounds is $\tilde{O}(\sqrt{n}kd)$, which does not significantly impact the overall runtime.

Question 4: In assumption 2, what does it mean that $P$ is average"? In addition, this assumption sounds very strong, in the sense of making the problem much easier. For example, sampling the points uniformly at rate $\log(n)/k^2$ will reduce the input size by factor almost $k^2$ and the optimal clusters/centers would probably (I have not verified it) change only a little bit. Why do you think this assumption is reasonable? Has it been used before?

Response: We thank the reviewer for this thoughtful question. In Assumption 2, $P$ is “average” means that each optimal cluster has size $\Omega(k^2)$, which implies the clustering instance is not heavily imbalanced.

This type of assumption is commonly used in the clustering literature. Prior works on clustering algorithm design (e.g., [2]–[4]) often assume that each optimal cluster has size $\Omega(n/k)$ to ensure meaningful sampling guarantees. Our assumption, which requires each cluster to have size $\Omega(k^2)$, is consistent with these when $k^3 = O(n)$. This is a reasonable setting in practice, as it typically holds that $n$ is much larger than $k$.

Regarding the reviewer’s observation about uniform sampling at rate $\log(n)/k^2$, we agree with the reviewer that this helps to reduce the input size to $\tilde{O}(n/k^2)$ while guaranteeing the theoretical loss (e.g., [2]). However, this approach often leads to worse clustering quality in practice.

To evaluate this strategy, we conducted additional experiments comparing our method with a baseline that uniformly samples $\beta$ points (ranging from 1,000 to 1,000,000) and applies a fast local search algorithm (LS++) to the subsample. Tables 1–3 give the comparison on large-scale datasets including SUSY (5M) and SIFT (100M). The results show that while using small samples (e.g., 1,000–10,000) yields faster runtimes than ours, it leads to at least a 5% increase in clustering cost. To match the clustering quality of ours, it requires to sample over 500,000 points. in which case our algorithm runs much faster than the uniform sampling baseline.

[2] Huang et al. The power of uniform sampling for k-median. Proc. ICML 2023, pp. 13933-13956.

[3] Meyerson et al. A k-median algorithm with running time independent of data size. Machine Learning, 2004, 56(1): 61-87.

[4] Huang et al. Bi-criteria Sublinear Time Algorithms for Clustering with Outliers in High Dimensions. Proc. COCOON 2024, pp. 91-103.

Table 1: Comparisons with uniform sampling on SUSY

Table 2: Comparisons with uniform sampling on SIFT

Thanks for the detailed rebuttal. It does not change my overall evaluation. but I appreciate the correction regarding Lemma 1 (a randomly shifted quadtree has an extra O(log n) factor in the expected distortion, which predates [12]) and explaining what is the source of the $\sqrt{n}$ factor.

We thank the reviewer for the thoughtful comments and constructive feedbacks. Below, we address the concerns.

Weakness 1: These assumptions are not standard or justified empirically, and no sensitivity analysis is provided to show how the algorithm performs when these assumptions are violated.

Response: We thank the reviewer for raising this concern. For the $k$-median problem, we obtain theoretical guarantees without assumptions by applying weighted sampling on $\ell_1$-norm objectives. However, this approach cannot be used for $k$-means due to the large approximation loss caused by squared distance definitions. To reduce the dependence on both $n$ and $k$ in the runtime, we apply bandit method for best swap pair selection. As shown in the literature (e.g., [6], [30], [31] in our ref list), to obtain theoretical guarantees in bandit settings, sub-Gaussianity is a standard assumption.

Previous assumptions for bandit methods are much stronger than ours. In [6], [30], and [31], it is assumed that the gain of each arm follows a fixed distribution with mean $\gamma$ and variance 1, or the variance-to-maximum-gain ratio is bounded by a constant. In contrast, we only require the variance to be bounded by a function of the average cost, making our assumption much weaker.

To further improve runtime for $k$-means, we design a practical algorithm (Alg 3) based on the Metropolis-Hastings method. To achieve theoretical guarantees, both exponential tail and non-degeneracy assumptions are required in previous work (e.g., [3] in the ref list). In contrast, our analysis only requires an exponential tail and a cluster size condition, making it much weaker than previous ones.

To validate the assumptions, we conduct the following experiments. For the sub-Gaussianity assumption, we use the fact that a variable with range $2\sigma$ is $\sigma$-sub-Gaussian, and measure the ratio between gain range and average clustering cost $\delta^2(P,C)/n$. The gain for a point $p$ under swap $(u, v)$ and center set $C$ is defined as $\delta^2(p, C \setminus u \cup v) - \delta^2(p, C)$. In Tables 1 and 2, we report the mean, variance, and maximum of the ratios calculated during the local search processes. The results show that this ratio is typically bounded by a constant (< 50) across most datasets, verifying the sub-Gaussianity assumption. Even on datasets like SYN, where the assumption may not hold (with large ratio), our algorithm still achieves over 20× speedup, demonstrating the empirical robustness.

For the exponential tail assumption, we evaluate it using approximate solutions from multi-swap local search, since the optimal clusters are unknown. Similar evaluation based on approximate solutions has also been adopted in prior work such as [1]. We rely on the fact that the log-CCDF (Complementary Cumulative Distribution Function) of an exponential distribution is upper bounded by a linear function of the variable. For each approximate cluster, we: (1) compute the log-CCDF of squared distances, and (2) apply linear regression using slope < 0 and p-value < 0.05 as criteria. As shown in Tables 3-4, 7 out of 12 datasets are likely to satisfy the assumption. All datasets with over 1M points meet the condition, while smaller ones (e.g., Glass) may not. Most datasets also satisfy the cluster size condition, though SYN ($k=100$) may not for its small clusters. However, our algorithm achieves over 20× speedup on datasets violating both sub-Gaussian and cluster size assumptions (e.g., SYN), demonstrating the empirical robustness.

[1] Huang et al. The power of uniform sampling for k-median. Proc. ICML 2023, pp. 13933-13956.

Table 1: Ratio of gain range to average cost on large datasets

Table 2: Ratio of gain range to average cost on small datasets

Table 3: Cluster size and tail tests on large datasets

Table 4: Cluster size and tail tests on small datasets

Weakness 2: Regarding the typos

Response: We apologize for the typos. In the revised version, we will carefully correct them to improve the presentation.

Weakness 3: The runtime comparison is also somewhat unfair, as MLS performs multi-swap operations while LS++ is based on single swaps. There is no clear reason why MLS should be faster than LS++

Response: We thank the reviewer for raising this concern.

Although MLS is a multi-swap method and LS++ is based on single-swap, as shown in prior work (e.g., [7], [21] in the ref list), both of them are compared together to evaluate the trade-offs in runtime and clustering quality.

While single-swap methods are generally expected to be faster than multi-swap methods, MLS is faster than LS++ because it employs a heuristic that estimates swap gains using a small sample of size $O(k \log k)$ instead of the full dataset. In contrast, LS++ evaluates exact gains on the entire dataset, incurring a cost of $O(ndk)$ per local search step. This makes MLS more efficient despite its multi-swap nature.

It can be seen that the sampling step is the key factor behind MLS’s speed advantage. Additional experiments show that when MLS is modified to compute exact gains over the full dataset (denoted as MLSP in Tables 5-7), it becomes much slower than LS++.

We will clarify this in our revised version.

Question 1: Can the authors report the performance of each method with fewer steps, such as 50 or 100, to provide a more realistic comparison of practical efficiency?

Response: We thank the reviewer for this insightful suggestion. In our experiments, 400 local search steps is set to ensure that all methods have sufficient iterations to converge. In the following (Tables 5-7), we conducted additional experiments on large datasets SUSY (5M), HIGGS (11M) and SIFT (100M) using 50 local search steps. The results show that our algorithm still outperforms other local search methods with 50x to 300x speedups and comparable clustering quality.

Table 5: Results on SUSY with 50 LS steps

Table 6: Results on HIGGS with 50 LS steps

Table 7: Results on SIFT with 50 LS steps

Comment: Thank you for the detailed rebuttal. My concern is partially resolved. The supplementary experiments with 50 and 100 iterations provide empirical support for the guess that the cost at 100 iterations is already very close to that at 400 iterations for most of the test instances. However, I still have a remaining concern: the validation of the exponential tail assumption is based on approximate solutions obtained via multi-swap local search. This introduces a form of circular reasoning—using the output of a heuristic algorithm to justify an assumption that is, in turn, necessary for the theoretical analysis of that same algorithm. Such empirical evidence, while suggestive, does not constitute a rigorous justification that the assumption holds for the optimal clusters, which is what the theoretical guarantees ultimately rely on.

Response: We thank the reviewer for raising this important concern. Since computing the optimal $k$-means solution is NP-hard in Euclidean space (even in the plane), it is generally infeasible to obtain the exact optimum for real-world datasets. This is why we rely on high-quality approximate solutions to evaluate the exponential tail and cluster size conditions. We agree with the reviewer that this approach may introduce a form of circular reasoning.

To address this concern, we construct synthetic datasets where the optimal clustering centers are fixed using the following methods. This eliminates potential circular reasoning and enables a direct evaluation of our algorithm's performance under both assumption-satisfying and assumption-violating scenarios.

Assumption-Satisfying Datasets (Tail-1 and Tail-2): Each dataset is generated by selecting cluster centers uniformly at random in range $[-100, 100]^d$ and adding zero-mean Laplace noise with scaling factor fixed as 1 (to ensure exponential tail). To guarantee that the given center is exactly the optimal 1-means center, we shift all points so that their empirical mean matches the predefined center. Additionally, we enforce a minimum separation of 20 units between clusters to guarantee that the constructed partition corresponds to the optimal ones. Tail-1 contains 1,250,000 points with $k=10$, $d=10$ and minimum cluster size $10k^2$. Tail-2 has 125,000 points with the same $k$, $d$, and minimum cluster size $k^2$.

Assumption-Violating Datasets (Non-Tail-1 and Non-Tail-2): These datasets are constructed with overlapping clusters, unbalanced sizes (ranging from 50 to 20,000), and random shifts to misalign the true cluster center with the optimal 1-means center. To violate exponential tails, uniform distributions are used to generate the points in the clusters. Non-Tail-1 has around 1,000,000 points and Non-Tail-2 has around 100,000 points, where both datasets are generated with $k=50$ and $d=10$.

Tables 1–4 present the comparison results. It can be seen that our algorithm consistently achieves comparable or better clustering quality across both types of datasets. On the assumption-satisfying Tail-1 and Tail-2 datasets, it matches or outperforms LS++ (the strongest baseline in clustering quality, while achieving at least 10× speedup. On the more challenging Non-Tail-1 and Non-Tail-2 datasets, which violate the theoretical assumptions, our method maintains similar clustering quality to LS++ with at least 7× speedup. These results demonstrate that our algorithm is not only theoretically grounded but also empirically robust, even when the assumptions are not fully satisfied.

Table 1: Comparisons on assumption-satisfying dataset Tail-1

Table 2: Comparisons on assumption-satisfying dataset Tail-2

Table 3: Comparisons on assumption-violating dataset Non-Tail-1

Table 4: Comparisons on assumption-violating dataset Non-Tail-2