Retrival-Augmented Generation for Knowledge-Intensive NLP Tasks

Piyush Chaudhari

Limitations in current methods

• Standard models (For e.g. BART) store their knowledge as model weights during training

• Limitations -

  • Difficult to update without retraining
  • Hallucinations- model interpolate from compressed parametric memory
  • There is no way to tell what part was looked at to produce output

• This paper treats retrieval as a learnable, differentiable element that can be plugged into a generative model

• A jointly trained system where

  • Model knows what to query
  • Generator knows how to use what is queried

Types of Memory

• Parametric Memory

  • Knowledge is frozen into model weights
  • Can’t update it without retraining
  • Fast to infer insights but limited and brittle

• Non-parametric Memory

  • Indexed text, external to the model, which the model looks up at runtime
  • The index can be swapped, updated or replaced independent of model weights

• RAG combines both types of memories - retriever fetches relevant documents from non-parametric memory, and generators (that stores parametric knowledge) synthesizes a response using both

• Prior work on the hybrid approach was done only on extractive tasks (REALM,ORQA) but RAG extends it to generative seq2seq tasks

Approach

The paper combines a pre-trained retriever (Query Encoder + Document Index) with a pre-trained seq2seq model (Generator) and fine-tune end-to-end. For query \(x\), we use Maximum Inner Product Search (MIPS) to find the top-K documents \(z\)_\(i\). For final prediction \(y\), we treat \(z\) as a latent variable and marginalize over seq2seq predictions given different documents

Methods

• The paper explores RAG models, that uses input sequence \(x\) to retrieve text documents \(z\) to use them as additional context when generating the target sequence \(y\)

• The model leverages two components

  • Retriever \(p\)_\(\eta\)\((z | x)\) with parameters \(\eta\) that returns top-K truncated distributions over the text given a query \(x\)
  • Generator \(p\)_\(\theta\)(\(y\)_i | \(x,z,y\)_1:i-1) parametrized by \(\theta\) that generates a current token based on the context of the previous \(i-1\) tokens \(y\)_1:i-1 , the original input \(x\) and a retrieved text \(z\)

• The retrieved documents are treated as latent variable - the model never directly supervises which document is correct

• The paper proposes two models - RAG-Sequence that uses the same document to predict each target token and RAG-Token which can predict each target token based on a different document

Models

1. RAG-Sequence Model

• uses the same retrived document to generate complete sequence
• treats retrieved document as a single latent variable marginalized to get seq2seq probability \(p(y|x)\) via a top-K approximation
• top K documents are retrived using the retriver and generator produces output sequence probability of each document, then marginalized

\(p\)_RAG-Sequence\((y|x)\) \(\approx\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x)p\)_\(\theta\)\((y|x,z)\) \(=\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x)\prod p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

• \(\prod p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\) is the standard autoregressive generator
• \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x)\) does it for top-K documents

Models

2. RAG-Token Model

• Draw different latent document for each target token
• Allows the generator to choose content from different documents when generating a response
• Once top K documents are retrieved, the generator produces a distribution for the next output token for each document, before marginalizing
• The process is repeated for following output token

\(p\)_RAG-Token\((y|x)\) \(\approx\) \(\prod_{i}^{N}\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x) p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

Models

2. RAG-Token Model

\(p\)_RAG-Token\((y|x)\) \(\approx\) \(\prod_{i}^{N}\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x) p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

• The product is outside while the sum is inside, compared to RAG Sequence. At each token position \(i\), the model considers all K documents simultaneously, weights each document’s contribution to this specific token by the retrieval probability, and sums.

• Then moves to position \(i+1\) and does it again — potentially with a different document dominating.

• RAG can be used for sequence classification tasks by considering target class as a target sequence of length one, where RAG-Sequence and RAG-Token are equivalent

Retriever

• The retriever \(p\)_\(\eta\)\((z|x)\) is based on Dense Passage Retrival (DPR) that follows a bi-encoder structure

