[論文レビュー] Distill-VQ: Learning Retrieval Oriented Vector Quantization By Distilling Knowledge from Dense Embeddings

Shitao Xiao, Zheng Liu+11 more|arXiv (Cornell University)|Apr 1, 2022

Advanced Image and Video Retrieval Techniques被引用数 4

ひとこと要約

Distill-VQは、密な埋め込みからの関連性予測を模倣することで、検索性能を維持する知識蒸留フレームワークを提案する。IVFおよびPQのベクトル量子化（VQ）モジュールを訓練し、ラベルなしのクエリを必要とせず、MS MARCOおよびNatural Questionsで最先端のリCALLおよびMRRを達成した。高圧縮および低レイテンシー下でも同様の性能を発揮する。

ABSTRACT

Vector quantization (VQ) based ANN indexes, such as Inverted File System (IVF) and Product Quantization (PQ), have been widely applied to embedding based document retrieval thanks to the competitive time and memory efficiency. Originally, VQ is learned to minimize the reconstruction loss, i.e., the distortions between the original dense embeddings and the reconstructed embeddings after quantization. Unfortunately, such an objective is inconsistent with the goal of selecting ground-truth documents for the input query, which may cause severe loss of retrieval quality. Recent works identify such a defect, and propose to minimize the retrieval loss through contrastive learning. However, these methods intensively rely on queries with ground-truth documents, whose performance is limited by the insufficiency of labeled data. In this paper, we propose Distill-VQ, which unifies the learning of IVF and PQ within a knowledge distillation framework. In Distill-VQ, the dense embeddings are leveraged as "teachers", which predict the query's relevance to the sampled documents. The VQ modules are treated as the "students", which are learned to reproduce the predicted relevance, such that the reconstructed embeddings may fully preserve the retrieval result of the dense embeddings. By doing so, Distill-VQ is able to derive substantial training signals from the massive unlabeled data, which significantly contributes to the retrieval quality. We perform comprehensive explorations for the optimal conduct of knowledge distillation, which may provide useful insights for the learning of VQ based ANN index. We also experimentally show that the labeled data is no longer a necessity for high-quality vector quantization, which indicates Distill-VQ's strong applicability in practice.

研究の動機と目的

• 従来のベクトル量子化（VQ）の目的関数（再構成誤差の最小化）と実際の検索性能との間の不一致を是正する。
• 対照学習に高価で限られたラベル付きクエリ-ドキュメントペairに依存する既存の検索最適化VQ手法の課題を克服する。
• 密な埋め込みを知識の教師として活用することで、大量のラベルなしデータを用いたVQモジュールの有効な訓練を可能にする。
• IVFおよびPQの両成分を同時に蒸留することで、高圧縮および低レイテンシー下での検索品質を向上させる。
• ラベル付きデータが不要であるにもかかわらず、高品質なベクトル量子化が実現可能であることを実証する。

提案手法

• 事前学習済みの密な埋め込みを「教師」として、クエリと候補ドキュメント間の関連性スコアを予測する。
• VQモジュール（IVFおよびPQ）を「生徒」として、教師の関連性予測を再現するように訓練する。
• 教師の予測順序と生徒の予測順序の不変性を保つために、ListNet損失を適用する。
• 真のラベルを必要とせず、Top-Kおよびバッチ内候補からドキュメントをサンプリングして、知識蒸留信号を構築する。
• 統一された蒸留フレームワーク内でIVFおよびPQの両成分を同時に最適化し、検索性能を維持する。
• 訓練済みのVQモジュールをFAISSのIVFPQインデックスに統合し、効率的な推論と評価を実現する。

実験結果

リサーチクエスチョン

• RQ1 ラベル付きクエリ-ドキュメントペアに依存せずに、密な埋め込みからの知識蒸留が、ベクトル量子化性能を効果的に向上させられるか？
• RQ2 VQ訓練において、再構成誤差または検索損失を直接最小化するのと比較して、密な埋め込みからの関連性予測を蒸留することは、どのように効果的か？
• RQ3 高圧縮および低レイテンシー下で、Distill-VQは既存のVQ手法に比べて、リCALLおよびMRRの面でどの程度優れているか？
• RQ4 蒸留はIVFおよびPQの両成分にどのような影響を及ぼし、効率-精度トレードオフにどのように寄与するか？
• RQ5 ラベル付きデータが乏しいあるいは存在しない状況でも、Distill-VQは最先端の性能を達成できるか？

主な発見

• Distill-VQは、MS MARCOおよびNatural Questionsの両ベンチマークにおいて、SOTAのVQ手法を顕著に上回るリCALLおよびMRRを達成した。
• 低ビットレート（例：192倍圧縮）でも、Distill-VQはリCALL@100およびMRR@10において顕著な優位性を維持した。特に高圧縮および低レイテンシー下で顕著である。
• コードブックサイズが小さい場合でも、FAISSのオリジナルIVFPQよりも優れた効率-リCALLトレードオフを達成した。
• 小さなコードブック（例：8×8）でも、より大きなコードブックを用いたベースラインを上回った。これは、強力な一般化能力とデータ効率性を示している。
• Top-K＋バッチ内サンプリングの活用により、ラベルなしで効果的な蒸留が可能となり、スケーラブルかつ実用的であることが示された。
• 異なる圧縮比や検索設定においても、Distill-VQは一貫した性能向上を示し、ロバストネスと一般化能力を確認した。