[논문 리뷰] BanditPAM: Almost Linear Time $k$-Medoids Clustering via Multi-Armed Bandits

Mohit Tiwari, Martin Jinye Zhang외 4명|arXiv (Cornell University)|2020. 06. 11.

Single-cell and spatial transcriptomics인용 수 15

한 줄 요약

BanditPAM은 다수의 랜덤 알고리즘으로, 다수의 랜덤 알고리즘을 사용하여 각 PAM 반복의 계산 복잡도를 $O(n^2)$에서 $O(n\log n)$으로 감소시켜 $k$-medoids 군집화를 가속화한다. 이는 대규모 실세계 데이터셋에서 기준 방법 대비 최대 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과 같은 최신 알고리즘에서 각 반복의 $O(n^2)$ 복잡도로 인한 $k$-medoids 군집화의 높은 계산 비용을 해결한다.
• 군집 중심의 해석 불가능성과 임의의 거리 척도와의 호환성 부족과 같은 $k$-means 군집화의 한계를 극복한다.
• 기존 방법이 너무 느려 대규모 데이터셋(예: 단일세포 RNA-seq 및 온라인 교육 데이터)에서 적용이 불가능한 상황에서 확장 가능한 $k$-medoids 군집화를 가능하게 한다.
• Bandit 기반 샘플링 전략을 통해 런타임과 거리 평가 횟수를 크게 줄이면서도 PAM 수준의 군집화 품질을 유지한다.

제안 방법

• PAM의 SWAP 단계에서 메도이드 교체를 평가하는 데 초점을 맞춰 다수의 랜덤 알고리즘을 $k$-medoids 문제에 적응시킨다.
• 상한 신뢰도(UCB) 랜덤 전략을 사용하여 잠재적 메도이드 교체의 손실 감소를 높은 확률로 샘플링하고 추정한다.
• 각 후보 메도이드와 대상 포인트에 대해 서브-구형 분포 매개변수를 적응적으로 추정하여 손실 추정치의 신뢰 구간을 보장한다.
• 이중 단계 접근법을 적용한다: 첫 번째로, 랜덤 샘플링을 통해 효율적으로 고보상 교체를 식별하고, 두 번째로 상위 후보들을 완전한 거리 계산으로 검증한다.
• 거리 계산이 런타임을 지배하기 때문에, 거리 평가 횟수는 알고리즘 복잡도의 핵심 지표가 된다.
• Bandit 프레임워크를 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), 트리 구조 코드 제출 등 다양한 데이터 유형에서 뛰어난 안정성을 보였다.