AlphaRouter: Token-level Routing Between SLM and LLM with Reinforcement Learning and Tree Search

SLM-LLM routing accelerates generation by strategically invoking LLMs for critical tokens. However, existing methods typically train routers to mimic the LLM, capping performance at the reference trajectory's limit. In this work, we demonstrate that the SLM-LLM collaborative inference space offers a richer solution set, yielding correct answers even when the LLM fails. To exploit this, we propose AlphaRouter, a routing framework learning optimal collaborative inference paths via a search and iterative update paradigm. Formulating routing as a Markov Decision Process, we introduce Collaborative Inference Tree Search (CITS) to explore the solution space. To tackle the severe credit assignment challenge posed by sparse rewards, we propose Tree-Advantage Policy Optimization (TAPO) to optimize the routing policy. By leveraging counterfactual advantages within the tree structure, TAPO effectively attributes the final reward to specific routing decisions, stabilizing training without dense supervision. Extensive experiments show that AlphaRouter advances the Pareto frontier of accuracy-efficiency trade-offs by exploiting better inference trajectories in the collaborative space. Code is available at https://github.com/StripeLife0217/AlphaRouter.

Large language models can solve difficult tasks, but using them for every part of every answer is slow and expensive. Smaller language models are much cheaper, yet they may make mistakes at a few key moments that determine the final result. This paper introduces AlphaRouter, a system that lets a small and a large language model work together while generating an answer. Most of the time, the small model produces the text; the large model is called only when its help is likely to matter. Instead of simply copying the large model's response, AlphaRouter learns by exploring alternative generation paths during training. When the two models disagree, it tries different continuations and checks which choices lead to better final answers. These comparisons teach the system which disagreements are harmless and which ones deserve the extra cost of using the large model. When the system is actually used, this expensive search is removed, allowing AlphaRouter to produce one answer efficiently while approaching the accuracy of much larger models and reducing computation, memory use, and latency.