Source-linked AI summary

Router-R1: Teaching LLMs Multi-Round Routing and Aggregation via Reinforcement Learning

Haozhen Zhang, Tao Feng, Jiaxuan You

arXiv:2506.09033v3cs.CLcs.AIcs.LG

TL;DR

Existing LLM routers typically make one-shot, single-model decisions, limiting coordination for complex tasks requiring complementary model strengths. Router-R1 uses reinforcement learning to train a capable LLM router that interleaves reasoning with multi-round model selection and cost-aware rewards. Across seven QA benchmarks, it outperforms strong baselines while maintaining robust generalization and performance–cost trade-offs.

  • Problem

    Existing routers usually map each query to one model in a single round, leaving limited support for complex tasks requiring coordinated use of multiple LLMs.

  • Method

    Router-R1 formulates multi-LLM routing as sequential decision-making, with an LLM router interleaving internal reasoning and model calls and optimizing format, outcome, and cost rewards.

  • Results

    Across seven diverse QA benchmarks, Router-R1 outperforms more than ten strong baselines and maintains robustness and generalization with distractor models.

  • Takeaways & Limitations

    The results support reinforcement-learning-based routing as a way to orchestrate multiple LLMs with flexible performance–cost trade-offs.

  • Takeaways & Limitations

    Evaluation focuses primarily on QA, while the reward may not capture nuanced objectives and multi-round routing can introduce inference latency.

Abstract

from arXiv · show

The rapid emergence of diverse large language models (LLMs) has spurred the development of LLM routers that assign user queries to the most suitable model. However, existing LLM routers typically perform a single-round, one-to-one mapping (\textit{i.e.}, assigning each query to a single model in isolation), which limits their capability to tackle complex tasks that demand the complementary strengths of multiple LLMs. In this paper, we present \textbf{Router-R1}, a reinforcement learning (RL)-based framework that formulates multi-LLM routing and aggregation as a sequential decision process. Router-R1 instantiates the router itself as a capable LLM, leveraging its reasoning ability to interleave "think" actions (internal deliberation) with "route" actions (dynamic model invocation), and integrates each response into its evolving context. To facilitate learning, we employ a lightweight rule-based reward comprising format rewards, final outcome rewards, and a novel cost reward for optimizing the balance between performance and cost, opening a pathway toward enhancing performance-cost trade-offs via RL. Router-R1 also conditions only on simple model descriptors such as pricing, latency, and example performance, enabling strong generalization to unseen model selection. Experiments on seven general and multi-hop QA benchmarks show that Router-R1 outperforms several strong baselines, achieving superior performance while maintaining robust generalization and cost management.

1 Introduction

Router-R1 addresses the limits of single-round routing by treating multi-LLM coordination as sequential reasoning and model selection. Its reinforcement-learning framework combines adaptive routing, rule-based rewards, descriptor-based generalization, and evaluation across seven QA benchmarks.

  • Motivation: Single-round routers select one LLM per query, but complex reasoning tasks often require coordinated interactions among multiple models.The challenge includes both non-differentiable multi-round selection and the limitations of one-off routing.
  • Router-R1: Router-R1 interleaves internal “think” actions with targeted “route” actions, integrating model responses into an evolving context.The router itself is instantiated as a capable LLM and incrementally constructs answers through sequential interaction.
  • Training Objective: Its rule-based reward combines format, final outcome, and cost rewards to navigate performance–cost trade-offs.The cost component penalizes excessive use of expensive routed models.
  • Generalization: Conditioning on simple model descriptors allows Router-R1 to generalize to unseen LLMs without retraining.This supports routing across a changing pool of candidate models.
  • Evaluation: Across seven question-answering benchmarks, Router-R1 outperforms competitive baselines while achieving robust generalization.The introduction reports superior performance across the evaluated benchmarks.

2 Related Work

Prior LLM routers primarily select models for individual queries, while Router-R1 extends routing into sequential multi-round coordination trained with reinforcement learning. Its distinguishing design combines interleaved reasoning, model routing, and cost-aware optimization.

  • Existing Routers: Existing query-based routers direct queries to appropriate models to improve response quality and efficiency, often using predicted difficulty or cost–performance estimates.Related methods include HybridLLM, GraphRouter, and FrugalGPT.
  • Router-R1: Router-R1 treats routing as a sequential process that interleaves internal “think” steps with multi-round model calls to refine answers.The router is itself a capable LLM that coordinates multiple candidates.
  • RL and Cost: Its reinforcement-learning training uses a cost reward to support flexible performance–cost trade-offs and resource-aware routing.This distinguishes the framework from prior routing approaches described in the related work.

3 Router-R1

