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단어 정도의 위키피디아 단락을 처리합니다. 이러한 설계는 검색기가 대규모 코퍼스에서 '바늘' 같은 단위를 찾도록 강제합니다. 반면, 리더는 짧게 검색된 단위에서 답변을 추출하기만 하면 됩니다. 이러한 불균형적인 '무거운' 검색기와 '가벼운' 리더 설계는 최적이 아닌 성능으로 이어질 수 있습니다. 이러한 불균형을 완화하기 위해, 우리는 '긴 검색기'와 '긴 리더'로 구성된 새로운 프레임워크인 LongRAG를 제안합니다. LongRAG는 전체 위키피디아를 4K 토큰 단위로 처리하며, 이는 이전보다 30배 더 깁니다. 단위 크기를 늘림으로써, 우리는 총 단위 수를 22M에서 700K로 크게 줄였습니다. 이는 검색기의 부담을 크게 낮추어, NQ에서 답변 recall@1=71%(이전 52%), HotpotQA(full-wiki)에서 답변 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.