# Predict the Retrieval! Test time adaptation for Retrieval Augmented Generation

arXiv: https://arxiv.org/abs/2601.11443 · HTML: https://mlsanity.com/paper/2601.11443/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RAG robustness and adaptive retrieval, one of 15 directions in Retrieval-augmented generation.
- Builds on 29 of its 38 references that are in the map.
- First posted on arXiv 16 Jan 2026.

## Placed in

- Retrieval-augmented generation / [RAG robustness and adaptive retrieval](https://mlsanity.com/rag/rag-robustness-and-adaptive-retrieval/)

## Builds on (most cited first, 29 of 29)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection](https://mlsanity.com/paper/2310.11511/) (2023)
- [From Local to Global: A Graph RAG Approach to Query-Focused Summarization](https://mlsanity.com/paper/2404.16130/) (2024)
- [PubMedQA: A Dataset for Biomedical Research Question Answering](https://mlsanity.com/paper/1909.06146/) (2019)
- [Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering](https://mlsanity.com/paper/2007.01282/) (2020)
- [Do We Really Need to Access the Source Data? Source Hypothesis Transfer for Unsupervised Domain Adaptation](https://mlsanity.com/paper/2002.08546/) (2020)
- [Active Retrieval Augmented Generation](https://mlsanity.com/paper/2305.06983/) (2023)
- [Adaptive-RAG: Learning to Adapt Retrieval-Augmented Large Language Models through Question Complexity](https://mlsanity.com/paper/2403.14403/) (2024)
- [Improving robustness against common corruptions by covariate shift adaptation](https://mlsanity.com/paper/2006.16971/) (2020)
- [Making Retrieval-Augmented Language Models Robust to Irrelevant Context](https://mlsanity.com/paper/2310.01558/) (2023)
- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs](https://mlsanity.com/paper/2407.02485/) (2024)
- [RA-DIT: Retrieval-Augmented Dual Instruction Tuning](https://mlsanity.com/paper/2310.01352/) (2023)
- [Parameter-free Online Test-time Adaptation](https://mlsanity.com/paper/2201.05718/) (2022)
- [Efficient Test-Time Adaptation of Vision-Language Models](https://mlsanity.com/paper/2403.18293/) (2024)
- [The Effect of Natural Distribution Shift on Question Answering Models](https://mlsanity.com/paper/2004.14444/) (2020)
- [CRAG -- Comprehensive RAG Benchmark](https://mlsanity.com/paper/2406.04744/) (2024)
- [ChatQA: Surpassing GPT-4 on Conversational QA and RAG](https://mlsanity.com/paper/2401.10225/) (2024)
- [Enhancing Noise Robustness of Retrieval-Augmented Language Models with Adaptive Adversarial Training](https://mlsanity.com/paper/2405.20978/) (2024)
- [MedAdapter: Efficient Test-Time Adaptation of Large Language Models towards Medical Reasoning](https://mlsanity.com/paper/2405.03000/) (2024)
- [Challenges in Generalization in Open Domain Question Answering](https://mlsanity.com/paper/2109.01156/) (2021)
- [SimRAG: Self-Improving Retrieval-Augmented Generation for Adapting Large Language Models to Specialized Domains](https://mlsanity.com/paper/2410.17952/) (2024)
- [UncertaintyRAG: Span-Level Uncertainty Enhanced Long-Context Modeling for Retrieval-Augmented Generation](https://mlsanity.com/paper/2410.02719/) (2024)
- [Divide-Then-Align: Honest Alignment based on the Knowledge Boundary of RAG](https://mlsanity.com/paper/2505.20871/) (2025)
- [DIVE: Subgraph Disagreement for Graph Out-of-Distribution Generalization](https://mlsanity.com/paper/2408.04400/) (2024)
- [KBQA-R1: Reinforcing Large Language Models for Knowledge Base Question Answering](https://mlsanity.com/paper/2512.10999/) (2025)

## Built on by (most cited first, 2 of 2)

- [Why AI systems don't learn and what to do about it: Lessons on autonomous learning from cognitive science](https://mlsanity.com/paper/2603.15381/) (2026)
- [HopRefusalBench: Diagnosing Refusal Failures in Search-Augmented Agents for Multi-Hop Reasoning](https://mlsanity.com/paper/2608.01358/) (2026)

## Abstract (arXiv)

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.
