REALM: Learning to Retrieve During Language-Model Pretraining

4 minute read

Published:

TL;DR: REALM trains a retriever through a language-modelling objective before task-specific adaptation. Retrieved documents help predict masked content, and the loss rewards documents that improve that prediction. Updating document representations creates an index-staleness problem that must be managed during learning.

Retrieval becomes part of pretraining

REALM’s contribution is to expose retrieval to the learning signal of language-model pretraining. Given masked input, the model retrieves external text and predicts the missing content conditioned on it. Retrieval is treated as a latent choice rather than requiring a human relevance label for each masked example.

The paper emphasises informative masking and efficient retrieval over a large collection. It also refreshes document representations while training so the searchable index does not remain permanently tied to obsolete encoder parameters. This differs from RAG’s task-specific mixture and from an inference-only prompt pipeline.

Marginalise the retrieved evidence

Let \(x\) be masked input, \(y\) the target content and \(z\) a retrieved document. The basic objective is

\[ \mathcal L=-\log\sum_z p_\eta(z\mid x)p_\theta(y\mid x,z). \]

The sum is approximated over retrieved candidates. Its posterior document weight is

\[ r(z)=\frac{p_\eta(z\mid x)p_\theta(y\mid x,z)}{\sum_{z'}p_\eta(z'\mid x)p_\theta(y\mid x,z')}. \]

This makes the training intuition concrete: a document earns responsibility when it was plausible to retrieve and helps predict the target.

Work through two documents

Suppose a retriever assigns 0.8 to a general overview and 0.2 to a release note. The masked product version has conditional probability 0.1 given the overview and 0.9 given the note. Their contributions are 0.08 and 0.18; the note receives posterior responsibility \(0.18/0.26\approx0.692\) despite its smaller retrieval prior.

The gradient can therefore favour evidence that the current retriever under-ranks. However, a document absent from the retrieved candidate set receives no such direct feedback in that truncated computation. Exploration, candidate quality and index freshness all affect what the learning process can discover.

Why some masks need retrieval and others do not

Masking a predictable function word can be solved from local syntax. Masking a fact-bearing entity or date is more likely to benefit from external information. If most examples are locally predictable, a reader may learn to ignore retrieval while the overall loss still improves.

A useful diagnostic compares performance with real evidence, shuffled evidence and no evidence. Include examples whose answer is genuinely absent from the visible local context. If the model performs identically under all three conditions, a retrieval improvement may not be the mechanism behind the reported gain.

The moving-index problem

A document encoder update changes the geometry of stored representations. Queries encoded by new parameters and documents encoded by old parameters can become poorly aligned. Rebuilding after every gradient step is expensive; never rebuilding leaves increasingly stale vectors.

REALM addresses this with asynchronous refreshes. In an implementation, log the encoder version associated with each index snapshot and measure how retrieval changes around refreshes. A training curve can jump because the search corpus representation changed, even if the number of source documents stayed constant.

Evaluate knowledge access separately from memorisation

A model may learn a fact in its parameters during pretraining, retrieve it from the collection, or use both. To investigate the contribution of retrieval, control corpus overlap, compare index updates and inspect evidence-dependent questions. Do not infer editable knowledge merely because documents can technically be replaced.

Contrast this with kNN-LM, which adds a token datastore to an existing model, and Atlas, which connects retrieval-aware learning to a multi-passage generative reader. These are distinct places to introduce external memory.

Key Insight: Latent retrieval learns through predictive usefulness. The important feedback is whether a retrieved document helps explain the target, not whether it merely resembles the masked query.
Common pitfall: A reader that solves most masks without external evidence can hide a weak retriever. Include evidence ablations and track the index version during training.

Recap

  • REALM brings retrieval into masked-language-model pretraining.
  • Posterior document responsibility connects prediction quality to retrieval learning.
  • Candidate truncation limits which documents receive feedback.
  • Encoder updates and index refreshes must remain coordinated.

References

  1. Guu et al. REALM: Retrieval-Augmented Language Model Pre-Training. ICML, 2020.
  2. Lewis et al. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS, 2020.
  3. Khandelwal et al. Generalization through Memorization: Nearest Neighbor Language Models. ICLR, 2020.