LongRAG: Melhorando a Geração Aumentada por Recuperação com LLMs de Contexto Longo

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

June 21, 2024

Autores: Ziyan Jiang, Xueguang Ma, Wenhu Chen

cs.AI

Resumo

No tradicional framework RAG, as unidades básicas de recuperação são normalmente curtas. Os recuperadores comuns, como DPR, normalmente trabalham com parágrafos de 100 palavras da Wikipedia. Tal design obriga o recuperador a buscar em um grande corpus para encontrar a unidade "agulha". Em contraste, os leitores só precisam extrair respostas das curtas unidades recuperadas. Tal design desequilibrado com um recuperador "pesado" e um leitor "leve" pode levar a um desempenho subótimo. Para aliviar o desequilíbrio, propomos um novo framework, LongRAG, composto por um "recuperador longo" e um "leitor longo". O LongRAG processa toda a Wikipedia em unidades de 4 mil tokens, o que é 30 vezes mais longo do que antes. Ao aumentar o tamanho da unidade, reduzimos significativamente o total de unidades de 22 milhões para 700 mil. Isso reduz consideravelmente o fardo do recuperador, o que resulta em uma pontuação de recuperação notável: recall de resposta@1=71% no NQ (anteriormente 52%) e recall de resposta@2=72% (anteriormente 47%) no HotpotQA (full-wiki). Em seguida, alimentamos as unidades recuperadas do top-k (aproximadamente 30 mil tokens) em um LLM de contexto longo existente para realizar extração de respostas sem treinamento. Sem exigir nenhum treinamento, o LongRAG alcança um EM de 62,7% no NQ, que é o melhor resultado conhecido. O LongRAG também alcança 64,3% no HotpotQA (full-wiki), o que está em pé de igualdade com o modelo SoTA. Nosso estudo oferece insights sobre o roadmap futuro para combinar RAG com LLMs de contexto longo.

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.