🤖 AI

Reason Before You Retrieve: Agentic Planning for Multi-modal RAG

The paper proposes MM-R2, a multimodal agentic framework that improves retrieval-augmented generation by explicitly reasoning about retrieval targets and structuring the search space via a KnowledgeMap, supported by a new trajectory dataset and two-stage post-training strategy to achieve superior accuracy and interpretability on benchmark datasets.

Original authors: Tianyu Yang, Shir Simon, Zhenzhen Li, Minhao Cheng, Xiangliang Zhang

Published 2026-07-28

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianyu Yang, Shir Simon, Zhenzhen Li, Minhao Cheng, Xiangliang Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a mystery, but instead of a single clue, you have a giant, messy pile of books, photos, and maps all thrown into one big box. This is the world of "Multimodal Retrieval-Augmented Generation" (mRAG). It's a fancy way of saying we want computers to answer questions about pictures by looking up facts in a huge library. Usually, when a computer tries to do this, it acts a bit like a frantic shopper in a massive supermarket: it grabs a question, shouts it into the whole store, and hopes the first few items it finds on the shelves are the right ones. The problem is, the computer often grabs the wrong things because it doesn't really understand what it's looking for or where in the store that item might be hidden. It's like asking for "the red car" and getting a picture of a red apple because the computer just saw the word "red" and grabbed the first red thing it could find.

This is where the new paper comes in. The researchers realized that before you even start searching, you need to stop and think. They call their new system MM-R2, which stands for "Reason Before You Retrieve." Instead of blindly shouting into the void, MM-R2 acts like a smart detective. First, it looks at the picture and the question together to figure out exactly what the "clue" is (the specific object in the photo) and what kind of information is needed. Then, instead of searching the whole messy box, it consults a special, organized map of the library called a "KnowledgeMap." It picks the right aisle, grabs the right book, reads it, and decides if it needs to look elsewhere. The paper suggests that by doing this "thinking first" step, the computer gets much better at finding the right answers and can even show you its work, proving exactly how it solved the mystery.

The Detective's New Strategy

The core idea of this paper is that computers are currently too eager to jump into action. When you ask a computer, "What kind of bird is in this picture?" while showing it a photo, current systems often just mash the words and the image together and start searching everything at once. The authors argue this is a mistake. They found that these systems get confused because they don't know which part of the image matters or what specific fact they need. It's like trying to find a specific recipe in a library by searching for the word "cake" across every single book in the building, rather than going straight to the "Baking" section.

To fix this, the team built MM-R2, a system that forces the computer to pause and plan before it searches. They break the process down into three fun steps:

The "What" and "Where" Plan: Before searching, the system creates a "retrieval state." Think of this as writing a perfect shopping list. It identifies the Information Need (what do I need to know?), the Grounded Referent (which specific thing in the photo am I talking about?), and the Constraints (do I need a date? a location? a scientific name?). This stops the computer from getting distracted by irrelevant details.
The Organized Map (KnowledgeMap): The researchers realized that searching a flat, messy list of facts is inefficient. So, they organized their library into a KnowledgeMap. Imagine the library isn't just a pile of books, but a set of neatly labeled bins: one bin for "Birds," one for "Cars," one for "History." When the computer has its shopping list, it doesn't search the whole library; it first decides which bin to open. This is called "routing."
The Detective Loop: Once the computer picks the right bin (the retrieval unit), it searches inside that bin for the specific answer. If it finds the answer, great! If not, it can decide to pick a different bin or ask a more specific question. It keeps doing this until it has enough evidence to solve the case.

How They Taught the Computer to Think

You might wonder, "How does a computer learn to plan like a detective?" The authors didn't just tell the computer to "be smart." Instead, they created a massive training dataset called MM-R2-Traj. This dataset contains 900,000 examples of "thought processes."

Imagine a teacher showing a student how to solve a puzzle. The teacher doesn't just give the answer; they show the steps: "First, look at the corner pieces. Then, sort by color. Then, try to fit these two together." The researchers used a "teacher" AI to generate these step-by-step stories for the student AI. They taught the student to:

