UR$^2$: Unify RAG and Reasoning through Reinforcement Learning

Technical Summary: UR2 (Unified RAG and Reasoning)

UR2 (Unified RAG and Reasoning) is a general reinforcement learning (RL) framework designed to bridge the gap between Retrieval-Augmented Generation (RAG) and Reinforcement Learning from Verifiable Rewards (RLVR). Accepted at ACL 2026, the paper addresses the fundamental tension between using external knowledge and maintaining internal reasoning capabilities in Large Language Models (LLMs).

---

1. The Problem: The Retrieval-Reasoning Imbalance

The authors identify a critical "imbalance" that occurs when attempting to unify RAG and reasoning through RL. Existing RAG-RL methods typically suffer from two failure modes:

• Over-reliance on Retrieval: Models become "lazy," issuing trivial or ill-posed queries (e.g., "which option is correct?") to offload the cognitive burden to the retriever, which weakens their internal reasoning.
• Collapse to Pure Reasoning: When faced with noisy or unstructured retrieved documents, models learn to ignore the retrieval tool entirely, degenerating into standalone Chain-of-Thought (CoT) reasoning and failing to utilize external knowledge when it is actually necessary.

2. Methodology: The UR2 Framework

To resolve this imbalance, UR2 introduces three core technical innovations:

A. LLM-Summarized Retrieval (Noise Reduction & Efficiency)
Instead of feeding raw, noisy, or long documents directly into the reasoning process, UR2 uses an LLM to transform retrieved content into compact, reasoning-compatible summaries.

• Benefit: This reduces the token count significantly (e.g., a 14.91× compression ratio in online settings), mitigating the impact of retrieval noise and reducing the quadratic computational cost of the attention mechanism during reasoning.

B. Difficulty-Aware Curriculum (Regulating Activation)
To prevent the model from defaulting to retrieval for every task, UR2 employs a curriculum based on task difficulty:

• Data Selection: Training instances are categorized into Easy, Medium, and Hard based on initial rollout performance.
• Task Mixing: For reasoning-intensive tasks (like Math), retrieval is only activated for "Hard" instances. For "Easy" and "Medium" instances, the model is trained using pure CoT. This forces the model to rely on its internal parametric knowledge for simpler problems while learning to use tools only when necessary.

C. Two-Stage Optimization (Skill Acquisition vs. Quality)
UR2 uses a decoupled training approach via REINFORCE++:

• Stage 1 (Retrieval Activation): The focus is on how to use the tool. The reward function prioritizes correct formatting, valid query generation, and successful tool invocation rather than just answer correctness.
• Stage 2 (Answer Optimization): Once the model knows how to retrieve, the focus shifts to accuracy. The reward function incorporates answer correctness while maintaining the formatting constraints learned in Stage 1.

3. Key Contributions

• First Unified Framework: Proposes a general RL framework that adapts to diverse domains (Math, Medicine, Open-domain QA) rather than being limited to Wikipedia-based QA.
• Dynamic Coordination: Successfully implements a mechanism that allows a model to decide when to retrieve and what to query.
• Robustness & Generalization: Demonstrates that the model can transition from offline corpora to real-time online web searches without retraining.

4. Experimental Results

The authors evaluated UR2 using Qwen-2.5 (3B/7B) and LLaMA-3.1-8B backbones across multiple benchmarks:

• Performance Gains: UR2 outperformed existing SOTA RAG and RL baselines by an average of 5.8% (7B) and 19.0% (3B).
• Competitive with Proprietary Models: The UR2 7B model achieved performance comparable to GPT-4o-mini and GPT-4.1-mini on several benchmarks.
• Domain Versatility: High performance was maintained across MMLU-Pro (general knowledge), MedQA (medical), and MATH500 (mathematical reasoning).
• Online Scalability: In open-domain QA, UR2 showed strong out-of-domain generalization, performing exceptionally well on benchmarks like Bamboogle and MusiQue when using real-time web search.

5. Significance

UR2 represents a significant milestone in the evolution of Agentic AI. It moves beyond "always-on" RAG (which is computationally expensive and noisy) and "pure reasoning" (which is limited by parametric knowledge). By treating retrieval as a learned, strategic behavior, UR2 enables smaller, open-source models (3B–8B) to exhibit sophisticated cognitive behaviors—such as self-verification, hypothesis revision, and query refinement—that were previously the domain of much larger, closed-source models.