LongRAG: 長文脈LLMによる検索拡張生成の強化

LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs

June 21, 2024

著者: Ziyan Jiang, Xueguang Ma, Wenhu Chen

cs.AI

要旨

従来のRAGフレームワークでは、基本的な検索単位は通常短いものです。DPRのような一般的な検索器は、通常100ワードのWikipedia段落を扱います。この設計では、検索器が大規模なコーパスから「針」となる単位を探すことを強制されます。一方、リーダーは短い検索された単位から答えを抽出するだけで済みます。このような不均衡な「重い」検索器と「軽い」リーダーの設計は、最適でないパフォーマンスを引き起こす可能性があります。この不均衡を緩和するために、我々は「長い検索器」と「長いリーダー」からなる新しいフレームワークLongRAGを提案します。LongRAGは、Wikipedia全体を4Kトークンの単位に処理します。これは以前の30倍の長さです。単位サイズを増やすことで、総単位数を22Mから700Kに大幅に削減します。これにより、検索器の負担が大幅に軽減され、驚異的な検索スコアが得られます：NQでのanswer recall@1=71%（以前は52%）、HotpotQA（full-wiki）でのanswer recall@2=72%（以前は47%）。次に、トップkの検索された単位（約30Kトークン）を既存の長文脈LLMに供給し、ゼロショットの答え抽出を実行します。LongRAGは、トレーニングを必要とせずに、NQで62.7%のEMを達成し、これは既知の最高の結果です。LongRAGはまた、HotpotQA（full-wiki）で64.3%を達成し、これはSoTAモデルと同等です。我々の研究は、RAGと長文脈LLMを組み合わせるための将来のロードマップに対する洞察を提供します。

English

In traditional RAG framework, the basic retrieval units are normally short. The common retrievers like DPR normally work with 100-word Wikipedia paragraphs. Such a design forces the retriever to search over a large corpus to find the `needle' unit. In contrast, the readers only need to extract answers from the short retrieved units. Such an imbalanced `heavy' retriever and `light' reader design can lead to sub-optimal performance. In order to alleviate the imbalance, we propose a new framework LongRAG, consisting of a `long retriever' and a `long reader'. LongRAG processes the entire Wikipedia into 4K-token units, which is 30x longer than before. By increasing the unit size, we significantly reduce the total units from 22M to 700K. This significantly lowers the burden of retriever, which leads to a remarkable retrieval score: answer recall@1=71% on NQ (previously 52%) and answer recall@2=72% (previously 47%) on HotpotQA (full-wiki). Then we feed the top-k retrieved units (approx 30K tokens) to an existing long-context LLM to perform zero-shot answer extraction. Without requiring any training, LongRAG achieves an EM of 62.7% on NQ, which is the best known result. LongRAG also achieves 64.3% on HotpotQA (full-wiki), which is on par of the SoTA model. Our study offers insights into the future roadmap for combining RAG with long-context LLMs.