Retrieval Built Into the Architecture

Ordinary RAG retrieves documents and pastes them into the prompt, so every retrieved token is an ordinary token paying ordinary attention costs. Architectural retrieval puts the lookup inside the network, where retrieved content can influence the output without ever entering the sequence.

The difference in one line

Prompt-level RAG:
  retrieved text -> tokens -> the same sequence as everything else
  5,000 retrieved tokens cost 5,000 tokens of context, of attention,
  of KV cache and of input billing.

Architectural retrieval:
  retrieved text -> a separate encoding -> read by a dedicated
                    mechanism inside the network
  the main sequence length does not change at all.

That is the whole structural claim, and it has one immediate consequence. Under quadratic attention, doubling the retrieved context quadruples the attention cost; under a design where retrieved content is read by a separate mechanism, the relationship can be linear or flat. It is the same motivation as keeping a source in an encoder rather than in the decoder’s own cache.

kNN-LM: interpolate the output distribution

The simplest member, and the one that requires no training at all (Khandelwal and colleagues, 2020).

1. Run the trained model over the whole training corpus once. At each position, store the hidden state as a key and the actual next token as the value. One entry per training token.
2. At generation time, take the current hidden state and find its k nearest neighbours in that datastore.
3. Turn the neighbours’ distances into a distribution over their stored next tokens.
4. Mix it with the model’s own: p = lambda * p_kNN + (1 - lambda) * p_LM.

The model itself is untouched; the retrieval attaches at the very last step. The appeal is that adding knowledge becomes adding rows to a table, with no gradient step anywhere.

The cost is severe and worth stating in numbers. A datastore over a 100-billion-token corpus holds 100 billion vectors; at 1,024 dimensions in 16-bit that is 200 KB per thousand entries and about 200 TB in total. Even heavily quantised and pruned it is enormous, and a nearest-neighbour query is needed for every generated token, which puts an index lookup on the critical path of the decode loop. This is why kNN-LM is a beautiful idea that is not in production.

RETRO: chunked cross-attention

The most complete architectural version (DeepMind, 2021), and the one whose shapes are worth knowing.

1. Split the input into chunks of 64 tokens.
2. For each chunk, embed it with a frozen BERT-style encoder and
   retrieve the nearest neighbours from a corpus index.
3. Encode each retrieved neighbour (and its continuation) with a
   small encoder.
4. Insert chunked cross-attention layers into the decoder: chunk i
   attends to the neighbours retrieved for chunk i-1.

The retrieved text NEVER enters the decoder's own sequence.

Two properties fall out. First, the decoder’s self-attention length is unchanged, so retrieving more neighbours does not touch the quadratic term. Second, the retrieval is chunk-local: neighbours are fetched for each 64-token chunk rather than once for the whole prompt, so relevance can shift as the text goes on — something prompt-level RAG does badly, since it retrieves once before generation begins and then commits.

The reported headline of that work was that a retrieval-augmented model could match the perplexity of a substantially larger one without the corpus. The framing is the interesting part: retrieval as a substitute for parameters. Memorising facts in weights is expensive; looking them up is cheap, if the lookup is cheap.

Training the retriever with the model

A separate strand, and the one with the sharpest idea in it. REALM and the original RAG paper (both 2020) do not change where retrieved text enters — it still goes into the reader’s input — but they make the retriever part of the trained system.

Treat the retrieved document z as a latent variable and
marginalise over the top-k:

  p(y | x) = sum over z in top-k of  p(z | x) * p(y | x, z)

  p(z | x) comes from the retriever: a softmax over
           embedding similarities, which is DIFFERENTIABLE.

So the gradient of the generation loss flows back into the
retriever's encoder.

The consequence is worth stating carefully because it is the whole argument. An off-the-shelf embedding model retrieves documents that are similar to the query. Similarity is a proxy for usefulness, and the two come apart constantly — the passage that answers a question often does not restate it, and the passage that restates it often does not answer it. A retriever trained through the reader optimises for the thing you actually want: documents that make the answer better. That is a different objective, not a better implementation of the same one.

The engineering cost is brutal and it is why this is rare. The document index is built by encoding the corpus with the retriever’s own encoder. The moment those weights update, every vector in the index is stale. REALM handled it by re-indexing asynchronously during training, which means running a second job that continuously re-embeds the corpus while the first one changes the encoder underneath it.

The practical compromise almost everyone makes is to freeze the document encoder and train only the query encoder, so the index stays valid. That recovers much of the benefit and is a normal embedding fine-tuning job rather than an architectural change — which is a fair summary of what happened to this whole line of work.

Memory layers: parameters you look up

A third variant does not retrieve text at all. A memory layer holds a very large table of learned key-value pairs and reads it with a learned query, returning a weighted sum of the top few values.

1,000,000 memory slots, value dimension 1,024

full table: 1e6 * 1,024 = 1.02 billion parameters
read cost:  top-k over the keys, then k value lookups
            with k = 32, that is 32 vectors read, not 1,000,000

Product-key trick: factor the keys into two halves of 1,000 each,
search each half separately, combine -> 2,000 comparisons instead
of 1,000,000 for the same 1,000,000-slot table.

The result is a layer with a billion parameters whose cost per token is that of reading thirty-two vectors. It is the same economics as mixture of experts — capacity that is stored but not activated — taken to its extreme, where the “expert” is a single row.

What it buys and what it costs

Bought:

• A knowledge source far larger than any context window, read at sub-quadratic cost.
• Updates without retraining: swap the datastore and the model’s effective knowledge changes.
• Retrieval that moves with the generation rather than being fixed before it starts.
• A smaller parameter budget for the same factual coverage, since facts live in an index rather than in weights.

Cost:

• You must own the pretraining. RETRO-style cross-attention layers cannot be bolted onto a finished model. That alone excludes almost everyone.
• The index becomes a serving dependency. Latency, memory, sharding and freshness of a vector index now sit inside the inference path rather than beside it.
• Batching gets harder. Two requests in a batch retrieve different neighbours, so the cross-attention operand differs per sequence — exactly the kind of irregularity that costs throughput.
• Attribution is worse. With prompt-level RAG you can show the user the retrieved passage. With interpolated distributions or cross-attended encodings there is no clean span to cite, which matters wherever citations are a requirement.
• No API surface exists for it. Nothing in a chat-completions request lets you attach a datastore.

Why prompt-level RAG won anyway

Not on elegance. On composability.

Prompt-level retrieval works with any model behind any API, changes when you change your index rather than when you retrain, is debuggable by printing the prompt, produces citable spans, and improves whenever the underlying models improve without any work on your side. Every one of those is an engineering property rather than a modelling one, and collectively they beat a better architecture that only its owner can deploy.

The economics also moved. Long context got cheaper and prefix caching made a large repeated retrieved block much less expensive than the naive token count suggests, which erodes the main cost argument for the architectural version.