Improving the Domain Adaptation of Retrieval Augmented Generation (RAG) Models for Open Domain Question Answering

In plain words

Some question-answering systems look up passages in a document store first and then write the answer from what they found. The original RAG design does this well for Wikipedia, because its search component was already trained on Wikipedia. For news or medical text, the search component is weaker, and the stored passage vectors are never updated, so the system cannot adapt.

This paper's variant, RAG-end2end, lets the search part and the answer writer learn together on the target domain. The passage vectors are recomputed in background jobs while training goes on, so the store stays current without stopping training. The authors also add a side task: rebuild a sentence from the passages that should support it, which pushes the system to learn facts from the domain text.

The authors test this on COVID-19 papers, news and conversations, and report clear gains over the original design on all three. The method costs more to train, because the whole passage store has to be re-encoded and re-indexed during training.

TL;DR

• RAG-end2end fine-tunes the DPR question and passage encoders and the BART generator together, and refreshes the passage index during training §3.3.
• An auxiliary task rebuilds a domain sentence from retrieved passages, marked with a control token <p> §3.4 eq. 3.
• Author-reported: exact match improves over RAG-original-QA+R on all three domains (Table 1, tab. 1).
• Costly for large knowledge bases: re-encoding takes GPU hours and re-indexing takes CPU hours §3.3.
• Results are single author-reported runs; training hyperparameters other than epochs are not stated.

Problem

The original RAG fine-tunes only the question encoder. Its passage encoder and the encoded knowledge base stay frozen after the Wikipedia-trained DPR checkpoint is used once §3.3. The authors argue this limits domain adaptation, because DPR was trained on Wikipedia-style data §1 §2.1. Updating only the question encoder without re-encoding the knowledge base "could degrade performance" §1. Full end-to-end updating is hard because each change to the passage encoder makes the stored passage vectors and the FAISS index stale; on a knowledge base of millions of passages, re-encoding takes several GPU hours and re-indexing several CPU hours §3.3.

Method

Retriever and generator (as in RAG). The retriever is a DPR model with two BERT towers, a question encoder E_Q and a passage encoder E_P. Relevance is the dot product of their CLS embeddings §3.1 eq. 1. The generator is a pre-trained BART seq2seq model. The loss marginalises each output token over the top-k retrieved passages, with the retrieval probability as weight §3.1 eq. 2.

\mathrm{sim}(p, q) \propto E_Q(q)^{\top} E_P(p) \qquad [\text{eq. 1}]

P_{\text{RAG-Token-Loss}}(y \mid x) = \prod_{i}^{n} \sum_{z \in \text{top-}k\,P(\cdot \mid x)} P_{\eta}(z \mid x)\, P_{\theta}(y_i \mid x, z, y_{1:i-1}) \qquad [\text{eq. 2}]

Indexing. All passages are encoded with E_P and stored in a FAISS index. The authors use HNSW (flat storage) for fast nearest-neighbour search §3.2 §4.2.

End-to-end retriever training (RAG-end2end). Both E_Q and E_P are trained through the generator loss. The gradients reach E_P through the passage selection probability in eq. 2 §3.3. For the trained passage encoder to matter, the stored passage vectors and the index must be rebuilt from it §3.3.

Asynchronous refresh. Three processes run: the training loop, a re-encoding job on GPUs outside the loop, and a re-indexing job that builds a FAISS index once re-encoding has finished and saved its output §3.3 fig. 1. The training loop loads each new index when it is ready, then starts the next re-encoding. Re-indexing must not start until all passage vectors are saved, or it fails §3.3. Gradients may be stale because training and indexing run asynchronously; the authors cite prior work saying this does not strongly hurt §3.3.

Statement reconstruction (auxiliary signal). A sentence from the domain corpus is encoded with E_Q, the top passages are retrieved from a knowledge base that does not contain that sentence, and the generator must rebuild the sentence from those passages §3.4. A control token <p> marks reconstruction inputs; QA inputs use the question instead §3.4 eq. 3. The paper does not give the loss weight between the QA and reconstruction signals, nor the order in which the two kinds of example are mixed §3.4.

\begin{aligned} \text{QA input:} &\quad \langle \text{Question} \rangle + \langle \text{Retrieved Passages} \rangle \\ \text{Reconstruction input:} &\quad \langle \texttt{<p>} \rangle + \langle \text{Retrieved Passages} \rangle \end{aligned} eq. 3

Pseudocode of the training loop. The paper gives no algorithm box; this sketch is written from §3.1 to §3.4 and fig. 1.

