Policy-Guided Stepwise Model Routing: RL-Based Step-Level Model Selection

> In May 2026, Si et al. proposed Policy-Guided Stepwise Model Routing, a method that dynamically selects model size at each step of a reasoning chain via a small control policy. The approach formalizes stepwise model routing as a constrained decision problem, trains a control policy with only a few hundred parameters using reinforcement learning, and combines threshold calibration to tune the performance-efficiency trade-off. Validation on GSM8K, MATH500, and OmniMath shows that the method consistently improves the accuracy-cost trade-off without relying on large process reward models (PRMs), and works for both open and closed models.

1. Background: The Efficiency Bottleneck of Inference-Time Compute

Test-time compute scaling improves LLM performance by allocating extra compute at inference, but brings significant cost growth:

| Strategy | Compute overhead | Use case | |:---|:---:|:---| | Self-Consistency | N× single inference | High confidence needed | | Best-of-N | N× inference + evaluation | Best answer needed | | Tree-of-Thoughts | branches × depth | Multi-path exploration | | Uniform large model | Fixed maximum cost | All scenarios |

Step-Level Heterogeneity

Compute demand within a reasoning chain is not uniform:

| Step type | Typical example | Required model capability | |:---|:---|:---:| | Basic calculation | "2 + 2 = 4" | Low | | Intermediate verification | "Verify this step is correct" | Medium | | Complex reasoning | "Prove this theorem" | High | | Synthesis | "Therefore the final answer is..." | Medium |

> Core observation: Using the largest model for every step is a serious waste of compute.

2. Method: A Constrained Decision Framework

Problem Formulation

Stepwise model routing is defined as a Markov decision process:

• State \(s_t\): CoT intermediate state at step \(t\)
• Action \(a_t\): select model \(M \in \{M_1, M_2, \ldots, M_k\}\)
• Transition: \(s_{t+1} = M(s_t)\)
• Reward \(r_t\): accuracy gain − cost penalty
• Constraint: total cost \(\leq B\) or accuracy \(\geq A_{\min}\)

Tiny Control Policy

| Component | Size | Function | |:---|:---:|:---| | Policy network \(\pi_\theta\) | Hundreds of parameters | \(s_t \rightarrow a_t\) mapping | | Threshold calibration \(\tau\) | Scalar | Performance-efficiency tuning | | Total | Negligible | Controls the entire inference cost structure |

The control policy's parameter count is far smaller than any candidate model; its training cost is negligible.

Reinforcement Learning Training

The policy is trained with policy-gradient methods:

\[\nabla_\theta J = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_t \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot G_t\right]\]

where \(G_t\) is the cumulative return combining:

| Return component | Weight | Role | |:---|:---:|:---| | Final answer correctness | \(\lambda_1\) | Accuracy-driven | | Model invocation cost | \(-\lambda_2\) | Cost constraint | | Constraint violation penalty | \(-\lambda_3\) | Hard constraint satisfaction |

Threshold Calibration

After training, a single threshold parameter \(\tau\) adjusts behavior:

\[\pi_\tau(a|s) = \mathbb{1}[Q(s, a) \geq \tau]\]
• High \(\tau\): favor large models (accuracy-sensitive)
• Low \(\tau\): favor small models (cost-sensitive)

No retraining is needed—a single-parameter slide adapts to different deployment needs.

3. Experimental Results

Benchmarks

| Benchmark | Domain | Difficulty | |:---|:---|:---:| | GSM8K | Grade-school math | Low–medium | | MATH500 | Competition math | High | | OmniMath | Mixed math | Mixed |

Accuracy–Cost Pareto Analysis

| Method | Pareto efficiency | PRM needed | Complexity | |:---|:---:|:---:|:---:| | Uniform large model | Low | No | Low | | Handcrafted routing | Medium | No | Medium | | Large PRM routing | High | Yes | High | | Policy-Guided routing | High | No | Low |

> Key advantage: Pareto efficiency comparable to large-PRM methods, while avoiding PRM training and deployment costs.

Cross-Model Validation

The method is effective for both open-weight models and closed API models.

4. Theoretical Analysis

Gap to Optimal Routing

Defining optimal routing as an omniscient decision-maker with full future information:

\[\text{Gap} = J(\pi^*) - J(\pi_\theta)\]

Experiments show the tiny control policy's gap is within an acceptable range, validating the hypothesis that step-level difficulty is predictable.

Cost Sensitivity

| Cost ratio (vs. uniform large model) | Accuracy retained | |:---:|:---:| | 100% | 100% | | 70% | ~98% | | 50% | ~95% | | 30% | ~90% |

> The first 50% cost reduction loses only ~5% accuracy; marginal losses grow after that.

5. Relation to Related Work

• Coupling Tax (Round 16): Coupling Tax concerns budget competition between reasoning chains and answers; this work concerns step-level budget allocation within the reasoning chain.
• DAST (Round 7): DAST allocates compute at the question level by difficulty; this work refines adaptivity to the step level.
• AutoTTS (Round 21): AutoTTS discovers TTS strategies automatically; this work shows another automation paradigm—RL-learned step-level model selection.
• 80/20 Rule (Round 14): Round 14 found 20% of tokens carry key gradient signal; step-level routing extends this idea to model selection: large models for critical steps, small models for the rest.

6. Limitations and Future Directions

• Model library granularity: currently 2–3 model options; scaling the action space and hierarchical routing (series first, then size) remain open.
• Cross-domain transfer: whether policies trained on math transfer to code generation, scientific reasoning, or creative writing.
• Parallel reasoning: how to route branches in Tree-of-Thoughts (independent routing vs. global coordination).
• Online adaptation: online RL fine-tuning of the control policy and user-feedback-driven threshold adaptation.

7. Conclusion

Policy-Guided Stepwise Model Routing offers a lightweight, scalable solution for inference-time compute optimization. Core contributions:

1. Constrained decision framework formalizing routing as a principled optimization problem 2. Minimal control policy—hundreds of parameters suffice for effective routing 3. No PRM dependency, avoiding the bottleneck of training large process reward models 4. Tunability via a single threshold parameter 5. Empirical effectiveness across multiple math benchmarks

As inference cost becomes a deployment bottleneck, step-level model routing represents a key step from uniform toward adaptive computation.

Paper Details

| Item | Content | |:---|:---| | Title | Policy-Guided Stepwise Model Routing for Cost-Effective Reasoning | | Authors | Wenwen Si, Insup Lee, Osbert Bastani | | Institution | University of Pennsylvania | | arXiv ID | 2605.06116 | | Date | 2026-05-07 | | Key contributions | Constrained decision framework; tiny control policy (hundreds of parameters); RL training + threshold calibration; no PRM; step-level model routing | | Key results | Consistent accuracy-cost improvements on GSM8K/MATH500/OmniMath; comparable to large-PRM methods; effective on open and closed models |