Papers
Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, Douwe Kiela · NeurIPS 2020, 2020
O problema: conhecimento preso nos pesos do modelo
Modelos de linguagem pré-treinados e grandes absorvem uma quantidade substancial de conhecimento factual nos seus pesos. O paper chama isso de memória paramétrica: tudo o que o modelo sabe está guardado nos parâmetros aprendidos. Os autores argumentam que isso tem limites. Esses modelos têm capacidade limitada de acessar e manipular o conhecimento com precisão, e dois problemas continuam em aberto: indicar de onde veio uma decisão (proveniência) e atualizar o que o modelo sabe sobre o mundo. Mudar um fato normalmente significa mudar os pesos, ou seja, treinar de novo. Por isso o paper mira tarefas intensivas em conhecimento, aquelas em que uma pessoa não conseguiria se sair bem sem consultar uma fonte externa. Ele propõe uma receita geral de fine-tuning chamada retrieval-augmented generation (RAG). A RAG combina um modelo sequência-para-sequência pré-treinado, que faz o papel de memória paramétrica, com um índice vetorial denso da Wikipedia, que faz o papel de memória não paramétrica. Um recuperador neural pré-treinado liga os dois. O gerador deixa de precisar memorizar todos os fatos: ele pode se condicionar a trechos buscados na hora da consulta, de modo que o conhecimento fica em parte nos pesos e em parte num índice que pode ser inspecionado e trocado. A ideia é testada em perguntas e respostas de domínio aberto, geração de perguntas e verificação de fatos.
A arquitetura: recuperador DPR alimentando um gerador BART
A RAG tem dois componentes. O recuperador é um Dense Passage Retriever (DPR), um bi-encoder: um encoder BERT-base transforma a consulta em vetor e outro faz o mesmo com cada documento, e a relevância é o produto interno dos dois vetores. Achar os top-k documentos vira, portanto, um problema de busca por produto interno máximo (MIPS), resolvido de forma aproximada com um índice FAISS hierárquico (HNSW). O recuperador parte de um DPR já treinado com supervisão de recuperação em TriviaQA e Natural Questions. O gerador é o BART-large, um transformer sequência-para-sequência pré-treinado de cerca de 400 milhões de parâmetros. Ele recebe a entrada concatenada com um trecho recuperado e produz o texto de saída. O índice de documentos vem de um dump da Wikipedia de dezembro de 2018, com artigos divididos em blocos disjuntos de 100 palavras, o que dá cerca de 21 milhões de documentos. A recuperação é tratada como variável latente. O modelo é treinado de ponta a ponta minimizando a log-verossimilhança negativa marginal da saída-alvo, somada sobre os documentos recuperados, então não são necessários rótulos dizendo qual documento era o correto. No fine-tuning, o encoder da consulta e o gerador BART são atualizados, enquanto o encoder de documentos e o índice ficam fixos. Os autores relatam que atualizar também o encoder de documentos não foi necessário para um bom desempenho.
RAG-Sequence versus RAG-Token
O paper define duas formas de combinar os documentos recuperados com o gerador. As duas tratam o documento como variável latente e marginalizam sobre os top-k documentos, mas diferem em quanto da saída um único documento pode influenciar. Na RAG-Sequence, o modelo usa o mesmo documento recuperado para a sequência de saída inteira. Ele pontua a sequência completa uma vez por documento e depois combina essas pontuações, ponderadas pela probabilidade que o recuperador dá a cada documento. Na RAG-Token, um documento diferente pode ser usado para cada token de saída. O gerador prevê cada token marginalizando sobre os documentos, de modo que a resposta pode tirar uma parte de um trecho e outra parte de outro trecho. Em tarefas de classificação, como a verificação de fatos, em que a saída é um rótulo de um único token, as duas formulações são equivalentes. A decodificação também difere. A RAG-Token pode usar beam search padrão sobre as probabilidades marginalizadas por token. A RAG-Sequence roda beam search separadamente para cada documento e, como algumas saídas candidatas não aparecem no beam de todos os documentos, ela então roda passes extras para pontuá-las (a Thorough Decoding do paper) ou aproxima as probabilidades ausentes como quase zero (Fast Decoding).
Trocar o índice sem retreinar, e o que os resultados mostram
Uma alegação central é que a memória não paramétrica pode ser substituída sem mexer no resto do modelo. Para testar isso, os autores montaram um índice a partir de um dump da Wikipedia de dezembro de 2016 e o compararam com o índice de dezembro de 2018 usado no restante do paper. Eles perguntaram sobre 82 líderes mundiais que mudaram entre as duas datas. Com índice e época coerentes, a acurácia foi de 70% usando o índice de 2016 para líderes de 2016 e de 68% usando o índice de 2018 para líderes de 2018. Com índice incompatível, a acurácia despencou para 12% (índice de 2018, líderes de 2016) e 4% (índice de 2016, líderes de 2018). O conhecimento do modelo acompanhou o índice que ele recebeu. Em outros experimentos, o paper reporta exact match em perguntas e respostas de domínio aberto. Em Natural Questions, a RAG-Sequence chegou a 44,5 e a RAG-Token a 44,1, contra 41,5 do DPR, 40,4 do REALM e 34,5 do T5-11B. Na avaliação humana de 452 pares na geração de perguntas do Jeopardy, os avaliadores julgaram a RAG-Token mais factual que o BART em 42,7% dos casos, contra 7,1% em que o BART foi mais factual. Os autores também reportam que aprender o recuperador ajudou em comparação com mantê-lo congelado.
Limites reconhecidos e outras ressalvas
O paper reconhece alguns limites, e os resultados trazem outras ressalvas. Primeiro, a recuperação pode colapsar em algumas tarefas. Em geração de histórias, por exemplo, nos experimentos preliminares o recuperador passou a trazer os mesmos documentos independentemente da entrada, o gerador aprendeu a ignorá-los e o modelo passou a se comportar como o BART puro. Segundo, os autores observam que a Wikipedia, ou qualquer outra fonte externa de conhecimento, provavelmente nunca será totalmente factual nem livre de viés. Terceiro, um resultado que o paper apresenta como vantagem da geração, e não como limite, mostra que a resposta não depende só do que foi recuperado: em Natural Questions, a RAG acertou 11,8% das vezes quando a resposta não estava em nenhum documento recuperado, caso em que um modelo extrativo marcaria 0%. Quarto, o melhor recuperador depende da tarefa: trocar o recuperador aprendido por BM25 piorou o desempenho em perguntas e respostas de domínio aberto, mas o BM25 foi o melhor na verificação de fatos (FEVER); os autores sugerem, como possibilidade, que isso ocorra porque as afirmações do FEVER são muito centradas em entidades, o que favorece a recuperação por sobreposição de palavras. Por fim, os autores discutem riscos sociais: como a RAG pode ser usada como modelo de linguagem, valem preocupações semelhantes às do GPT-2, embora talvez em menor grau, como gerar conteúdo abusivo, falso ou enganoso, se passar por outras pessoas ou automatizar spam e phishing.
Veja as ideias em ação
Pratique na escola de System Design os desenhos que usam este paper.