\(d(z)\) = BERT_\(d\)\((z)\) document encoder

\(q(x)\) = BERT_\(q\)\((x)\) query encoder

with a retrival probability of

\(p\)_\(\eta\)\((z|x)\) \(\propto\) exp \((d(z)^T.q(x))\)

• Higher dot product equals higher probability of retrival
• Finding top-K is a Maximum Inner Product Search problem
• Solved approximately using FAISS + HNSW graphs - sub-linear time over 21 million documents

Retriever

• One performace tradeoff is that the Document Encoder \(d(z)\) is frozen during fine-tuning

  • Re-embedding 21 million documents every training step is expensive
  • The initialization DPR retriever was pre-trained specifically to retrieve passages relevant to Natural Questions and TriviaQA. This means the document embeddings it produces are already semantically well-organized
  • So \(d(z)\) already has a meaningful embedding space before fine-tuning begins
  • Doesn’t hurt the performance significantly
  • Prevents catastrophic indexing - Updating the encoded during training changes the embedding, clashing with FAISS which was build over old embeddings

Generator

• 400M parameter seq2seq transformer, pre-trained with a denoising objective
• Input is simply the concatenation of query and retrieved document: [\(x\) ; \(z\)]
• Generates answer autoregressively, one token at a time:

\(p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

• At each step: conditioned on the full query, the retrieved document, and all prior output tokens
• BART over T5:
  • BART outperforms comparably-sized T5 on generation tasks
  • Bidirectional encoder + left-to-right decoder is well-suited for this hybrid input structure

Marginalizing Over Documents

• Since the “correct” document is unknown, both RAG variants sum (marginalize) over top-K documents, weighted by retrieval probability

• RAG-Sequence has the same document for the whole output
  • Generates a complete answer for each document, then weight by how likely it was to be retrieved

\(p(y|x)=\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x)\prod p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

• RAG-Token potentially has different document per output token
  • At each token step, the model can draw from a different document

\(p(y|x)\) \(\approx\) \(\prod_{i}^{N}\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z|x) . p\)_\(\theta\)\((y\)_i\(|x,z,y\)_1:i-1 \()\)

RAG-Sequence vs RAG-Token

Cohort	RAG-Sequence	RAG-Token
Document Usage	One Document for full answer	Different document per token
Flexibility	Lower	Higher
Applications	Coherent, Single-Source Outputs	Outputs synthesizing multiple sources
Decoding	Beam Search per document + re-scoring	Standard Beam Decoder

Training

• Trained end-to-end on input/output pairs (\(x\)_\(j\),\(y\)_\(j\) ) only — no supervision on which document to retrieve
• Loss function: negative marginal log-likelihood of the target output \(Loss\) = - \(\sum\)_\(j\) \(log\) \(p(y\)_\(j\)|\(x\)_\(j\) \()\)
• Gradient flows through BART (generator) and the query encoder
• The system learns to retrieve by learning what helps it generate correct answers
• Only the query encoder and BART are updated; document encoder + index are frozen
• Optimizer: Adam with stochastic gradient descent

Decoding

