[論文レビュー] PLAID: An Efficient Engine for Late Interaction Retrieval

Keshav Santhanam, Omar Khattab+2 more|arXiv (Cornell University)|May 19, 2022

Topic Modeling被引用数 8

ひとこと要約

PLAID は、ColBERTv2 を用いて後段相互作用検索を高速化する高度に最適化された検索エンジンであり、重心相互作用と重心プルーニングを導入することで、GPU で最大 7 倍、CPU で最大 45 倍の検索遅延低減を達成した。140M パassage までのスケールのデータセットにおいても、最先端の検索品質を維持している。

ABSTRACT

Pre-trained language models are increasingly important components across multiple information retrieval (IR) paradigms. Late interaction, introduced with the ColBERT model and recently refined in ColBERTv2, is a popular paradigm that holds state-of-the-art status across many benchmarks. To dramatically speed up the search latency of late interaction, we introduce the Performance-optimized Late Interaction Driver (PLAID). Without impacting quality, PLAID swiftly eliminates low-scoring passages using a novel centroid interaction mechanism that treats every passage as a lightweight bag of centroids. PLAID uses centroid interaction as well as centroid pruning, a mechanism for sparsifying the bag of centroids, within a highly-optimized engine to reduce late interaction search latency by up to 7$ imes$ on a GPU and 45$ imes$ on a CPU against vanilla ColBERTv2, while continuing to deliver state-of-the-art retrieval quality. This allows the PLAID engine with ColBERTv2 to achieve latency of tens of milliseconds on a GPU and tens or just few hundreds of milliseconds on a CPU at large scale, even at the largest scales we evaluate with 140M passages.

研究の動機と目的

• 大規模な情報検索システムにおいて、最先端のパフォーマンスを示すものの、高遅延である後段相互作用検索の課題を解決すること。
• 残留表現と重心ベースの圧縮を用いる ColBERTv2 のエンドツーエンド検索パイプラインを最適化すること。
• 特に 100M パassage 以上のプロダクション環境向けに、CPU および GPU での低遅延推論を可能にすること。
• 候補スコアリング中の計算およびメモリのオーバーヘッドを著しく削減しながら、検索品質を維持すること。
• 重心 ID の離散性を活用した効率的なフィルタリングとスコアリングを可能にするモジュラで高最適化されたエンジンの開発

提案手法

• 重心相互作用の導入：各パassage を整数の重心 ID の軽量なバッグとして扱い、完全な残留ベクトルをロードせずに高速で近似スコアリングを可能にする。
• 事前計算された重心-クエリ間の距離を用いることで、検索中に全重心に対して類似度スコアを再計算する必要を回避し、複数のパassage 間で再利用可能にする。
• 重心プルーニングの実装：スコアが低い重心 ID のスパarsity を活用し、初期フィルタリング段階で遠く離れた重心 ID をスキップすることで、候補セットのサイズを削減する。
• GPU 向けのパディングフリーな MaxSim カーネルと、CPU 向けの最適化された C++ カーネルを設計し、残留ベクトルの再圧縮およびスコアリング処理を高速化する。
• 完全な残留スコアリングの前段階で重心相互作用とプルーニングを適用するマルチステージ検索パイプラインを構築し、低品質な候補に対する高コストな処理を最小限に抑える。
• 重心の固定で離散的な語彙を利用することで、CPU および GPU 実行における効率的なベクトル化演算およびメモリアクセスパターンを実現する。

実験結果

リサーチクエスチョン

• RQ1 重心ベースの表現のみで、後段相互作用検索において高いリCALLを達成できるか。これにより、完全な残留スコアリングの前段階での効率的なフィルタリングが可能になるか？
• RQ2 重心プルーニングは、検索品質を劣化させることなく、候補パassage の数をどれほど削減できるか？
• RQ3 アルゴリズム最適化（重心相互作用とプルーニング）と低レベルのカーネル最適化を組み合わせることで、CPU および GPU におけるエンドツーエンド遅延をどの程度改善できるか？
• RQ4 異なるハードウェア構成において、コーパスサイズおよび検索深さ（k）の増加に伴い、PLAID はどのようにスケーリングするか？
• RQ5 初期段階で完全な残留スコアリングを重心のみのスコアリングに置き換えた場合、遅延低減と検索品質のトレードオフはどの程度か？

主な発見

• PLAID ColBERTv2 は、MS MARCO v1, v2, Wikipedia, LoTTE ベンチマークにおいて、vanilla ColBERTv2 と比較して GPU で最大 7 倍、CPU で最大 45 倍の高速化を達成し、品質に顕著な損失は認められない。
• 重心相互作用のみで、GPU で 5.2 倍、CPU で 8.6 倍の高速化が達成され、重心のみのスコアリングが高リCALLを維持し、効果的な初期フィルタリングを可能にしていることが示された。
• アルゴリズム最適化（重心相互作用とプルーニング）と低レベルカーネル最適化の組み合わせにより、CPU で k=1000 時に 42.4 倍の高速化が達成された。これは、カーネルのみで得られる 3 倍の高速化と比較して顕著な向上である。
• エンドツーエンド遅延は、データセットサイズの平方根にほぼ比例してスケーリングされ、これは ColBERTv2 における埋め込み数に対する重心数のスケーリングと整合的である。
• 140M パassage でも、CPU で 100ms 未満、GPU で 50ms 未満のエンドツーエンド遅延を達成し、大規模なリアルタイム推論を可能にした。
• CPU スレッド数の増加に伴い、16 スレッドで 4.9 倍の高速化を達成しており、非理想的なロードバランシングにもかかわらず、並列処理の効果が顕著に現れている。