Router-R1 formulates coordination with an LLM routing pool as sequential policy optimization, where a capable router alternates reasoning and candidate-model calls. It uses structured interaction and hierarchical rule-based rewards to balance correctness, format, and cost.

  • 3.1 Reinforcement Learning: Router-R1 optimizes a policy LLM over a pool of candidate models, with KL regularization keeping the updated policy close to a reference LLM.The objective supports regularized reinforcement-learning algorithms such as PPO and GRPO.
  • 3.1 Reinforcement Learning: The policy LLM can select candidate models multiple times and provide relevant context, allowing several models to jointly solve complex tasks.This turns the policy LLM into a coordinator rather than a one-shot dispatcher.
  • 3.2 Reward Curation: The final outcome reward uses Exact Match between the extracted predicted answer and the ground truth.EM is used as the only final outcome reward for task correctness.
  • 3.2 Reward Curation: The cost reward decreases with model size and output-token count, using a predefined model-size cost function to support performance–cost balance.The reward is normalized between 0 and 1 during training.
  • 3.2 Reward Curation: The reward function combines format, final outcome, and cost rewards, with hierarchical priorities that nullify later rewards when format validation fails.Format correctness is enforced before optimizing task performance or computational efficiency.
  • 3.3 Multi-round Interaction: The structured prompt alternates internal analysis with selective candidate-LLM querying, inserts returned information into the sequence, and may repeat this process.The prompt uses <think>, <search>, <info>, and <answer> blocks for the interaction format.

4 Experimental Setup

The experiments evaluate Router-R1 on seven general and multi-hop QA datasets against basic, retrieval, search, and query-based routing baselines. Training uses small supervised resources with controlled routing and evaluation settings for in-domain and out-of-domain testing.

  • Datasets: The evaluation covers seven QA datasets: NQ, TriviaQA, PopQA, HotpotQA, 2WikiMultiHopQA, Musique, and Bamboogle.The set includes three general QA and four multi-hop QA benchmarks.
  • Baselines: Router-R1 is compared with Direct, CoT, SFT, RAG, Search-R1, and query-based LLM routers.The baselines include both basic inference methods and routing systems using one or more candidate models.
  • Baseline Training: Query-based router training labels each candidate LLM with question-specific Exact Match scores obtained from repeated temperature-sampled responses.Each training question is independently evaluated across the routing pool.
  • Training Configuration: Router-R1 uses Qwen2.5-3B-Instruct and LLaMA-3.2-3B-Instruct as base models, with a maximum of 4 routing steps per query and PPO as the default RL algorithm.The reported setup uses batch size 64, up to 225 training steps, and α = 0.0 in the main experiment.
  • Training Data: Training uses 7K NQ samples and 7K HotpotQA samples, producing a 14K-sample joint dataset for routing and aggregation.The authors report that this modest dataset is sufficient to induce effective routing strategies.
  • Evaluation Protocol: In-domain evaluation uses NQ and HotpotQA, while five additional QA datasets assess out-of-domain generalization under consistent baseline settings.Each dataset uses 500 randomly sampled test instances except Bamboogle, which has about 120 total examples.

5 Experimental Analysis

Router-R1 is evaluated across seven QA benchmarks for performance, cost-sensitive routing, unseen-model generalization, adaptive API usage, and training convergence. It consistently outperforms baselines, adapts routing to task difficulty and cost constraints, generalizes to new candidates, and converges rapidly.

  • 5.1 Main Results: Router-R1 consistently outperforms basic and LLM-router baselines across all seven QA datasets.It achieves SOTA performance and surpasses both single-round and advanced routing baselines.
  • 5.1 Main Results: 0.416 average exact match is achieved by Router-R1-Qwen, exceeding Router-R1-Llama’s 0.409.These are the highest reported average exact match scores for the two Router-R1 variants.
  • 5.3 Generalization to Unseen LLMs: Router-R1 maintains robust performance on five out-of-domain datasets after training only on NQ and HotpotQA.The evaluation covers general and multi-hop QA, with 500 sampled instances per dataset except Bamboogle.
  • 5.2 Analysis of Cost Rewards: Increasing the cost coefficient gradually lowers overall performance while raising the cost reward.The learned policy increasingly avoids expensive calls, querying smaller models first and escalating only when necessary.
  • 5.3 Generalization to Unseen LLMs: Adding previously unseen candidate LLMs at inference preserves and slightly improves Router-R1-Qwen’s performance across all four reported benchmarks.The original checkpoint uses textual model descriptors without additional fine-tuning and reaches new best scores on several datasets and overall averages.
  • 5.4 Discussion: Router-R1-Qwen makes more API calls on multi-hop than general QA benchmarks and converges within 100 training steps.Rising rewards and decreasing policy entropy indicate convergence; format rewards reduce instability from malformed outputs.

6 Conclusion

The paper concludes that Router-R1 uses reinforcement learning to coordinate multiple LLMs through sequential reasoning and routing. Its rule-based rewards support performance–cost trade-offs, while descriptor-based prompting supplies model capability priors and enables integration of new models without retraining.

  • 6 Conclusion: Router-R1 formulates multi-LLM routing and aggregation as a sequential decision process.The router interleaves internal reasoning with targeted model selection and incrementally builds responses through multi-round interaction.
  • 6 Conclusion: Model descriptors provide cold-start priors that Router-R1 refines through interaction and reward feedback.The descriptors are generated from public model cards and summarize candidate capabilities for the training prompt.
  • 6 Conclusion: Format rewards penalize malformed outputs and enforce disciplined, interpretable reasoning trajectories.The expected structure includes think, search, info, and answer tags with specific pairing and ordering requirements.
  • 6 Conclusion: Cost rewards penalize high-cost LLM use and guide cost-sensitive routing without hard constraints.The reward uses normalized cost, including square-root smoothing and percentile-based normalization, to adapt to resource profiles.
  • 6 Conclusion: Cost-reward normalization may become inconsistent across heterogeneous domains because response-length variation produces differing cost scales.The paper identifies task-level normalization as a future direction for improving multi-task training stability.

