Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

Guia de leitura

Problema e motivação

Modelos de linguagem pré-treinados armazenam substancial conhecimento factual dentro de seus parâmetros (memória paramétrica). No entanto, apresentam desvantagens significativas: têm dificuldade para expandir ou atualizar seus conhecimentos sem retreinamento, não conseguem fornecer facilmente proveniência ou interpretabilidade para suas saídas e são propensos a alucinações.

Embora modelos híbridos paramétricos/não paramétricos como REALM e ORQA tenham combinado recuperação diferenciável com modelagem de linguagem mascarada, eles ficaram restritos ao question answering extrativo. Arquiteturas sequence-to-sequence (seq2seq) padrão como BART e T5 permaneceram amplamente closed-book. Este trabalho preenche essa lacuna propondo o Retrieval-Augmented Generation (RAG), um framework geral de fine-tuning que acopla um recuperador neural denso pré-treinado a um gerador seq2seq autorregressivo, tanto para geração de domínio aberto quanto para classificação.

Ideias principais

• Memória Híbrida Paramétrica e Não Paramétrica: A memória não paramétrica consiste em um índice vetorial denso sobre 21 milhões de passagens de 100 palavras da Wikipedia, acessadas via Maximum Inner Product Search (MIPS). A memória paramétrica é um gerador BART pré-treinado.
• Fine-Tuning de Ponta a Ponta: O codificador de consultas e o gerador são treinados conjuntamente, tratando as passagens recuperadas como variáveis latentes e minimizando a log-verossimilhança marginal negativa dos tokens-alvo sem supervisão explícita de recuperação.
• Duas Formulações de Condicionamento:
  • RAG-Sequence: Condiciona no mesmo documento recuperado ao longo de toda a sequência gerada.
  • RAG-Token: Permite que cada token gerado atenda e marginalize sobre diferentes passagens recuperadas.
• Conhecimento de Mundo Intercambiável (Hot-Swappable): O índice externo pode ser trocado em tempo de inferência (por exemplo, trocando um índice da Wikipedia de 2016 por um de 2018) para atualizar o conhecimento de mundo sem modificar os pesos do modelo.

Método

Dada uma sequência de entrada $x$, o recuperador $p_{\eta }(z|x)$ produz uma distribuição sobre documentos $z$, e o gerador $p_{\theta }(y_i|x,z,y_{1:i-1})$ gera os tokens-alvo condicionados tanto em $x$ quanto em $z$.

O recuperador utiliza o Dense Passage Retriever (DPR): $$p_{\eta }(z|x)\propto \exp \left(\mathbf{d}(z{)}^{\top }\mathbf{q}(x)\right)$$ onde $\mathbf{d}(z)={\text{BERT}}_d(z)$ e $\mathbf{q}(x)={\text{BERT}}_q(x)$. Os $k$ principais candidatos são recuperados usando FAISS MIPS. O índice de documentos e o codificador de documentos ${\text{BERT}}_d$ são mantidos fixos, enquanto ${\text{BERT}}_q$ e os parâmetros do gerador BART $\theta$ são ajustados via fine-tuning.

RAG-Sequence

O modelo trata o documento recuperado como uma única variável latente ao longo de toda a sequência $y$: $$p_{\text{RAG-Sequence}}(y|x)\approx \sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)\prod _i^N{p}_{\theta }(y_i|x,z,y_{1:i-1})$$ A decodificação pode ser avaliada usando Thorough Decoding (executando beam search por documento e computando passagens forward para hipóteses candidatas em todos os documentos) ou Fast Decoding (aproximando probabilidades ausentes do beam como 0).

RAG-Token

O modelo marginaliza sobre os documentos a cada etapa de token, permitindo a mistura de conteúdo em nível de token entre fontes: $$p_{\text{RAG-Token}}(y|x)\approx \prod _i^N\sum _{z\in \text{top-}k(p(\cdot |x))}{p}_{\eta }(z|x)p_{\theta }(y_i|x,z,y_{1:i-1})$$ Isso gera uma probabilidade de transição autorregressiva padrão que se conecta diretamente ao beam search.

Resultados

• QA de Domínio Aberto: O RAG estabelece novos resultados de estado da arte em benchmarks abertos usando correspondência exata (EM):
  • Natural Questions: 44,5 (RAG-Seq) vs. 41,5 (DPR) e 36,6 (T5-11B+SSM).
  • TriviaQA (divisão aberta padrão): 56,8 (RAG-Seq); no split de teste TQA-Wiki: 68,0 (RAG-Seq) vs. 60,5 (T5-11B+SSM).
  • WebQuestions: 45,5 (RAG-Token) vs. 41,1 (DPR).
  • CuratedTrec: 52,2 (RAG-Seq) vs. 50,6 (DPR).
  • Notavelmente, no Natural Questions, o RAG alcançou 11,8% de EM quando a resposta correta estava completamente ausente do contexto recuperado (onde modelos extrativos obtêm 0%).
• Geração Abstrativa: No Open MS-MARCO, o RAG-Sequence atinge 40,8 de Rouge-L e 44,2 de Bleu-1, superando o BART (38,2 Rouge-L, 41,6 Bleu-1).
• Geração de Perguntas do Jeopardy: Avaliadores humanos classificaram o RAG-Token como mais factual que o BART em 42,7% dos casos, contra apenas 7,1% em que o BART foi considerado superior. Os avaliadores também consideraram o RAG mais específico (37,4% vs. 16,8%).
• Verificação de Fatos: No FEVER 3-way, o RAG atinge 72,5% de acurácia de rótulo (ficando a 4,3% de sistemas em pipeline que usam supervisão com evidência padrão-ouro); no 2-way, o RAG alcança 89,5% de acurácia.
• Atualização de Conhecimento via Hot-Swapping: Usando um teste temporal de 82 líderes mundiais entre 2016 e 2018, o RAG obteve 70% de acurácia para líderes de 2016 usando o índice de 2016 e 68% para líderes de 2018 com o índice de 2018, caindo drasticamente (12% e 4%) quando os índices foram invertidos.

Limitações e questões em aberto

• Colapso da Recuperação: Para tarefas que exigem um condicionamento factual explícito mais fraco (como geração de histórias abertas), o recuperador pode colapsar e recuperar documentos estáticos e invariantes, enquanto o gerador aprende a ignorá-los.
• Indexação de Documentos Fixa: O codificador de documentos é mantido estático durante o fine-tuning da tarefa devido ao custo computacional de reindexar continuamente 21M de vetores, o que limita a adaptação do recuperador.
• Fundamentação em Fontes e Viés: A memória não paramétrica permanece vulnerável a desinformação e vieses presentes no corpus de conhecimento subjacente (Wikipedia).
• Pré-treinamento Conjunto: Os componentes do recuperador e do gerador foram pré-treinados separadamente antes do fine-tuning nas tarefas downstream, deixando o pré-treinamento conjunto a partir do zero como uma direção aberta de pesquisa.