Retrieval as a Latent-Variable Model

RAG written as marginalisation over a latent document, the exact decomposition of its error, and why retrieval failures propagate rather than average out.

RAGResearch⏱ 1 day13 min readraglatent variablemarginalisationerror propagationreranking

5 problems →

AssumesVector Search — Exact and Approximate Nearest Neighbours · Language Modelling, Cross Entropy and Perplexity

The model

DefinitionRAG as marginalisation over a latent document

Let $x$ be a query, $y$ a response, and $z$ a document from a corpus $\Zcal$ . Treat $z$ as a latent variable:

p(y\mid x) = \sum_{z\in\Zcal} \underbrace{p_\eta(z\mid x)}_{\text{retriever}}\ \underbrace{p_\theta(y\mid x,z)}_{\text{generator}} .

In practice the sum is truncated to the top- $k$ retrieved documents, with the retriever posterior

p_\eta(z\mid x) = \frac{\exp\big(\ip{f_\eta(x)}{g_\eta(z)}\big)}{\sum_{z'\in\mathrm{top}\text{-}k}\exp\big(\ip{f_\eta(x)}{g_\eta(z')}\big)} .

PropositionRAG-sequence and RAG-token are different models

p_{\text{seq}}(y\mid x) = \sum_z p_\eta(z\mid x)\prod_{t=1}^{|y|}p_\theta(y_t\mid x,z,y_{<t}), \qquad p_{\text{tok}}(y\mid x) = \prod_{t=1}^{|y|}\sum_z p_\eta(z\mid x)\,p_\theta(y_t\mid x,z,y_{<t}).

These are not equal: the first is a mixture of products and the second a product of mixtures. $p_{\text{seq}}$ commits to a single document for the whole response; $p_{\text{tok}}$ can draw different tokens from different documents.

A variational view

TheoremA lower bound on the RAG log-likelihood

For any distribution $q(z)$ over documents,

\log p(y\mid x) \;\ge\; \E_{z\sim q}\big[\log p_\theta(y\mid x,z)\big] - \mathrm{KL}\big(q(z)\,\|\,p_\eta(z\mid x)\big),

with equality if and only if $q(z) = p(z\mid x,y)$ , the true posterior given both query and answer.

Key idea

The bound identifies what a retriever should approximate: not $p(z\mid x)$ , the documents relevant to the query, but $p(z\mid x,y)$ , the documents that explain the answer. Those differ, and the gap is the reason two-stage training helps — the retriever trained on query–document relevance is optimising the wrong posterior, and methods that distil the generator's likelihood back into the retriever (REALM, and the retriever-update schemes that follow it) are closing exactly this gap.

Error propagation

TheoremRetrieval error propagates into generation

Fix a query $x$ and suppose there is a set $Z^\star$ of documents sufficient to answer it. Let

\gamma = p_\eta\big(Z^\star \mid x\big)

be the retriever's probability of supplying a sufficient document (its recall, in probability form). Let $p_{\mathrm{RAG}}(\cdot\mid x)$ be the RAG output distribution and $p_{\mathrm{oracle}}(\cdot\mid x)$ the distribution obtained by conditioning on $Z^\star$ with certainty. Then

\big\| p_{\mathrm{RAG}}(\cdot\mid x) - p_{\mathrm{oracle}}(\cdot\mid x)\big\|_{\mathrm{TV}} \;\le\; 1-\gamma .

PropositionThe bound is attained

There are instances with equality. Take a binary answer space $\{y_0,y_1\}$ , a generator that outputs $y_1$ with probability $1$ given any $z\in Z^\star$ and $y_0$ with probability $1$ given any $z\notin Z^\star$ . Then $\norm{\tilde p_{A^c}-\tilde p_A}_{\mathrm{TV}}=1$ and the bound holds with equality.

Key idea

Read the bound as: retrieval error is a hard ceiling, not a soft degradation. A retriever that fails $10\%$ of the time permits the output distribution to be wrong on $10\%$ of the probability mass, and the tightness proposition shows a generator can realise the whole of that error.

There is no averaging-out. One might hope that the $k-1$ irrelevant documents in the context act as neutral noise the model ignores; the evidence is against that. Shi et al. show that irrelevant context measurably degrades reasoning accuracy, and Liu et al. show that relevant context placed in the middle of a long prompt is used substantially less than context at either end. So the failure modes compound rather than cancel: missing documents cap the achievable quality, and extra documents actively hurt.

Combined with the multiplicative recall of the pipeline, the arithmetic is unforgiving. A chunker that splits the answer across boundaries ( $\gamma_1 = 0.95$ ), an ANN index at $\gamma_2 = 0.9$ , and a top- $k$ truncation at $\gamma_3 = 0.95$ give $\gamma \approx 0.81$ — so nearly a fifth of the output distribution is unconstrained by the evidence, before the generator makes a single mistake of its own.