Analyze the picture and question to make a plan.
Choose the right "bin" (KnowledgeMap unit).
Ask a specific question inside that bin.
Summarize what it found.
Decide if it's done or needs to look again.

They trained the computer in two stages. First, they used Supervised Fine-Tuning (SFT), which is like having the student copy the teacher's notes to learn the basic rules. Then, they used Group Relative Policy Optimization (GRPO), which is like a game where the student tries different strategies, and if it gets the right answer, it gets a high score. If it gets the answer wrong, it learns to try a different path next time. This helped the computer get really good at making the right choices.

The Results: Smarter and Clearer

When they tested MM-R2 on two big challenges—InfoSeek (a dataset with 1.3 million image-question pairs) and Encyclopedic VQA (a test of encyclopedia-style knowledge)—the results were impressive.

On the InfoSeek test, MM-R2 got the right answer 54.3% of the time overall. This is a big jump compared to other smart systems, which usually hovered around 46-47%. On the specific "Unseen-Question" part of the test (where the computer has to handle questions it has never seen before), MM-R2 scored 56.0%, beating the previous best by a solid margin. Even on the Encyclopedic VQA test, it reached 39.4%, showing it works well even when the questions are very specific and hard.

But the paper suggests something even more important than just getting the right score: transparency. Because MM-R2 plans its steps, we can actually see how it solved the problem.

Better Routing: The system correctly picked the right "bin" (KnowledgeMap unit) 74% of the time, compared to only 23% for the older systems. It stopped wasting time looking in the wrong places.
Clearer Explanations: When the system explained its answer, it was more concise and focused. The "Fidelity" score (which checks if the answer really came from the evidence found) was 0.90, meaning the answer was almost entirely based on the facts it found, not just guessing.

What This Means

The paper suggests that the old way of just "searching everything" is hitting a wall. By forcing the computer to reason before it retrieves, we can build systems that are not only smarter but also more trustworthy. They don't just guess; they follow a logical path, check their work, and can show you exactly how they found the answer. It's the difference between a chaotic scavenger hunt and a well-planned detective mission. While the researchers note that this is a specific approach to multimodal AI and doesn't solve every problem in the world, the results suggest that taking a moment to think before acting is a powerful strategy for the future of smart computers.

Technical Summary: Reason Before You Retrieve: Agentic Planning for Multi-modal RAG

Problem Statement

Multimodal Retrieval-Augmented Generation (mRAG) aims to answer image-text queries by leveraging external knowledge. However, existing systems typically employ a "retrieve-then-generate" pipeline that operates over a flat, weakly structured evidence space. The authors identify two critical limitations in this prevailing design:

Retrieval-Intent Ambiguity: In multimodal settings, the text question defines the information need, while the image grounds the specific visual referent. Existing systems often fail to explicitly ground the query to the correct visual instance, leading to retrieval drift toward textual priors or generic caption-like matching rather than evidence specific to the image-question pair.
Weakly Structured Search Space: Most systems embed diverse evidence into a single representation space and rely on global similarity functions. This forces semantically distinct evidence to compete in a single ranking step, often favoring generic content over precise, domain-specific evidence.

The paper argues that multimodal retrieval should be reframed as "reason before you retrieve," drawing an analogy to Schema Theory in cognitive science, where structured knowledge frameworks guide attention and memory access.

Methodology: The MM-R2 Framework

The authors propose MM-R2 (MultiModal Reason-to-Retrieve), an agentic retrieval framework that explicitly models what to retrieve and where to search before executing the retrieval. The framework operates through the following components:

1. Intent-Referent Binding Module

Before retrieval, the system constructs a structured retrieval state from the image-question pair $(I, Q)$ . This module, denoted as $\Phi$ , outputs a triplet $z = (i, r, c)$ :

$i$ (Information Need): The specific knowledge required (e.g., taxonomy, location, function).
$r$ (Grounded Referent): The specific object, entity, or region in the image the query refers to.
$c$ (Constraints): Task-specific conditions such as answer type, granularity, or temporal scope.

This step resolves retrieval intent prior to search, reducing query drift.

2. Structured KnowledgeMap

Instead of a flat corpus, the external knowledge base is pre-structured into a KnowledgeMap ( $K$ ), consisting of semantically coherent retrieval units ( $u_1, \dots, u_N$ ).

