Source-linked AI summary
Improving the Domain Adaptation of Retrieval Augmented Generation (RAG) Models for Open Domain Question Answering
Shamane Siriwardhana, Rivindu Weerasekera, Elliott Wen, Tharindu Kaluarachchi, Rajib Rana, Suranga Nanayakkara
TL;DR
RAG was not well optimized for specialized-domain ODQA because its passage encodings and knowledge base remained fixed after training. The paper proposes RAG-end2end, jointly adapting retrieval and generation while updating the knowledge base, and reports significant improvements over original RAG across COVID-19, News, and Conversations.
Problem
RAG had mainly been trained with Wikipedia-based knowledge bases, so its feasibility for specialized-domain ODQA such as research papers and news was not well understood.
Method
RAG-end2end jointly fine-tunes the retriever and generator while asynchronously updating passage encodings and the external knowledge-base index, with an auxiliary reconstruction signal.
Results
RAG-end2end significantly improves performance over original RAG across COVID-19, News, and Conversations, while the auxiliary signal further improves retrieval and answer generation accuracy.
Takeaways & Limitations
Retriever adaptation is critical for domain adaptation, and joint RAG-end2end training can outperform independently fine-tuned DPR for end-to-end QA.
Takeaways & Limitations
RAG-end2end fine-tuning can be expensive for large knowledge bases because passages must be re-encoded and re-indexed during training.
Abstract
from arXiv · showhide
Retrieval Augment Generation (RAG) is a recent advancement in Open-Domain Question Answering (ODQA). RAG has only been trained and explored with a Wikipedia-based external knowledge base and is not optimized for use in other specialized domains such as healthcare and news. In this paper, we evaluate the impact of joint training of the retriever and generator components of RAG for the task of domain adaptation in ODQA. We propose \textit{RAG-end2end}, an extension to RAG, that can adapt to a domain-specific knowledge base by updating all components of the external knowledge base during training. In addition, we introduce an auxiliary training signal to inject more domain-specific knowledge. This auxiliary signal forces \textit{RAG-end2end} to reconstruct a given sentence by accessing the relevant information from the external knowledge base. Our novel contribution is unlike RAG, RAG-end2end does joint training of the retriever and generator for the end QA task and domain adaptation. We evaluate our approach with datasets from three domains: COVID-19, News, and Conversations, and achieve significant performance improvements compared to the original RAG model. Our work has been open-sourced through the Huggingface Transformers library, attesting to our work's credibility and technical consistency.
1 Introduction
The paper targets domain adaptation for retrieval-augmented ODQA beyond Wikipedia, proposing RAG-end2end to jointly adapt retrieval, generation, and the external knowledge base. Experiments across COVID-19, Conversations, and News report improved performance over original RAG.
- RAG has mainly been explored with Wikipedia-based knowledge bases, leaving adaptation to specialized domains such as research papers and news insufficiently understood.
- RAG-end2end updates the external knowledge base, DPR retriever, and BART generator during training for domain adaptation.The method updates knowledge encodings asynchronously and jointly trains retrieval and generation for end-to-end QA.
- An auxiliary signal trains the model to generate concise, factual statements about documents from self-retrieved domain-specific passages.This signal is intended to inject additional domain-specific knowledge during training.
- RAG-end2end is evaluated on COVID-19 research, Conversations, and News datasets.
- The major finding is that adapting the retriever is critical, while jointly fine-tuning it within RAG-end2end and using the auxiliary signal improves overall accuracy.Updating only the question encoder without updating knowledge-base encodings could degrade performance.
- The implementation is open-sourced through the HuggingFace Transformers library.
2 Background and Related Work
The background distinguishes conventional two-stage ODQA systems from retrieval-augmented architectures such as RAG and REALM. The paper extends RAG because prior work largely emphasized general Wikipedia-based QA rather than domain adaptation.
- Conventional ODQA separates passage retrieval from answer generation, traditionally using sparse retrieval methods such as TF-IDF and BM25.Dense retrievers instead model semantic similarity between questions and passages.
- RAG combines retrieval and answer generation differentiably, using a BART generator and dense Wikipedia passage representations indexed with FAISS.Its loss can fine-tune the generator and question encoder together while passage encodings remain part of the external memory.
- REALM introduced an end-to-end trainable retriever during masked-language pre-training but keeps it frozen during downstream QA fine-tuning.
- Compared with REALM, RAG uses pretrained models, avoids a heavy pre-training stage, is less computationally expensive, and has open-source code.
- Prior retrieval-augmented work explored end-to-end retriever training or general QA, but did not focus on domain adaptation for retrieval-augmented QA models.
3 Model Architecture and Training Procedure
RAG-end2end extends RAG by jointly training its retriever and generator while asynchronously updating knowledge-base embeddings and indexes. It also adds statement reconstruction alongside QA training to inject domain-specific information.
- Training Signals: The training architecture combines QA pairs with reconstruction signals and uses a control token to distinguish answer generation from statement reconstruction.The BART reader receives retrieved passages for both tasks, with reconstruction inputs prefixed by a special <p> token.
- RAG Retriever and Generator: The retriever uses DPR’s question and passage encoders, while the generator uses a pre-trained BART sequence-to-sequence model.DPR represents questions and passages with CLS-token embeddings and scores their similarity using a dot product.
- RAG Retriever and Generator: The original RAG architecture fixes passage encodings and the passage encoder during training, motivating RAG-end2end’s broader parameter updates.This fixed-index design is effective for Wikipedia-like datasets because DPR was pretrained on Wikipedia-based question-answering data.
- End-to-End Retriever Training: RAG-end2end fine-tunes both DPR encoders and updates the external knowledge-base index during training.Unlike the original RAG setup, passage and question encoders are both optimized, and the index is refreshed with updated passage representations.
- End-to-End Retriever Training: Asynchronous re-encoding and re-indexing prevent expensive knowledge-base updates from stalling the main training loop.Separate GPU re-encoding and FAISS re-indexing processes run independently of gradient updates, while preserving the required process order.
4 Experiments & Results
Experiments across COVID-19, News, and Conversation domains show that jointly adapting RAG’s retriever and external knowledge base improves domain-specific ODQA, with further gains from statement reconstruction.
- End-to-end retriever training: RAG-end2end-QA significantly outperforms RAG-original-QA across EM, F1, Top-5, and Top-20 metrics in all three domains.EM improvements range from 1.13 points in News to 12.16 points in Conversation.
- End-to-end retriever training: Around 25 points of retrieval improvement occur in Conversation, while the other domains improve by 4.7 to 6.6 points on Top-5 and Top-20 scores.
- Statement-reconstruction auxiliary task: Adding statement reconstruction makes RAG-end2end-QA + R outperform RAG-end2end-QA across all three domains.EM increases range from 1.7 points in Conversation to 8.39 points in News, while Top-20 retrieval accuracy rises by 3.2 to 8 points.
- Overall comparison: Compared with RAG-original-QA, RAG-end2end-QA + R produces a 13-point EM improvement and 27-point Top-5 improvement in Conversation.The highlighted Conversation Top-20 improvement is 16 points.
- Statement-reconstruction auxiliary task: Reconstructed statements generally capture retrieved-document context and resemble the input, but they are not always factually 100% correct.
- Retriever adaptation trade-offs: Updating the entire DPR retriever improves domain adaptation but can be computationally costly because external passages must be re-encoded and re-indexed.
5 Discussion
The discussion identifies retriever adaptation as central to RAG domain adaptation, while noting that end-to-end training improves results but can be computationally expensive for large knowledge bases.
- RAG-end2end improves performance across all evaluated datasets compared with original RAG finetuning.
- Updating knowledge-base embeddings and indexing enables RAG-end2end to improve domain adaptation beyond independently finetuning DPR.
- End-to-end retriever adaptation can be expensive when the external knowledge base contains millions of passages.Re-encoding and FAISS re-indexing require substantial GPU and CPU resources.
- The authors contrast their findings with RETRO by arguing that domain-adapted retrievers benefit domain-specific retrieval-augmented models.
- RAG-end2end also outperforms original RAG on an adapted Wikipedia-based SQuAD ODQA dataset.
6 Conclusion and Future Work
The paper concludes that RAG-end2end jointly trains retrieval and generation while updating the knowledge base, improving performance across three specialized domains. It also identifies retriever adaptation and auxiliary training signals as important directions for future work.
- RAG-end2end jointly trains the retriever and generator for end-to-end question answering and domain adaptation.
- RAG-end2end significantly improves performance over original RAG across COVID-19, News, and Conversations datasets.
- The reconstruction auxiliary signal further improves retriever and final answer-generation accuracy.
- Future Work: Future work includes evaluating RAG-end2end on fact checking, summarisation, and conversational response generation.
- Future Work: The authors propose exploring retriever improvements and additional auxiliary signals for RAG models.
A Appendix
The appendix illustrates conversational QA retrieval and answer prediction, including examples where retrieved passages contain the information needed to answer questions.
- One conversational example asks where a Kiwi girl commutes from and gives “what seems like N Africa to get to work” as the answer.
- Figure 2 presents predicted answers and retrieved passages for conversational-domain questions.