Policy-Guided Stepwise Model Routing for Cost-Effective Reasoning
For practitioners deploying LLMs on reasoning tasks, this work provides a cost-effective alternative to expensive process reward models without sacrificing performance.
The authors formulate stepwise model routing as a constrained decision-making problem and train a small control policy using reinforcement learning with threshold calibration to balance accuracy and cost. Their method consistently improves the accuracy-cost tradeoff on three math benchmarks compared to handcrafted approaches and achieves comparable tradeoff to methods requiring large process reward models.
Inference-time computation has greatly enhanced the performance of large language models (LLMs) on challenging reasoning tasks, but this strategy can incur high inference costs. One solution is to route intermediate chain-of-thought (CoT) states to language models of different sizes; however, existing approaches rely on handcrafted routing strategies that limit performance, or on training large process reward models that may be infeasible in many applications. We formulate stepwise model routing as a constrained decision-making problem, which we solve by training a small control policy using reinforcement learning in conjunction with threshold calibration to tune the performance-efficiency tradeoff. We validate our method on three math benchmarks (GSM8K, MATH500, and OmniMath) on both open and closed models. Our method consistently improves the accuracy-cost tradeoff compared to handcrafted approaches, while achieving a comparable tradeoff to methods that require training large process reward models.