[Paper Review] REALM: Retrieval-Augmented Language Model Pre-Training

Kelvin Guu, Kenton Lee+3 more|arXiv (Cornell University)|Feb 10, 2020

Topic Modeling38 references513 citations

TL;DR

REALM augments language model pre-training with a learned neural retriever over a large text corpus, enabling end-to-end unsupervised pre-training and improved open-domain QA performance. It outperforms prior retrieval-based and generation-based methods on three Open-QA benchmarks while maintaining moderate model size.

ABSTRACT

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.

Motivation & Objective

• Motivate integrating a latent knowledge retriever into pre-training to capture world knowledge more modularly.
• Develop an end-to-end retriever and knowledge-augmented encoder trained with unsupervised MLM signals.
• Enable retrieval-then-predict during pre-training and fine-tuning for Open-QA tasks.
• Demonstrate scalability with MIPS-based retrieval and asynchronous index refreshing.
• Show interpretable retrieval behavior and ablations to inform design choices.

Proposed method

• Model architecture combines a dense retriever with inner-product scoring to select top documents for each query.
• Retriever embeddings are learned via a Transformer-based encoder for input x and a separate document encoder for z, with f(x,z) = Embed_input(x) dot Embed_doc(z).
• Knowledge-augmented encoder conditions on retrieved documents to predict y, using cross-attention between x and z.
• Training optimizes p(y|x) by marginalizing over top-k retrieved documents: p(y|x) = sum_z p(y|x,z) p(z|x).
• Pre-training uses masked language modeling with retrieved documents to predict masked tokens.
• Fine-tuning for Open-QA retrieves top-k documents (top 5) and marginalizes over them to predict answer spans or tokens.
• To handle large corpora, employs Maximum Inner Product Search (MIPS) with asynchronous index refreshes to keep retrieval index up-to-date.

Experimental results

Research questions

• RQ1 Can a latent knowledge retriever trained with unsupervised signals improve language model pre-training for knowledge-intensive tasks?
• RQ2 Does end-to-end RETRIEVAL-augmented pre-training improve Open-QA benchmarks beyond prior retrieval-based and generation-based methods?
• RQ3 How does the quality of retrieved documents impact downstream prediction and QA accuracy?
• RQ4 What are effective inductive biases (e.g., salient span masking) to guide retrieval during pre-training?
• RQ5 Is the MIPS-based retrieval scalable and stable when integrated into large-scale pre-training?

Key findings

• REALM achieves 4-16 percentage points absolute improvement over prior Open-QA methods on three benchmarks.
• REALM outperforms the largest T5-11B model while being 30x smaller.
• Both encoder and retriever benefit from REALM pre-training, with best results from joint optimization.
• Salient span masking and including a null document improve retrieval usefulness and model performance.
• Asynchronous MIPS index refreshing yields stable optimization and better retrieval quality.
• REALM demonstrates interpretable retrieval behavior by showing how retrieved documents influence predictions.