Source-linked AI summary

REALM: Retrieval-Augmented Language Model Pre-Training

Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, Ming-Wei Chang

arXiv:2002.08909v1cs.CLcs.LG

TL;DR

REALM addresses the difficulty of storing and inspecting growing amounts of world knowledge in language-model parameters by augmenting pre-training with a learned retriever over external documents. It achieves new state-of-the-art results on three Open-QA benchmarks, outperforming previous systems by 4–16% absolute accuracy while offering interpretability and modularity.

  • Problem

    Language models store world knowledge implicitly in network parameters, making it difficult to inspect and requiring increasingly large, costly networks to capture more facts.

  • Method

    REALM augments language-model pre-training with a learned textual retriever that uses masked-language-modeling signals and efficient maximum inner-product search over document embeddings.

  • Results

    4–16% absolute accuracy: REALM achieves new state-of-the-art results on all three Open-QA benchmarks, outperforming previous systems.

  • Takeaways & Limitations

    REALM provides a more interpretable and modular way to access world knowledge than storing it solely in model parameters.

  • Takeaways & Limitations

    The work is a minimal instantiation, leaving structured, multilingual, and multimodal extensions for future work.

Abstract

from arXiv · show

Language model pre-training has been shown to capture a surprising amount of world knowledge, crucial for NLP tasks such as question answering. However, this knowledge is stored implicitly in the parameters of a neural network, requiring ever-larger networks to cover more facts. To capture knowledge in a more modular and interpretable way, we augment language model pre-training with a latent knowledge retriever, which allows the model to retrieve and attend over documents from a large corpus such as Wikipedia, used during pre-training, fine-tuning and inference. For the first time, we show how to pre-train such a knowledge retriever in an unsupervised manner, using masked language modeling as the learning signal and backpropagating through a retrieval step that considers millions of documents. We demonstrate the effectiveness of Retrieval-Augmented Language Model pre-training (REALM) by fine-tuning on the challenging task of Open-domain Question Answering (Open-QA). We compare against state-of-the-art models for both explicit and implicit knowledge storage on three popular Open-QA benchmarks, and find that we outperform all previous methods by a significant margin (4-16% absolute accuracy), while also providing qualitative benefits such as interpretability and modularity.

1. Introduction

REALM addresses the opacity and scalability limits of parameter-stored world knowledge by augmenting language-model pre-training with a learned retriever over external textual knowledge. It trains this retrieve-then-predict system end-to-end and achieves state-of-the-art Open-QA results on three benchmarks.

  • Motivation: Parameter-stored knowledge is difficult to inspect and requires increasingly large networks to capture more facts.The storage limitation can make larger models prohibitively slow or expensive.
  • Method: REALM explicitly retrieves documents from a corpus such as Wikipedia before each prediction and attends over them to use external knowledge.The retriever exposes which knowledge the model selects during inference.
  • Training: REALM trains retrieval from unsupervised text by rewarding documents that improve language-model perplexity and penalizing uninformative retrievals.The retrieve-then-predict model is optimized as a latent-variable language model using marginal likelihood.
  • Scalability: The retriever addresses million-document computation through cached, asynchronously updated document representations and Maximum Inner Product Search.The system must backpropagate through retrieval decisions during pre-training.
  • Results: 4-16% absolute accuracy: REALM significantly outperforms previous systems on NATURALQUESTIONS-OPEN, WEBQUESTIONS, and CURATEDTREC.The evaluation fine-tunes REALM-pre-trained models on Open-domain Question Answering and compares them with implicit-knowledge and retrieval-based systems.

2. Background

Language model pre-training learns useful representations from unlabeled text for downstream tasks, while masked language modeling also captures syntactic, semantic, and some world knowledge. Open-domain question answering tests knowledge-intensive reasoning over potentially millions of documents, motivating retrieval-based systems such as REALM.

  • Language model pre-training: Pre-training learns language representations from unlabeled text and can improve downstream generalization after fine-tuning.The paper focuses on Open-QA as the downstream task.
  • Language model pre-training: Masked language modeling trains models to predict randomly masked tokens in unlabeled text passages.The example masks tokens in a Wikipedia-like passage and predicts the missing words.
  • Language model pre-training: A good masked language model encodes syntactic and semantic information as well as some world knowledge.The passage illustrates semantic prediction with “of” and world-knowledge prediction with “pound.”
  • Open-domain question answering: Open-domain question answering requires producing an answer to a question without receiving a single preidentified document, potentially drawing on millions of documents.The task uses questions such as “What is the currency of the UK?” with “pound” as the answer.
  • Open-domain question answering: Textual-corpus Open-QA systems commonly retrieve potentially relevant documents from corpus Z and extract an answer, a paradigm REALM extends to language model pre-training.The retrieval-based approach is contrasted with generation-based work in the passage.

