[논문 리뷰] TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s

Felix Chern, Blake A. Hechtman외 4명|arXiv (Cornell University)|2022. 06. 28.

Algorithms and Data Compression인용 수 5

한 줄 요약

이 논문은 메모리 및 명령어 수준의 병목 현상을 최적화하기 위해 정교화된 룻라인 모델을 활용하여 Google의 TPU에서 최고의 FLOP/s 성능을 달성하는 TPU-KNN이라는 K-최근접 이웃 검색 알고리즘을 제시한다. 이 방법은 복잡한 색인 구조나 하이퍼파라미터 튜닝 없이도 상태 기반의 속도-재현율 트레이드오프를 제공하며, Sift1M 및 Glove1.2M와 같은 표준 벤치마크에서 유사한 재현율 수준에서 GPU 기반 기준선을 능가한다.

ABSTRACT

This paper presents a novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-the-art GPU algorithms with similar level of recall. The design of the proposed algorithm is motivated by an accurate accelerator performance model that takes into account both the memory and instruction bottlenecks. Our algorithm comes with an analytical guarantee of recall in expectation and does not require maintaining sophisticated index data structure or tuning, making it suitable for applications with frequent updates. Our work is available in the open-source package of Jax and Tensorflow on TPU.

연구 동기 및 목표

• 메모리 및 명령어 수준의 병목 현상을 모두 모델링하여 TPU 하드웨어에서 최고 성능을 달성하는 K-NN 알고리즘을 설계하기.
• 복잡한 색인 구조나 하이퍼파라미터 튜닝에 의존하지 않고 기대치 기반의 재현율 보장을 갖는 방법을 개발하기.
• 유사한 재현율 조건에서 TPU 기반 K-NN가 최신 기술 수준의 GPU 기반 알고리즘을 능가할 수 있음을 입증하기.
• 명령어 처리량의 차이를 고려하여 룻라인 성능 모델을 확장하여 가속기에서 알고리즘 성능을 정확하게 예측할 수 있도록 하기.
• TPU 배포를 위한 JAX 및 텐서플로우 기반의 프로덕션 준비 완료된 오픈소스 구현을 제공하기.

제안 방법

• 저자들은 명령어 처리량의 다양성을 반영하기 위해 룻라인 모델을 확장하여 가속기 성능 병목 현상을 정확하게 예측할 수 있도록 한다.
• 데이터 이동을 최소화하고 계산 및 메모리 대역폭을 균형 있게 조절함으로써 FLOP/s 활용도를 극대화하는 K-NN 알고리즘을 설계한다.
• 모든 데이터 포인트에 대해 브루트포스 거리 계산을 수행함으로써 재현율이 압축 아티팩트가 아닌 순서 통계에 기반함을 보장한다.
• 성능 모델에서 유도된 효율적인 벡터화, 배치 처리, 메모리 접근 패턴을 활용한 TPU 전용 최적화를 적용한다.
• 낮은 처리량을 가지는 연산에서 발생하는 대역폭 병목 현상을 식별하고 완화하기 위해 계수별 명령어 모델을 사용한다.
• JAX 및 텐서플로우에서 오픈소스로 제공되며, TPU에서 효율적인 상위-K 선택을 위한 새로운 연산 'approx_max_k'를 노출한다.

실험 결과

연구 질문

• RQ1 메모리 및 명령어 수준의 병목 현상을 모두 모델링함으로써 TPU에서 최고의 FLOP/s 활용도를 달성할 수 있는 K-NN 알고리즘을 설계할 수 있는가?
• RQ2 명령어 처리량의 차이를 포함하도록 확장된 룻라인 기반 성능 모델이 K-NN 워크로드에 대한 가속기 성능을 정확하게 예측할 수 있는가?
• RQ3 비틀림된 도메인 검색 방법(예: IVF-Flat 및 IVF-PQ)보다 더 복잡한 구조를 갖지 않은 단순한 브루트포스 K-NN 접근 방식이 유사한 재현율 수준에서 TPU에서 더 뛰어난 성능을 낼 수 있는가?
• RQ4 복잡한 색인 구조나 하이퍼파라미터 튜닝 없이도 분석적으로 보장된 높은 재현율을 달성할 수 있는가?
• RQ5 유사한 재현율 조건에서 TPU 최적화된 K-NN 알고리즘의 성능은 최신 기술 수준의 GPU 기반 K-NN 방법과 비교해 어떻게 되는가?

주요 결과

• TPU-KNN은 Sift1M 및 Glove1.2M 벤치마크에서 GPU 기반 IVF-Flat, IVF-PQ 및 플랫 검색 방법과 유사한 재현율 수준에서 최고의 속도-재현율 트레이드오프를 달성한다.
• TPU 구현은 확장된 룻라인 모델의 성능 예측과 매우 유사하게 성능을 보이며, 성능 모델링 프레임워크의 정확성을 검증한다.
• 데이터 이동을 최소화하고 계산 및 메모리 대역폭을 효율적으로 균형 있게 조절함으로써 TPU에서 최고의 FLOP/s 활용도를 달성한다.
• 알고리즘이 순서 통계에 기반하기 때문에 데이터 분포의 변화에 대해 강건하며, 다양한 데이터셋에서 일관된 재현율을 유지한다.
• JAX 및 텐서플로우에서의 오픈소스 배포로 인해 TPU 하드웨어에서의 프로덕션 배포 및 기존 머신러닝 파ip라인과의 통합이 가능해졌다.
• 알고리즘이 차원의 저주에 영향을 받지 않으며, 차원 축소나 압축 없이 모든 쌍의 거리를 계산하기 때문이다.