[論文レビュー] Fine-Tuning LLaMA for Multi-Stage Text Retrieval

Xueguang Ma, Liang Wang+3 more|arXiv (Cornell University)|Oct 12, 2023

Topic Modeling被引用数 7

ひとこと要約

本稿では、LLaMA-2 large language modelを微調整して、多段階テキスト検索用の濃密なリトリーバー（RepLLaMA）およびポイントワイズ再ランク機構（RankLLaMA）を提案し、MS MARCOおよびBEIRベンチマークで最先端の性能を達成した。このアプローチは、LLMの長文脈能力を活用して文書全体を包括的に表現することで、従来のセグメンテーションおよびプーリング手法を不要にし、オープンソースのモデルチェックポイントを用いたゼロショット効果が顕著である。

ABSTRACT

The effectiveness of multi-stage text retrieval has been solidly demonstrated since before the era of pre-trained language models. However, most existing studies utilize models that predate recent advances in large language models (LLMs). This study seeks to explore potential improvements that state-of-the-art LLMs can bring. We conduct a comprehensive study, fine-tuning the latest LLaMA model both as a dense retriever (RepLLaMA) and as a pointwise reranker (RankLLaMA) for both passage retrieval and document retrieval using the MS MARCO datasets. Our findings demonstrate that the effectiveness of large language models indeed surpasses that of smaller models. Additionally, since LLMs can inherently handle longer contexts, they can represent entire documents holistically, obviating the need for traditional segmenting and pooling strategies. Furthermore, evaluations on BEIR demonstrate that our RepLLaMA-RankLLaMA pipeline exhibits strong zero-shot effectiveness. Model checkpoints from this study are available on HuggingFace.

研究の動機と目的

• 大規模言語モデル（LLM）であるLLaMA-2が、より小さな事前学習モデルよりも多段階テキスト検索において優れた性能を示すかどうかを調査すること。
• プロンプトベースのLLM再ランク機構の限界（非並列化された推論およびラベル付きデータの活用不足）を、エンドツーエンドのリトリーバーおよび再ランク機構としてのLLMの微調整により解決すること。
• LLMの長文脈能力に起因する、文書全体を包括的に表現する可能性を検討することにより、文書のセグメンテーションおよびプーリング戦略の必要性を排除すること。
• LLaMA-2をバックボーンとして用いた完全にオープンソースで最適化された多段階リトリーブパイプラインを確立し、より優れた効果性および推論効率を実現すること。
• 提案されたRepLLaMA–RankLLaMAパイプラインのゼロショット一般化能力が、多様なリトリーブベンチマークでどの程度有効であるかを評価すること。

提案手法

• 対照的損失を用いて、LLaMA-2をバイエンコーダー形式の濃密なリトリーバー（RepLLaMA）として微調整し、密度空間におけるクエリ-ドキュメント類似度を学習する。
• ラベル付き関連性判断を用いて、LLaMA-2をポイントワイズ再ランク機構（RankLLaMA）として微調整し、クエリ-ドキュメントペアの関連性スコアを直接予測する。
• 標準的な評価指標を用いて、パassageおよびドキュメントリトリーブタスクの両方で、MS MARCOデータセットを用いてドメイン内での微調整を実施する。
• LLaMA-2の長文脈アテンションを活用し、セグメンテーションを施さずに文書全体をエンコードすることで、包括的な文書表現を可能にする。
• 対照的損失およびポイントワイズランク損失を別々に用いて、標準的な教師あり微調整の枠組みで両モデルを訓練する。
• HuggingFaceにモデルチェックポイントをリリースし、オープンアクセスおよび再現可能性を確保する。

Figure 1: Cumulative distribution function of document lengths in the MS MARCO document corpus, showing the proportion of documents that has a length less than a specific value (determined by the LLaMA tokenizer). For clarity, we exclude 3% of documents with a length exceeding 10,000 tokens.

実験結果

リサーチクエスチョン

• RQ1 微調整されたLLaMA-2モデルは、多段階テキスト検索において、より小さな事前学習モデルを上回る性能を示せるか？
• RQ2 ラベル付きデータで微調整されたLLaMA-2は、効果的な濃密なリトリーバーおよびポイントワイズ再ランク機構として機能できるか？
• RQ3 LLaMA-2の長文脈能力により、文書リトリーブにおける文書セグメンテーションおよびプーリング戦略の必要性が排除されるか？
• RQ4 RepLLaMA–RankLLaMAパイプラインは、多様なBEIRベンチマークにおいてゼロショットリトリーブでどの程度効果的か？
• RQ5 LLMのエンドツーエンド微調整は、プロンプトベースの生成アプローチと比較して、リトリーブの効果性を向上させられるか？

主な発見

• RepLLaMAおよびRankLLaMAは、MS MARCOデータセットを用いたパassageおよびドキュメントリトリーブタスクで、以前のより小さなモデルを上回る最先端の性能を達成した。
• RepLLaMA–RankLLaMAパイプラインは、BEIRベンチマークにおいて強いゼロショット効果を示し、多様なドメインにわたる強力な一般化能力を示している。
• LLaMA-2の長文脈アテンションにより、包括的な文書エンコードが可能となり、従来のセグメンテーションおよびプーリング戦略の必要性が排除された。
• LLaMA-2をリトリーバーおよび再ランク機構として微調整することで、非並列的で複数回のデコードを要するプロンプトベースのLLM再ランク機構の欠点を克服し、並列スコアリングが可能になった。
• ラベル付きデータ（例：MS MARCOから得たもの）を効果的に活用することで、このようなアノテーションを無視するゼロショットプロンプト手法よりも優れた性能が得られた。
• HuggingFaceにリリースされたオープンソースのモデルチェックポイントにより、LLMベースのリトリーブシステムに関する再現性およびさらなる研究が可能になった。

Figure 2: Comparison of document ranking MRR@100 scores for RankLLaMA trained with different maximum input lengths and evaluated using different maximum input lengths. Each line represents a model trained with a specific maximum length, while points along the line indicate the effectiveness when var