Source-linked AI summary
Self-Correcting RAG: Enhancing Faithfulness via MMKP Context Selection and NLI-Guided MCTS
Shijia Xu, Zhou Wu, Xiaolong Jia, Yu Wang, Kai Liu, April Xiaowen Dong
TL;DR
Self-Correcting RAG addresses low context utilization and hallucinations in complex reasoning. It combines MMKP-based context selection with NLI-guided MCTS generation, and reports stronger reasoning, retrieval, and faithfulness outcomes across six benchmarks. The approach increases inference latency and depends on limitations of its auxiliary NLI and redundancy assumptions.
Problem
Complex reasoning tasks suffer from low context utilization and hallucinations, while greedy generation cannot guarantee feasible or globally optimal decisions.
Method
The framework selects context with a constrained MMKP optimizer and explores generation trajectories with NLI-guided MCTS.
Results
Across six datasets, the framework improves reasoning accuracy and reduces hallucinations; MMKP raises Recall@5 from 49.6% to 71.8%.
Takeaways & Limitations
The framework jointly optimizes evidence selection and faithfulness-oriented reasoning for complex multi-hop and fact-checking tasks.
Takeaways & Limitations
MCTS increases computational overhead and inference latency, limiting applicability in ultra-low-latency real-time scenarios.
Abstract
from arXiv · showhide
Retrieval-augmented generation (RAG) substantially extends the knowledge boundary of large language models. However, it still faces two major challenges when handling complex reasoning tasks: low context utilization and frequent hallucinations. To address these issues, we propose Self-Correcting RAG, a unified framework that reformulates retrieval and generation as constrained optimization and path planning. On the input side, we move beyond traditional greedy retrieval and, for the first time, formalize context selection as a multi-dimensional multiple-choice knapsack problem (MMKP), thereby maximizing information density and removing redundancy under a strict token budget. On the output side, we introduce a natural language inference (NLI)-guided Monte Carlo Tree Search (MCTS) mechanism, which leverages test-time compute to dynamically explore reasoning trajectories and validate the faithfulness of generated answers. Experiments on six multi-hop question answering and fact-checking datasets demonstrate that our method significantly improves reasoning accuracy on complex queries while effectively reducing hallucinations, outperforming strong existing baselines.Our code is available at https://github.com/xjiacs/Self-Correcting-RAG .
1 Introduction
Self-Correcting RAG targets hallucinations and inefficient context use in complex reasoning by combining constrained context selection with search-based generation. It reports improved reasoning accuracy and reduced hallucinations across six datasets.
- Motivation: Complex reasoning exposes hallucinations and makes greedy token generation unable to guarantee feasible or globally optimal decisions.Errors in early decisions can propagate through the exponentially large search space.
- Framework: The framework unifies retrieval and generation by grounding local feasibility checks in evidence and organizing decisions into a backtrackable tree.This design synergizes RAG with MCTS for MMKP solving.
- Contributions: The MMKP Context Selector maximizes information density under token budgets, minimizes redundancy, and outperforms greedy ranking strategies.It models document selection as a constrained knapsack problem.
- Contributions: The NLI-Guided MCTS Generator uses test-time compute to explore reasoning paths, penalize contradictions, and improve faithfulness to retrieved context.Natural Language Inference serves as the reward model.
- Results: The framework achieves strong performance across six multi-hop question answering and fact verification datasets while reducing hallucinations and improving reasoning accuracy.The comparison is against strong agentic baselines.
2 Related Work
Related work spans prompting, agentic reasoning, retrieval robustness, LLM-based optimization, and non-linear search. The paper positions MMKP as an under-explored constrained optimization setting requiring verifiable evidence.
- Reasoning and Agents: Prompting and agentic frameworks improve intermediate reasoning and enable models to execute actions through APIs.Examples include CoT, Self-Consistency, ReAct, and Toolformer.
- Retrieval-Augmented Generation: RAG research addresses hallucinations through external grounding, corrective mechanisms, hierarchical indexing, and improved retrieval components.These developments include Self-RAG, CRAG, and instruction-tuned embeddings.
- LLMs for Optimization: LLMs have been used as iterative optimizers, evolving heuristics, and generators of mathematical constructions for hard problems.The cited examples include OPRO, ReEvo, and FunSearch.
- Combinatorial Optimization: MMKP remains under-explored because maintaining feasibility is difficult and a hallucinated constraint can invalidate a solution.The paper addresses this by aligning optimization steps with verifiable evidence.
- Search-Based Reasoning: Tree- and graph-based methods maintain multiple reasoning paths, while MCTS trades inference compute for solution quality and lookahead.These approaches include ToT, GoT, Reflexion, LATS, and reasoning-via-planning.
3 Methodology
The method separates RAG optimization into MMKP-based context selection and logic-guided MCTS reasoning. It combines semantic diversity, multidimensional constraints, retrieval augmentation, and NLI-based faithfulness rewards.
- Framework Overview: The framework bifurcates RAG optimization into pre-generation MMKP context optimization and inference-time logic-guided MCTS reasoning.These are the two main methodological phases.
- MMKP Context Selection: Conventional top-k selection is treated as suboptimal because it selects from a flat document list despite high inter-document redundancy.The proposed formulation instead selects a subset maximizing information density under multidimensional constraints.
- MMKP Context Selection: Semantic clustering groups highly similar documents so that selecting multiple members provides diminishing returns and candidates become mutually exclusive.Groups are defined using embedding cosine similarity above threshold τ.
- MMKP Optimization: Binary selection variables represent documents, while the MMKP objective maximizes utility subject to token and redundancy capacities.The optimization selects at most one representative from each semantic group.
- MCTS Generation: MCTS models generation as an MDP in which the LLM either generates a continuation or retrieves additional context based on uncertainty.States include the query, selected documents, and partial answer; augmentation updates the context.
- NLI Reward: The NLI reward evaluates entailment between answer sentences and evidence snippets, assigning a severe contradiction penalty to prune hallucinated branches.UCT and PUCT procedures are supplied separately, while the reward uses entailment, neutrality, and contradiction probabilities.
- Optimization Solver: Exact MMKP solution is computationally prohibitive, so the method uses a single-dimensional FPTAS and Pareto-pruned dynamic programming for multidimensional cases.For D = 2, the practical implementation uses dynamic programming with Pareto pruning.
4 Experiments
The experiments evaluate Self-Correcting RAG across six challenging datasets and compare it with standard, query-optimized, and advanced retrieval pipelines. Evaluation covers generation, retrieval, and faithfulness.
- Evaluation Setup: The evaluation spans six challenging datasets and three task types, ranging from single-hop retrieval to complex multi-step reasoning.Dataset statistics are summarized in Table 1.
- Evaluation Setup: Long-tail knowledge evaluation tests whether the system filters red-herring documents with lexical overlap but no answer-bearing evidence.This setting emphasizes context filtering.
- Baselines: Comparisons include Naive RAG, HyDE, RRR, and advanced selection or reranking methods spanning standard and agentic pipelines.Naive RAG uses BGE-Large embeddings and top-k truncation.
- Metrics: The multi-dimensional protocol reports EM and F1 for generation, Recall@5 for retrieval, and Citation Precision and Contradiction Rate for faithfulness.Faithfulness is verified with RoBERTa-large-mnli entailment judgments.
- Implementation: The implementation uses Qwen2.5-7B-Instruct, dense and sparse retrieval fused by RRF, Pareto-pruned MMKP dynamic programming, and RoBERTa-large-mnli for MCTS penalties.Experiments ran on 8× NVIDIA A100 (80GB) GPUs.
5 Results
Self-Correcting RAG achieves strong overall QA and retrieval results across diverse benchmarks, with especially notable gains on complex multi-hop reasoning, noisy contexts, and context selection.
- 37.1 average EM and 45.8 average F1 surpass the strongest baselines across six diverse datasets.
- Complex Reasoning Tasks: 4.5% absolute EM improvement over CRAG on MuSiQue demonstrates the strongest gains in complex multi-hop reasoning.
- Robustness to Noise: 35.3 EM on MultiHop-RAG exceeds RAG + MMR at 32.1 and Filco at 29.8 under noisy-context conditions.
- Retrieval Quality and Context Optimization: 72.0% average recall exceeds RAG+MMR's 63.3% by 8.7% across all datasets.
- Retrieval Quality and Context Optimization: The MMKP selector jointly optimizes relevance and diversity within a token budget, retaining evidence that greedy methods may discard.
- Handling Information Scarcity: 93.6% recall on HotpotQA exceeds CRAG's 91.5%, supporting selection of complementary document pairs for multi-hop inference.
6 Discussion
Ablations separate the contributions of MMKP retrieval and NLI-guided MCTS generation: MMKP improves retrieval and answer matching, while MCTS improves faithfulness. Sensitivity analyses identify useful budget and search settings, and failure analysis links context crowding to reasoning shortcuts.
- Impact of MMKP Context Selection: Recall@5 improves from 49.6% to 71.8% when Top-k retrieval is replaced by MMKP selection.
- Impact of MMKP Context Selection: MMKP raises EM to 34.5 and reduces redundancy within a 1500-token budget, but Attribution Precision remains 0.58.
- Impact of MMKP Context Selection: Improved retrieval alone does not prevent hallucinated ungrounded claims, motivating a separate faithfulness mechanism.
- Impact of NLI-Guided MCTS: NLI-guided MCTS increases Attribution Precision to 0.85 and reduces Contradiction Rate to 0.04 by penalizing contradictory reasoning paths.
- Sensitivity Analysis: A scaled redundancy budget of approximately 120 yields optimal recall, while stricter budgets can discard relevant lexically similar evidence.
- Sensitivity Analysis: A branching factor of k = 3 and maximum depth of 3 yield optimal reasoning accuracy, while weaker contradiction penalties increase errors.
- Failure Analysis: Context crowding can displace key evidence and trigger hallucinated connections based on parametric memory.
7 Conclusion
Self-Correcting RAG combines MMKP-based context selection with NLI-guided MCTS to improve retrieval density and answer faithfulness. Across six benchmarks, it outperforms strong agentic baselines, especially on complex multi-hop tasks, but incurs higher inference latency.
- MMKP context selection maximizes information density under strict token budgets, while NLI-guided MCTS prunes hallucinatory reasoning paths.
- Experiments across six benchmarks show stronger performance than strong agentic baselines, especially on complex multi-hop tasks.
- The planner's iterative search increases inference latency, motivating future work on sample efficiency and test-time search cost.
Limitations
Self-Correcting RAG improves faithfulness and reasoning but incurs substantial computational and latency costs. Its robustness also depends on NLI quality and assumptions about redundancy in retrieved groups.
- Computational overhead: Multiple forward passes and NLI verifications increase computational overhead and inference latency compared with single-pass RAG.Test-time search increases time-to-first-token and limits applicability in ultra-low-latency real-time scenarios.
- Computational overhead: Pareto-pruning keeps the MMKP selector polynomial-time, but does not remove the latency introduced by test-time search.
- Robustness: The reward mechanism relies on an off-the-shelf NLI model whose quality may be insufficient for subtle contradictions in specialized domains.The paper identifies law and medicine as examples requiring domain-specific fine-tuning.
- Selection assumptions: The MMKP selector assumes redundancy within retrieved groups is semantically uniform, which can discard complementary minority opinions in highly complex contexts.
Ethical Considerations
The framework is intended to reduce hallucinations and improve reliability, but its test-time computation increases energy use. Its outputs remain dependent on retrieved sources and may reproduce their biases or toxicity.
- Environmental impact: Monte Carlo Tree Search increases GPU utilization and energy consumption compared with standard decoding methods.The paper identifies planner sample efficiency as a future direction for reducing this computational footprint.
- Source dependence: RAG outputs are constrained by the retrieved corpora, so their quality and potential biases reflect the source documents.
- Source risks: If the knowledge base contains biased or toxic content, the system may reproduce those issues.The paper distinguishes faithfulness to retrieved context from absolute objective truth.
Appendices
The supplementary material covers dataset details, implementation, theory, experiments, sensitivity analysis, prompts, and qualitative case studies. These appendices provide supporting methodological, analytical, and empirical material.
- Supplementary contents: Appendix A provides dataset details for the evaluated tasks.
- Supplementary contents: Appendices B and F cover implementation details, hyperparameters, and LLM prompts.
- Supplementary contents: Appendix C presents theoretical proofs and analysis.
- Supplementary contents: Appendices D and E provide detailed experimental results and sensitivity analysis results.
- Supplementary contents: Appendix G contains a qualitative analysis case study.
A Dataset Details
The evaluation covers six datasets spanning simple QA, multi-hop QA, and multi-document QA, while the framework combines constrained MMKP retrieval with NLI-guided MCTS generation.
- Datasets: Six datasets cover Simple QA, Multi-Hop QA, and Multi-Document QA scenarios.The evaluation includes NQ, PopQA, HotpotQA, MuSiQue, 2WikiMultiHopQA, and MultiHop-RAG.
- Simple QA: NQ and PopQA evaluate open-domain and long-tail knowledge retrieval using 1,000 queries each.NQ uses 9,633 passages, while PopQA uses 8,676 passages.
- Multi-Hop QA: MuSiQue and 2WikiMultiHopQA test reasoning across multiple documents, typically requiring 2–4 hops.MuSiQue uses 1,000 queries and 11,656 passages; 2Wiki uses 1,000 queries and 6,119 passages.
- Multi-Document QA: HotpotQA uses a distractor setting to test irrelevant-information filtering with 1,000 queries and 9,811 passages.
- Multi-Document QA: MultiHop-RAG evaluates noise robustness with 2,556 queries over 609 passages using F1 score and retrieval accuracy.
- Framework: The MMKP selector assigns relevance-diversity value and two costs: token length and semantic redundancy penalty.Candidates are grouped by cosine similarity, and dynamic programming with Pareto pruning solves the resulting NP-hard problem.
- Framework: NLI-guided MCTS chooses between answering and evidence augmentation, rewarding entailment and strongly penalizing contradiction.The search uses UCT to balance exploration and exploitation, while the NLI reward assigns wc = −2.0 to contradictions.
- Theory: Theoretical results establish NP-hardness for RAG-MMKP, an FPTAS guarantee for the single-dimensional variant, and asymptotic MCTS consistency.The FPTAS returns V(S) ≥ (1−ϵ)OPT in time polynomial in n and 1/ϵ, while MCTS converges under an approximate truth oracle.
D Detailed Experimental Results
Experiments show that MMKP retrieval remains effective under tight token budgets and that performance gains grow on complex reasoning tasks; analyses also examine redundancy, MCTS settings, and failure correction.
- Token Budget Robustness: At 500 tokens, MMKP achieves 61.5% Recall@5 versus 38.0% for Top-K, a 23.5% margin, and 11.5% over RAG+MMR.The result supports MMKP’s resilience when context space is scarce.
- Reasoning on Hard Queries: MMKP improves over Naive by 19.0% on Simple QA, 73.4% on Multi-Hop QA, and 57.6% on Multi-Doc QA.On Multi-Hop QA, it also exceeds CRAG by 2.5 absolute points.
- Sensitivity Analysis: Recall@5 is most favorable near a redundancy budget of approximately 120, balancing diversity with retention of relevant evidence.
- MCTS Sensitivity: Increasing MCTS simulations beyond 24 yields diminishing F1 returns, while depth beyond 3 provides marginal faithfulness gains and complicates traces.The selected configuration is N = 24 and D = 3.
- Failure Analysis: A baseline fails when redundant chunks crowd out evidence for a second entity, causing a hallucinated founding date from parametric memory.
- Failure Analysis: MMKP removes semantically duplicate chunks to reserve budget for diverse evidence, while MCTS backtracks from low-entailment hallucinations to verified reasoning paths.In the case study, entailment rises from 0.12 on the hallucinated branch to 0.98 on the correct branch.
- Case Studies: A second case study examines attribute comparison under information noise, including a distractor-dense query about DeepMind and Google executives.