Policy-Guided Stepwise Model Routing for Cost-Effective Reasoning

• URL: http://arxiv.org/abs/2605.06116v1
• Date: Thu, 07 May 2026 12:26:56 GMT
• Title: Policy-Guided Stepwise Model Routing for Cost-Effective Reasoning
• Abstract summary: Inference-time computation has greatly enhanced the performance of large language models (LLMs) on challenging reasoning tasks.<n>One solution is to route intermediate chain-of-thought (CoT) states to language models of different sizes.<n>We formulate stepwise model routing as a constrained decision-making problem.
• Score: 29.183252659668696
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• TRIM: Hybrid Inference via Targeted Stepwise Routing in Multi-Step Reasoning Tasks [26.198066761026297]
  Current methods assign entire queries to one model, treating all reasoning as equal to one model.<n>We propose a new model that handles all multi-step reasoning tasks.<n>We develop several strategies within ranging from a simple threshold to more expressive routing policies.
  arXiv  Detail & Related papers  (2026-01-15T10:06:06Z)
• DiRL: An Efficient Post-Training Framework for Diffusion Language Models [54.405206032785706]
  Diffusion Language Models (dLLMs) have emerged as promising alternatives to Auto-Regressive (AR) models.<n>Existing methods suffer from computational inefficiency and objective mismatches between training and inference.<n>We introduce DiRL, an efficient post-training framework that tightly integrates FlexAttention-accelerated blockwise training with LMDeploy-optimized inference.
  arXiv  Detail & Related papers  (2025-12-23T08:33:19Z)
• CONCUR: A Framework for Continual Constrained and Unconstrained Routing [79.85419373937765]
  AI tasks differ in complexity and are best addressed with different computation strategies.<n>Most prior methods build the routing framework by training a single model across all strategies.<n>We propose CONCUR, a continual routing framework that supports both constrained and unconstrained routing.
  arXiv  Detail & Related papers  (2025-12-10T07:30:13Z)
• Ring-lite: Scalable Reasoning via C3PO-Stabilized Reinforcement Learning for LLMs [51.21041884010009]
  Ring-lite is a Mixture-of-Experts (MoE)-based large language model optimized via reinforcement learning (RL)<n>Our approach matches the performance of state-of-the-art (SOTA) small-scale reasoning models on challenging benchmarks.
  arXiv  Detail & Related papers  (2025-06-17T17:12:34Z)
• Route to Reason: Adaptive Routing for LLM and Reasoning Strategy Selection [7.045509749924679]
  Route-To-Reason (RTR) is a novel unified routing framework that dynamically allocates both LMs and reasoning strategies according to task difficulty under budget constraints.<n>RTR learns compressed representations of both expert models and reasoning strategies, enabling their joint and adaptive selection at inference time.
  arXiv  Detail & Related papers  (2025-05-26T02:53:17Z)
• Reinforcement Learning for Machine Learning Model Deployment: Evaluating Multi-Armed Bandits in ML Ops Environments [0.0]
  We investigate whether reinforcement learning (RL)-based model management can manage deployment decisions more effectively.<n>Our approach enables more adaptive production environments by continuously evaluating deployed models and rolling back underperforming ones in real-time.<n>Our findings suggest that RL-based model management can improve automation, reduce reliance on manual interventions, and mitigate risks associated with post-deployment model failures.
  arXiv  Detail & Related papers  (2025-03-28T16:42:21Z)
• FTP: A Fine-grained Token-wise Pruner for Large Language Models via Token Routing [17.01412432658081]
  Large language models (LLMs) have demonstrated superior performance across various tasks by adhering to scaling laws.<n>We propose a fine-grained token-wise pruning approach for the LLMs, which presents a learnable router to adaptively identify the less important tokens.<n>Our approach achieves state-of-the-art (SOTA) pruning results, surpassing other existing pruning methods.
  arXiv  Detail & Related papers  (2024-12-16T07:09:46Z)
• Enhancing Multi-Step Reasoning Abilities of Language Models through Direct Q-Function Optimization [49.362750475706235]
  Reinforcement Learning (RL) plays a crucial role in aligning large language models with human preferences and improving their ability to perform complex tasks.<n>We introduce Direct Q-function Optimization (DQO), which formulates the response generation process as a Markov Decision Process (MDP) and utilizes the soft actor-critic (SAC) framework to optimize a Q-function directly parameterized by the language model.<n> Experimental results on two math problem-solving datasets, GSM8K and MATH, demonstrate that DQO outperforms previous methods, establishing it as a promising offline reinforcement learning approach for aligning language models.
  arXiv  Detail & Related papers  (2024-10-11T23:29:20Z)
• Towards Efficient Pareto Set Approximation via Mixture of Experts Based Model Fusion [53.33473557562837]
  Solving multi-objective optimization problems for large deep neural networks is a challenging task due to the complexity of the loss landscape and the expensive computational cost. We propose a practical and scalable approach to solve this problem via mixture of experts (MoE) based model fusion. By ensembling the weights of specialized single-task models, the MoE module can effectively capture the trade-offs between multiple objectives.
  arXiv  Detail & Related papers  (2024-06-14T07:16:18Z)
• Task-optimal data-driven surrogate models for eNMPC via differentiable simulation and optimization [42.72938925647165]
  We present a method for end-to-end learning of Koopman surrogate models for optimal performance in a specific control task.<n>We evaluate the performance of our method by comparing it to that of other training algorithms on an existing economic nonlinear model predictive control (eNMPC) case study.
  arXiv  Detail & Related papers  (2024-03-21T14:28:43Z)
• Evaluating model-based planning and planner amortization for continuous control [79.49319308600228]
  We take a hybrid approach, combining model predictive control (MPC) with a learned model and model-free policy learning. We find that well-tuned model-free agents are strong baselines even for high DoF control problems. We show that it is possible to distil a model-based planner into a policy that amortizes the planning without any loss of performance.
  arXiv  Detail & Related papers  (2021-10-07T12:00:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.