A.2.1 Main Results w.r.t. F1-Score

Router-R1 consistently outperforms all baselines across seven general and multi-hop QA datasets, achieving the highest average F1-scores for both Qwen and LLaMA backbones.

  • Router-R1 consistently outperforms all baselines across seven general and multi-hop QA datasets.
  • Router-R1 achieves the highest average F1-scores for both Qwen and LLaMA backbones.
  • The results are reported with respect to F1-Score in Table 3.

A.2.2 Extensive Analysis of Cost Rewards

Router-R1 uses cost-aware routing to navigate a controllable accuracy–cost frontier, matching strong-model accuracy at lower cost or prioritizing efficiency while retaining competitive accuracy.

  • Varying α shifts Router-R1 from performance-oriented routing toward more cost-efficient decisions.At α = 0.0, Router-R1 achieves the highest EM across nearly all datasets; increasing α substantially reduces average cost with a moderate accuracy effect.
  • Router-R1 spans a wide cost–performance frontier by tuning α across different deployment preferences.The router can approximate larger, stronger models’ accuracy at a fraction of their cost or prioritize efficiency while retaining competitive accuracy.
  • Mixtral-8x22B combines high EM with relatively lower cost among individual LLMs.Its lower cost is attributed to generating shorter responses under identical prompts.
  • Router-R1 achieves competitive or superior efficiency–accuracy trade-offs relative to individual models and learned routers.Its cost-aware design enables smooth traversal along the efficiency–accuracy continuum for flexible, resource-adaptive deployment.
  • Router-R1 consistently outperforms FrugalGPT across general and multi-hop QA tasks.

A.2.4 Sensitivity Analysis of Model Descriptors

The descriptor-sensitivity experiment compresses detailed model descriptions to retain key attribute information while removing stylistic and secondary details.

  • Router-R1 is evaluated with compressed model descriptors that retain key attributes while removing stylistic and secondary details.The modified descriptors preserve the overall format while reducing descriptive richness.
  • Table 6 reports sensitivity analysis of model descriptors across seven QA datasets using Exact Match.

A.2.5 New Candidate Incorporation During Training

Router-R1 remains robust when the routing pool evolves during training: incorporating a new model slightly improves average exact match while maintaining stable dataset-level performance.

  • Adding LLaMA3-ChatQA-1.5-8B during training tests Router-R1’s adaptability to an expanding model pool.The main results instead focus on models introduced only at test time.
  • 0.416 to 0.425: average exact-match score slightly increases after incorporating the new model.
  • Router-R1 maintains stable performance across all datasets when the routing pool evolves during training.
  • Table 7 reports new candidate incorporation during training across seven QA datasets using Exact Match.
  • The case study examines Router-R1 under zero and non-zero cost coefficients α.It illustrates routing behavior with and without cost-reward constraints.

B.1 Case Study without Cost Reward Constraints

The case study shows Router-R1 using external LLMs to resolve knowledge gaps in multi-step questions, including cases where initial answers are uncertain or incomplete.

  • Multi-round routing: For the film-release comparison, the routed answer selects “Ek Haseena Thi Ek Deewana Tha” (2017) over “Sacred Silence” (2016).The answer is based on the reported release years.
  • Multi-round routing: Router-R1 recognizes when it lacks sufficient information and issues targeted sub-questions to an external LLM.For the film-director comparison, the policy first identifies the missing director information before routing a follow-up query.
  • Multi-round routing: The routed response identifies Biren Nag as the director of “To See or Not to See” (1969) and reports his death date as August 5, 2018.The response also notes uncertainty about the exact title match.

E Limitations

Router-R1’s evaluation is concentrated on QA, while its reward design, latency, and descriptor-based generalization impose additional boundaries on deployment and broader applicability.

  • Scope and deployment: The evaluation focuses primarily on QA, leaving generalization to dialogue, summarization, and code generation unresolved.These domains may involve different routing dynamics.
  • Scope and deployment: The rule-based reward may not capture nuanced objectives such as factual consistency or long-term dialogue coherence.The paper suggests learned or human-in-the-loop rewards as a possible extension.
  • Scope and deployment: Multi-round routing introduces inference latency, especially when reasoning steps are interleaved with multiple model calls, limiting suitability for time-sensitive applications.
  • Scope and deployment: Generalization to unseen LLMs relies on simple descriptors such as pricing and latency that may miss deeper capabilities when performance history is limited.
  • Future directions: Future work includes improving reward modeling, reducing latency through pruning or routing heuristics, and expanding beyond the current language-task scope.
Loading 2506.09033v3…