Predict the Retrieval! Test time adaptation for Retrieval Augmented Generation

TL;DR

This work tackles domain shift in Retrieval-Augmented Generation by introducing TTARAG, a test-time adaptation method that updates model parameters during inference via a self-supervised objective grounded in retrieved passages. The method splits retrieved texts into prefix-suffix pairs and trains the model to predict suffix content from the prefix and the query, enabling real-time domain specialization without labeled data, formalized as $\mathcal{L}_{adapt} = -\sum_{i=1}^k \log P(p_i^{suffix}|p_i^{prefix}, q; \theta)$. Evaluations across CRAG and medical datasets show TTARAG consistently improves over naive RAG and several baselines, with pronounced gains in biomedical domains and competitive efficiency. The results suggest test-time adaptation is a practical route to enhance RAG performance in specialized domains, supported by segmentation benefits and robust hyperparameter behavior, with the accompanying code available at the project repository.

Abstract

Retrieval-Augmented Generation (RAG) has emerged as a powerful approach for enhancing large language models' question-answering capabilities through the integration of external knowledge. However, when adapting RAG systems to specialized domains, challenges arise from distribution shifts, resulting in suboptimal generalization performance. In this work, we propose TTARAG, a test-time adaptation method that dynamically updates the language model's parameters during inference to improve RAG system performance in specialized domains. Our method introduces a simple yet effective approach where the model learns to predict retrieved content, enabling automatic parameter adjustment to the target domain. Through extensive experiments across six specialized domains, we demonstrate that TTARAG achieves substantial performance improvements over baseline RAG systems. Code available at https://github.com/sunxin000/TTARAG.

Figures (3)

• Figure 1: Comparison between Standard RAG and TTARAG systems.
• Figure 2: Accuracy vs. Learning Rate
• Figure 3: Accuracy vs. Number of Adaptation Pairs