BanditPAM: Almost Linear-Time k-medoids Clustering via Multi-Armed Bandits

**Abstract: BanditPAM - An Almost Linear-Time k-medoids Clustering Algorithm via Multi-Armed Bandits** **Introduction:** Clustering algorithms are fundamental tools in machine learning (ML) and data analysis, used to group similar data points into clusters. Two prominent clustering methods are \(k\)-means and \(k\)-medoids. While \(k\)-means is widely used due to its linear time complexity and simplicity, \(k\)-medoids offers several advantages, including greater interpretability of cluster centers and robustness to outliers. However, the traditional \(k\)-medoids algorithm, known as Partitioning Around Medoids (PAM), has a quadratic time complexity (\(O(n^2)\)), which has historically limited its practicality for large datasets. In a recent breakthrough, the BanditPAM algorithm, introduced at NeurIPS 2021, significantly reduces the complexity of \(k\)-medoids clustering to almost linear time (\(O(n\log n)\)), making it a viable alternative to \(k\)-means. **Key Concepts:** - **k-medoids vs. k-means:** - **k-medoids:** This clustering method requires that the cluster centers (medoids) be actual data points from the dataset, enhancing the interpretability of the clusters. It supports arbitrary distance metrics, making it more versatile and robust to outliers compared to \(k\)-means. - **k-means:** Typically uses the \(L_2\) distance metric and computes the mean of data points within each cluster to determine the center. This method is computationally efficient with \(O(n)\) complexity but can produce less interpretable cluster centers and is less robust to outliers. **Advantages of k-medoids:** - **Interpretability:** Medoids are actual data points, which can provide clearer insights into the characteristics of each cluster. For instance, in image clustering, the medoid is a real image, whereas the \(k\)-means center is often a nondescript blob. - **Versatility:** \(k\)-medoids supports any pairwise dissimilarity function, allowing for the clustering of complex data types such as strings, natural language, trees, and graphs without the need to embed them in a vector space. - **Robustness:** Using robust distance metrics like \(L_1\) (median) can make \(k\)-medoids more resilient to outliers compared to \(L_2\) (mean) used in \(k\)-means. **BanditPAM Algorithm:** - **Problem Reformulation:** The core insight of BanditPAM is to reformulate the PAM algorithm's steps as multi-armed bandit problems. This allows for efficient sampling of distances rather than computing every pairwise distance, significantly reducing computational complexity. - **BUILD Step:** In the BUILD step, BanditPAM samples distances to \(O(\log n)\) points for each candidate medoid, rather than computing all \(O(n^2)\) distances. This step initializes the \(k\) medoids. - **SWAP Step:** In the SWAP step, BanditPAM samples distances to \(O(\log n)\) points for each (medoid, non-medoid) pair, reducing the complexity from \(O(kn^2)\) to \(O(kn \log n)\). This step iteratively improves the medoid set by considering potential swaps. **Performance and Implementation:** - **Complexity Reduction:** BanditPAM reduces the time complexity of \(k\)-medoids from \(O(n^2)\) to \(O(n \log n)\), making it almost as efficient as \(k\)-means. - **High-Performance Implementation:** The BanditPAM algorithm is implemented in C++ for speed and supports parallelization and intelligent caching. It is available via Python bindings, making it easy to install and use (\(\texttt{pip install banditpam}\)). - **User-Friendly Interface:** The interface of BanditPAM matches that of \(\texttt{sklearn.cluster.KMeans}\), allowing users to integrate it into existing code with minimal changes. - **Customization:** Users can implement their own distance metrics, further extending the algorithm's flexibility and applicability to various data types. **Conclusion:** BanditPAM represents a significant advancement in \(k\)-medoids clustering, offering the benefits of interpretability, versatility, and robustness with a time complexity that is almost linear. This makes \(k\)-medoids a more practical and powerful tool for clustering large datasets. The algorithm's efficient implementation and user-friendly interface ensure that it can be easily adopted by ML practitioners, potentially leading to more accurate and meaningful clustering results in a wide range of applications. **References:** - **Useful Links:** - 3-minute video summary - PyPI - Github Repository - Full Paper **Acknowledgments:** This summary is based on the paper "BanditPAM: Almost Linear Time \(k\)-medoids Clustering via Multi-Armed Bandits" presented at NeurIPS 2021. The authors, Martin Jinye Zhang, James Mayclin, Sebastian Thrun, Chris Piech, and Ilan Shomorony, thank Drew A. Hudson and Sidd Karamcheti for their feedback on this blog post.