[論文レビュー] A Multi-Task Embedder For Retrieval Augmented LLMs

Peitian Zhang, Shitao Xiao+3 more|arXiv (Cornell University)|Oct 11, 2023

Topic Modeling被引用数 18

ひとこと要約

本論文は LLM-Embedder を提案する。これは multi-task、報酬 guided、蒸留駆動のフレームワークで訓練された統一埋め込みモデルで、知識強化、長文脈のモデリング、ツール学習、対話検索にわたる retrieval augmentation を強化する。

ABSTRACT

LLMs confront inherent limitations in terms of its knowledge, memory, and action. The retrieval augmentation stands as a vital mechanism to address these limitations, which brings in useful information from external sources to augment the LLM. However, existing retrieval methods encounter two pressing issues. On one hand, the general retrievers are not properly optimized for retrieval augmentation hence exhibit limited effectiveness; on the other hand, the task-specific retrievers excel in the targeted retrieval augmentation scenario, while lack the versatility to handle diverse scenarios. In this work, we propose extbf{LLM-Embedder} for the unified support of diverse retrieval augmentation scenarios. Our method presents three technical contributions. Firstly, we introduce a new extit{reward formulation}, namely {rank-aware reward}. It exploits the ranking position of the desired output among $N$ sampled outputs from the LLM, which leads to fine-grained and robust computation of reward from the LLM's feedback. Secondly, we design a novel extit{distillation objective}, called graded distillation. It incorporates both the absolute value and the relative order of the reward for more sufficient utilization of the LLM's feedback. Thirdly, we systematically optimize the extit{multi-task learning}, which effectively unifies the multiple retrieval functionalities into one model. In our experiment, LLM-Embedder notably improves the LLM's performances in various downstream tasks, and outperforms both general and task-specific retrievers with a substantial advantage.

研究の動機と目的

• 知識、記憶、能力という LLM の三つの intrinsic boundaries を retrieval augmentation を活用して動機づけ、対処する。
• 多様な retrieval augmentation タスクを一つのバックボーンでサポートする統一埋め込みモデル（LLM-Embedder）を提案する。
• LLM ベースの報酬、安定化した蒸留、指示ベースのファインチューニング、同質なインバッチネガティブを組み合わせた頑健なトレーニング手法を開発する。
• 複数の retrieval シナリオで、一般目的およびタスク特化のリトリーバを上回ることを示す。

提案手法

• 多様なデータソースを用いて統一埋め込みモデル（LLM-Embedder）を訓練する。データ源は QA (MSMARCO, Natural Questions)、対話検索 (QReCC)、ツール学習 (ToolLLM)、指示調整 (FLAN, UPRISE)、生成/履歴データ (Books3, ArXiv, CodeParrot, Multi-Session Chat)。
• トレーニングを導くために、ハードラベルと LLM由来の報酬を含む複合的な監督信号を使用する（報酬は Equation 1）。
• 指示ベースのファインチューニング（タスク固有のプロンプト I_t）と同質なインバッチネガティブサンプリングを用いたコントラスト学習を適用し、識別力を高める。
• 報酬の変動を抑制するため、ソフト報酬ベースの重みとハードランキングラベルを組み合わせて安定化した蒸留を実施する（Equation 3）。
• KL-divergence ベースの蒸留を用いて、埋め込みベースの候補分布を LLM が生成した報酬と整合させ、シグナルを安定化させるトップダウンの再ランキングを導入する。
• 知識強化、長文脈のモデリング、インコンテキスト学習、ツール学習、対話検索のための retrieval augmentation ワークフローを示す。

実験結果

リサーチクエスチョン

• RQ1 RQ1: LLM-Embedder は LLMs の多様な retrieval augmentation シナリオを総合的にサポートできるか？
• RQ2 RQ2: 各個別の retrieval シナリオに対する LLM-Embedder の影響は？
• RQ3 RQ3: LLM-Embedder の実証的性能に影響を与える要因は？
• RQ4 RQ4: 統一埋め込みモデルは、ベンチマーク全体で一般目的およびタスク特化のリトリーバとどのように比較されるか？

主な発見

• LLM-Embedder は、None および多くのベースラインと比較して、知識、記憶、指示遵守タスク全般で LLM の性能を一貫して向上させる。
• 一般的な埋め込みモデル（例: Contriever, RetroMAE-BEIR, BGE）および複数のタスク特化リトリーバ（例: AAR, LLM-R, API-Retriever）を、複数のシナリオで上回る。
• タスク特化のリトリーバは最適化されたシナリオ外で性能が低下することがある一方、LLM-Embedder は多様なタスクで競争力のある性能を維持する。
• 知識強化はより高い検索精度によって利益を得ており、PopQA のような知識集約型データセットでは特に顕著な向上が見られる。
• インコンテキスト学習と長文脈モデリングは、それぞれ retrieved demonstrations と memory chunks を使用することで一貫した向上を示す。
• 長文脈シナリオ（対話と言語モデリング）は、困惑度と検索支援指標の改善を示し、ツール学習と対話検索はランキング指標（NDCG）で評価される。