SlowGuess/NeurIPS_2023_OCR huggingface dataset:
BanditPAM++: Faster  $k$ -medoids Clustering
Mo Tiwari
motiwari@stanford.edu
Ryan Kang*
txryank@stanford.edu
Donghyun Lee*
donghyun.lee.21@ucl.ac.uk
thrun@stanford.edu
Chris Piech
piech@cs.stanford.edu
Ilan Shomorony
ilans@illinois.edu
Martin Jinye Zhang
martinzh@andrew.cmu.edu
Clustering is a fundamental task in data science with wide-ranging applications. In  $k$ -medoids clustering, cluster centers must be actual datapoints and arbitrary distance metrics may be used; these features allow for greater interpretability of the cluster centers and the clustering of exotic objects in  $k$ -medoids clustering, respectively.  $k$ -medoids clustering has recently grown in popularity due to the discovery of more efficient  $k$ -medoids algorithms. In particular, recent research has proposed BanditPAM, a randomized  $k$ -medoids algorithm with state-of-the-art complexity and clustering accuracy. In this paper, we present BanditPAM++, which accelerates BanditPAM via two algorithmic improvements, and is  $O(k)$  faster than BanditPAM in complexity and substantially faster than BanditPAM in wall-clock runtime. First, we demonstrate that BanditPAM has a special structure that allows the reuse of clustering information within each iteration. Second, we demonstrate that BanditPAM has additional structure that permits the reuse of information across different iterations. These observations inspire our proposed algorithm, BanditPAM++, which returns the same clustering solutions as BanditPAM but often several times faster. For example, on the CIFAR10 dataset, BanditPAM++ returns the same results as BanditPAM but runs over  $10 \times$  faster. Finally, we provide a high-performance  $C++$  implementation of BanditPAM++, callable from Python and R, that may be of interest to practitioners at https://github.com/motiwari/BanditPAM. Auxiliary code to reproduce all of our experiments via a one-line script is available at https://github.com/ThrunGroup/BanditPAM_plusplus_experiments/.
Clustering is a critical and ubiquitous task in data science and machine learning. Clustering aims to separate a dataset  $\mathcal{X}$  of  $n$  datapoints into  $k$  disjoint sets that form a partition of the original dataset. Intuitively, datapoints within a cluster are similar and datapoints in different clusters are dissimilar. Clustering problems and algorithms have found numerous applications in textual data [20], social network analysis [18], biology [28], and education [28].
A common objective function used in clustering is Equation 1:
$$
L (\mathcal {C}) = \sum_ {j = 1} ^ {n} \min  _ {c \in \mathcal {C}} d (c, x _ {j}). \tag {1}
$$
Under this loss function, the goal becomes to minimize the distance, defined by the distance function  $d$ , from each datapoint to its nearest cluster center  $c$  among the set of cluster centers  $\mathcal{C}$ . Note that this formulation is general and does not require the datapoints to be vectors or assume a specific functional form of the distance function  $d$ .
Specific choices of  $\mathcal{C}$ , dataset  $\mathcal{X}$ , and distance function  $d$  give rise to different clustering problems. Perhaps one of the most commonly used clustering methods is  $k$ -means clustering [17, 16]. In  $k$ -means clustering, each datapoint is a vector in  $\mathbb{R}^p$  and the distance function  $d$  is usually taken to be squared  $L_2$  distance; there are no constraints on  $\mathcal{C}$  other than that it is a subset of  $\mathbb{R}^p$ . Mathematically, the  $k$ -means objective function is
$$
L (\mathcal {C}) = \sum_ {j = 1} ^ {n} \min  _ {c \in \mathcal {C}} \| x _ {j} - c \| _ {2} ^ {2} \tag {2}
$$
subject to  $|\mathcal{C}| = k$
The most common algorithm for the  $k$ -means problem is Lloyd iteration [16], which has been improved by other algorithms such as  $k$ -means++ [2]. These algorithms are widely used in practice due to their simplicity and computational efficiency. Despite widespread use in practice,  $k$ -means clustering suffers from several limitations. Perhaps its most significant restriction is the choice of  $d$  as the squared  $L_{2}$  distance. This choice of  $d$  is for computational efficiency – as the mean of many points can be efficiently computed under squared  $L_{2}$  distance – but prevents clustering with other distance metrics that may be preferred in other contexts [22, 8, 5]. For example,  $k$ -means is difficult to use on textual data that necessitates string edit distance [20], social network analyses using graph metrics [18], or sparse datasets (such as those found in recommendation systems [15]) that lend themselves to other distance functions. While  $k$ -means algorithms have been adapted to specific other metrics, e.g., cosine distance [23], these methods are bespoke to the metric and not readily generalizable. Another limitation of  $k$ -means clustering is that the set of cluster centers  $C$  may often be uninterpretable, as each cluster center is (generally) a linear combination of datapoints. This limitation can be especially problematic when dealing with structured data, such as parse trees in context-free grammars, where the mean of trees is not necessarily well-defined, or images in computer vision, where the mean image can appear as random noise [28, 15].
In contrast with  $k$ -means clustering,  $k$ -medoids clustering [10, 11] requires the cluster centers to be actual datapoints, i.e., requires  $\mathcal{C} \subset \mathcal{X}$ . More formally, the objective is to find a set of medoids  $\mathcal{M} \subset \mathcal{X}$  (versus  $\mathbb{R}^p$  in  $k$ -means) that minimizes
$$
L (\mathcal {M}) = \sum_ {j = 1} ^ {n} \min  _ {m \in \mathcal {M}} d (m, x _ {j}) \tag {3}
$$
subject to  $|\mathcal{M}| = k$ . Note that there is no restriction on the distance function  $d$ .
$k$ -medoids clustering has several advantages over  $k$ -means. Crucially, the requirement that each cluster center is a datapoint leads to greater interpretability of the cluster centers because each cluster center can be inspected. Furthermore,  $k$ -medoids supports arbitrary dissimilarity measures; the distance function  $d$  in Equation (3) need not be a proper metric, i.e., may be negative, asymmetric, or violate the triangle inequality. Because  $k$ -medoids supports arbitrary dissimilarity measures, it can also be used to cluster "exotic" objects that are not vectors in  $\mathbb{R}^p$ , such as trees and strings [28], without embedding them in  $\mathbb{R}^p$ .
The  $k$ -medoids clustering problem in Equation (3) is a combinatorial optimization algorithm that is NP-hard in general [25]; as such, algorithms for  $k$ -medoids clustering are restricted to heuristic solutions. A popular early heuristic solution for the  $k$ -medoids problem was the Partitioning Around Medoids (PAM) algorithm [11]; however, PAM is quadratic in dataset size  $n$  in each iteration, which
is prohibitively expensive on large dataset. Improving the computational efficiency of these heuristic solutions is an active area of research. Recently, [28] proposed BanditPAM, the first subquadratic algorithm for the  $k$ -medoids problem that matched prior state-of-the-art solutions in clustering quality. In this work, we propose BanditPAM++, which improves the computational efficiency of BanditPAM while maintaining the same results. We anticipate these computational improvements will be important in the era of big data, when  $k$ -medoids clustering is used on huge datasets.
Contributions: We propose a new algorithm, BanditPAM++, that is significantly more computationally efficient than PAM and BanditPAM but returns the same clustering results with high probability. BanditPAM++ is  $O(k)$  faster than BanditPAM in complexity and substantially faster than BanditPAM in actual runtime wall-clock runtime. Consequently, BanditPAM++ is faster than prior state-of-the-art  $k$ -medoids algorithms while maintaining the same clustering quality. BanditPAM++ is based on two observations about the structure of BanditPAM and the  $k$ -medoids problem, described in Section 4. The first observation leads to a technique that we call Virtual Arms (VA). The second observation leads to a technique that we refer to as Permutation-Invariant Caching (PIC). We combine these techniques in BanditPAM++ and prove (in Section 5) and experimentally validate (in Section 6) that BanditPAM++ returns the same solution to the  $k$ -medoids problem as PAM and BanditPAM with high probability, but is more computationally efficient. In some instances, BanditPAM++ is over  $10 \times$  faster than BanditPAM. Additionally, we provide a highly optimized implementation of BanditPAM++ in  $C++$  that is callable from Python and R and may be of interest to practitioners.
As discussed in Section 1, global optimization of the  $k$ -medoids problem (Equation (3)) is NP-hard in general [25]. Recent work attempts to perform attempts global optimization and is able to achieve an optimality gap of  $0.1\%$  on one million datapoints, but is restricted to  $L_{2}$  distance and takes several hours to run on commodity hardware [24].
Because of the difficulty of global optimization of the  $k$ -medoids problem, many heuristic algorithms have been developed for the  $k$ -medoids problem that scale polynomially with the dataset size and number of clusters. The complexity of these algorithms is measured by their sample complexity, i.e., the number of pairwise distance computations that are computed; these computations have been observed to dominate runtime costs and, as such, sample complexity translates to wall-clock time via an approximately constant factor [28] (this is also consistent with our experiments in Section 6 and Appendix 3).
Among the heuristic solutions for the  $k$ -medoids problem, the algorithm with the best clustering loss is Partitioning Around Medoids (PAM) [10, 11], which consists of two phases: the BUILD phase and the SWAP phase. However, the BUILD phase and each SWAP iteration of PAM perform  $O(kn^{2})$  distance computations, which can be impractical for large datasets or when distance computations are expensive. We provide greater details about the PAM algorithm in Section 3 because it is an important baseline against which we assess the clustering quality of new algorithms.
Though PAM achieves the best clustering loss among heuristic algorithms, the era of huge data has necessitated the development of faster  $k$ -medoids algorithms in recent years. These algorithms have typically been divided into two categories: those that agree with PAM and recover the same solution to the  $k$ -medoids problem but scale quadratically in  $n$ , and those that sacrifice clustering quality for runtime improvements. In the former category, [25] proposed a deterministic algorithm called FastPAM1, which maintains the same output as PAM but reduces the computational complexity of each SWAP iteration from  $O(kn^2)$  to  $O(n^2)$ . However, this algorithm still scales quadratically in  $n$  in every iteration, which is prohibitively expensive on large datasets.
Faster heuristic algorithms have been proposed but these usually sacrifice clustering quality; such algorithms include CLARA [11], CLARANS [21], and FastPAM [25]. While these algorithms scale subquadratically in  $n$ , they return substantially worse solutions than PAM [28]. Other algorithms with better sample complexity, such as optimizations for Euclidean space and those based on tabu search heuristics [6] also return worse solutions. Finally, [1] attempts to minimize the number of unique pairwise distances or adaptively estimate these distances or coordinate-wise distances in specific settings [14, 3], but all these approaches sacrifice clustering quality for runtime.
Recently, [28] proposed BanditPAM, a state-of-the-art  $k$ -medoids algorithm that arrives at the same solution as PAM with high probability in  $O(kn \log n)$  time. BanditPAM borrows techniques from the multi-armed bandit literature to sample pairwise distance computations rather than compute all  $O(n^2)$ . In this work, we show that BanditPAM can be made more efficient by reusing distance computations both within iterations and across iterations.
We note that the use of adaptive sampling techniques and multi-armed bandits to accelerate algorithms has also had recent successes in other work, e.g., to accelerate the training of Random Forests [27], solve the Maximum Inner Product Search problem [27], and more [26].
3 Preliminaries and Background
Notation: We consider a dataset  $\mathcal{X}$  of size  $n$  (that may contain vectors in  $\mathbb{R}^p$  or other objects). Our goal is to find a solution to the  $k$ -medoids problem, Equation (3). We are also given a dissimilarity function  $d$  that measures the dissimilarity between two objects in  $\mathcal{X}$ . Note that we do not assume a specific functional form of  $d$ . We use  $[n]$  to denote the set  $\{1, \dots, n\}$ , and  $a \wedge b$  (respectively,  $a \vee b$ ) to denote the minimum (respectively, maximum) of  $a$  and  $b$ .
Partitioning Around Medoids (PAM): The original Partitioning Around Medoids (PAM) algorithm [10, 11] consists of two main phases: BUILD and SWAP. In the BUILD phase, PAM iteratively initializes each medoid in a greedy, one-by-one fashion: in each iteration, it selects the next medoid that would reduce the  $k$ -medoids clustering loss (Equation (3)) the most, given the prior choices of medoids. More precisely, given the current set of  $l$  medoids  $\mathcal{M}_l = \{m_1,\dots ,m_l\}$ , the next point to add as a medoid is:
$$
m ^ {*} = \underset {x \in \mathcal {X} \backslash \mathcal {M} _ {l}} {\arg \min } \sum_ {j = 1} ^ {n} \left[ d \left(x, x _ {j}\right) \wedge \underset {m ^ {\prime} \in \mathcal {M} _ {l}} {\min } d \left(m ^ {\prime}, x _ {j}\right) \right] \tag {4}
$$
The output of the BUILD step is an initial set of the  $k$  medoids, around which a local search is performed by the SWAP phase. The SWAP phase involves iteratively examining all  $k(n - k)$  medoid-nonmedoid pairs and performs the swap that would lower the total loss the most. More precisely, with  $\mathcal{M}$  the current set of  $k$  medoids, PAM finds the best medoid-nonmedoid pair to swap:
$$
\left(m ^ {*}, x ^ {*}\right) = \underset {(m, x) \in \mathcal {M} \times \left(\mathcal {X} \backslash \mathcal {M}\right)} {\arg \min } \sum_ {j = 1} ^ {n} \left[ d \left(x, x _ {j}\right) \wedge \underset {m ^ {\prime} \in \mathcal {M} \backslash \{m \}} {\min } d \left(m ^ {\prime}, x _ {j}\right) \right] \tag {5}
$$
PAM requires  $O(kn^2)$  distance computations for the  $k$  greedy searches in the BUILD step and  $O(kn^2)$  distance computations for each SWAP iteration [28]. The quadratic complexity of PAM makes it prohibitively expensive to run on large datasets. Nonetheless, we describe the PAM algorithm because it has been observed to have the best clustering loss among heuristic solutions to the  $k$ -medoids problem. More recent algorithms, such as BanditPAM [28], achieve the same clustering loss but have a significantly improved complexity of  $O(kn\log n)$  in each step. Our proposed algorithm, BanditPAM++, improves upon the computational complexity of BanditPAM by a factor of  $O(k)$ .
Sequential Multi-Armed Bandits: BanditPAM [28] improves the computational complexity of PAM by converting each step of PAM to a multi-armed bandit problem. A multi-armed bandit problem (MAB) is defined as a collection of random variables  $\{R_1, \dots, R_n\}$ , called actions or arms. We are commonly interested in the best-arm identification problem, which is to identify the arm with the highest mean, i.e.,  $\arg \max_i \mathbb{E}[R_i]$ , with a given probability of possible error  $\delta$ . Many algorithms for this problem exist, each of which make distributional assumptions about the random variables  $\{R_1, \dots, R_n\}$ ; popular ones include the upper confidence bound (UCB) algorithm and successive elimination. For an overview of common algorithms, we refer the reader to [9].
We define a sequential multi-armed bandit problem to be an ordered sequence of multi-armed bandit problems  $Q = \{B_1, \ldots, B_T\}$  where each individual multi-armed bandit problem  $B_t = \{R_1^t, \ldots, R_n^t\}$  has the same number of arms  $n$ , with respective, timestep-dependent means  $\mu_1^t, \ldots, \mu_n^t$ . At each timestep  $t$ , our goal is to determine (and take) the best action  $a_t = \arg \max_i \mathbb{E}[R_t^i]$ . Crucially, our choices of  $a_t$  will affect the rewards at future timesteps, i.e., the  $R_i^{t'}$  for  $t' > t$ . Our definition of a sequential multi-armed bandit problem is similar to non-stationary multi-armed bandit problems, with the added restriction that the only non-stationarity in the problem comes from our previous actions.
We now make a few assumptions for tractability. We assume that each  $R_{i}^{t}$  is observed by sampling an element from a set  $S$  with  $S$  possible values. We refer to the values of  $S$  as the reference points, where each possible reference point is sampled with equal probability and determines the observed reward. With some abuse of notation, we write  $R_{i}^{t}(x_{s})$  for the reward observed from arm  $R_{i}^{t}$  when the latent variable from  $S$  is observed to be  $x_{s}$ . We refer to a sequential multi-armed bandit as permutation-invariant, or as a SPIMAB (for Sequential, Permutation-Invariant Multi-Armed Bandit), if the following conditions hold:
1. For every arm  $i$  and timestep  $t$ ,  $R_i^t = f(D_i, \{a_0, a_1, \ldots, a_{t-1}\})$  for some known function  $f$  and some random variable  $D_i$  with mean  $\mu_i \coloneqq \mathbb{E}[D_i] = \frac{1}{S} \sum_{s=1}^{S} D_i(x_s)$ ,
2. There exists a common set of reference points,  $S$ , shared amongst each  $D_{i}$
3. It is possible to sample each  $D_{i}$  in  $O(1)$  time by drawing from the points in  $\mathcal{S}$  without replacement, and
4.  $f$  is computable in  $O(1)$  time given its inputs.
Intuitively, the conditions above require that at each timestep, each random variable  $R_{i}^{t}$  is expressible as a known function of another random variable  $D_{i}$  and the prior actions taken in the sequential multi-armed bandit problem. Crucially,  $D_{i}$  does not depend on the timestep;  $R_{i}^{t}$  is only permitted to depend on the timestep through the agent's previously taken actions  $\{a_{0}, a_{1}, \ldots, a_{t - 1}\}$ . The SPIMAB conditions imply that if  $\mathbb{E}[D_i] = \mu_i$  is known for each  $i$ , then  $\mu_i^t \coloneqq \mathbb{E}[R_i^t]$  is also computable in  $O(1)$  time for each  $i$  and  $t$ , i.e., for each arm and timestep.
BanditPAM: BanditPAM [28] reduces the scaling with  $n$  of each step of the PAM algorithm by reformulating each step as a best-arm identification problem. In PAM, each of the  $k$  BUILD steps has complexity  $O(n^{2})$  and each SWAP iteration has complexity  $O(kn^{2})$ . In contrast, the complexity of BanditPAM is  $O(n\log n)$  for each of the  $k$  BUILD steps and  $O(kn\log n)$  for each SWAP iteration. Fundamentally, BanditPAM achieves this reduction in complexity by sampling distance computations instead of using all  $O(n^{2})$  pairwise distances in each iteration. We note that all  $k$  BUILD steps of BanditPAM (respectively, PAM) have the same complexity as each SWAP iteration of BanditPAM (respectively, PAM). Since the number of SWAP iterations is usually  $O(k)$  ([28]; see also our experiments in the Appendix 3), most of BanditPAM's runtime is spent in the SWAP iterations; this suggests improvements to BanditPAM should focus on expediting its SWAP phase.
4 BanditPAM++: Algorithmic Improvements to BanditPAM
In this section, we discuss two improvements to the BanditPAM algorithm. We first show how each SWAP iteration of BanditPAM can be improved via a technique we call Virtual Arms (VA). With this improvement, the modified algorithm can be cast as a SPIMAB. The conversion to a SPIMAB permits a second improvement via a technique we call the Permutation-Invariant Cache (PIC). Whereas the VA technique improves only the SWAP phase, the PIC technique improves both the BUILD and SWAP phases. The VA technique improves the complexity of each SWAP iteration by a factor of  $O(k)$ , whereas the PIC improves the wall-clock runtime of both the BUILD and SWAP phases.
4.1 Virtual Arms (VA)
As discussed in Section 3, most of the runtime of BanditPAM is spent in the SWAP iterations. When evaluating a medoid-nonmedoid pair  $(m,x_{i})$  to potentially swap, BanditPAM estimates the quantity:
$$
\Delta L _ {m, x _ {i}} = \frac {1}{n} \sum_ {s = 1} ^ {n} \Delta l _ {m, x _ {i}} \left(x _ {s}\right), \tag {6}
$$
for each medoid  $m$  and candidate nonmedoid  $x_{i}$ , where
$$
\Delta l _ {m, x _ {i}} \left(x _ {s}\right) = \left(d \left(x _ {i}, x _ {s}\right) - \min  _ {m ^ {\prime} \in \mathcal {M} \backslash \{m \}} d \left(m ^ {\prime}, x _ {s}\right)\right) \wedge 0 \tag {7}
$$
is the change in clustering loss (Equation (3)) induced on point  $x_{s}$  for swapping medoid  $m$  with nonmedoid  $x_{i}$  in the set of medoids  $\mathcal{M}$ . Crucially, we will find that for a given  $x_{s}$ , each  $\Delta l_{m,x_i}(x_s)$  for  $m = 1,\dots ,k$ , except possibly one, is equal. We state this observation formally in Theorem 1.
General SPIMAB Term	BanditPAM++ BUILD step	BanditPAM++ SWAP step
Arms, {Rt}n j=1 Reference points, S Timestep, t Di(xs) f(Di, {a0, a1, ..., at-1})	Candidate points for medoids Points of dataset X t-th medoid to be added Distance between xi and xs Equation 8	Points to swap in as medoids Points of dataset X (t-k)-th swap to be performed Distance between xi and xs Equation 8
Table 1: BanditPAM++'s two phases can each be cast in the SPIMAB framework.
Theorem 1. Let  $\Delta l_{m,x_i}(x_s)$  be the change in clustering loss induced on point  $x_s$  by swapping medoid  $m$  with nonmedoid  $x_i$ , given in Equation (7), with  $x_s$  and  $x_i$  fixed. Then the values  $\Delta l_{m,x_i}(x_s)$  for  $m = 1,\ldots ,k$  are equal, except possibly where  $m$  is the medoid for reference point  $x_s$ .
Theorem 1 is proven in Appendix 2. Crucially, Theorem 1 tells us that when estimating  $\Delta L_{m,x_i}$  in Equation (6) for fixed  $x_{i}$  and various values of  $m$ , we can reuse a significant number of the summands across different indices  $m$  (across  $k$  medoids). We note that Theorem 1 has been observed in alternate forms, e.g., as the "FastPAM1" trick, in prior work [28]. However, to the best of our knowledge, we are the first to provide a formal statement and proof of Theorem 1 and demonstrate its use in an adaptive sampling scheme inspired by multi-armed bandits.
Motivated by Theorem 1, we present the SWAP step of our algorithm, BanditPAM++, in Algorithm 1. BanditPAM++ uses the VA technique to improve the complexity of each SWAP iteration by a factor of  $O(k)$ . We call the technique "virtual arms" because it uses only a constant number of distance computations to update each of the  $k$  "virtual" arms for each of the "real" arms, where a "real" arm corresponds to a datapoint.
4.2 Permutation-Invariant Caching (PIC):
The original BanditPAM algorithms consider each of the  $k(n - k)$  medoid-nonmedoid pairs as arms. With the VA technique described in Section 4.1, BanditPAM++ instead considers each of the  $n$  datapoints (including existing medoids) as arms in each SWAP iteration. Crucially, this implies that the BUILD phase and each SWAP iteration of BanditPAM++ consider the same set of arms. It is this observation, induced by the VA technique, that allows us to cast BanditPAM++ as a SPIMAB and implement a second improvement. We call this second improvement the Permutation-Invariant Cache (PIC).