# Joint training with an asynchronous index refresh  [§3.3] [fig. 1]
# Input:  qa_pairs, statements (text); passages (knowledge base text)
# Output: trained question encoder, passage encoder and generator
# Not stated: how QA and statement examples are mixed [§3.4]

def train(qa_pairs, statements, passages):
    # Index built from the pretrained DPR passage encoder  [§4.2]
    index = build_index(encode_all(passages, passage_encoder))
    for example in training_examples(qa_pairs, statements):
        # Question or statement -> dense vector  [§3.1] [eq. 1]
        query_vector = question_encoder(example.input_text)
        # Top-k passages from the current index; k not stated  [eq. 2]
        retrieved = index.search(query_vector, k)
        # QA: question + passages; statement: <p> + passages  [eq. 3]
        loss = generator_loss(example, retrieved)  # marginalised [eq. 2]
        update_parameters(loss)  # gradients reach both encoders  [§3.3]
        if refresh_due():  # about every 750 updates on average [§3.3]
            # Runs on other GPUs while training continues  [§3.3]
            start_refresh(passages)
        if refresh_finished():
            # Swap in only after all vectors are saved  [§3.3]
            index = refreshed_index()

Domain datasets. Three domains, each with its own knowledge base, QA data and reconstruction statements §4.1:

• COVID-19: 5,000 CORD-19 full-text articles split into 250,000 passages of 100 words, each prefixed with its paper title. Reconstruction statements come from abstract sentences of 10,000 papers, filtered to 15 to 35 words, about 50,000 statements. Training and validation QA are synthetic: a BART-large model trained two epochs on SQuAD writes questions for passages, with round-trip consistency filtering, giving 225,000 pairs split 90/10. The test set is 2,000 human-labelled COVID-QA pairs §4.1.
• News: 85,000 passages of 100 words from 10,000 NewsQA articles. Reconstruction statements are the CNN/DM summary sentences after the first, about 35,000. NewsQA gives 90,000 train, 5,000 validation and 5,000 test QA pairs §4.1.
• Conversation: 110,000 passages of at most 100 words from 10,000 QAConv conversations, each passage prefixed with the conversation id and the speaker name. Reconstruction statements are one-sentence summaries of conversations longer than 45 words, about 35,000. QAConv gives 25,000 train, 5,000 validation and 5,000 test QA pairs §4.1.

Standalone DPR fine-tuning (comparison, §4.5). The authors also fine-tune DPR alone with its original loss, using hard negatives chosen by BM25 §4.5. For COVID-19 they use synthetic QA pairs. For news and conversations they report that fine-tuning on the supervised data did not help, and attribute this to answers sitting in passages longer than the 100 words DPR was trained on §4.5. They then generate synthetic QA pairs for those domains as well.

Hyperparameters

Results

Author-reported. Setup: three domain QA sets (COVID-19, News, Conversations), RAG-token base model, each model fine-tuned on that domain §4.2. "RAG-original" is the public NQ checkpoint, zero-shot §4.2. "QA" means trained on domain QA pairs only; "+R" adds the reconstruction signal §4.2.

Table 1: domain adaptation results tab. 1

Table 3: DPR retrieval accuracy after domain adaptation tab. 3

Table 4: RAG initialised from a domain-adapted DPR tab. 4

Table 5: open-domain SQuAD (in-domain check) tab. 5

Headline comparisons from the text §4.3 §4.4:

• RAG-end2end-QA over RAG-original-QA: EM rises by 1.13 points (News) to 12.16 points (Conversations) §4.3.
• Adding the reconstruction signal to RAG-end2end-QA: EM rises by 1.7 points (Conversations) to 5.69 points (News, from Table 1) §4.4.
• Full method (5) over RAG-original-QA (2): EM rises 13.86 points on Conversations tab. 1.
• The authors also say RAG-end2end's DPR (DPR-RAG-end2end) beats DPR-domain-adapted on top-5 and top-20 in all three domains tab. 3 §4.5.

Paper vs code

Code location. The paper's footnote links the implementation at examples/research_projects/rag-end2end-retriever in huggingface/transformers §1. That folder was removed from main in March 2025 (commit 1e4286fd5962438d15af0257fa5feb9f2b2e58be, "Remove research projects"; found through the GitHub API, not a local file). The last commit containing it is 37508816d650a8074bc31c761e20872c2e5eaec4. The README there names Shamane Siriwardhana as maintainer, and two other contributors README.md. The repository itself belongs to Hugging Face, not the authors, and the folder is gone from main. So this card has no code field and no code_commit. The code below was read from the snapshot at 37508816 through GitHub, not cloned into raw/code/, so no line-numbered citations are possible.

Discrepancies found in that snapshot. The snapshot is the authors' example code. Its training configuration is a SQuAD example, not the paper's domain setup.

