Mini-Batch Kernel $k$-means

We present the first mini-batch kernel $k$-means algorithm, offering an order of magnitude improvement in running time compared to the full batch algorithm. A single iteration of our algorithm takes $\widetilde{O}(kb^2)$ time, significantly faster than the $O(n^2)$ time required by the full batch kernel $k$-means, where $n$ is the dataset size and $b$ is the batch size. Extensive experiments demonstrate that our algorithm consistently achieves a 10-100x speedup with minimal loss in quality, addressing the slow runtime that has limited kernel $k$-means adoption in practice. We further complement these results with a theoretical analysis under an early stopping condition, proving that with a batch size of $\widetilde{\Omega}(\max \{\gamma^{4}, \gamma^{2}\} \cdot \epsilon^{-2})$, the algorithm terminates in $O(\gamma^2/\epsilon)$ iterations with high probability, where $\gamma$ bounds the norm of points in feature space and $\epsilon$ is a termination threshold. Our analysis holds for any reasonable center initialization, and when using $k$-means++ initialization, the algorithm achieves an approximation ratio of $O(\log k)$ in expectation. For normalized kernels, such as Gaussian or Laplacian it holds that $\gamma=1$. Taking $\epsilon = O(1)$ and $b=\Theta(\log n)$, the algorithm terminates in $O(1)$ iterations, with each iteration running in $\widetilde{O}(k)$ time.

Paper

References (31)

Scroll for more · 19 remaining

Similar papers

Reviewer c6yY6/10 · confidence 3/52024-11-03

Summary

This paper introduces a mini-batch kernel k-means algorithm designed to accelerate the clustering process by operating on small, randomly selected data subsets rather than the entire dataset in each iteration. A detailed theoretical analysis of the per-iteration computational complexity is provided, along with an examination of the algorithm's convergence behavior under early stopping criteria, batch size considerations, and kernel selection. The findings demonstrate that the proposed method achieves accuracy comparable to that of full-batch kernel k-means, particularly when initialized with k-means++. Empirical results across multiple datasets validate the significant performance gains in terms of execution speed.

Strengths

1\. The paper is well-structured and easy to follow. 2\. The work represents a natural extension of mini-batch k-means to its kernelized version, incorporating tailored analysis for infinite-dimensional feature spaces.

Weaknesses

