BanditPAM：基于多臂老虎机的近线性时间 $k$-Medoids 聚类

聚类是数据科学中普遍存在的任务。与常用的 $k$-means 聚类相比，$k$-medoids 聚类要求簇中心为实际数据点并支持任意距离度量，这带来了更强的可解释性以及对结构化对象的聚类能力。当前最先进的 $k$-medoids 聚类算法，如围绕 medoids 划分（PAM），是迭代式的，且每次迭代在数据集大小 $n$ 上呈二次复杂度，对大型数据集而言极其昂贵。我们提出 BanditPAM，一种受多臂老虎机技术启发的随机化算法，在数据上满足实践中常成立的假设下，将每次 PAM 迭代的复杂度从 $O(n^2)$ 降至 $O(n \log n)$，并以高概率返回相同结果。因此，BanditPAM 在匹配最先进聚类损失的同时更快地得到解。我们在多个大型真实数据集上实证验证结果，包括编程练习提交数据集、10x Genomics 68k PBMC 单细胞 RNA 测序数据集和 MNIST 手写数字数据集。在这些实验中，我们观察到 BanditPAM 返回与最先进类 PAM 算法相同的结果，最高快 4 倍，同时执行最多少 200 倍的距离计算。BanditPAM 所展现的改进使得 $k$-medoids 聚类能应用于广泛场景，包括大规模单细胞数据中的细胞类型识别，以及为在线学习计算机科学的学生提供可扩展反馈。我们还发布了高度优化的 Python 与 C++ 算法实现。