HierRouter: Coordinated Routing of Specialized Large Language Models via Reinforcement Learning

TL;DR

HierRouter addresses the high inference cost of large language models by coordinating a pool of specialized lightweight LLMs through a multi-hop, RL-driven router. Formulated as a finite-horizon MDP and trained with Proximal Policy Optimization, the router selects models at each hop to maximize a final quality score while penalizing cumulative compute via a cost term. Across six diverse benchmarks, HierRouter yields substantial quality gains over static small models and often matches or exceeds larger LLMs, with a modest, predictable increase in compute that is justified by the improved results. This approach demonstrates a scalable, cost-aware paradigm for deploying efficient, high-performance LLM inference in resource-constrained settings.

Abstract

Large Language Models (LLMs) deliver state-of-the-art performance across many tasks but impose high computational and memory costs, limiting their deployment in resource-constrained or real-time settings. To address this, we propose HierRouter, a hierarchical routing approach that dynamically assembles inference pipelines from a pool of specialized, lightweight language models. Formulated as a finite-horizon Markov Decision Process (MDP), our approach trains a Proximal Policy Optimization (PPO)-based reinforcement learning agent to iteratively select which models to invoke at each stage of multi-hop inference. The agent conditions on the evolving context and accumulated cost to make context-aware routing decisions. Experiments with three open-source candidate LLMs across six benchmarks, including QA, code generation, and mathematical reasoning, show that HierRouter improves response quality by up to 2.4x compared to using individual models independently, while incurring only a minimal additional inference cost on average. These results highlight the promise of hierarchical routing for cost-efficient, high-performance LLM inference. All codes can be found here https://github.com/ Nikunj-Gupta/hierouter.

Figures (3)

• Figure 1: Specialized LLMs are smaller models fine-tuned for specific domains or task types, such as code generation, math reasoning, or biomedical QA. Unlike general-purpose LLMs, these models exhibit high efficiency and competitive accuracy on aligned tasks. By coordinating such models through routing rather than relying on monolithic architectures, systems like HierRouter can support adaptive, low-cost inference while preserving task performance.
• Figure 2: Overview of the HierRouter architecture. At each of the $L$ inference stages, a PPO-based (RL) router selects one model from a pool of specialized LLMs to respond to the current context. The model's output is appended to the evolving context, which is passed to the next router layer. This multi-hop process enables compositional reasoning: earlier hops handle simpler sub-tasks, while later hops refine the output. The routing policy is trained to optimize response quality under inference cost constraints, balancing accuracy and efficiency through sequential model coordination.
• Figure 3: Comparing met score or reward comparison ($R = Q - \alpha C$). $Q$ denotes average F1 quality score, $C$ denotes average cost per query (token-based), and $\alpha$ is the cost penalty coefficient used during training. Despite slightly higher costs, HierRouter achieves superior or comparable net reward in all datasets, demonstrating efficient tradeoffs between performance and compute usage.