3. Approach

REALM models prediction as retrieving documents from a knowledge corpus and then generating outputs conditioned on the input and retrieved evidence. Its jointly trained retriever and knowledge-augmented encoder support masked-language-model pre-training and Open-QA fine-tuning, with approximate retrieval enabling scalable optimization.

  • Training objectives: REALM pre-training predicts masked tokens, whereas Open-QA fine-tuning predicts answer spans assumed to occur contiguously in retrieved documents.Both tasks optimize the likelihood of the correct output, but use task-specific output formulations.
  • Retrieve-then-predict: REALM decomposes p(y | x) into retrieving documents z with p(z | x) and predicting y with p(y | z, x), marginalizing over possible documents.The latent-document formulation applies to both pre-training and fine-tuning.
  • Model architecture: The knowledge retriever uses BERT-style embeddings and a dense inner product, while the knowledge-augmented encoder joins x and z in a distinct Transformer with cross-attention.The retriever’s softmax over relevance scores defines p(z | x), and the encoder defines p(y | z, x).
  • Scalable retrieval: To scale marginalization over millions of documents, REALM sums over the top k retrieved documents and uses Maximum Inner Product Search with sub-linear time and storage scaling.The method refreshes its document-embedding index asynchronously because retriever updates otherwise make the index stale.
  • Learning dynamics: A document receives a positive retriever update exactly when its conditional prediction probability exceeds the marginal prediction probability p(y | x).This gradient signal increases relevance for documents that help predict the correct output and decreases it for documents that perform worse than expectation.

4. Experiments

REALM is evaluated on three Open-QA benchmarks using exact-match answer evaluation and is reported to outperform existing systems, including strong generation-based and retrieval-based baselines. Ablations show that both the encoder and retriever benefit from REALM pre-training, while salient span masking is crucial to its effectiveness.

  • Benchmarks: REALM is evaluated on NaturalQuestions-Open, WebQuestions, and CuratedTrec, with predictions scored by exact match against reference answers.NaturalQuestions-Open retains short-answer questions of at most five tokens; CuratedTrec uses regular expressions for multiple correct answers and spelling variants.
  • Results: REALM outperforms all existing systems on the three Open-QA datasets.Table 1 reports accuracy results and model parameter counts for the compared approaches.
  • Results: 50 times larger T5-11B models gain roughly 5 points in accuracy from Base, whereas REALM outperforms T5-11B while being 30 times smaller.The comparison highlights a tradeoff between generative model scaling and REALM’s parameter efficiency.
  • Results: REALM’s improvement over ORQA is attributable to better pre-training because the fine-tuning setup, hyperparameters, and training data are identical.The pre-training method works in both single-corpus and separate-corpus settings.
  • Results: REALM achieves the best overall performance while retrieving only 5 documents, compared with 20 to 80 documents for other retrieval-based systems.This comparison concerns systems cited as Asai et al. (2019) and Min et al. (2019a;b).
  • Ablations: Both the encoder and retriever benefit separately from REALM training, but the best result requires both components.Resetting both components reduces the system to the ORQA baseline.
  • Ablations: Salient span masking is crucial for REALM, unlike its lack of demonstrated impact in previous standard BERT training.The ablation compares salient span masking with random token masking and random span masking.

5. Discussion and Related Work

REALM connects corpus-conditioned language modeling to broader ideas in learned retrieval, scalable neural memory, and unsupervised alignment. It extends context scope to entire corpora while learning useful documents and producing document-level alignments.

  • Language modeling with corpus as context: REALM generalizes language representation models from word-, sentence-, and paragraph-level context to conditioning on an entire text corpus.The cited progression includes Mikolov et al. (2013a;b), Kiros et al. (2015), Peters et al. (2018), Radford et al. (2018), and Devlin et al. (2018).
  • Retrieve-and-edit with learned retrieval: Unlike retrieve-and-edit models that use high lexical overlap, REALM learns which texts are most useful for reducing perplexity.The comparison is with Guu et al. (2018) and the retrieve-and-edit framework of Hashimoto et al. (2018).
  • Scalable grounded neural memory: REALM’s document index can be viewed as a memory whose keys are document embeddings, connecting it to scalable neural memory methods.The passage relates this view to product key memory and memory networks that enable sub-linear memory access.
  • Unsupervised Corpus Alignment: REALM’s latent document selection analogously provides model-centric unsupervised alignments between pre-training-corpus text X and knowledge-corpus text Z.This parallels latent selection of relevant tokens in sequence-to-sequence models with attention.

