SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Moses Charikar
Stanford, CA
moses@cs.stanford.edu
Monika Henzinger
monika.henzinger@ist.ac.at
Lunjia Hu
lunjia@stanford.edu
Maximilian Votsch
Doctoral School of Computer Science DoCS Vienna
maximilian.voetsch@univie.ac.at
Department of Computer and Information Sciences
ewaingar@seas.upenn.edu
Clustering is a fundamental problem in unsupervised machine learning with many applications in data analysis. Popular clustering algorithms such as Lloyd's algorithm and  $k$ -means++ can take  $\Omega (ndk)$  time when clustering  $n$  points in a  $d$ -dimensional space (represented by an  $n \times d$  matrix  $X$ ) into  $k$  clusters. In applications with moderate to large  $k$ , the multiplicative  $k$  factor can become very expensive. We introduce a simple randomized clustering algorithm that provably runs in expected time  $O(\mathsf{nnz}(X) + n\log n)$  for arbitrary  $k$ . Here  $\mathsf{nnz}(X)$  is the total number of non-zero entries in the input dataset  $X$ , which is upper bounded by  $nd$  and can be significantly smaller for sparse datasets. We prove that our algorithm achieves approximation ratio  $\tilde{O} (k^4)$  on any input dataset for the  $k$ -means objective. We also believe that our theoretical analysis is of independent interest, as we show that the approximation ratio of a  $k$ -means algorithm is approximately preserved under a class of projections and that  $k$ -means++ seeding can be implemented in expected  $O(n\log n)$  time in one dimension. Finally, we show experimentally that our clustering algorithm gives a new tradeoff between running time and cluster quality compared to previous state-of-the-art methods for these tasks.
Clustering is an essential and powerful tool for data analysis with broad applications in computer vision and computational biology, and it is one of the fundamental problems in unsupervised machine learning. In large-scale applications, datasets often contain billions of high-dimensional points. Grouping similar data points into clusters is crucial for understanding and organizing datasets. Because of its practical importance, the problem of designing efficient and effective clustering algorithms has attracted the attention of numerous researchers for many decades.
One of the most popular algorithms for the  $k$ -means clustering problem is Lloyd's algorithm [52], which seeks to locate  $k$  centers in the space that minimize the sum of squared distances from the points of the dataset to their closest center (we call this the "  $k$ -means cost"). While finding the centers minimizing the objective is NP-hard [3], in practice we can find high-quality sets of centers using Lloyd's iterative algorithm. Lloyd's algorithm maintains a set of  $k$  centers. It iteratively updates them by assigning points to one of  $k$  clusters (according to their closest center), then redefining the center as the points' center of mass. It needs a good initial set of centers to obtain a high-quality clustering and fast convergence. In practice, the  $k$ -means++ algorithm [6], a randomized seeding procedure, is used to choose the initial  $k$  centers.  $k$ -means++ achieves an  $O(\log k)$ -approximation ratio in expectation, upon which each iteration of Lloyd's algorithm improves. Beyond their effectiveness, these algorithms are simple to describe and implement, contributing to their popularity.
The downside of these algorithms is that they do not scale well to massive datasets. A standard implementation of an iteration of Lloyd's algorithm needs to calculate the distance to each center for each point in the dataset, leading to a  $\Theta(ndk)$  running time. Similarly, the standard implementation of the  $k$ -means++ seeding procedure produces  $k$  samples from the so-called  $D^2$  distribution (see Appendix A for details). Maintaining the distribution requires making a pass over the entire dataset after choosing each sample. Generating  $k$  centers leads to a  $\Theta(ndk)$  running time. Even for moderate values of  $k$ , making  $k$  passes over the entire dataset can be prohibitively expensive.
One particularly relevant application of large-scale  $k$ -means clustering is in approximate nearest neighbor search [66] (for example, in product quantization [44] and building inverted file indices [16]). There,  $k$ -means clustering is used to compress entire datasets by mapping vectors to their nearest centers, leading to billion-scale clustering problems with large  $k$  (on the order of hundreds or thousands). Other applications on large datasets requiring a large number of centers may be spam filtering [61, 65], near-duplicate detection [42], and compression or reconciliation tasks [62]. New algorithmic ideas are needed for these massive scales, and this motivates the following challenge:
Can we design a simple, practical algorithm for  $k$ -means that runs in time roughly  $O(nd)$ , independent of  $k$ , and produces high-quality clusters?
Given its importance in theory and practice, a significant amount of effort has been devoted to algorithms for fast  $k$ -means clustering. We summarize a few of the approaches below with the pros and cons of each so that we may highlight our work's position within the literature:
A. Standard  $k$ -means++: This is our standard benchmark. Plus: Guaranteed to be an  $O(\log k)$ -approximation [6]; outputs centers, as well as the assignments of dataset points to centers. Minus: The running time is  $O(ndk)$ , which is prohibitively expensive in large-scale applications.
B. Using Approximate Nearest Neighbor Search: One may implement  $k$ -means++ faster using techniques from approximate nearest neighbor search (instead of a brute force search each iteration). Plus: The algorithms with provable guarantees, like [26], obtain an  $O_{\varepsilon}(\log k)$ -approximation. Minus: The running time is  $\widetilde{O}(nd + (n\log(\Delta))^{1+\varepsilon})$ , depending on a dataset dependent parameter  $\Delta$ , the ratio between the maximum and minimum distances between input points. The techniques are algorithmically sophisticated and incur extra poly-logarithmic factors (hidden in  $\widetilde{O}(\cdot)$ ), making the implementation significantly more complicated.
C. Approximating the  $D^2$ -Distribution: Algorithms that speed up the seeding procedure for Lloyd's algorithm or generate fast coresets (we expand on this below) have been proposed in [8, 7, 10]. Plus: These algorithms are fast, making only one pass over the dataset in time  $O(nd)$ . (For [8, 7], there is an additional additive  $O(k^2 d)$  term in the running time). Minus: The approximation guarantees are qualitatively weaker than the approximation of  $k$ -means clustering. They incur an additional additive approximation error that grows with the entire dataset's variance (which can lead to an arbitrarily large error; see Section 3). These algorithms output a set of  $k$  centers but not the cluster assignments. Naively producing the assignments would take time  $O(ndk)$ .<sup>2</sup>
1 Approximation is with respect to the  $k$ -means cost. A  $c$ -approximation has  $k$ -means cost, which is at most  $c$  times larger than the optimal  $k$ -means cost.
One may use approximate nearest neighbor search techniques to improve on the  $O(ndk)$  running time. However, as discussed above, approximate nearest neighbor search adds a significant layer of complexity (and approximation).
Coresets. At a high level, coresets are a dataset-reduction mechanism. A large dataset  $X$  of  $n$  points in  $\mathbb{R}^d$  is distilled into a significantly smaller (weighted) dataset  $Y$  of  $m$  points in  $\mathbb{R}^d$ , called a "coreset" which serves as a good proxy for  $X$ , i.e., the clustering cost of any  $k$  centers on  $Y$  is approximately the cost of the same centers on  $X$ . We point the reader to [9, 35] for a recent survey on coresets. Importantly, coreset constructions (with provable multiplicative-approximation guarantees) require an initial approximate clustering of the original dataset  $X$ . Therefore, any fast algorithm for  $k$ -means clustering automatically speeds up any algorithmic pipeline that uses coresets for clustering — looking forward, we will show how our algorithm can significantly speed up coreset constructions without sacrificing approximation.
Beyond those mentioned above, many works seek to speed up  $k$ -means++ or Lloyd iterations by maintaining some nearest neighbor search data structures [58, 54, 46, 45, 33, 60, 40, 59, 71, 30, 29, 13, 55, 28, 18], or by running some first-order methods [64]. These techniques do not give provable guarantees on the quality of the  $k$ -means clustering or on the running time of their algorithms.
Theoretical Results. We give a simple randomized clustering algorithm with provable guarantees on its running time and approximation ratio without making any assumptions about the data. It has the benefit of being fast (like the algorithms in Category C above) while achieving a multiplicative error guarantee without additional additive error (like the algorithms in Category B above).
- The algorithm runs in time  $O(nd + n \log n)$  irrespective of  $k$ . It passes over the dataset once to perform data reduction, which gives the  $nd$  factor plus an additive  $O(n \log n)$  term to solve  $k$ -means on the reduced data, producing  $k$  centers and cluster assignments. On sparse input datasets, the  $nd$  term becomes  $\mathsf{nnz}(X)$ , where  $\mathsf{nnz}(X)$  is the number of non-zero entries in the dataset. Thus, our algorithm runs in  $O(\mathsf{nnz}(X) + n \log n)$  time on sparse matrices.
- The algorithm is as simple as the  $k$ -means++ algorithm while significantly more efficient. The approximation ratio we prove is poly(k), which is worse than the  $O(\log k)$ -approximation achieved by  $k$ -means++ but it is purely multiplicative (see the remark below on improving this to  $O(\log k)$ ). It does not incur the additional additive errors from the fast algorithms in [8, 7, 10].
Our algorithm projects the input points to a random one-dimensional space and runs an efficient  $k$ -means++ seeding after the projection. For the approximation guarantee, we analyze how the approximation ratio achieved after the projection can be transferred to the original points (Lemma 2.5). We bound the running time of our algorithm by efficiently implementing the  $k$ -means++ seeding in one dimension and analyzing the running time via a potential function argument (Lemma 2.4). Our algorithm applies beyond  $k$ -means to other clustering objectives that sum up the  $z$ -th power of the distances for general  $z \geq 1$ , and our guarantees on its running time and approximation ratio extend smoothly to these settings.
Improving the Approximation from  $\mathrm{poly}(k)$  to  $O(\log k)$ . The approximation ratio of  $\mathrm{poly}(k)$  may seem significantly worse than the  $O(\log k)$  approximations achievable with  $k$ -means++. However, we can improve this to  $O(\log k)$  with an additional, additive  $O(\mathrm{poly}(kd) \cdot \log n)$  term in the running time. Using previous results discussed in Appendix A.2 (specifically Theorem A.2), a multiplicative  $\mathrm{poly}(k)$ -approximation suffices to construct a coreset of size  $\mathrm{poly}(kd)$  and run  $k$ -means++ on the coreset. Constructing the coreset is simple and takes time  $\mathrm{poly}(kd) \cdot \log n$  (by sampling from an appropriate distribution); running  $k$ -means++ on the coreset takes  $\mathrm{poly}(kd)$  time (with no dependence on  $n$ ). Combining our algorithm with coresets, we get a  $O(\log k)$ -approximation in  $O(\mathsf{nnz}(X)) + O(n \log n) + \mathrm{poly}(kd) \cdot \log n$  time. Notably, these guarantees cannot be achieved with the additive approximations of [8, 7, 10].
Experimental Results. We implemented our algorithm, as well as the lightweight coreset of [10] and  $k$ -means++ with sensitivity sampling [15]. We ran two types of experiments, highlighting various aspects of our algorithm. We provide our code in the supplementary material. The two types of experiments are:
- Coreset Construction Comparison: First, we evaluate the performance of our clustering algorithm when we use it to construct coresets. We compare the performance of our algorithm to  $k$ -means++ with sensitivity sampling [9] and lightweight coresets [10]. In real-world, high-dimensional data, the cost of the resulting clusters from the three algorithms is roughly the same. However, ours and the lightweight coresets can be significantly faster (ours is up to 190x faster than  $k$ -means++,
see Figure 2 and Table 1). The lightweight coresets can be faster than our algorithm (between  $3 - 5\mathrm{x}$ ); however, our algorithm is "robust" (achieving multiplicative approximation guarantees).<sup>3</sup> Additionally, we show that the clustering from lightweight coresets can have an arbitrarily high cost for a synthetic dataset. On the other hand, our algorithm achieves provable (multiplicative) approximation guarantees irrespective of the dataset (this is demonstrated in the right-most column of Figure 2).
- Direct k-means++ comparison: Second, we compare the speed and cost of our algorithm to k-means++[6] as a stand-alone clustering algorithm (we also compare two other natural variants of our algorithm). Our algorithm can be up to 800x faster than  $k$ -means++ for  $k = 5000$  and our slowest variant up to 100x faster (Table 1). The cost of the cluster assignments can be significantly worse than that of  $k$ -means++ (see Figure 3). Such a result is expected since our theoretical results show a poly  $(k)$ -approximation. The other (similarly) fast algorithms (based on approximating the  $D^2$ -distribution) which run in time  $O(nd)$  [8, 7] do not produce the cluster assignments (they only output  $k$  centers). These algorithms would take  $O(ndk)$  time to find the cluster assignments — this is precisely the computational cost our algorithm avoids.
We do not compare our algorithm with [26] nor implement approximate nearest neighbor search to speed up  $k$ -means++ for the following reasons. The algorithm in [26] is significantly more complicated, and there is no publicly available implementation. In addition, both [26] and approximate nearest neighbor search incur additional poly-logarithmic (or even  $n^{o(1)}$ -factors for nearest neighbor search over  $\ell_2$  [4]) which add significant layers of complexity to the implementation and make a thorough evaluation of the algorithm significantly more complicated. Instead, our current implementation demonstrates that a simple, one-dimensional projection and  $k$ -means++ on the line enables dramatic speedups to coreset constructions without sacrificing approximation quality.
Related Work. Efficient algorithms for clustering problems with provable approximation guarantees have been studied extensively, with a few approaches in the literature. There are polynomial-time (constant) approximation algorithms (an exponential dependence on  $k$  is not allowed) (see [50, 17, 2, 39] for some of the most recent and strongest results), nearly linear time  $(1 \pm \varepsilon)$ -approximations with running time exponential in  $k$  which proceed via coresets (see [41, 19, 36, 37, 15, 9, 27, 25] and references therein, as well as the surveys [1, 35]), and nearly-linear time  $(1 \pm \varepsilon)$ -approximations in fixed/low-dimensional spaces [5, 48, 68, 38, 24, 22, 23]. Our  $O(n \log n)$ -expected-time implementation of  $k$ -means++ seeding achieves an  $O(\log k)$  expected approximation ratio for  $k$ -median and  $k$ -means in one dimension. We are unaware of previous work on clustering algorithms running in time  $O(n \log n)$ .
Another line of research has been on dimensionality reduction techniques for  $k$ -means clustering. Dimensionality reduction can be achieved via PCA based methods [31, 37, 21, 67], or random projection [21, 11, 53]. For random projection methods, it has been shown that the  $k$ -means objective is preserved up to small multiplicative factors when projecting onto  $O_{\varepsilon}(\log (k))$  dimensional space. Additional work has shown that dimensionality reduction can be performed in  $O(\mathsf{nnz}(A))$  time [51]. To the best of our knowledge, we are the first to show that clustering objectives such as  $k$ -median and  $k$ -means are preserved up to a poly(  $k$  ) factor by one-dimensional projections.
Some works show that the  $O(\log k)$  expected approximation ratio for  $k$ -means++ can be improved by adding local search steps after the seeding procedure [49, 20]. In particular, Choo et al. [20] showed that adding  $\varepsilon k$  local search steps achieves an  $O(1 / \varepsilon^3)$  approximation ratio with high probability.
Several other algorithmic approaches exist for fast clustering of points in metric spaces. These include density-based methods like DBSCAN [34] and DBSCAN++ [43] and the line of heuristics based on the Partitioning Around Medoids (PAM) approach, such as FastPAM [63], Clarans [56], and BanditPAM [69]. While these algorithms can produce high-quality clustering, their running time is at least linear in the number of clusters (DBSCAN++ and BanditPAM) or superlinear in the number of points (DBSCAN, FastPAM, Clarans).
3Recall that the lightweight coresets incur an additional additive error which can be arbitrarily large.
2 Overview of Our Algorithm and Proof Techniques
Our algorithm, which we call PRONE (PProjected ONE-dimensional clustering), takes a random projection onto a one-dimensional space, sorts the projected (scalar) numbers, and runs the  $k$ -means++ seeding strategy on the projected numbers. By virtue of its simplicity, the algorithm is scalable and effective at clustering massive datasets. More formally, PRONE receives as input a dataset of  $n$  points in  $\mathbb{R}^d$ , a parameter  $k \in \mathbb{N}$  (the number of desired clusters), and proceeds as follows:
1. Sample a random vector  $v \in \mathbb{R}^d$  from the standard Gaussian distribution and project the data points to one dimension along the direction of  $v$ . That is, we compute  $x_i' = \langle x_i, v \rangle \in \mathbb{R}$  in time  $O(\mathrm{nnz}(X))$  by making a single pass over the data, effectively reducing our dataset to the collection of one-dimensional points  $x_1', \ldots, x_n' \in \mathbb{R}$ .
2. Run  $k$ -means++ seeding on  $x_1', \ldots, x_n'$  to obtain  $k$  indices  $j_1, \ldots, j_k \in [n]$  indicating the chosen centers  $x_{j_1}', \ldots, x_{j_k}'$  and an assignment  $\sigma : [n] \to [k]$  assigning point  $x_i'$  to center  $x_{j_{\sigma(i)}}'$ . Even though  $k$ -means++ seeding generally takes  $O(nk)$  time in one dimension, we give an efficient implementation, leveraging the fact that points are one-dimensional, which runs in  $O(n \log n)$  expected time, independent of  $k$ . A detailed algorithm description is in the appendix.
3. The one-dimensional  $k$ -means++ algorithm produces a collection of  $k$  centers  $x_{j_1},\ldots ,x_{j_k}$ , as well as the assignment  $\sigma$  mapping each point  $x_{i}$  to the center  $x_{j_{\sigma (i)}}$ . For each  $\ell \in [k]$ , we update the cluster center for cluster  $\ell$  to be the center of mass of all points assigned to  $x_{j_\ell}$ .
While the algorithm is straightforward, the main technical difficulty lies in the analysis. In particular, our analysis (1) bounds the approximation loss incurred from the one-dimensional projection in Step 1 and (2) shows that we can implement Step 2 in  $O(n \log n)$  expected time, as opposed to  $O(nk)$  time. We summarize the theoretical contributions in the following theorems.
Theorem 2.1. The algorithm PRONE has expected running time  $O(\mathfrak{nnz}(X) + n\log n)$  on any dataset  $X = \{x_{1},\ldots ,x_{n}\} \subset \mathbb{R}^{d}$ . Moreover, for any  $\delta \in (0,1 / 2)$  and any dataset  $X$ , with probability at least  $1 - \delta$ , the algorithm runs in time  $O(\mathfrak{nnz}(X) + n\log (n / \delta))$ .
Theorem 2.2. The algorithm PRONE achieves an  $\widetilde{O}(k^4)$  approximation ratio for the  $k$ -means objective with probability at least 0.9.
To our knowledge, PRONE is the first algorithm for  $k$ -means running in time  $O(nd + n\log n)$  for arbitrary  $k$ . As mentioned in the paragraph on improving the competitive ratio, we obtain the following corollary of the previous two theorems using a two-stage approach with a coreset:
Corollary 2.3. By using PRONE as the  $\alpha$ -approximation algorithm in Theorem A.2 and running  $k$ -means++ on the resulting coreset, we obtain an algorithm with an approximation ratio of  $O(\log k)$  that runs in time  $O(\mathsf{nnz}(X) + n\log n + \mathrm{poly}(kd)\log n)$ , with constant success probability.
Due to space constraints, all proofs of our theoretical results are deferred to the appendix, where we also generalize them beyond  $k$ -means to clustering objectives that sum up the  $z$ -th power of Euclidean distances for general  $z \geq 1$ . The following subsections give a high-level overview of the main techniques we develop to prove our main theorems above.
2.1 Efficient Seeding in One Dimension
The  $k$ -means++ seeding procedure has  $k$  iterations, where a new center is sampled in each iteration. Since a new center may need to update  $\Omega(n)$  distances to maintain the  $D^2$  distribution, which samples each point with probability proportional to its distance to its closest center, a naive analysis leads to a running time of  $O(nk)$ . A key ingredient in the proof of Theorem 2.1 is showing that, for one-dimensional datasets,  $k$ -means++ only needs to make  $O(n \log n)$  updates, irrespective of  $k$ .
Lemma 2.4. The  $k$ -means++ seeding procedure can be implemented in expected time  $O(n \log n)$  in one dimension. Moreover, for any  $\delta \in (0,1/2)$ , with probability at least  $1 - \delta$ , the implementation runs in time  $O(n \log (n / \delta))$ .
The intuition of the proof is as follows: Since points are one-dimensional, we always maintain them in sorted order. In addition, each data point  $x_{i}$  will maintain its center assignment and distance  $p_i$  to the closest center. By building a binary tree over the sorted points (where internal nodes maintain sums of  $p_i^{2}$ 's), it is easy to sample a new center from the  $D^2$  distribution in  $O(\log n)$  time. The difficulty is
that adding a new center may result in changes to  $p_i$ 's of multiple points  $x_i$ , so the challenge is to bound the number of times these values are updated (see Figure 1 below).
Figure 1: From the top to the bottom, a new center (black circle) is chosen. Every point has an arrow pointing to its closest center. The points in the dashed box are the ones that require updates.
To bound the total running time, we leverage the one-dimensional structure. Observe that, for a new center, the updated points lie in a contiguous interval around the newly chosen center. Once a center is chosen, the algorithm scans the points (to the left and the right) until we reach a point that does not need to be updated. This point identifies that points to the other side of it need not be updated, so we can get away without necessarily checking all  $n$  points (see Figure 1). Somewhat surprisingly, when sampling centers from the  $D^2$ -distribution, the expected number of times that each point will be updated is only  $O(\log n)$ , which implies a bound of  $O(n\log n)$  on the total number of updates in expectation. The analysis of the fact that each point is updated  $O(\log n)$  times is non-trivial and uses a carefully designed potential function (Lemma C.5).
2.2 Approximation Guarantees from One-Dimensional Projections
Our proof of Theorem 2.2 builds on a line of work studying randomized dimension reduction for clustering problems [14, 21, 11, 53]. Prior work studied randomized dimension reduction for accurate  $(1\pm \epsilon)$ -approximations. Our perspective is slightly different; we restrict ourselves to one-dimensional projections and give an upper bound on the distortion.