BanditPAM++: Faster $k$-medoids Clustering (2310.18844v1)

Abstract: 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 $\textit{within}$ each iteration. Second, we demonstrate that BanditPAM has additional structure that permits the reuse of information $\textit{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.

• The paper introduces two key enhancements—Virtual Arms and Permutation-Invariant Caching—that significantly reduce unnecessary distance computations.
• It demonstrates over 10x speedup on datasets like CIFAR-10 while maintaining the same clustering accuracy as BanditPAM.
• The method reformulates the SWAP phase as an SPIMAB problem, ensuring effective selection of optimal swaps under sub-Gaussian data assumptions.

BanditPAM++: Faster $k$-medoids Clustering

This paper introduces BanditPAM++, a novel algorithm for $k$-medoids clustering that builds upon the existing BanditPAM framework to enhance computational efficiency. The $k$-medoids problem involves selecting representative data points (medoids) from a dataset, which provides better interpretability compared to $k$-means clustering where centers can be arbitrary points in space. BanditPAM++ achieves significant improvements in computational complexity, making it particularly suited for handling large datasets.

Key Contributions

The authors propose two primary algorithmic enhancements to the BanditPAM method:

1. Virtual Arms (VA): This technique optimizes the SWAP phase of BanditPAM by reducing the number of distance computations necessary. By recognizing that many swap evaluations yield redundant information, the algorithm can reuse computed distances. This leads to substantial gains in efficiency, improving the complexity by a factor of $O(k)$.
2. Permutation-Invariant Caching (PIC): This caching strategy allows the algorithm to reuse previously computed information across different iterations. By sampling reference points in a predetermined order, computational reuse becomes feasible, further reducing runtime.

Collectively, these improvements enable BanditPAM++ to maintain the accuracy of clustering achieved by BanditPAM while significantly reducing computational costs, particularly notable in wall-clock runtimes.

Numerical Results

The authors demonstrate that BanditPAM++ maintains the clustering loss of BanditPAM, achieving the same clustering accuracy. On the CIFAR-10 dataset, BanditPAM++ is observed to be over 10 times faster than BanditPAM while producing identical clustering results. These impressive speedups are consistent across diverse datasets, including MNIST and 20 Newsgroups, highlighting the robustness of the proposed enhancements.

Theoretical Considerations

The theoretical validation is structured around the formulation of the SWAP phase as a Sequential Permutation-Invariant Multi-Armed Bandit (SPIMAB) problem, which underpins the efficiency gains. Assuming a sub-Gaussian distribution of data, the authors prove that the likelihood of selecting the optimal swap remains high, thus ensuring the reliability of the results.

Implications and Future Work

The introduction of BanditPAM++ signifies a substantial leap forward in $k$-medoids clustering performance, especially important in the current era where data is plentiful and computational resources are strained. By reducing the computational burden, this algorithm opens avenues for applying $k$-medoids clustering to larger and more complex datasets—enabling new research opportunities and practical applications.

Practically, the implementation in C++ with interfaces for Python and R increases accessibility for practitioners, promoting its adoption in real-world scenarios. The authors acknowledge assumptions such as sub-Gaussian data distributions and typical swap iterations being $O(k)$, which may not hold universally. Future developments could focus on refining these assumptions, potentially broadening the algorithm's applicability.

In summary, BanditPAM++ represents an important step in clustering methodologies, striking an effective balance between interpretability and computational efficiency, and setting a new benchmark for $k$-medoids clustering algorithms.