[論文レビュー] 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$-メディオイドクラスタリングを高速化する確率的アルゴリズムである。大規模な実世界のデータセットにおいて、ベースライン手法と比較して最大 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.

研究の動機と目的

• PAM や FastPAM1 などの最先端のアルゴリズムで見られる、1 イテレーションあたり $O(n^2)$ の計算コストに起因する $k$-メディオイドクラスタリングの高コストを軽減すること。
• クラスタ中心の解釈不能性や任意の距離尺度との非適合性といった $k$-メディオイドクラスタリングの限界を克服すること。
• 従来の手法が著しく遅いために実行不可能である、単細胞 RNA-Seq やオンライン教育データなど大規模データセットにおけるスケーラブルな $k$-メディオイドクラスタリングを可能にすること。
• PAM と同等のクラスタリング品質を維持しながら、バンディットに基づくサンプリング戦略によって実行時間と距離評価回数を著しく削減すること。

提案手法

• 多腕バンディットアルゴリズムを $k$-メディオイド問題に適応し、特に PAM の SWAP ステップにおけるメディオイド交換の評価を対象とする。
• 上界信頼区間（UCB）バンディット戦略を用い、高い確率で損失低減を推定するためのサンプリングを行う。
• 各候補メディオイドとターゲットポイントに対して、サブガウス型パrameter を動的に推定し、損失推定値の信頼区間を保証する。
• 2段階アプローチを採用：まずバンディットサンプリングを用いて高報酬の交換を効率的に同定し、次に上位候補を完全な距離計算で検証する。
• 距離計算が実行時間の大部分を占めるため、距離評価回数をアルゴリズムの複雑性の主要な代理指標とする。
• PAM の SWAP ステップにバンディットフレームワークを統合し、$O(kn^2)$ の全探索的比較を、やや緩いデータ仮定のもとで $O(n\log n)$ の期待評価回数に置き換える。

実験結果

リサーチクエスチョン

• RQ1 多腕バンディット手法を、クラスタリング品質を損なわせることなく、$k$-メディオイドの SWAP ステップの計算コストを低減するために効果的に適用できるか？
• RQ2 提案されたバンディットベースのサンプリング戦略が、データセットサイズ $n$ に対して準二次的（sub-quadratic）にスケーリングされ、実際には $O(n\log n)$ の複雑性を達成できるか？
• RQ3 PAM と同一の最終メディオイド割り当てを維持しつつ、距離評価回数をどの程度削減できるか？
• RQ4 損失推定値とアームパラメータの経験的分布が、サブガウス型仮定と分布モデリングの理論的仮定をどの程度裏付けているか？
• RQ5 BanditPAM は、多様なデータタイプと距離尺度を有する実世界のデータセットにおいて、PAM や FastPAM1 と同等またはそれ以上のクラスタリングパフォーマンスを達成できるか？

主な発見

• BanditPAM は 1 イテレーションあたりの距離評価回数を $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）、木構造のコード提出など、多様なデータタイプにおいても頑健であった。