• Re-encoding batch size: the paper uses 32 per GPU §3.3. The passage encoding step in kb_encode_utils.py uses batch_size=16 kb_encode_utils.py.
• Index refresh interval: the paper's average is about 750 updates §3.3. The example shell script sets --indexing_freq 500 [code: examples/research_projects/rag-end2end-retriever/finetune_rag_ray_end2end.sh].
• Training hyperparameters: the paper gives 10 epochs §4.2 and nothing else. The example script sets learning rate 3e-5, training batch size 4 with 8 gradient accumulation steps, 500 warmup steps, label smoothing 0.1, dropout 0.1, weight decay 0.001, max gradient norm 0.1 and a polynomial schedule [code: examples/research_projects/rag-end2end-retriever/finetune_rag_ray_end2end.sh]. These belong to the SQuAD example and are not the paper's values.
• Passage-encoder checkpoint: the index build uses facebook/dpr-ctx_encoder-multiset-base [code: examples/research_projects/rag-end2end-retriever/kb_encode_utils.py]. The paper names no checkpoint §3.1.
• Index structure: the paper says HNSW with 128 links §4.2. The code builds an IndexHNSWFlat with dimension 768, 128 links and inner-product metric kb_encode_utils.py. It sets no efConstruction or efSearch. This matches the paper.
• Passage length: the paper says 100 words §4.2. The chunking function defaults to 100 words kb_encode_utils.py. This matches the paper.
• Reconstruction signal: the paper describes the statement reconstruction task and the <p> control token §3.4. The files I read (finetune_rag.py, the Ray script, the callbacks, the README) contain no reconstruction data or <p> token. I did not search the whole folder, so treat this as "not found in the files read", not as proof of absence.
• Joint training: the code enables end-to-end training with a flag and sets the context encoder for training when it is on [code: examples/research_projects/rag-end2end-retriever/finetune_rag.py]. This matches the paper's joint training.

Pitfalls

• Cost scales with the knowledge base. Re-encoding takes several GPU hours and re-indexing several CPU hours on large bases §3.3. The authors say this is the main cost §5.2.
• Ordering matters. Re-indexing must start only after all passage vectors are saved, or it fails §3.3. An implementation that refreshes in place will produce mixed-version indices.
• Stale gradients are accepted. The index lags behind the encoder during training §3.3. The paper gives no measurement of how much lag is tolerable.
• Top-k and loss weighting are missing. The number of retrieved passages k is not stated §3.1. Neither is the weight of the reconstruction loss against the QA loss §3.4. The fine-tuning optimiser, learning rate, batch size and warmup are also not stated §4.2.
• Fixed passage length. Passages are 100 words §4.2. DPR was trained on 100-word passages, and the authors found that supervised DPR fine-tuning did not help news or conversations because the answers sit in longer passages §4.5.
• COVID-19 training is synthetic. Its training and validation pairs are generated, and only the test set is human-labelled §4.1. Checkpoint selection uses validation accuracy on synthetic data §4.2, so the choice of checkpoint is never checked against human questions.
• Eq. 2 is typeset as a probability. The text calls it a loss §3.1. An implementer should take the negative log, which is how the loss is normally defined. The paper does not say this.
• Inconsistent numbers in the paper.
  • Table 1 gives F1 12.20 for RAG-original-QA+R on COVID-19, but Table 4 gives 12.12 for the same model tab. 1 tab. 4.
  • §4.4 says the EM increase from the reconstruction signal reaches 8.39 points on News. Table 1 gives 14.08 minus 8.39, which is 5.69 §4.4 tab. 1.
  • §4.4 gives a top-20 range of 3.2 to 8 points. Table 1 gives 3.20 to 9.75 across the three domains.
• Single runs, no variance. The paper reports one run per configuration, with no seeds or confidence intervals §4.2.
• Dataset release. The authors say they released the domain knowledge bases and QA pairs, through a Google Drive folder linked in §4.1 §4.1. I did not check that link or any license on it, so do not assume the data can be reused.

Lineage

• Builds on: RAG (Lewis et al.), with arXiv:2005.11401 as its id in the code README §2.1 README.md. Adds joint retriever training and an asynchronous index refresh.
• Builds on: DPR (Karpukhin et al.) for the retriever encoders §3.1. The paper gives no arXiv id.
• Compared against: REALM (Guu et al.), which trains its retriever with masked language modelling and freezes it for fine-tuning §2.2; Sachan et al. and Singh et al., who train the retriever end-to-end for QA §2.2; RETRO, which keeps frozen BERT retrieval §5.2.
• Superseded by: not stated.