[論文レビュー] Efficiently Teaching an Effective Dense Retriever with Balanced Topic Aware Sampling

Sebastian Hofstätter, Sheng-Chieh Lin+3 more|arXiv (Cornell University)|Apr 14, 2021

Topic Modeling被引用数 4

ひとこと要約

本稿では、ペairwiseおよびin-batch教師モデルからの二重教師付き学習を用いて、効率的なドキュメント検索モデルを訓練するためのトピックに配慮したマージンバランス型サンプリング手法TAS-Balancedを提案する。1台のコンsumer GPUで48時間未満で訓練可能であり、2つのTREC-DLデータセットにおいて、NDCG@10でBM25を44%上回り、低遅延性能（1クエリあたり64ms）を達成し、最先端の性能を実現した。

ABSTRACT

A vital step towards the widespread adoption of neural retrieval models is their resource efficiency throughout the training, indexing and query workflows. The neural IR community made great advancements in training effective dual-encoder dense retrieval (DR) models recently. A dense text retrieval model uses a single vector representation per query and passage to score a match, which enables low-latency first stage retrieval with a nearest neighbor search. Increasingly common, training approaches require enormous compute power, as they either conduct negative passage sampling out of a continuously updating refreshing index or require very large batch sizes for in-batch negative sampling. Instead of relying on more compute capability, we introduce an efficient topic-aware query and balanced margin sampling technique, called TAS-Balanced. We cluster queries once before training and sample queries out of a cluster per batch. We train our lightweight 6-layer DR model with a novel dual-teacher supervision that combines pairwise and in-batch negative teachers. Our method is trainable on a single consumer-grade GPU in under 48 hours (as opposed to a common configuration of 8x V100s). We show that our TAS-Balanced training method achieves state-of-the-art low-latency (64ms per query) results on two TREC Deep Learning Track query sets. Evaluated on NDCG@10, we outperform BM25 by 44%, a plainly trained DR by 19%, docT5query by 11%, and the previous best DR model by 5%. Additionally, TAS-Balanced produces the first dense retriever that outperforms every other method on recall at any cutoff on TREC-DL and allows more resource intensive re-ranking models to operate on fewer passages to improve results further.

研究の動機と目的

• 高コストなGPUクラスターや大規模バッチサイズを必要とする有効なドキュメント検索モデルのトレーニングの計算コストを低減すること。
• 各トレーニングバッチあたりの情報量を最大化するサンプリング戦略を導入することで、高価なインフラ依存を軽減すること。
• ハードウェア要件を増加させることなく、トレーニングの効率性とモデルの有効性を向上させ、強力なドキュメント検索モデルの広範な利用を可能にすること。
• 最初の段階でのリコール性能を向上させ、再ランク付きモデルと互換性を持つことで、低遅延推論を実現する高パフォーマンスな検索パイプラインを提供すること。
• 異なるランダムシードや教師付き学習設定に対しても、リソース効率的かつ頑健なトレーニング手法を設計すること。

提案手法

• トレーニング前一度だけ、意味的ドット積類似度を用いてクエリをトピックグループにクラスタリングし、トピックに配慮したバッチ構成を可能にする。
• 各バッチから1つのトピッククラスタからクエリをサンプリングすることで、バッチ内に集中的な意味的情報を集中させ、バッチ内ネガティブ教師効果を向上させる。
• ペアワイズ教師スコアのマージンをバランスさせることで、トレーニングの安定化と一般化性能の向上を図る。
• ペアワイズ連結BERT（BERT-CAT）とin-batchネガティブに基づくColBERT型教師モデルの両方を用いて、二重教師付き蒸留を実施し、信号品質を向上させる。
• 両方の教師信号を用いてMargin-MSE損失を適用し、6層の軽量なデュアルエンコーダーBERT-DOTモデルを訓練することで、ランキング品質を向上させる。
• 大規模バッチや継続的更新が必要なインデックス戦略を回避するため、1台のコンsumerグレードGPU（11GB VRAM）を用いてエンドツーエンドのトレーニングを実施する。

実験結果

リサーチクエスチョン

• RQ1 ランダムサンプリングと比較して、トピックに配慮したサンプリング（TAS）は、トレーニング効率性および検索パフォーマンスの向上にどの程度有効であるか？
• RQ2 ペアワイズ教師スコアのマージンをバランスさせたTAS-Balancedは、さまざまなトレーニング設定においてさらにモデル性能を向上させるか？
• RQ3 ペアワイズおよびin-batchネガティブ教師モデルからの二重教師付き学習は、モデルの有効性と頑健性にどのように影響を与えるか？
• RQ4 TAS-Balancedは、大規模GPUクラスタでトレーニングされたモデルと比較して、1台のコンsumer GPUで高性能なドキュメント検索モデルを訓練可能か？
• RQ5 再ランク付きモデルを含むエンドツーエンドの検索パイプラインにおいて、TAS-Balancedは特に低遅延および中程度遅延の状況でどの程度の性能を示すか？

主な発見

• TAS-Balancedは2つのTREC-DLデータセットにおいて、NDCG@10で最先端の性能を達成し、BM25を44%、単純にトレーニングされたDRモデルを19%、以前の最良DRモデルを5%上回った。
• 1クエリあたり64msという低遅延推論時間を達成しており、リアルタイム応用に適している。
• TAS-Balancedは、TREC-DLにおいて、すべてのカットオフですべてのベースラインを上回るリコール性能を達成した最初のドキュメント検索モデルであり、最初の段階でのリコール性能を顕著に向上させた。
• TAS-BalancedとdocT5queryを融合させることで、Recall@1Kが向上し、2～6倍の遅延を持つシステムを上回った。これは、パイプラインとの高い互換性を示している。
• 1台のコンsumerグレードGPU（11GB VRAM）で48時間未満でトレーニング可能であり、マルチGPUや大規模バッチトレーニングの必要性がなくなった。
• Margin-MSE損失を用いた二重教師付き学習が最も高いパフォーマンスを示し、特にリコール性能で顕著な向上を示した。また、異なるランダムシードや教師設定に対しても一貫した向上が見られた。