[Paper Review] Fine-tune the Entire RAG Architecture (including DPR retriever) for Question-Answering

Shamane Siriwardhana, Rivindu Weerasekera+2 more|arXiv (Cornell University)|Jun 22, 2021

Topic Modeling7 references4 citations

TL;DR

This paper proposes end-to-end fine-tuning of the entire Retrieval-Augmented Generation (RAG) architecture, including the Dense Passage Retrieval (DPR) retriever, by jointly training the question encoder and passage encoder. The method addresses engineering challenges like re-encoding and re-indexing the knowledge base during training using parallel processing, achieving a 12% improvement in exact match on SQuAD (40.02% vs. 28.12%) compared to the original RAG.

ABSTRACT

In this paper, we illustrate how to fine-tune the entire Retrieval Augment Generation (RAG) architecture in an end-to-end manner. We highlighted the main engineering challenges that needed to be addressed to achieve this objective. We also compare how end-to-end RAG architecture outperforms the original RAG architecture for the task of question answering. We have open-sourced our implementation in the HuggingFace Transformers library.

Motivation & Objective

• To enable end-to-end training of the full RAG architecture, including the DPR retriever, which was previously not supported in the original Hugging Face implementation.
• To address the engineering challenges of updating the knowledge base embeddings and re-indexing during training, which are computationally expensive and time-consuming.
• To evaluate whether joint fine-tuning of both question and passage encoders improves question-answering performance compared to frozen passage encoders.
• To open-source the implementation for broader research use and reproducibility.

Proposed method

• Extended the Hugging Face Transformers library to support end-to-end training of the RAG model, including the DPR passage encoder.
• Used a dual BERT architecture: one for encoding questions and one for encoding passages, both fine-tuned jointly.
• Implemented parallel re-encoding and re-indexing of the knowledge base using Python multiprocessing and HuggingFace Datasets to avoid blocking the main training loop.
• Leveraged FAISS for efficient approximate nearest-neighbor search and updated the vector index after each re-encoding step.
• Applied gradient backpropagation through the document retrieval probability p(z|x) to update both the question and passage encoders.
• Used a periodic update strategy—re-encoding and re-indexing every N training steps—to balance performance and computational cost.

Experimental results

Research questions

• RQ1 Can end-to-end fine-tuning of the DPR retriever in RAG improve question-answering performance compared to frozen passage encoders?
• RQ2 What are the key engineering challenges in maintaining a dynamically updated knowledge base during RAG training?
• RQ3 How does the frequency of knowledge base re-encoding and re-indexing affect model performance and training efficiency?
• RQ4 Does joint training of both question and passage encoders improve domain adaptation and generalization?

Key findings

• End-to-end fine-tuning of the RAG architecture, including the DPR retriever, improved exact match on the SQuAD dataset from 28.12% (original RAG) to 40.02%.
• The performance gain demonstrates that joint optimization of both the question and passage encoders enhances retrieval and generation quality.
• Re-encoding and re-indexing the knowledge base during training is computationally expensive, but parallel processing makes it feasible within a reasonable training timeline.
• The study confirms that updating the passage encoder during training leads to better alignment between queries and relevant passages, especially in domain-specific or non-Wikipedia-like settings.
• The open-sourced implementation enables researchers to replicate and extend the method using Hugging Face Transformers and PyTorch Lightning.