[论文解读] BanditPAM: Almost Linear Time $k$-Medoids Clustering via Multi-Armed Bandits

Mohit Tiwari, Martin Jinye Zhang+4 more|arXiv (Cornell University)|Jun 11, 2020

Single-cell and spatial transcriptomics被引用 15

一句话总结

BanditPAM 是一种随机化算法，通过使用多臂赌博机技术，将每次 PAM 迭代的计算复杂度从 $O(n^2)$ 降低至 $O(n\log n)$，从而加速 $k$-medoids 聚类。该方法在实现 SOTA 聚类质量的同时，相比基线方法在大规模真实数据集上减少了高达 200 倍的距离计算量，并实现了最高 4 倍的运行速度提升。

ABSTRACT

Clustering is a ubiquitous task in data science. Compared to the commonly used $k$-means clustering, $k$-medoids clustering requires the cluster centers to be actual data points and support arbitrary distance metrics, which permits greater interpretability and the clustering of structured objects. Current state-of-the-art $k$-medoids clustering algorithms, such as Partitioning Around Medoids (PAM), are iterative and are quadratic in the dataset size $n$ for each iteration, being prohibitively expensive for large datasets. We propose BanditPAM, a randomized algorithm inspired by techniques from multi-armed bandits, that reduces the complexity of each PAM iteration from $O(n^2)$ to $O(n \\log n)$ and returns the same results with high probability, under assumptions on the data that often hold in practice. As such, BanditPAM matches state-of-the-art clustering loss while reaching solutions much faster. We empirically validate our results on several large real-world datasets, including a coding exercise submissions dataset, the 10x Genomics 68k PBMC single-cell RNA sequencing dataset, and the MNIST handwritten digits dataset. In these experiments, we observe that BanditPAM returns the same results as state-of-the-art PAM-like algorithms up to 4x faster while performing up to 200x fewer distance computations. The improvements demonstrated by BanditPAM enable $k$-medoids clustering on a wide range of applications, including identifying cell types in large-scale single-cell data and providing scalable feedback for students learning computer science online. We also release highly optimized Python and C++ implementations of our algorithm.

研究动机与目标

• 解决 $k$-medoids 聚类的高计算成本问题，特别是像 PAM 和 FastPAM1 这类先进算法中每次迭代的 $O(n^2)$ 复杂度。
• 克服 $k$-means 聚类的局限性，例如聚类中心不可解释以及与任意距离度量不兼容的问题。
• 实现在大规模数据集（如单细胞 RNA-seq 和在线教育数据）上的可扩展 $k$-medoids 聚类，这些数据集此前的算法因速度过慢而难以应用。
• 在保持与 PAM 相当的聚类质量的同时，通过基于赌博机的采样策略显著减少运行时间和距离评估次数。

提出的方法

• 将多臂赌博机算法适配至 $k$-medoids 问题，特别针对 PAM 中的 SWAP 步骤，即评估潜在的聚类中心替换。
• 使用上界置信区间（UCB）赌博机策略，以高概率采样并估计潜在聚类中心替换带来的损失减少量。
• 为每个候选聚类中心和目标点自适应估计次高斯参数，以确保损失估计的置信区间。
• 采用两阶段方法：首先，利用赌博机采样高效识别高回报的替换；其次，通过完整距离计算验证排名靠前的候选。
• 利用距离计算主导运行时间的事实，使距离计算次数成为算法复杂度的关键代理指标。
• 将赌博机框架集成至 PAM 的 SWAP 阶段，将原本 $O(kn^2)$ 的穷举比较替换为在温和数据假设下的 $O(n\log n)$ 期望计算量。

实验结果

研究问题

• RQ1 多臂赌博机技术能否有效应用于降低 $k$-medoids 的 SWAP 步骤计算成本，同时不牺牲聚类质量？
• RQ2 所提出的基于赌博机的采样策略在数据集规模 $n$ 上是否实现亚二次方扩展，实际中达到 $O(n\log n)$ 的复杂度？
• RQ3 在保持与 PAM 相同的最终聚类中心分配的前提下，距离计算次数最多可减少多少？
• RQ4 损失估计的实证分布与臂参数如何支持次高斯性与分布建模的理论假设？
• RQ5 BanditPAM 在包含多种数据类型和距离度量的真实数据集上，能否实现与 PAM 和 FastPAM1 相当或更优的聚类性能？

主要发现

• BanditPAM 将每次迭代的距离计算次数减少至 $O(n\log n)$，在双对数图上测得斜率为 1.046，证实了近似线性扩展。
• 在 10x Genomics 68k PBMC 数据集上，BanditPAM 的距离计算量相比 PAM 最多减少了 200 倍，同时达到了完全相同的聚类结果。
• 在 Code.org 的 HOC4 数据集实验中，BanditPAM 实现了与 PAM 相同的聚类中心分配，但运行速度最高提升了 4 倍。
• 损失估计的实证分布与真实臂参数紧密匹配次高斯假设，验证了理论建模选择的合理性。
• BanditPAM 在聚类质量上匹配或超越了 FastPAM1 和 PAM，同时在运行时间和距离评估效率方面显著优于二者。
• 该算法在多种数据类型上表现稳健，包括单细胞 RNA-seq、手写数字（MNIST）以及树状结构的代码提交数据。