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

TL;DR

This work tackles the rigidity of fixed retrieval in multimodal retrieval-augmented reasoning by introducing R1-Router, which enables LLMs to decide when and where to retrieve knowledge across diverse KBs in a stepwise fashion. It couples this with Step-GRPO, a reinforcement learning objective that provides step-specific rewards to optimize both intermediate querying and final answers within a reasoning trajectory. Empirical results across Text, Visual, and Table QA show about a 7% average improvement over strong baselines, with notable gains from adaptive KB routing and more efficient use of diverse knowledge sources. The findings demonstrate the potential for generalizable, reasoning-driven RAG systems that dynamically coordinate information from heterogeneous knowledge bases to tackle complex, multimodal questions.

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.

Figures (22)

• Figure 1: Comparison of Different RAG Architectures.
• Figure 2: Performance of Step-GRPO Training Method.
• Figure 3: Query Routing Perfromance on VQA During the Reasoning Process of R1-Router. Three training strategeies are compared, including Prompt, SFT, and Step-GRPO.
• Figure 4: An Example from the Visual QA Task Demonstrating the Performance of R1-Router.
• Figure 5: Case Studies of Different Retrievers.