- 1\. Inconsistencies in the Empirical Results:   - For the \`mnist_784\` dataset, the Adjusted Rand Index (ARI) and Normalized Mutual Information (NMI) scores of the mini-batch kernel method surpass those of the full-batch kernel method, which seems counterintuitive. Is there any explanation for this anomaly?   - In the case of the \`letter\` dataset, the truncated \(\beta\)-mini-batch kernel method did not achieve significant acceleration in runtime as shown in Figure 1. Moreover, it appears slower in Figure 9 when \(\tau = 300\). Could you provide an explanation for this inconsistency? - 2\. The theoretical analysis on batch size for termination guarantees is well-constructed. However, given the wide range of batch sizes available for practical use, do you have any recommendations for selecting an appropriate batch size in real-world implementations?

Questions

See the weakness.

Rating

6

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Authorsrebuttal2024-11-14

Response to questions

Thank you for your review. Below are our responses to your questions. 1) Thank you for taking a careful look at the results, however, these are not inconsistencies. - It is common for mini-batch algorithms to outperform full batch algorithms (e.g., stochastic gradient descent outperforms non-stochastic gradient descent). This is because the stochastic nature of the algorithm allows it to escape local minima. Recall that (kernel) k-means only arrives at a local minima, which means noise can lead to improvements. - The runtime improvement of using truncation is more visible with larger datasets. Compared to the other datasets, the ratio of the number of points to the number of classes is smallest for Letter. This means that with even moderately large values of tau such as 300, we end up doing more work per iteration than the untruncated version. 2) A batch size of 1024 seems reasonable across our experiments and aligns with the default batch size for sklearn's implementation of mini-batch k-means. We didn't see any significant improvements with tau above 300 and note that raising it further can have a significant effect on the runtime due to the quadratic dependence on tau.

Reviewer Ae525/10 · confidence 4/52024-11-07

Summary

This paper studies efficient implementations of the kernel k-means problems. In this problem, we are given $n$ datapoints $x_1, \ldots, x_n \in X$ and a kernel function $K : X \times X \to \mathbb{R}$ such that there exists a function $\varphi: X \rightarrow \mathcal{H}$, where $\mathcal{H}$ is a Hilbert space, satisfying $K(x, y) = \langle \varphi(x), \varphi(y)\rangle$ for all $x, y \in X$. The goal is to cluster the set of points $\varphi(x_1), \ldots, \varphi(x_n)$ using the usual $k$-means objective. One key issue here is that the outputs of $\varphi(\cdot)$ can be infinite dimensional. So, we would ideally want to work without ever needing to map the input points into the Hilbert space $\mathcal{H}$, The authors argue that, the $k$-means++ initialization and further iterations of Lloyd's can be implemented without ever needing to apply the map $\varphi(\cdot)$ on the input points. A key issue is that each iteration of the Lloyd's algorithm can take $O(n^2)$ time, since the centers in the intermediate iterations are possibly linear combinations of all the input points. Thus, the authors propose using mini-batches in each iteration. Instead of using all the points in every iteration, sample only $b$ points in each iteration and update the centers slightly in the centers determined by the Lloyd's update rule. The learning rate parameters set as $\sqrt{b^j/b}$, where $b^j$ is the number of points sampled from the cluster $j$. This was proposed in an earlier work by Schwartzman et al. Using a dynamic programming algorithm, the authors show that this update rule can then be implemented in $O(n \cdot b)$ time. To improve the run time further, the authors propose truncation, whereby a truncated representation of centers is stored instead of storing them all as linear combinations of (possibly all the $n$ points). They show that the quality hit with truncation is not large.

Strengths

The main contributions of the work: 1. A simple dp algorithm to decrease the cost of mini-batch lloyd's iterations. 2. Truncation to further decrease the cost of each iteration. 3. Experiments implementing their algorithm.

Weaknesses

Utility and correctness issues I describe in the section below,

Questions

1. I do not think the proof of Lemma 15 is correct. The way $\bar{\mathcal{C}}\_{{i+1}}$ is defined in line 379 uses, $\alpha_j^i = \sqrt{b^i_j / b}$ values which themselves are functions of the samples $B_i$. So how can you claim that $\bar{\mathcal{C}}_{i+1}$ is independent of $B_i$? Please clarify if I misunderstood anything. I'm willing to update my evaluation after this clarification. 2. While the experimental results do show that the algorithm generates reasonable clusters, I am unconvinced that the iteration bound for minibatch algorithms is super illuminating. Consider an algorithm which outputs the cluster centers unchanged. That algorithm will always terminate in a single iteration by this stopping criteria. So, why should someone (theoretically) care about iteration bound with respect to this stopping criteria (i.e., the value decrease is small) though I understand that the sklearn library uses this stopping criteria and since we are sure that the Lloyd's algorithm only makes progress in terms of the loss, it is a useful criteria to stop the algorithm when the progress is small. 3. In Lemma 13, where is probability coming from? 4. Definition of $\alpha^j_i$ must be highlighted. I will update my evaluation based on the answers for these questions.

Rating

5

Confidence

4

Soundness

2

Presentation

3

Contribution

2

Authorsrebuttal2024-11-14

Response to questions

Thank you for your review. Here are our responses to your questions. Q1) Great observation! Indeed $\bar{C}_{i+1}$ is not independent of $C_i$ due to the learning rate. We've fixed Lemma 15 in the updated paper (changes are in red). This results in an extra multiplicative k-factor in the batch size. Q2) The motivation for this analysis was indeed because the stopping condition is used in practice. However, our main contribution is not the theoretical guarantees, but the actual algorithm which reduced the running time per iteration from $O(n^2)$ to $\tilde{O}(bk^2)$. We consider our result to be an analogue to Sculley's classical mini-batch k-means algorithm which is widely used in practice. The truncation was essential for breaking the linear dependence on $n$, bringing the per iteration cost closer to that of classical k-means - $O(kb)$. Q3) You are right. We remove "w.h.p" from the statement. Q4) We've added the definition explicitly in the algorithm. Please let us know if you would like it to be highlighted elsewhere.

Reviewer Ae522024-11-23

More comments

I thank the authors for correcting the proofs! I have another question: What's the need for a martingale inequality in the proof of Lemma 14? I don't understand why a Hoeffding bound is not sufficient? Another comment is that in the current analysis, the truncation is done right at the end but it is not quite correct since C_{i+1} would be computed using truncated C_i and not simply truncating all the C_i's at the end. This is a minor issue but would be good to correct the proof.

Authorsrebuttal2024-11-24

Response to further comments

Thank you for your further comments. We assume you mean a Hilbert space-valued Hoeffding bound here (a simple Hoeffding bound won't work as our random variables take values in a Hilbert space). Indeed, a Hilbert space-valued Hoeffding bound would suffice (and is a special case of the martingale bound), however, we could not find a reference for such a Hoeffding bound that we could easily cite. The truncated centers are not computed recursively, but using the relation above lemma 3. Recall that we show that they are not too far from the non-truncated centers (which are computed recursively but do not depend on the truncated centers). The guarantees for the truncated centers follow from the guarantees from the non-truncated centers, together with the fact that they are close. In other words, we never use truncated centers for computing non-truncated centers or vice versa. We may have misunderstood your concern here, so please let us know.

Reviewer Ae522024-11-24

Truncated Centers

I didn't mean recursive here. Your analysis computes C_1,..., C_t where C_2 depends on C_1, C_3 depends on C_2, .... Your final truncation analysis directly truncates C_1,..., C_t to get \hat C_1,..., \hat C_t respectively. Wouldn't the implemented algorithm compute C_2 using \hat C_1 and then truncate it to get \hat C_2, then compute C_3 based on \hat C_2 and then truncate to get \hat C_3 and so on? This is what I meant.

Authorsrebuttal2024-11-24

On the Truncated Centers

The key point you might be missing is in the paragraph just before Lemma 3 as well as Lemma 3 itself. To simplify things*, assume that truncated centers only remember the $\tau$ most recent points assigned to cluster $j$ from previous batches. If the total number of points seen by cluster $j$ is less than $\tau$, then $\hat C_{i+1}^j$ will be identical to $C_{i+1}^j$. That is, both centers are identical linear combinations of at most $\tau$ points. After $\tau$ points have been assigned to cluster $j$, the truncated version starts to forget the oldest points and their coefficients, diverging from the untruncated center. However, Lemma 3 tells us that so long as $\tau$ is suitably large, then for every iteration, the truncated and untruncated centers will be close. This is all we need for the termination guarantee. The untruncated version maintains the full linear combination of all points assigned to cluster $j$ while the truncated one only maintains a linear combination of the $\tau$ most recently seen points. To update either one, we only need to know which points are in the next batch and we can update the coefficients of the truncated linear combination according to definition (1). *$Q_i^j$ is defined so that points from the same batch enter and leave the truncated linear combination simultaneously. If we only remembered the $\tau$ most recent points, the indexing would become a lot more awkward.

Reviewer Ae522024-11-28

Score update

I updated my score based on the back-and-forth. My main remaining concerns are: 1. The (uncorrected) analysis is extremely similar to earlier work which makes the results in this work seem only a marginal advancement over previous work. 2. All the experiments seem to use a fixed number of 200 iterations. Does the full batch k-means converge to the loss quicker? Or the mini-batch k-means remains close to full-batch throughout? Are there instances where there is a clear gap between the loss attained by full-batch vs mini-batch? I would have liked to see much more extensive analysis here. Please also indicate to the authors of "Mini-batch k-means terminates within O(d/ε) iterations" about the mistake in their analysis and your fix so they can post a corrected version of their results.

Authorsrebuttal2024-11-28

Response

We'd like to thank you for your time during this discussion period. We would greatly appreciate your input on the following points you raised. 1) You are correct, the original analysis is very similar to earlier work and represents only a marginal advancement. However, we feel strongly that the truncation procedure is sufficiently novel and will be very useful in practice. Without the truncation, it's easy to argue that the algorithm could not be called a true minibatch algorithm, in the sense that the time complexity of each iteration would still have a linear dependence on the number of points. This is prohibitively expensive for large datasets. By reducing this dependence to being only logarithmic in $n$, we hope that our truncated algorithm will do for kernel k-means what Sculley's paper (https://dl.acm.org/doi/10.1145/1772690.1772862) did for k-means. What more could we have done to convince you of this? 2) We thought running each algorithm for 200 iterations would be the fairest and most intuitive approach as this was almost certainly long enough for each method to converge. We did consider using early stopping for comparison instead, but controlling the effect of early stopping across the full batch and mini-batch algorithms with various batch sizes is not obvious. Another approach would be to keep the number of times we read the entire dataset constant instead of the number of iterations. When you say "Does the full batch k-means converge to the loss quicker? Or the mini-batch k-means remains close to full-batch throughout?", there are multiple ways to measure this. We could measure convergence in terms of iterations, wall clock time, or by the number of times we read the input data before reaching convergence (which will be affected by batch size). While we didn't report the loss (inertia), clear gaps in ARI/NMI typically correspond to clear gaps in inertia. An example of a clear gap can be found in Figure 3, where we see that for small $\tau$, the truncated method struggles significantly to meet the performance of the full batch method. Which experimental approaches would you like to have seen? With more time, we would have gladly performed the experiments again with different hyperparameters. We will contact the authors of the previous paper.

Reviewer Ae522024-11-28

Response

Thanks for taking time to reply. I didn't mean for you to run extensive experiments in such a short time during the rebuttal period. It's just that it'd have added quite a bit to the strength of the paper with a more extensive set of experiments. - I very much understand and appreciate that truncation is important for the iterations to be truly sublinear and that it is the main contribution of the work. My main concern was that the algorithm itself (along with the learning rate) is exactly the same as the previous work. We will definitely consider your response in the reviewer discussion. Thanks again.

Reviewer JASs5/10 · confidence 2/52024-11-10

Summary

The paper proposes a kernel mini-batch $k$-means algorithm. The authors provide theoretical analysis that shows that the number of iterations of their algorithm is dimension-independent. In addition, they provide experimantal results that suggest that their algorithm might be useful in practice.

Strengths

The dimension-independent result seem to be interesting, though I'm not an expert in the field.

Weaknesses

I'm not an expert in the area, and I might get some things wrong. So far it seems to me that the main difference of your analysis and the analysis of [Schwartzman (2023)] is your Lemma 12, that is very similar to their Lemma 10. These lemmas have almost identical proofs, except the last step, where both proofs use the Cauchy-Schwartz inequality, but in different spaces. This corresponds to your $\gamma^2$ and their $d$ in the number of iterations. Could you tell what are other important differences between your proof and their proof? If there are no other crucial differences, the technical contribution of this paper is not strong enough for ICLR.

Questions

see Weaknesses above

Rating

5

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Authorsrebuttal2024-11-14

Response to questions

Thank you for your review. Although the theoretical guarantees are nice to have, our main contribution is introducing the first mini-batch algorithm for kernel k-means. We would like to emphasize that before our algorithm, kernel k-means was simply too slow to execute on large datasets. Using our algorithm we get an order of magnitude speed up, which is our main contribution. You can see our algorithm as an analogue to the widely used (non-kernel) mini-batch algorithm (the one implemented in sklearn). We believe that our result is a step towards making kernel k-means applicable in practice. As for the analysis, we needed to overcome the fact that we cannot represent centers explicitly in kernel spaces (as Schwartzman does), we do this by introducing $\gamma$ and using Hilbert space-valued martingales (while Schwartzman uses Hoeffding bounds). But again, our main contribution is in the design of the algorithm and runtime analysis. Please also note that the running time per iteration only has a logarithmic dependence on $n$, which is very difficult to achieve in kernel spaces and wouldn't be possible without our truncation procedure.

Authorsrebuttal2024-11-23

Response

Dear reviewers, the discussion period is ending soon. We would greatly appreciate it if you could reply to our responses and confirm that we properly answered your questions, or clarify what we need to improve. Thank you for your time.

Area Chair E26dmeta-review2024-12-23

Meta-review

This paper combines mini batches with kernel k means. Naturally, this greatly improves the run time compared to a full sweep of the data at each iteration. There are some nice contributions and the theoretical analysis seems correct. The reviewers are mixed and borderline. The reviewers unfortunately did not engage in discussion after the rebuttal period. Overall, the methods used to derive the main results are not sufficiently novel and the result is somewhat unsurprising.

Additional comments on reviewer discussion

Reviews were very shallow and reviewers did not engage in discussion after the rebuttal nor after I commented

© 2026 NYSGPT2525 LLC