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
The problem: knowledge trapped in the model weights
Large pre-trained language models absorb a substantial amount of factual knowledge in their weights. The paper calls this parametric memory: everything the model knows is stored in the learned parameters. The authors argue that this has limits. These models have limited ability to access and manipulate knowledge precisely, and two problems remain open: indicating where a decision came from (provenance) and updating what the model knows about the world. Changing a fact usually means changing the weights, that is, training again. That is why the paper targets knowledge-intensive tasks, those in which a person could not do well without consulting an external source. It proposes a general fine-tuning recipe called retrieval-augmented generation (RAG). RAG combines a pre-trained sequence-to-sequence model, which plays the role of parametric memory, with a dense vector index of Wikipedia, which plays the role of non-parametric memory. A pre-trained neural retriever connects the two. The generator no longer needs to memorize every fact: it can condition on passages fetched at query time, so the knowledge lives partly in the weights and partly in an index that can be inspected and swapped. The idea is tested on open-domain question answering, question generation, and fact verification.
The architecture: a DPR retriever feeding a BART generator
RAG has two components. The retriever is a Dense Passage Retriever (DPR), a bi-encoder: one BERT-base encoder turns the query into a vector and another does the same for each document, and relevance is the inner product of the two vectors. Finding the top-k documents therefore becomes a maximum inner product search (MIPS) problem, solved approximately with a hierarchical FAISS index (HNSW). The retriever starts from a DPR already trained with retrieval supervision on TriviaQA and Natural Questions. The generator is BART-large, a pretrained sequence-to-sequence transformer with about 400 million parameters. It receives the input concatenated with a retrieved passage and produces the output text. The document index comes from a December 2018 Wikipedia dump, with articles split into disjoint 100-word chunks, which yields about 21 million documents. Retrieval is treated as a latent variable. The model is trained end to end by minimizing the marginal negative log-likelihood of the target output, summed over the retrieved documents, so no labels saying which document was the correct one are needed. During fine-tuning, the query encoder and the BART generator are updated, while the document encoder and the index stay fixed. The authors report that also updating the document encoder was not necessary for good performance.
RAG-Sequence versus RAG-Token
The paper defines two ways of combining the retrieved documents with the generator. Both treat the document as a latent variable and marginalize over the top-k documents, but they differ in how much of the output a single document can influence. In RAG-Sequence, the model uses the same retrieved document for the entire output sequence. It scores the complete sequence once per document and then combines those scores, weighted by the probability the retriever assigns to each document. In RAG-Token, a different document can be used for each output token. The generator predicts each token by marginalizing over the documents, so the answer can take one part from one passage and another part from another passage. In classification tasks, such as fact verification, where the output is a single-token label, the two formulations are equivalent. Decoding also differs. RAG-Token can use standard beam search over the per-token marginalized probabilities. RAG-Sequence runs beam search separately for each document and, since some candidate outputs do not appear in the beam of every document, it then runs extra passes to score them (the paper's Thorough Decoding) or approximates the missing probabilities as nearly zero (Fast Decoding).
Swapping the index without retraining, and what the results show
A central claim is that the non-parametric memory can be swapped out without touching the rest of the model. To test this, the authors built an index from a December 2016 Wikipedia dump and compared it with the December 2018 index used in the rest of the paper. They asked about 82 world leaders who changed between the two dates. With a matching index and era, accuracy was 70% using the 2016 index for 2016 leaders and 68% using the 2018 index for 2018 leaders. With a mismatched index, accuracy plummeted to 12% (2018 index, 2016 leaders) and 4% (2016 index, 2018 leaders). The model's knowledge followed the index it was given. In other experiments, the paper reports exact match on open-domain question answering. On Natural Questions, RAG-Sequence reached 44.5 and RAG-Token 44.1, versus 41.5 for DPR, 40.4 for REALM and 34.5 for T5-11B. In the human evaluation of 452 pairs on Jeopardy question generation, raters judged RAG-Token more factual than BART in 42.7% of cases, versus 7.1% in which BART was more factual. The authors also report that learning the retriever helped compared with keeping it frozen.
Acknowledged limits and other caveats
The paper acknowledges some limits, and the results bring other caveats. First, retrieval can collapse on some tasks. In story generation, for example, in preliminary experiments the retriever started returning the same documents regardless of the input, the generator learned to ignore them, and the model began to behave like plain BART. Second, the authors note that Wikipedia, or any other external knowledge source, will probably never be fully factual or free of bias. Third, a result that the paper presents as an advantage of generation, rather than as a limit, shows that the answer does not depend only on what was retrieved: on Natural Questions, RAG got 11.8% right when the answer was in none of the retrieved documents, a case where an extractive model would score 0%. Fourth, the best retriever depends on the task: replacing the learned retriever with BM25 worsened performance on open-domain question answering, but BM25 was the best on fact verification (FEVER); the authors suggest, as a possibility, that this happens because FEVER claims are highly entity-centric, which favors retrieval by word overlap. Finally, the authors discuss societal risks: since RAG can be used as a language model, concerns similar to those about GPT-2 apply, although perhaps to a lesser degree, such as generating abusive, false, or misleading content, impersonating other people, or automating spam and phishing.
See the ideas at work
Practice the designs that use this paper in the System Design school.