6. Future Work

The work is presented as a minimal REALM-like instantiation for reasoning over large knowledge corpora during inference. Future directions include structured knowledge and multilingual settings.

  • 6. Future Work: REALM is framed as a minimal member of a broader family of approaches that reason over large knowledge corpora on the fly during inference.The representation is pre-trained for this inference-time reasoning.
  • 6. Future Work: Structured knowledge could generalize Peters et al. (2019) by learning which entities are informative.This is identified as a particularly promising extension of the approach.
  • 6. Future Work: The authors also identify multilingual retrieval as a future direction.The passage introduces multilingual settings as a possible generalization but does not specify its implementation.

A. Derivation of the gradient with respect to the knowledge retriever

The section derives the gradient of the REALM log-likelihood objective with respect to the knowledge retriever parameters θ. The derivation applies conditional Bayes’ rule and rewrites the gradient as an expectation under p(z | y, x).

  • The derivation computes the gradient of the REALM pre-training objective, a log-likelihood, with respect to the knowledge retriever parameters θ.
  • Conditional Bayes’ rule is used to rewrite p(z | y, x) before expanding ∇log p(z | x).
  • The resulting expression is interpreted as an expectation with respect to p(z | y, x), allowing terms that depend on z′ but not z to move outside the expectation.

B. Connection between REALM and supervised learning

Under an idealized setting with one document yielding perfect predictions and all others yielding zero accuracy, REALM’s gradient descent objective becomes equivalent to supervised maximum-likelihood training on the gold document.

  • Connection between REALM and supervised learning: When one document z∗ gives perfect prediction accuracy and every other document gives zero accuracy, p(z∗ | y, x) = 1 if p(z∗ | x) is non-zero.The setup assumes p(y | z∗, x) = 1 and p(y | z′, x) = 0 for all other documents.
  • Connection between REALM and supervised learning: Gradient descent on the REALM objective is equivalent to gradient descent on log p(z∗ | x), the supervised maximum-likelihood objective for the “gold” document.This identifies the latent retrieved document z∗ with the document used as the supervised training target.

C. Adapting to new knowledge

REALM adapts to new world knowledge by replacing its document corpus after pre-training, enabling predictions to reflect updated information when corpora disagree. However, this adaptation is imperfect for facts whose answers remain strongly associated with outdated information in Wikipedia.

  • Corpus adaptation: REALM can adapt to new world knowledge by replacing its corpus documents with a more recent Wikipedia version after pre-training.The retrieval corpus is modified without retraining the model.
  • Corpus adaptation: When the original and updated corpora disagree about a fact, REALM can change its prediction to reflect the updated information.The paper demonstrates this behavior in Table 4.
  • Limitations: For “is the prime minister of United Kingdom,” the model predicts “Thatcher” on both corpora, likely because her name is frequently mentioned in Wikipedia articles.This example illustrates that explicit retrieval does not guarantee adaptation when a previously associated answer remains highly frequent.

D. Retrieval Utility

Retrieval utility (RU) measures how much a retrieved document improves masked-token log-likelihood over a null document. RU rises during pre-training, reflects retrieval usefulness, and supports adaptation when the knowledge corpus is updated.

  • Definition: RU is the difference between log p(y | z, x) and log p(y | ∅, x), comparing a retrieved document with the null document.The null document provides the reference condition for measuring a retrieved document’s importance.
  • Interpretation: Negative RU means the retrieved document is less useful than the null document for predicting y.This may indicate irrelevance, no need for world knowledge, or knowledge already encoded in the model’s parameters.
  • Training dynamics: RU increases steadily during pre-training and is more predictive of downstream Open-QA performance than the overall [truncated].The passage explicitly reports this relationship but ends before naming the comparison measure.
  • Corpus adaptation: The same model trained on the 2018 corpus retrieves the newly added “Excellent Cadaver” page in the 2020 corpus and generates “Lawrence”.With the outdated 2018 corpus, the model could not recover the word; updating the corpus enables retrieval of the new knowledge.
  • Training dynamics: RU estimates retrieval usefulness and is affected by masking choices and the number of pre-training steps.Figure 4 plots RU against pre-training steps.
Loading 2002.08909v1…