Learning to Route Queries Across Knowledge Bases for Step-wise Retrieval-Augmented Reasoning

2025-05-28

AI Summary

Current multimodal retrieval-augmented generation (MRAG) methods rely on static retrieval pipelines and fail to leverage the dynamic reasoning capabilities of multimodal large language models (MLLMs) and their interactive potential with knowledge bases. To address this, we propose R1-Router, a dynamic stepwise retrieval-augmented reasoning framework. R1-Router introduces, for the first time, an adaptive knowledge-base routing mechanism conditioned on real-time reasoning states, enabling the model to autonomously decide *when* to retrieve, *which* knowledge source to consult, and *what* to query—while generating intermediate queries to orchestrate multi-source knowledge. We further design Step-wise GRPO, a reinforcement learning algorithm that optimizes rewards at the granularity of individual reasoning steps. Evaluated on multimodal open-domain question answering benchmarks, R1-Router outperforms strong baselines by over 7% in answer accuracy, while reducing redundant retrieval, improving reasoning fidelity, and enhancing computational efficiency.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Multimodal Retrieval-Augmented Generation (MRAG) has shown promise in mitigating hallucinations in Multimodal Large Language Models (MLLMs) by incorporating external knowledge during generation. Existing MRAG methods typically adopt a static retrieval pipeline that fetches relevant information from multiple Knowledge Bases (KBs), followed by a refinement step. However, these approaches overlook the reasoning and planning capabilities of MLLMs to dynamically determine how to interact with different KBs during the reasoning process. To address this limitation, we propose R1-Router, a novel MRAG framework that learns to decide when and where to retrieve knowledge based on the evolving reasoning state. Specifically, R1-Router can generate follow-up queries according to the current reasoning step, routing these intermediate queries to the most suitable KB, and integrating external knowledge into a coherent reasoning trajectory to answer the original query. Furthermore, we introduce Step-wise Group Relative Policy Optimization (Step-GRPO), a tailored reinforcement learning algorithm that assigns step-specific rewards to optimize the reasoning behavior of MLLMs. Experimental results on various open-domain QA benchmarks across multiple modalities demonstrate that R1-Router outperforms baseline models by over 7%. Further analysis shows that R1-Router can adaptively and effectively leverage diverse KBs, reducing unnecessary retrievals and improving both efficiency and accuracy.

Problem

• Dynamic routing of queries across knowledge bases for reasoning
• Optimizing retrieval-augmented generation with step-specific knowledge integration
• Reducing unnecessary retrievals while improving efficiency and accuracy

Innovation

• Dynamic query routing across knowledge bases
• Step-specific reinforcement learning optimization
• Adaptive retrieval reducing unnecessary queries