SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We formalize the reduction of BanditPAM++ to a SPIMAB in Table 1. In the SPIMAB formulation of BanditPAM++, the set of reference points  $S$  is the same as the set of datapoints  $\mathcal{X}$ . Each  $D_{i}$  is a random variable representing the distance from point  $x_{i}$  to one of the sampled reference points and can be sampled in  $O(1)$  time without replacement. Each  $\mu_{i} = \mathbb{E}[D_{i}]$  corresponds to the average distance from point  $x_{i}$  to all the points in the dataset  $\mathcal{X}$ . Each arm  $R_{i}^{t}$  corresponds to the point that we would add to the set of medoids (for  $t \leq k$ ) or swap in to the set of medoids (for  $t > k$ ), of which there are  $n$  at each possible timestep. Similarly, the actions  $\{a_0, \dots, a_p, \dots, a_t\}$  correspond to points added to the set of medoids (for  $t \leq k$ ) or swaps performed (for  $t > k$ ). Equation (8) provides the functional forms of  $f$  for each of the BUILD and SWAP steps.
$$
f \left(D _ {i} \left(x _ {s}\right), \mathcal {A}\right) = \left(d \left(x _ {i}, x _ {s}\right) - \min  _ {m ^ {\prime} \in \mathcal {M}} d \left(m ^ {\prime}, x _ {s}\right)\right) \wedge 0. \tag {8}
$$
where  $\mathcal{M}$  is a set of medoids. For the BUILD step,  $\mathcal{A}$  is a sequence of  $t$  actions that results in a set of medoids of size  $t \leq k$ , and, for the SWAP step,  $\mathcal{A}$  is a set of actions that results in a set of medoids  $\mathcal{M}$  of size  $k$ .
The observation that BanditPAM++ is a SPIMAB allows us to develop an intelligent cache design, which we call a permutation-invariant cache (PIC). We may choose a permutation  $\pi$  of the reference points  $S = \mathcal{X}$  and sample distance computations to these reference points in the order of the permutation. Since we only need to sample some of the reference points, and not all of them, we do not need to compute all  $O(n^2)$  pairwise distances. Crucially, we can also reuse distance computations across different steps of BanditPAM++ to save on computational cost and runtime.
Algorithm 1 BanditPAM++ SWAP Step  $(f_{j}(D_{j},\{a_{1},\ldots ,a_{t}\}),\delta ,\sigma_{x},$  permutation  $\pi$  of [n])   
1:  $S_{solution}\gets [n]$  ▷ Set of potential solutions to MAB   
2:  $t^\prime \leftarrow 0$  ▷ Number of reference points evaluated   
3: For all  $(i,j)\in [n]\times [k],$  set  $\hat{\mu}_{i,j}\gets 0,C_{i,j}\gets \infty$  ▷ Initial means and CIs for all swaps   
4: while  $t^\prime < n$  and  $|\mathcal{S}_{\mathrm{solution}}| > 1$  do   
5:  $s\gets \pi (t^{\prime})$  ▷ Uses PIC   
6: for all  $i\in S_{\mathrm{solution}}$  do   
7: Let  $c(s)$  and  $c^{(2)}(s)$  be the indices of  $x_{s}$  's closest and second closest medoids ▷ Cached   
8: Compute distance to  $x_{s}$  's closest medoid  $d_1\coloneqq d(m_{c(s)},x_s)$  ▷ Cached   
9: Compute distance to  $x_{s}$  's second closest medoid  $d_2\coloneqq d(m_{c^{(2)}(s)},x_s)$  ▷ Cached   
10: Compute  $d_i\coloneqq d(x_i,x_s)$  ▷ Reusing  $x_{s}$  's across calls leads to more cache hits   
11:  $\hat{\mu}_{i,c(s)}\gets \frac{t'\hat{\mu}_{i,c(s)} - d_1 + \min(d_2,d_i)}{t' + 1}$  ▷ Update running mean for  $x_{s}$  's medoid   
12:  $C_{i,c(s)}\gets \sigma_i\sqrt{\frac{\log(\frac{1}{\delta})}{t' + 1}}$  ▷ Update confidence interval for  $x_{s}$  's medoid   
13: for all  $j\in \{1,\dots ,k\} \setminus \{c(s)\}$  do   
14:  $\hat{\mu}_{i,j}\gets \frac{t'\hat{\mu}_{i,j} + f(D_i(x_s),a_1,\dots,a_k)}{t' + 1}$  ▷ Update running means; does not depend on j   
15:  $C_{i,j}\gets \sigma_i\sqrt{\frac{\log(\frac{1}{\delta})}{t' + 1}}$  ▷ Update confidence intervals; does not depend on j   
16:  $S_{solution}\gets \{i:\exists j s.t.\hat{\mu}_{i,j} - C_{i,j}\leq \min_{i,j}(\hat{\mu}_{i,j} + C_{i,j})\}$  ▷ Filter suboptimal arms   
17:  $t^\prime \gets t^\prime +1$    
18: if  $|\mathcal{S}_{\mathrm{solution}}| = 1$  then   
19: return  $i^{*}\in S_{\mathrm{solution}}$  and  $j^{*} = \arg \min_{j}\mu_{i^{*},j}$    
20: else   
21: Compute  $\mu_{i,j}$  exactly for all  $i\in S_{\mathrm{solution}}$  ▷ At most 3n distance computations   
22: return  $(i^{*},j^{*}) = \arg \min_{(i,j):i\in S_{\mathrm{solution}}}~\mu_{i,j}$
The full BanditPAM++ algorithm is given in Algorithm 1. Crucially, for each candidate point  $x_{i}$  to swap into the set of medoids on Line 6, we only perform 3 distance computations (not  $k$ ) to update all  $k$  arms, each of which has a mean and confidence interval (CI), on Lines 13-15. This is permitted by Theorem 1 and the VA technique which says that  $k - 1$  "virtual" arms for a fixed  $i$  will get the same update. The PIC technique allows us to choose a permutation of reference points (the  $x_{s}$ 's) and reuse those  $x_{s}$ 's across the BUILD and SWAP steps; as such, many values of  $d(x_{i},x_{s})$  can be cached.
We emphasize that BanditPAM++ uses the same BUILD step as the original BanditPAM algorithm, but with the PIC. The PIC is also used during the SWAP step of BanditPAM++, as is the VA technique. We prove that the full BanditPAM++ algorithm returns the same results as BanditPAM and PAM in Section 5 and demonstrate the empirical benefits of both the PIC and VA techniques in Section 6.
5 Analysis of the Algorithm
In this section, we demonstrate that, with high probability, BanditPAM++ returns the same answer to the  $k$ -medoids clustering problem as PAM and BanditPAM while improving the SWAP complexity of BanditPAM by  $O(k)$  and substantially decreasing its runtime. Since the BUILD step of BanditPAM is the same as the BUILD step of BanditPAM++, it is sufficient to show that each SWAP step of BanditPAM++ returns the same swap as the corresponding step of BanditPAM (and PAM). All of the following theorems are proven in Appendix 2.
First, we demonstrate that PIC does not affect the results of BanditPAM++ in Theorem 2:
Theorem 2. Let  $\mathcal{X} = \{x_1, \ldots, x_S\}$  be the reference points of  $D_i$ , and let  $\pi$  be a random permutation of  $\{1, \ldots, S\}$ . Then for any  $c \leq S$ ,  $\sum_{q=1}^{c} D_i(x_{\pi(p_q)})$  has the same distribution as  $\sum_{q=1}^{c} D_i(x_{p_q})$ , where each  $p_q$  is drawn uniformly without replacement from  $\{1, \ldots, S\}$ .
Intuitively, Theorem 2 says that instead of randomly sampling new reference points at each iteration of BanditPAM++, we may choose a fixed permutation  $\pi$  in advance and sample in permutation order at each step of the algorithm. This allows us to reuse computation across different steps of the algorithm.
We now show that BanditPAM++ returns the same result as BanditPAM (and PAM) in every SWAP iteration and has the same complexity in  $n$  as BanditPAM. First, we consider a single call to Algorithm 1. Let  $\mu_{i} \coloneqq \min_{j \in [k]} \mu_{i,j}$  and let  $i^{*} \coloneqq \arg \min_{i \in [n]} \mu_{i}$  be the optimal point to swap in to the set of medoids, so that the medoid to swap out is  $j^{*} \coloneqq \arg \min_{j \in [k]} \mu_{i^{*},j}$ . For another candidate point  $i \in [n]$  with  $i \neq i^{*}$ , let  $\Delta_{i} \coloneqq \mu_{i} - \mu_{i^{*}}$ , and for  $i = i^{*}$ , let  $\Delta_{i} \coloneqq \min_{j}^{(2)} \mu_{i,j} - \min_{j} \mu_{i,j}$  where  $\min_{j}^{(2)}$  denotes the second smallest value over the indices  $j$ . To state the following results, we will assume that, for a fixed candidate point  $i$  and a randomly sampled reference point  $x_{s}$ , the random variable  $f(D_{i}(x_{s}), A)$  is  $\sigma_{i}$ -sub-Gaussian for some known parameter  $\sigma_{i}$  (which, in practice, can be estimated from the data [28]):
Theorem 3. For  $\delta = 1 / kn^3$ , with probability at least  $1 - \frac{2}{n}$ , Algorithm 1 returns the optimal swap to perform using a total of  $M$  distance computations, where
$$
E [ M ] \leq 6 n + \sum_ {i \in [ n ]} \min \left[ \frac {1 2}{\Delta_ {i} ^ {2}} \left(\sigma_ {i} + \sigma_ {i ^ {*}}\right) ^ {2} \log k n + B, 3 n \right].
$$
Intuitively, Theorem 3 states that with high probability, each SWAP iteration of BanditPAM++ returns the same result as BanditPAM and PAM. Since the BUILD step of BanditPAM++ is the same as the BUILD step of BanditPAM, this implies that BanditPAM++ follows the exact same optimization trajectories as BanditPAM and PAM over the entire course of the algorithm with high probability. We formalize this observation in Theorem 4:
Theorem 4. If BanditPAM++ is run on a dataset  $\mathcal{X}$  with  $\delta = 1 / kn^3$ , then it returns the same set of  $k$  medoids as PAM with probability  $1 - o(1)$ . Furthermore, the total number of distance computations  $M_{\mathrm{total}}$  required satisfies
$$
E [ M _ {\mathrm {t o t a l}} ] = O \left(n \log k n\right).
$$
Note on assumptions: For Theorem 3, we assumed that the data is generated in a way such that the observations  $f(D_{i}(x_{s}),\mathcal{A})$  follow a sub-Gaussian distribution. Furthermore, for Theorem 4, we assume that the  $\Delta_i$ 's are not all close to 0, i.e., that we are not in the degenerate arm setting where many of the swaps are equally optimal, and assume that the  $\sigma_{i}$ 's are bounded (we formalize these assumptions in Appendix 2). These assumptions have been found to hold in many real-world datasets [28]; see Section 7 and Appendices 1.1, and 2 for more formal discussions.
Additionally, we assume that both BanditPAM and BanditPAM++ place a hard constraint  $T$  on the maximum number of SWAP iterations that are allowed. While the limit on the maximum number of swap steps  $T$  may seem restrictive, it is not uncommon to place a maximum number of iterations on iterative algorithms. Furthermore,  $T$  has been observed empirically to be  $O(k)$  [25], consistent with our experiments in Section 6 and Appendix 3.
We note that statements similar to Theorems 3 and 4 can be proven for other values of  $\delta$ . We provide additional experiments to understand the effects of the hyperparameters  $T$  and  $\delta$  in Appendix 3.
Complexity in  $k$ : The original BanditPAM algorithm scales as  $O(kc(k)n\log kn)$ , where  $c(k)$  is a problem-dependent function of  $k$ . Intuitively,  $c(k)$  governs the "hardness" of the problem as a function of  $k$ ; as more medoids are added, the average distance from each point to its closest medoid will decrease and the arm gaps (the  $\Delta_i$ 's) will decrease, increasing the sample complexity in Theorem 3. With the VA technique, BanditPAM++ removes the explicit factor of  $k$ : each SWAP iteration has complexity  $O(c(k)n\log kn)$ . The implicit dependence on  $k$  may still enter through the term  $c(k)$  for a fixed dataset, which we observe in our experiments in Section 6.
Setup: BanditPAM++ consists of two improvements upon the original BanditPAM algorithm: the VA and PIC techniques. We measure the gains of each technique by presenting an ablation study in which we compare the original BanditPAM algorithm (BP), BanditPAM with only the VA technique (BP+VA), BanditPAM with only the PIC (BP+PIC), and BanditPAM with both the VA and PIC techniques (BP++, the final BanditPAM++ algorithm).
First, we demonstrate that all algorithms achieve the same clustering solution and loss as BanditPAM across a variety of datasets and dataset sizes. In particular, this implies that BanditPAM++ matches
Dataset Size (n):	10,000	15,000	20,000	25,000	30,000
MNIST (L2,k=10)	1.00	1.00	1.00	1.00	1.00
CIFAR10 (L1,k=10)	1.00	1.00	1.00	1.00	1.00
20 Newsgroups (cosine,k=5)	1.00	1.00	1.00	1.00	1.00
Table 2: Clustering loss of BanditPAM++, normalized to clustering loss of BanditPAM, across a variety of datasets, metrics, and dataset sizes. In all scenarios, BanditPAM++ matches the loss of BanditPAM (in fact, returns the exact same solution).
prior state-of-the-art in clustering quality. Next, we investigate the scaling of all four algorithms in both  $n$  and  $k$  across a variety of datasets and metrics. We present our results in both sample-complexity and wall-clock runtime. BanditPAM++ outperforms BanditPAM by up to  $10 \times$ . Furthermore, our results demonstrate that each of the VA and PIC techniques improves the runtime of the original BanditPAM algorithm.
For an experiment on a dataset of size  $n$ , we sampled  $n$  datapoints from the original dataset with replacement. In all experiments using the PIC technique, we allowed the algorithm to store up to 1,000 distance computations per point. For the wall-clock runtime and sample complexity metrics, we divide the result of each experiment by the number of swap iterations  $+1$ , where the  $+1$  accounts for the complexity of the BUILD step.
Datasets: We conduct experiments on several public, real-world datasets to evaluate BanditPAM++'s performance: the MNIST dataset, the CIFAR10 dataset, and the 20 Newsgroups dataset. The MNIST dataset [13] contains 70,000 black-and-white images of handwritten digits. The CIFAR10 dataset [12] comprises 60,000 images, where each image consists of  $32 \times 32$  pixels and each pixel has 3 colors. The 20 Newsgroups dataset [19] consist of approximately 18,000 posts on 20 topics split in two subsets: train and test. We used a fixed subsample of 10,000 training posts and embedding them into 385-dimensional vectors using a sentence transformer from HuggingFace [7]. We use the  $L_{2}$ ,  $L_{1}$ , and cosine distances across the MNIST, CIFAR10, and 20 Newsgroups datasets, respectively.
6.1 Clustering/loss quality
First, we assess the solution quality all four algorithms across various datasets, metrics, and dataset sizes. Table 2 shows the relative losses of BanditPAM++ with respect to the loss of BanditPAM; the results for BP+PIC and BP+VA are identical and omitted for clarity. All four algorithms return identical solutions; this demonstrates that neither the VA nor the PIC technique affect solution quality. In particular, this implies that BanditPAM++ matches the prior state-of-the-art algorithms, BanditPAM and PAM, in clustering quality.
6.2 Scaling with  $k$
Figure 1 compares the wall-clock runtime scaling with  $k$  of BP, BP+PIC, BP+VA, and BP++ on same datasets as Figure 1. Across all data subset sizes, metrics, and values of  $k$ , BP++ outperforms each of BP+VA and BP+PIC, both of which in turn outperform BP. As  $k$  increases, the performance gap between algorithms using the VA technique and the other algorithms increases. For example, on the CIFAR10 dataset with  $k = 15$ , BanditPAM++ is over  $10 \times$  faster than BanditPAM. This provides empirical evidence for our claims in Section 5 that the VA technique improves the scaling of the BanditPAM algorithm with  $k$ .
We provide similar experiments that demonstrate the scaling with  $n$  of BanditPAM++ and each of the baseline algorithms in Appendix 3. The results are qualitatively similar to those shown here; in particular, BanditPAM++ outperforms BP+PIC and BP+VA, which both outperform the original BanditPAM algorithm.
7 Conclusions and Limitations
We proposed BanditPAM++, an improvement upon BanditPAM that produces state-of-the-art results for the  $k$ -medoids problem. BanditPAM++ improves upon BanditPAM using the Virtual Arms
(c)
Figure 1: Average wall-clock runtime versus  $k$  for various dataset sizes, metrics, and subsample sizes  $n$  BP++ outperforms BP+PIC and BP+VA, both of which outperform BP. Negligible error bars are omitted for clarity.
technique, which improves the complexity of each SWAP iteration by  $O(k)$ , and the Permutation-Invariant Cache, which allows the reuse of computation across different phases of the algorithm. We prove that BanditPAM++ returns the same results as BanditPAM (and therefore PAM) with high probability. Furthermore, our experimental evidence demonstrates the superiority of BanditPAM++ over baseline algorithms; across a variety of datasets, BanditPAM++ is up to  $10 \times$  faster than prior state-of-the-art while returning the same results. While the assumptions of BanditPAM and BanditPAM++ are likely to hold in many practical scenarios, it is important to acknowledge that these assumptions can impose limitations on our approach. Specifically, when numerous arm gaps are narrow, BanditPAM++ might employ a naive and fall back to brute force computation. Similarly, if the distributional assumptions on arm returns are violated, the complexity of BanditPAM++ may be no better than PAM. We discuss these settings in greater detail in Appendix 1.1.
We would like to thank the anonymous Reviewers, PCs, and ACs for their helpful feedback on our paper. Mo Tiwari was supported by a Stanford Interdisciplinary Graduate Fellowship (SIGF) and a Standard Data Science Scholarship. The work of Ilan Shomorony was supported in part by the National Science Foundation (NSF) under grant CCF-2046991. Martin Zhang is supported by NIH R01 MH115676.
[1] Amin Aghaee, Mehrdad Ghadiri, and Mahdieh Soleymani Baghshah. Active distance-based clustering using k-medoids. Pacific-Asia Conference on Knowledge Discovery and Data Mining, 9651:253–264, 2016.
[2] David Arthur and Sergei Vassilvitskii. K-means++ the advantages of careful seeding. In Proceedings of the eighteenth annual ACM-SIAM symposium on Discrete algorithms, pages 1027-1035, 2007.
[3] Vivek Bagaria, Tavor Z. Baharav, Govinda M. Kamath, and David Tse. Bandit-based monte carlo optimization for nearest neighbors. In Advances in Neural Information Processing Systems, pages 3650-3659, 2019.
[4] Vivek Bagaria, Govinda M. Kamath, Vasilis Ntranos, Martin J. Zhang, and David Tse. Medoids in almost-linear time via multi-armed bandits. In International Conference on Artificial Intelligence and Statistics, pages 500–509, 2018.
[5] Paul S. Bradley, Olvi L. Mangasarian, and W. N. Street. Clustering via concave minimization. In Advances in Neural Information Processing Systems, pages 368-374, 1997.
[6] Vladimir Estivill-Castro and Michael E. Houle. Robust distance-based clustering with applications to spatial data mining. Algorithmica, 30(2):216-242, 2001.
[7] HuggingFace. all-minilm-l6-v2 model card. https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2. Accessed: 2023-10-28.
[8] Anil K. Jain and Richard C. Dubes. Algorithms for clustering data. Prentice-Hall, 1988.
[9] Kevin Jamieson and Robert Nowak. Best-arm identification algorithms for multi-armed bandits in the fixed confidence setting. In Annual Conference on Information Sciences and Systems, pages 1–6, 2014.
[10] Leonard Kaufman and Peter J. Rousseeuw. Clustering by means of medoids. Statistical Data Analysis based on the L1 Norm and Related Methods, pages 405-416, 1987.
[11] Leonard Kaufman and Peter J. Rousseeuw. Partitioning around medoids (program pam). Finding groups in data: an introduction to cluster analysis, pages 68-125, 1990.
[12] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, Citeseer, 2009.
[13] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278-2324, 1998.
[14] Daniel LeJeune, Richard G. Baraniuk, and Reinhard Heckel. Adaptive estimation for approximate k-nearest-neighbor computations. In International Conference on Artificial Intelligence and Statistics, 2019.
[15] Jure Leskovec, Anand Rajaraman, and Jeffrey D. Ullman. Mining of massive data sets. Cambridge university press, 2020.
[16] Stuart Lloyd. Least squares quantization in pmc. IEEE transactions on information theory, 28(2):129-137, 1982.
[17] James MacQueen. Some methods for classification and analysis of multivariate observations. In *Berkeley Symposium on Mathematical Statistics and Probability*, volume 1, pages 281-297, 1967.
[18] Nina Mishra, Robert Schreiber, Isabelle Stanton, and Robert E. Tarjan. Clustering social networks. In International Workshop on Algorithms and Models for the Web-Graph, pages 56-67. Springer, 2007.
[19] Tom Mitchell. Twenty Newsgroups. UCI Machine Learning Repository, 1999. DOI: https://doi.org/10.24432/C5C323.
[20] Gonzalo Navarro. A guided tour to approximate string matching. ACM computing surveys, 33(1):31-88, 2001.
[21] Raymond T. Ng and Jiawei Han. Clarans: A method for clustering objects for spatial data mining. IEEE transactions on knowledge and data engineering, 14(5):1003-1016, 2002.
[22] Michael L Overton. A quadratically convergent method for minimizing a sum of euclidean norms. Mathematical Programming, 27(1):34-63, 1983.
[23] Rameshwar Pratap, Anup Deshmukh, Pratheeksha Nair, and Tarun Dutt. A faster sampling algorithm for spherical  $k$ -means. In Asian Conference on Machine Learning, pages 343-358. PMLR, 2018.
[24] Jiayang Ren, Kaixun Hua, and Yankai Cao. Global optimal k-medoids clustering of one million samples. Advances in Neural Information Processing Systems, 35:982-994, 2022.
[25] Erich Schubert and Peter J Rousseeuw. Faster k-medoids clustering: improving the pam, clara, and clarans algorithms. In International Conference on Similarity Search and Applications, pages 171-187. Springer, 2019.
[26] Mo Tiwari. Accelerating machine learning algorithms with adaptive sampling. arXiv preprint arXiv:2309.14221, 2023.
[27] Mo Tiwari, Ryan Kang, Jaeyong Lee, Chris Piech, Ilan Shomorony, Sebastian Thrun, and Martin J Zhang. Mabsplit: Faster forest training using multi-armed bandits. Advances in Neural Information Processing Systems, 35:1223–1237, 2022.
[28] Mo Tiwari, Martin Jinye Zhang, James Mayclin, Sebastian Thrun, Chris Piech, and Ilan Shomorony. Banditpam: Almost linear time k-medoids clustering via multi-armed bandits. In Advances in Neural Information Processing Systems, 2020.