Construction: Passages are encoded, clustered using HDBSCAN (to handle varying densities without pre-fixing cluster numbers), and assigned lightweight metadata (titles and descriptions) via an LLM.
Function: This acts as a routing scaffold. The agent first selects relevant units ( $a_t$ ) and then executes grounded queries ( $q_t$ ) within those units, rather than searching the entire corpus globally.

3. Agentic Planning Loop

The agent $A$ (a multimodal language model) performs a sequential decision-making process:

Routing: Selects one or more retrieval units from the KnowledgeMap based on the current state (intent $z$ and history $h_t$ ).
Query Generation: Formulates a grounded query specific to the selected unit.
Retrieval & Update: Retrieves evidence $e_t$ and updates the history $h_{t+1}$ .
Termination Decision: The agent iteratively decides whether to continue retrieving or to generate the final answer $\hat{y}$ based on the accumulated evidence.

4. Training Strategy

To train the agent to perform this multi-step planning, the authors introduce MM-R2-Traj, a large-scale dataset of 900K grounded multi-step retrieval trajectories synthesized using a teacher model. The training follows a two-stage post-training pipeline:

Stage 1 (Supervised Fine-Tuning - SFT): The agent learns the basic structure of retrieval behaviors (unit selection, query generation, summarization, answer generation) by imitating the structured trajectories.
Stage 2 (Group Relative Policy Optimization - GRPO): The agent is further optimized using reinforcement learning with a rollout-level reward function ( $R_i$ ) that jointly evaluates routing quality ( $R_{route}$ ) and final answer correctness ( $R_{ans}$ ). This strengthens the agent's ability to plan retrieval paths under task-level feedback.

Key Contributions

Novel Framework: MM-R2 is a multimodal agentic framework that explicitly reasons about both the retrieval target and the search space before retrieval, making the "what" and "where" of retrieval interpretable.
Architectural Innovations: The design of the Intent-Referent Binding Module and the KnowledgeMap provides an explicit decision state and a structured search space, addressing the ambiguity and flatness of prior mRAG systems.
Dataset and Training: The construction of MM-R2-Traj and the adoption of a two-stage SFT+GRPO training strategy enable agents to learn complex, multi-step retrieval planning.
Process-Level Analysis: The framework enables lightweight process-level analysis, allowing for the inspection of inferred intent, selected units, and evidence separately.

Experimental Results

The authors evaluated MM-R2 on two benchmarks: InfoSeek and Encyclopedic VQA.

Task-Level Effectiveness: MM-R2 substantially outperforms strong baselines. On InfoSeek, it achieved 54.3% accuracy (All), improving by 7.8 points over the strongest prior baseline (Wiki-R1-7B at 46.5%). On Encyclopedic VQA, it reached 39.4%. The performance gains were consistent across different retrieval budgets (Top-1 to Top-5), with improvements increasing as the budget expanded, suggesting effective utilization of additional evidence.
Process-Level Verifiability:
- Routing Accuracy: Route Recall@1 improved significantly from 0.2339 (baseline) to 0.7400, indicating the agent successfully identifies the correct KnowledgeMap regions.
- Trace Quality: LLM-as-a-Judge scores for trace-answer consistency and evidence coverage improved.
- Fidelity: The model achieved a fidelity score of 0.9011, indicating that the final prediction is highly reproducible from the revealed evidence summary, unlike baselines which often rely on hidden parametric knowledge.

Significance and Claims

The paper claims that MM-R2 demonstrates that reframing multimodal retrieval as a reasoning task ("reason before you retrieve") yields substantial benefits. By explicitly grounding the retrieval intent and structuring the search space, the system achieves:

Higher Accuracy: More precise evidence retrieval leads to better answer generation.
Interpretability: The retrieval trajectory is transparent, allowing users to inspect the agent's reasoning, routing decisions, and evidence usage.
Verifiability: The exposed evidence summaries are sufficient to justify the final answer, reducing reliance on "black box" parametric knowledge.

The authors conclude that MM-R2 offers a more robust and trustworthy approach for complex multimodal queries compared to static, flat retrieval pipelines.