LongRAG : Amélioration de la Génération Augmentée par Récupération avec des LLMs à Contexte Long

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

June 21, 2024

Auteurs: Ziyan Jiang, Xueguang Ma, Wenhu Chen

cs.AI

Résumé

Dans le cadre traditionnel de RAG, les unités de récupération de base sont généralement courtes. Les récupérateurs courants comme DPR fonctionnent normalement avec des paragraphes de Wikipédia de 100 mots. Une telle conception oblige le récupérateur à parcourir un large corpus pour trouver l'unité "aiguille". En revanche, les lecteurs n'ont besoin d'extraire les réponses que des courtes unités récupérées. Cette conception déséquilibrée, avec un récupérateur "lourd" et un lecteur "léger", peut entraîner des performances sous-optimales. Pour atténuer ce déséquilibre, nous proposons un nouveau cadre appelé LongRAG, composé d'un "récupérateur long" et d'un "lecteur long". LongRAG traite l'intégralité de Wikipédia en unités de 4 000 tokens, soit 30 fois plus longues qu'auparavant. En augmentant la taille des unités, nous réduisons considérablement le nombre total d'unités, passant de 22 millions à 700 000. Cela allège significativement la charge du récupérateur, ce qui conduit à un score de récupération remarquable : un rappel de réponse@1 de 71 % sur NQ (contre 52 % précédemment) et un rappel de réponse@2 de 72 % (contre 47 % précédemment) sur HotpotQA (full-wiki). Ensuite, nous alimentons les k premières unités récupérées (environ 30 000 tokens) dans un LLM à contexte long existant pour effectuer une extraction de réponse en zero-shot. Sans nécessiter d'entraînement, LongRAG atteint un EM de 62,7 % sur NQ, ce qui constitue le meilleur résultat connu. LongRAG atteint également 64,3 % sur HotpotQA (full-wiki), ce qui est comparable au modèle SoTA. Notre étude offre des perspectives sur la feuille de route future pour combiner RAG avec des LLMs à contexte long.

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.