MIRAGE: Scaling Test-Time Inference with Parallel Graph-Retrieval-Augmented Reasoning Chains

Abstract: Large reasoning models (LRMs) have shown significant progress in test-time scaling through chain-of-thought prompting. Current approaches like search-o1 integrate retrieval augmented generation (RAG) into multi-step reasoning processes but rely on a single, linear reasoning chain while incorporating unstructured textual information in a flat, context-agnostic manner. As a result, these approaches can lead to error accumulation throughout the reasoning chain, which significantly limits its effectiveness in medical question-answering (QA) tasks where both accuracy and traceability are critical requirements. To address these challenges, we propose MIRAGE (Multi-chain Inference with Retrieval-Augmented Graph Exploration), a novel test-time scalable reasoning framework that performs dynamic multi-chain inference over structured medical knowledge graphs. Specifically, MIRAGE 1) decomposes complex queries into entity-grounded sub-questions, 2) executes parallel inference chains, 3) retrieves evidence adaptively via neighbor expansion and multi-hop traversal, and 4) integrates answers using cross-chain verification to resolve contradictions. Experiments on three medical QA benchmarks (GenMedGPT-5k, CMCQA, and ExplainCPE) show that MIRAGE consistently outperforms GPT-4o, Tree-of-Thought variants, and other retrieval-augmented baselines in both automatic and human evaluations. Additionally, MIRAGE improves interpretability by generating explicit reasoning chains that trace each factual claim to concrete chains within the knowledge graph, making it well-suited for complex medical reasoning scenarios. The code will be available for further research.

• The paper introduces the MIRAGE framework that decomposes complex medical queries into entity-grounded sub-questions, using parallel reasoning chains for robust and accurate inference.
• It employs advanced graph-retrieval methods, including neighbor expansion and multi-hop traversal, to dynamically access and aggregate evidence from structured knowledge graphs.
• It implements a cross-chain verification mechanism that resolves inconsistencies, significantly outperforming baseline models in rigorous evaluations on medical QA datasets.

MIRAGE: Scaling Test-Time Inference with Parallel Graph-Retrieval-Augmented Reasoning Chains

Introduction

The paper presents MIRAGE, an innovative framework aimed at enhancing test-time inference by leveraging parallel reasoning chains complemented by graph-retrieval-augmented methods. The framework is meticulously designed to address shortcomings in current reasoning approaches which typically rely on linear reasoning models susceptible to error propagation. Specifically, MIRAGE resolves these deficiencies by integrating structured knowledge graphs, thus enabling efficient and accurate dynamic multi-hop retrieval processes for more reliable inference results in medical QA tasks.

Methodology

Framework Overview

MIRAGE introduces a comprehensive system of interconnected modules designed to elevate inference capabilities. This involves decomposing medical queries into structured, entity-grounded sub-questions and then conducting parallel inference chains for each. This method is instrumental in dynamically retrieving pertinent evidence through adaptive strategies, subsequently synthesizing multiple reasoning chains for coherent and correct responses (Figure 1).

Figure 1: Overview of the proposed MIRAGE framework. Given a clinical query, the system decomposes it into sub-questions, each initiating a reasoning chain. For each chain, the system iteratively retrieves knowledge graph evidence using either Anchor mode or Bridge mode. Retrieved results are coordinated and aggregated to generate the final answer.

Graph-Retrieval-Augmented Reasoning

MIRAGE's critical advancement lies in its strategy for graph-based retrieval. It uses knowledge graphs to expand retrieval capabilities beyond flat text integration, fostering nuanced understanding and solving complex, multi-dimensional inquiries. Through neighbor expansion and multi-hop traversal processes, the system identifies structurally relevant entities and relations, organizing them into coherent reasoning chains that enhance interpretability and trace evidence back to explicit, grounded knowledge paths.

Cross-Chain Verification

An essential feature of MIRAGE is its cross-chain verification mechanism, which compiles the outputs of parallel reasoning chains, examines inter-chain agreements, and resolves disparities. This verification ensures the final product's consistency and reliability, effectively navigating the inherent variability within medical question-answering contexts. The decomposition and assembly strategy minimizes the noise prevalent in oversimplified question handling, fostering a more precise, domain-specific application.

Evaluation

Benchmark Performance

In rigorous evaluations on challenging medical QA datasets, including GenMedGPT-5k, CMCQA, and ExplainCPE, MIRAGE achieved substantial performance surpassing well-recognized baselines such as GPT-4o and Search-o1. The framework's success is evident in both automatic and human evaluations, where it demonstrated significant improvements in accuracy and inference reliability for complex medical scenarios (Figure 2).

Figure 2: Human evaluation results on GenMedGPT-5k.

Parameter Influence

The experimental analysis underscored the importance of decomposition and retrieval thresholds on MIRAGE's efficiency. Adjustments to these parameters markedly influenced the model's ability to refine its response accuracy, showcasing robust adaptability and resilience against the variable complexities of medical contexts (Figure 3).

Figure 3: Effect of the decomposition threshold $N_q$ (a) and retrieval threshold $N_r$ (b) on GPT-4o Ranking and accuracy.

Implications and Future Directions

MIRAGE embodies a formidable advancement in reasoning methodology, particularly relevant in domains that prioritize accuracy and traceability, such as medicine. The framework's use of structured knowledge graphs marks a significant stride towards improved scalability and interpretability in AI reasoning tasks.

Future work can explore the expansion of MIRAGE's graph-retrieval frameworks to other domains, further fine-tuning the model's parameter calibration to optimize its application scope. Enhanced integration with diverse LLMs and deeper investigation into error-proofing mechanisms might yield even richer domain capabilities and inference efficacy.

Conclusion

MIRAGE successfully introduces a structured and parallelized approach to enhancing test-time inference using knowledge graphs, marking notable advancements over existing models. By strategically aligning reasoning chains with graph-retrieval processes, MIRAGE positions itself as a powerful tool for complex problem-solving and in-depth domain analysis in AI. The methodology's robust performance in medical QA scenarios demonstrates its potential to significantly influence AI's capacity for nuanced and scalable reasoning across diverse applications.