Adaptive Retrieval and Scalable Indexing for k-NN Search with Cross-Encoders

要約

クエリ-項目のペアを共同で符号化することで類似度を計算するクロスエンコーダ（CE）モデルは、クエリ-項目の関連性を推定する際に、埋め込みベースのモデル（デュアルエンコーダ）よりも優れた性能を発揮する。既存のアプローチは、CEの類似度を、デュアルエンコーダ（DE）またはCUR行列分解に適合するベクトル埋め込み空間で近似することにより、CEを用いてk-NN検索を行う。DEに基づく検索と再ランクのアプローチは、新しいドメインでの再現性の低さに悩まされており、DEによる検索はCEから切り離されている。CURベースのアプローチはDEベースのアプローチよりも精度が高いが、項目の埋め込みを計算するために法外な数のCEコールを必要とするため、大規模な展開には非現実的である。本論文では、CEスコアを近似するために潜在的なクエリとアイテムの埋め込みを効率的に計算し、近似CE類似度を用いてk-NN探索を実行する、我々の提案するスパース行列分解に基づく手法により、これらの欠点に対処する。我々は、訓練クエリ集合に対するクエリ-項目CEスコアを含む疎行列を因数分解することにより、項目埋め込みをオフラインで計算する。本手法は、CURベースの手法と比較して、CE呼び出しの数分の1しか必要とせず、埋め込み空間を初期化するためにDEを活用することができる一方で、蒸留によるDEの計算量とリソース集約的な微調整を回避しながら、高品質な近似を生成する。テスト時には、項目の埋め込みは固定されたままであり、検索は、a)これまでに検索された項目のCEスコアの近似誤差を最小化することによりテストクエリの埋め込みを推定する、b)更新されたテストクエリの埋め込みを用いてより多くの項目を検索する、ということを交互に繰り返しながら、ラウンドを経て行われる。我々のk-NN検索手法は、DEベースのアプローチと比較して、最大5% (k=1)、54% (k=100)までリコールを改善する。さらに、我々のインデックス作成手法は、CURベースの手法に対して100倍、DE蒸留手法に対して5倍のスピードアップを達成し、同時にk-NN検索リコールはベースラインと同等か向上する。

要約(オリジナル)

Cross-encoder (CE) models which compute similarity by jointly encoding a query-item pair perform better than embedding-based models (dual-encoders) at estimating query-item relevance. Existing approaches perform k-NN search with CE by approximating the CE similarity with a vector embedding space fit either with dual-encoders (DE) or CUR matrix factorization. DE-based retrieve-and-rerank approaches suffer from poor recall on new domains and the retrieval with DE is decoupled from the CE. While CUR-based approaches can be more accurate than the DE-based approach, they require a prohibitively large number of CE calls to compute item embeddings, thus making it impractical for deployment at scale. In this paper, we address these shortcomings with our proposed sparse-matrix factorization based method that efficiently computes latent query and item embeddings to approximate CE scores and performs k-NN search with the approximate CE similarity. We compute item embeddings offline by factorizing a sparse matrix containing query-item CE scores for a set of train queries. Our method produces a high-quality approximation while requiring only a fraction of CE calls as compared to CUR-based methods, and allows for leveraging DE to initialize the embedding space while avoiding compute- and resource-intensive finetuning of DE via distillation. At test time, the item embeddings remain fixed and retrieval occurs over rounds, alternating between a) estimating the test query embedding by minimizing error in approximating CE scores of items retrieved thus far, and b) using the updated test query embedding for retrieving more items. Our k-NN search method improves recall by up to 5% (k=1) and 54% (k=100) over DE-based approaches. Additionally, our indexing approach achieves a speedup of up to 100x over CUR-based and 5x over DE distillation methods, while matching or improving k-NN search recall over baselines.

arxiv情報

提供元, 利用サービス

arxiv.jp, DeepL