Reranking

PropositionWhat reranking can and cannot do

Let the first stage return a candidate set $C$ with $\P(Z^\star\cap C\ne\emptyset) = \gamma_1$ , and let a reranker select the top $k'<|C|$ from $C$ . Then the post-rerank recall satisfies

\gamma_{\text{rerank}} = \gamma_1\cdot\P\big(\text{relevant document in top }k' \mid \text{it is in } C\big) \le \gamma_1 .

Reranking can only improve the precision of the final set, never the recall of the pipeline.

The variant landscape

HeuristicRAG variants as choices about the latent structure

Heuristic — a reasoning shortcut that predicts the right answer in the regimes discussed here. It is not a proof and can fail outside them.

The named RAG variants are all interventions on the posterior $p_\eta(z\mid x)$ or on the latent space $\Zcal$ , within the same probabilistic model.

Variant	What it changes	In the model
Naive RAG	—	one-shot $p_\eta(z\mid x)$ , top- $k$ , concatenate
Hybrid retrieval	scoring function	$p_\eta \propto \exp(\lambda s_{\text{dense}} + (1-\lambda)s_{\text{BM25}})$
Reranking	truncation of the support	two-stage construction of $p_\eta$
Query expansion / HyDE	the conditioning variable	replace $x$ by $\tilde x = h(x)$ closer to document space
Corrective RAG	conditional re-retrieval	evaluate sufficiency, re-draw $z$ if inadequate
Agentic / iterative RAG	sequential latent	$z_1,\dots,z_m$ drawn adaptively, a latent sequence not a latent variable
Graph RAG	the latent space itself	$\Zcal$ becomes a graph; retrieval returns a subgraph, so multi-hop paths are single latents
Fusion-in-Decoder	where mixing happens	attention over separately encoded documents, not an explicit sum

What this does not say

Two honest caveats about this table.

First, the fact that all of these fit one framework does not mean they are equally good or that the framework predicts which will win. The probabilistic view organises the design space; it does not supply a theory of retrieval quality, and the comparative evidence between variants is empirical and benchmark-dependent.

Second, the iterative and agentic variants strictly leave the single-latent model: with $m$ adaptive retrieval rounds the object is a latent trajectory, the marginalisation is over sequences, and the error analysis above does not apply directly — each round's failure probability compounds, and there is no clean analogue of the total-variation bound. That is a genuine gap: the systems being built are ahead of the analysis, and the error propagation for multi-round retrieval has not, to this author's knowledge, been characterised.

Problems

RAG-sequence and RAG-token are different modelsProof · Intermediate · 30 min The RAG evidence lower boundProof · Advanced · 30 min Retrieval error propagates into generationProof · Graduate · 1 hour End-to-end recall arithmetic for a RAG systemNumerical · Beginner · 15 min Build and instrument a minimal RAG pipelineCoding · Advanced · 1 day

References

PaperPatrick Lewis and collaborators (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
PaperKelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat and Ming-Wei Chang (2020). REALM: Retrieval-Augmented Language Model Pre-Training
PaperGautier Izacard and Edouard Grave (2021). Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering
PaperFreda Shi and collaborators (2023). Large Language Models Can Be Easily Distracted by Irrelevant Context
PaperNelson F. Liu and collaborators (2024). Lost in the Middle: How Language Models Use Long Contexts

These works are cited as sources to read, not reproduced. All exposition, examples and problems on this site are written independently; where an argument is standard, the citation points to a canonical treatment.

Results stated on this page

DefinitionRAG as marginalisation over a latent documentThe generation probability is a mixture over retrieved documents weighted by the retriever's posterior.
PropositionRAG-sequence and RAG-token are different modelsproof includedMarginalising once per sequence and once per token give different distributions, and only the second can switch sources mid-generation.
TheoremA lower bound on the RAG log-likelihoodproof includedJensen's inequality gives an evidence lower bound whose gap is the KL divergence between the retriever posterior and the true posterior.
TheoremRetrieval error propagates into generationproof includedThe total variation distance between the RAG output and the oracle output is bounded by the retriever's probability of missing the relevant document.
PropositionThe error bound is attainedproof includedThere are instances where a retrieval failure changes the output distribution by exactly the failure probability, so the bound cannot be improved.
PropositionWhat reranking can and cannot doproof includedA reranker improves precision within the retrieved set and cannot exceed the first-stage recall.
HeuristicRAG variants as choices about the latent structureCorrective, agentic, hybrid and graph RAG are different ways of shaping the posterior over documents, not different probabilistic frameworks.

Problems (5)

All problems