• RAG-Token
  • Per-token probability factors are plugged directly into standard beam decoder
  • Transition probability
  \(p'\)_\(\theta\)\((y\)_i\(|x,y\)_1:i-1 \()\) \(=\) \(\sum_{z\in top-k(p(.|x))}\) \(p\)_\(\eta\)\((z\)_\(i\)\(|x)p\)_\(\theta\)\((y\)_i\(|x,z\)_\(i\)\(,y\)_1:i-1 \()\)
  • To decode, \(p'\)_\(\theta\)\((y\)_i\(|x,y\)_1:i-1 \()\) can be plugged into a standard beam decoder

Decoding

• RAG-Sequence
  • Full sequence probability \(p(y|x)\) doesn’t factor per-token, hence can’t do a single beam search
  • Thorough Decoding:
    1. Get \(K\) sets of candidate sequences by running beam search independently for each of the \(K\) retrieved documents
       
    2. Take the union of all candidates across all K beams which is the hypothesis set \(Y\)
    3. For any candidate \(y\) that didn’t appear in a particular document’s beam, run an additional forward pass through BART to score it against that document anyway
    4. Final probability \(p(y|x)=\) \(\sum_{z}\) \(p\)_\(\eta\)\((z|x) . p\)_\(\theta\)\((y|x,z)\) summed across all \(K\) documents
    5. Pick the highest scoring \(y\)

Decoding

• RAG-Sequence
  • Fast Decoding
    • Similar to Thorough, but skips Step 3
    • If a candidate was generated from document 1’s beam but not document 2’s, Thorough Decoding still need to score it against document 2. With K=10 and a large hypothesis set, this is computinally expensive
    • In Fast Decoding, if \(y\) didn’t appear in document \(z\)’s beam, \(p\)_\(\theta\)\((y|x,z)\) is assumed as 0

Experiments

RAG was tested across five distinct task types to prove generality:

• Open-Domain QA (Natural Questions, TriviaQA, WebQuestions, CuratedTrec) — exact match
• Abstractive QA (MS-MARCO) — free-form answer generation, BLEU + Rouge-L
• Jeopardy Question Generation — given an entity, generate the Jeopardy clue — Q-BLEU-1 + human evaluation
• Fact Verification (FEVER) — classify claims as supported / refuted / not enough information

Results

Open-Domain Q&A

Model	Natural Questions	TriviaQA	WebQ	CuratedTrec
T5-11B (closed book)	34.5	—/50.1	37.4	—
T5-11B+SSM (closed book)	36.6	—/60.5	44.7	—
REALM (open book)	40.4	—	40.7	46.8
DPR (open book)	41.5	57.9	41.1	50.6
RAG-Token	44.1	55.2/66.1	45.5	50.0
RAG-Sequence	44.5	56.8/68.0	45.2	52.2

• RAG beats T5-11B (27x larger, expensive pre-training) on NQ and WebQ
• RAG beats DPR’s extractive pipeline — without a re-ranker or cross-encoder
• RAG scores 11.8% accuracy even when the correct answer is absent from all retrieved documents — extractive systems score 0% in this case

Results

Generation Tasks

Model	Jeopardy Q-BLEU-1	MS-MARCO Rouge-L	FEVER-3 Accuracy
BART baseline	19.7	38.2	64.0
RAG-Token	22.2	40.1	72.5
RAG-Sequence	21.4	40.8	—
State of the art*	—	49.8*	76.8

*SotA uses gold context/evidence passages; RAG retrieves its own

Results

Jeopardy

Human evaluation on Jeopardy (452 pairs):

• RAG more factual than BART: 42.7% of cases
• BART more factual than RAG: 7.1% of cases
• RAG more specific: 37.4% vs. 16.8%

On FEVER, RAG comes within 4.3% of complex pipeline systems that use domain-specific architectures and supervised retrieval labels — while using none of that.

Results

Generation Diversity

Measured as ratio of distinct tri-grams to total tri-grams:

Model	MS-MARCO	Jeopardy
Gold (human)	89.6%	90.0%
BART	70.7%	32.4%
RAG-Token	77.8%	46.8%
RAG-Sequence	83.5%	53.8%

• RAG generates significantly more diverse outputs than BART — without any diversity-promoting decoding tricks

• The retrieval of varied evidence naturally pushes the generator toward more specific, varied language

• BART on Jeopardy at 32.4% - a pure parametric model tends to repeat generic patterns

• RAG is anchored to specific retrieved facts

Ablations

• Frozen vs. learned retrieval:

  • RAG-Token with frozen retriever: 37.8 on NQ and 43.5 with learned
  • Learning to retrieve is worth ~6 EM points on NQ alone

• Dense retrieval (DPR) vs. sparse retrieval (BM25):

  • BM25 replacement tanks QA performance significantly (29.7 vs. 43.5 on NQ for RAG-Token)
  • Exception: FEVER — BM25 performs comparably or better, because fact-checking claims are entity-heavy and keyword overlap is a reliable signal
  • For generation tasks: dense retrieval is crucial

• Number of retrieved documents:

  • RAG-Sequence: more docs monotonically helps QA
  • RAG-Token: performance peaks at K=10, then levels off
  • Flexibility to adjust K at inference time without retraining is a practical advantage

Index Hot-Swapping

• Same trained RAG model, two different Wikipedia indexes (Dec 2016 vs. Dec 2018)
• 82 world leaders who changed roles between those dates

Index	2016 leaders query	2018 leaders query
2016 index	70% correct	4% correct
2018 index	12% correct	68% correct

• Mismatched index has a near-zero accuracy
• Matched index has about 70% accuracy with same model weights, but different index
• RAG’s world knowledge lives in the index, not the weights, and hence a deplyed RAG system can be updated on factual knowledge by replacing its document store
  • without retraining
  • without fine-tuning

Generation vs Extraction

• All four QA benchmarks in the paper are classified as extractive tasks — the answer exists verbatim in some document.

• Still, RAG’s generative approach beats purpose-built extractive systems.

  • Documents that contain clues but not the verbatim answer can still guide generation — extraction can’t use them at all
  • Marginalization over \(K\) documents lets multiple partial-evidence documents contribute to a single correct answer
  • Generation can synthesize across documents while extraction is limited to one span in one document
  • RAG can produce correct answers even when no retrieved document contains the answer

Parametric vs Non-Parametric Memory

Property	Parametric (BART weights)	Non-Parametric (FAISS index)
Update mechanism	Retraining	Index replacement
Interpretability	Opaque	Human-readable raw text
Writability	No	Yes — edit documents directly
Access speed	Instant (forward pass)	Fast (MIPS, sub-linear)
Knowledge source	Training corpus	Any document collection
Hallucination risk	High	Lower (grounded in text)

• The non-parametric memory being raw text (not embeddings) is a design choice
• Keeps the memory human-readable and human-writable
• Gives you authority to correct - a factual error can be corrected by editing a document in the index, the same thing cannot be done with weights

Limitations

• Single knowledge source (Wikipedia was the one source tested)
• Joint pre-training wasn’t explored, only fine-tuning. was performed
• RAG-Sequencing Decoding doesn’t scale optimally
• Knowledge spanning multiple documents can’t be multi-hopped while reasoning natively
• No inference regarding when to retrieve vs. when to rely purely on parametric memory

Legacy of the Paper

While some of the specific components have been replaced by stronger alternatives, the core ideas from this paper propel the field even today

• Treat retrieval as a learnable, differentiable function
• Marginalize over uncertainty in which document is relevant
• Separate updatable knowledge (index) from reasoning capability (weights)
• Train end-to-end on (input, output) pairs without document-level supervision
• Hot-Swapping has saved retraining costs across the industry- The observation that a 400M parameter model with a good index beats an 11B parameter model without one reshaped how to think about the scaling-vs-retrieval tradeoff

Current Advances

• Retrievers
  • Sparse + Dense Hybrid Retrievals outperform and have replaced DPR
  • Both encoders can be fine-tuned using. contrastive learning, solving the frozen document-encoder compromise
• Generators
  • BART has been replaced by stronger LLMs like LLaMA, Mistral, GPT-4
  • Capability to synthesize retrieved context has improved
• Multi-Hop RAG
  • Modern systems reads retrieved docs, generates a sub-question, retrieves again, and chains reasoning across multiple retrieval steps
• Production RAG
  • The index hot-swap insight became the architectural backbone of virtually every enterprise AI system — LlamaIndex, LangChain, and similar frameworks are essentially RAG pipelines with better chunking, re-ranking, and hybrid retrieval layered on top