TRIM: Hybrid Inference via Targeted Stepwise Routing in Multi-Step Reasoning Tasks

TL;DR

TRIM (Targeted routing in multi-step reasoning tasks), which routes only critical steps, those likely to derail the solution, to larger models while letting smaller models handle routine continuations, demonstrating that step-level difficulty represents fundamental characteristics of reasoning.

Abstract

Multi-step reasoning tasks like mathematical problem solving are vulnerable to cascading failures, where a single incorrect step leads to complete solution breakdown. Current LLM routing methods assign entire queries to one model, treating all reasoning steps as equal. We propose TRIM (Targeted routing in multi-step reasoning tasks), which routes only critical steps$\unicode{x2013}$those likely to derail the solution$\unicode{x2013}$to larger models while letting smaller models handle routine continuations. Our key insight is that targeted step-level interventions can fundamentally transform inference efficiency by confining expensive calls to precisely those steps where stronger models prevent cascading errors. TRIM operates at the step-level: it uses process reward models to identify erroneous steps and makes routing decisions based on step-level uncertainty and budget constraints. We develop several routing strategies within TRIM, ranging from a simple threshold-based policy to more expressive policies that reason about long-horizon accuracy-cost trade-offs and uncertainty in step-level correctness estimates. On MATH-500, even the simplest thresholding strategy surpasses prior routing methods with 5x higher cost efficiency, while more advanced policies match the strong, expensive model's performance using 80% fewer expensive model tokens. On harder benchmarks such as AIME, TRIM achieves up to 6x higher cost efficiency. All methods generalize effectively across math reasoning tasks, demonstrating that step-level difficulty represents fundamental characteristics of reasoning.

Figures (7)

• Figure 1: Schematic Overview of a Two-Model Setup for TRIM
• Figure 2: Comparison of task performance–cost trade-offs for Qwen2.5-3B-Instruct ($M_w$) and Claude 3.7 Sonnet ($M_s$), under the myopic thresholding policy (with Qwen2.5-Math-PRM-7B) versus the Idealized Oracle Query-level Router, across multiple math benchmarks. The Oracle Router is evaluated by incrementally varying the number of queries routed to $M_s$, selecting those solvable by $M_s$ but not by $M_w$. Since multiple query subsets can achieve the same accuracy with different token costs, the oracle’s performance–cost curve forms a shaded region rather than a single line, reflecting the full trade-off frontier.
• Figure 3: Step-Wise Router architecture for TRIM using process rewards to evaluate partial solutions and uses RL-based policies or POMDP-based solvers for making routing decisions.
• Figure 4: POMDP state space $S$ classes and observation space $\Omega$ in TRIM-POMDP. The latent state space $S$ consists of three correctness classes: $S_0$ (trajectory correct so far), $S_1$ (trajectory irrecoverably incorrect due to an earlier error), and $S_2$ (most recent step incorrect but prior steps correct and thus potentially recoverable), augmented with current step index and token cost. The observation space $\Omega$ comprises the PRM-based cumulative correctness score of the prior steps, PRM score of the current step, and auxiliary features, which together serve as noisy observations of the underlying latent correctness state.
• Figure 5: Learned observation function illustrating PRM noise. Heatmaps show the empirical probability density of PRM-based observations conditioned on each latent correctness state, estimated from ProcessBench (Omni-MATH). While density concentrates around state-consistent regions (e.g., near $(1,0)$ for $S_2$, where prior steps are correct but the current step is incorrect), the substantial spread reflects noise in PRM scores. If PRM estimates were perfectly accurate, each distribution would collapse to a single point; the observed variance motivates treating PRM outputs as noisy observations in TRIM-POMDP.