Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Retrieval is Not Enough: Enhancing RAG Reasoning through Test-Time Critique and Optimization

Jiaqi Wei, Hao Zhou, Xiang Zhang, Di Zhang, Zijie Qiu, Wei Wei, Jinzhe Li, Wanli Ouyang, Siqi Sun

TL;DR

This work reframes Retrieval-Augmented Generation as Retrieval-Augmented Reasoning and tackles reasoning misalignment between a model's internal reasoning and retrieved evidence. It introduces Critique-Driven Alignment (CDA) and a retrieval-augmented Critic Language Model (CLM) trained via Contrastive Critique Synthesis to iteratively align reasoning with evidence. AlignRAG and its autonomous variant AlignRAG-auto dynamically refine responses, achieving state-of-the-art results across seven QA benchmarks and three model families, including strong robustness to noisy retrieval. The approach is plug-and-play, scalable at test time, and improves out-of-domain generalization while reducing computation through dynamic stopping in AlignRAG-auto.

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.

Retrieval is Not Enough: Enhancing RAG Reasoning through Test-Time Critique and Optimization

TL;DR

This work reframes Retrieval-Augmented Generation as Retrieval-Augmented Reasoning and tackles reasoning misalignment between a model's internal reasoning and retrieved evidence. It introduces Critique-Driven Alignment (CDA) and a retrieval-augmented Critic Language Model (CLM) trained via Contrastive Critique Synthesis to iteratively align reasoning with evidence. AlignRAG and its autonomous variant AlignRAG-auto dynamically refine responses, achieving state-of-the-art results across seven QA benchmarks and three model families, including strong robustness to noisy retrieval. The approach is plug-and-play, scalable at test time, and improves out-of-domain generalization while reducing computation through dynamic stopping in AlignRAG-auto.

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.

Paper Structure

This paper contains 46 sections, 12 equations, 9 figures, 18 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Reasoning Misalignment in RAG
  3. Critique-Driven Alignment for Retrieval-Augmented Reasoning
  4. Problem Setting
  5. Critic Training
  6. Training Corpus Construction
  7. Contrastive Critique Synthesis
  8. Critic LLM Training
  9. Critique-Driven Alignment Systems
  10. AlignRAG-auto: Dynamic and Domain-General Alignment
  11. Training for Dynamic Refinement
  12. Dynamic Inference-Time Stopping
  13. Experiments
  14. Experiment Setup
  15. Main Result
  16. ...and 31 more sections

Figures (9)

  • Figure 1: Overview of our AlignRAG Framework.
  • Figure 2: Performance Comparison: AlignRAG-fixed (1 iter.) vs. AlignRAG-auto
  • Figure 3: Drop in average OOD performance compared to average In-Domain performance.
  • Figure 4: Performance comparison of AlignRAG and Self-Refine across five refinement iterations.
  • Figure 5: AlignRAG performance under Noisy Retrieval (a) and Informative Retrieval (b) scenarios.
  • ...and 4 more figures