Retrieval is Not Enough: Enhancing RAG through Test-Time Critique and Optimization
Jiaqi Wei, Hao Zhou, Xiang Zhang, Di Zhang, Zijie Qiu, Noah Wei, Jinzhe Li, Wanli Ouyang, Siqi Sun
Abstract
Retrieval-augmented generation (RAG) has become a widely adopted paradigm for enabling knowledge-grounded large language models (LLMs). However, standard RAG pipelines often fail to ensure that model reasoning remains consistent with the evidence retrieved, leading to factual inconsistencies or unsupported conclusions. In this work, we reinterpret RAG as Retrieval-Augmented Reasoning and identify a central but underexplored problem: Reasoning Misalignment-the divergence between an LLM's internal reasoning trajectory and the evidential constraints provided by retrieval. To address this issue, we propose ALIGNRAG, a novel iterative framework grounded in Critique-Driven Alignment (CDA). We further introduce ALIGNRAG-AUTO, an autonomous variant that dynamically terminates refinement, removing the need to pre-specify the number of critique iterations. At the heart of ALIGNRAG lies a contrastive critique synthesis mechanism that generates retrieval-sensitive critiques while mitigating self-bias. This mechanism trains a dedicated retrieval-augmented Critic Language Model (CLM) using labeled critiques that distinguish between evidence-aligned and misaligned reasoning. Empirical evaluations show that our approach significantly improves reasoning fidelity. Our 8B-parameter CLM improves performance over the Self-Refine baseline by 12.1% on out-of-domain tasks and outperforms a standard 72B-parameter CLM by 2.2%. Furthermore, ALIGNRAG-AUTO achieves this state-of-the-art performance while dynamically determining the optimal number of refinement steps, enhancing efficiency and usability. ALIGNRAG remains compatible with existing RAG architectures as a plug-and-play module and demonstrates strong robustness under both informative and noisy retrieval scenarios. Overall, ALIGNRAG offers a principled solution for aligning model reasoning with retrieved evidence, substantially improving the factual reliability and robustness of RAG systems. Our source code is provided at Github.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext d3935f15-d769-4954-b326-26bc3f8296eaCited by top-tier papers2
- When to Think, When to Speak: Learning Disclosure Policies for LLM ReasoningJiaqi Wei, Xuehang Guo, Pengfei Yu, Xiang Zhang et al.ICML 2026 · 2 citations
- Optimizing Agentic Reasoning with Retrieval via Synthetic Semantic Information Gain RewardSenkang Hu, Yong Dai, Yuzhi Zhao, Yihang Tao et al.ICML 2026
Builds on34
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Retrieval-Augmented Generation for Knowledge-Intensive NLP TasksPatrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni et al.NeurIPS 2020 · 19,162 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- Toolformer: Language Models Can Teach Themselves to Use ToolsTimo Schick, Jane Dwivedi-Yu, Roberto Dessì, Roberta Raileanu et al.NeurIPS 2023 · 5,989 citations
- Reflexion: language agents with verbal reinforcement learningNoah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan et al.NeurIPS 2023 · 5,828 citations
Related papers
- InstructRAG: Instructing Retrieval-Augmented Generation via Self-Synthesized RationalesZhepei Wei, Wei-Lin Chen, Yu MengICLR 2025
- Understand What LLM Needs: Dual Preference Alignment for Retrieval-Augmented GenerationGuanting Dong, Yutao Zhu, Chenghao Zhang, Zechen Wang et al.WWW 2025 · 44 citations
- FaithfulRAG: Fact-Level Conflict Modeling for Context-Faithful Retrieval-Augmented GenerationQinggang Zhang, Zhishang Xiang, Yilin Xiao, Le Wang et al.ACL 2025 · 18 citations
- RAG+: Enhancing Retrieval-Augmented Generation with Application-Aware ReasoningYu Wang, Shiwan Zhao, Zhihu Wang, Ming Fan et al.EMNLP 2025 · 3 citations
- Iterative Multi-Granular RAG with Contextual Hierarchical GraphYanli Hu, Teng Liu, Zhuangyi Zhou, Weixin Zeng et al.AAAI 2026
