[論文レビュー] TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s

Felix Chern, Blake A. Hechtman+4 more|arXiv (Cornell University)|Jun 28, 2022

Algorithms and Data Compression被引用数 5

ひとこと要約

本論文では、メモリおよび命令レベルのボトル neck をモデル化することで、Google の TPU でピーク FLOP/s 性能を達成する TPU-KNN と呼ばれる K 近傍探索アルゴリズムを提示する。この手法は、複雑なインデックス構造やハイパーパramータチューニングを必要とせず、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.

研究の動機と目的

• 両方のメモリおよび命令レベルのボトル neck をモデル化することで、TPU ハードウェアでピーク性能を達成する K-NN アルゴリズムの設計。
• 複雑なインデックス構造やハイパーパramータチューニングに依存せず、期待値における再現率保証を有する手法の開発。
• 類似の再現率条件下で、TPU ベースの K-NN が最先端の GPU ベースのアルゴリズムを上回ることの実証。
• 命令スループットの差を考慮するように屋根線性能モデルを拡張し、アクセcelレータ上でアルゴリズムの性能を正確に予測可能にする。
• TPU デプロイメントを想定した生産準備完了のオープンソース実装（JAX および TensorFlow で提供）の提供。

提案手法

• 著者らは、異なる命令スループットを考慮するように屋根線モデルを拡張し、アクセcelレータの性能ボトル neck を正確に予測可能にする。
• データ移動を最小限に抑え、計算とメモリ帯域幅をバランスさせることで、FLOP/s 利用率を最大化する K-NN アルゴリズムを設計する。
• すべてのデータポイントに対してブルートフォースな距離計算を実行し、再現率を圧縮アーティファクトではなく順序統計に基づくものにする。
• パフォーマンスモデルから導出された効率的なベクトル化、バッチ処理、メモリアクセスパターンを活用した TPU 特有の最適化を採用する。
• 低スループット演算からの帯域幅ボトル neck を特定・緩和するため、係数単位の命令モデルを用いる。
• JAX および TensorFlow でオープンソース化され、TPU での効率的なトップ-K 選択を可能にする新しい演算 'approx_max_k' を公開する。

実験結果

リサーチクエスチョン

• RQ1 両方のメモリおよび命令レベルのボトル neck をモデル化することで、TPU でピーク FLOP/s 利用率を達成する K-NN アルゴリズムを設計可能か？
• RQ2 命令スループットの差を含めるように拡張された屋根線ベースのパフォーマンスモデルは、K-NN ワークロードのアクセcelレータ性能を正確に予測可能か？
• RQ3 IVF-Flat や IVF-PQ などの複雑な圧縮ドメイン探索手法と比較して、TPU で類似の再現率レベルで単純なブルートフォース K-NN アプローチが優れた性能を発揮可能か？
• RQ4 複雑なインデックス構造やハイパーパramータチューニングを回避しつつ、解析的保証付きの高い再現率を達成可能か？
• RQ5 類似の再現率条件下で、TPU 最適化 K-NN アルゴリズムのパフォーマンスは、最先端の GPU ベースの K-NN 方法と比較してどうか？

主な発見

• TPU-KNN は、Sift1M および Glove1.2M ベンチマークで最先端の速度-再現率トレードオフを達成し、類似の再現率レベルで GPU ベースの IVF-Flat、IVF-PQ、フラット探索手法を上回る。
• TPU 実装は、拡張された屋根線モデルからのパフォーマンス予測と非常に近い結果を示し、性能モデリングフレームワークの正確性を裏付けた。
• データ移動を最小限に抑え、計算帯域幅とメモリ帯域幅を効率的にバランスさせることで、TPU でピーク FLOP/s 利用率を達成した。
• 順序統計に基づくため、データ分布のシフトに強く、再現率がデータセット間で一貫性を保つ。
• JAX および TensorFlow でのオープンソース公開により、生産環境でのデプロイメントおよび既存の ML パイプラインへの統合が可能になった。
• 次元の呪いに影響されない。次元削減や圧縮を伴わず、すべてのペアワイズ距離を計算するためである。