BanditPAM: K-Medoids Finally Catches Up to K-Means

K-means clustering has long dominated the field, not because it's ideal, but because it's fast. Its lesser-known cousin, k-medoids, offers compelling advantages, if only it weren't so slow. Stanford researchers have changed that calculus with BanditPAM, an algorithm that makes k-medoids clustering nearly as fast as k-means while preserving its unique benefits.

K-medoids forces cluster centers to be actual data points, not artificial averages. This seemingly small distinction yields major practical benefits:

1. Interpretability: When clustering ImageNet images with k=1, k-means produces a nondescript blob. K-medoids selects an actual image, in this case, a bee. This real-world interpretability is crucial in many domains.

2. Flexibility: K-medoids works with any dissimilarity function, not just L2 distance. This allows clustering of non-vector data like strings, graphs, or trees without embedding them first.

3. Robustness: Using L1 distance instead of L2 can make clustering more resistant to outliers. K-medoids supports this natively.

Despite these advantages, k-medoids has seen limited adoption due to its prohibitive O(n^2) time complexity compared to k-means' O(n). BanditPAM, introduced in a NeurIPS paper, slashes that complexity to O(n log n), nearly matching k-means' speed.

The key insight? You don't need to compute every pairwise distance to find good medoids. By cleverly sampling distances using multi-armed bandit techniques, BanditPAM achieves the same quality as exhaustive approaches with far less computation.

The algorithm builds on the classic Partitioning Around Medoids (PAM) approach but replaces its exhaustive distance calculations with strategic sampling. This sampling is framed as a multi-armed bandit problem, where each potential medoid is an 'arm' to be pulled.

Crucially, BanditPAM recovers the exact same solutions as PAM, just much faster. This isn't a heuristic approximation, it's a more efficient path to the optimal answer.

The researchers have released a high-performance C++ implementation with Python bindings, installable via pip. It mimics the scikit-learn KMeans interface, allowing easy drop-in replacement:

from banditpam import KMedoids
kmeds = KMedoids(n_clusters=5, metric='manhattan')
kmeds.fit(X)

BanditPAM isn't a universal replacement for k-means. For very large datasets where O(n) beats O(n log n), k-means can still be faster. And when interpretable centers aren't needed and L2 distance works well, k-means remains solid.

But for tasks requiring robust, interpretable clustering, especially with non-Euclidean data, BanditPAM removes the primary barrier to k-medoids adoption. It opens up new possibilities in anomaly detection, customer segmentation, and pattern recognition in complex data types.

The broader impact of BanditPAM may be in demonstrating how classical algorithms can be dramatically improved through modern machine learning techniques. By reframing parts of PAM as bandit problems, the researchers achieved an algorithmic speedup that had eluded the field for decades. This success suggests there may be other algorithmic gems waiting to be polished by similar techniques.