TRIM: Hybrid Inference via Targeted Stepwise Routing in Multi-Step Reasoning Tasks
Authors: Vansh Kapoor
Aman Gupta
Hao Chen
Anurag Beniwal
Jing Huang
Aviral Kumar
Venue: ICLR 2026 Poster (2026)
Keywords: Large Language Models
Model routing
Reinforcement Learning
Partially Observable MDPs (POMDPs)
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce TRIM, a routing approach for multi-step reasoning that selectively routes only critical steps to larger LLMs based on uncertainty and budget.

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.

Published as a conference paper at ICLR 2026
TRIM: HYBRID INFERENCE VIA TARGETED STEPWISE ROUTING IN MULTI-STEP REASONING TASKS
Vansh Kapoor't, Aman Gupta?, Hao Chen", Anurag Beniwal", Jing Huang", Aviral Kumar! Carnegie Mellon University AAmazon
ABSTRACT
Multi-step reasoning tasks like mathematical problem solving vulnerable cascading failures where a single incorrect step leads to complete solution break- down. Current L.M routing methods assign entire queries one model, treating all reasoning steps as equal. We propose 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. 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: 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- long-horizon accuracy-coss trade-offs and uncertainty in step-level correctness estimates. On MATH-500, even the sim- plest thresholding strategy surpasses prior routing methods with 5x higher cost efficiency, while more advanced policies match the strong, expensive model's per- formance 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.
INTRODUCTION
The rapid progress in large language models (LLMs) has led to a diverse ecosystem of models, spanning wide spectrum of sizes, capabilities, and computational demands. Larger models typically achieve stronger performance but incur substantial serving costs, rendering them impractical for many routine applications. In contrast, smaller models are more affordable to deploy but often produce lower-quality responses. This trade-off poses a fundamental dilemma for the practical deployment of LLMs: routing all queries to the largest available model ensures high-quality outputs but is prohibitively expensive, whereas relying solely on smaller models reduces serving costs at the expense of degraded response quality, especially on challenging queries.
Contemporary routing strategies attempt to mitigate this dilemma by assigning each query to a single model, which then responsible for the entire generation. However, not tokens in an LLM's response are equally challenging to generate: some represent critical decision points that can dramatically alter the solution path (Setlur et al., 2024a; Wang et al., 2025; Qu et al., 2025), while others routine continuations that easier generate. By committing over a smaller one for a given query, existing routing methods implicitly assume that the stronger model's intervention is equally necessary at every token produce high-quality response. This inefficiency is particularly pronounced in multi-step reasoning tasks such as step-by-step reasoning or code generation, where mistakes early on can snowball into complete failure (Zhang et al., 2024). It is precisely at these erroneous steps that the intervention of a stronger model is most valuable. Yet, contemporary routing methods often incur substantial inefficiency by defaulting to full generations from the larger model, even when targeted interventions at these erroneous steps would suffice. To address this inefficiency, we introduce TRIM- TRIM-Targeted Stepwise Routing for Inference in Multi-step
larger LLM
Work primarily conducted during an internship at Amazon. Corresponding author: vanahkêca emv.edu

---

Published as a conference paper at ICLR 2026
Problem Statement
Cheap
Pur liell Genurated Solution
Generated Stop
Step 1 Step 2
x
Step N
.*..
Replace with Expensive LLM Step present
Step N-1
Regenerate Step with Expensive LLM?
Expansivee.. LLM Regensrater Step
Stop- Wise Router
Expensive LLM
Step N
Figure 1: Schematic Overview of a Two-Model Setup for TRIM
Reasoning Tasks: an approach that selectively routes only the most critical steps to larger LLMs. Unlike prior prompt-level routers, TRIM operates at the granularity of individual reasoning steps, generating solutions one step at a time. As illustrated in Figure 1, stepwise router evaluates each intermediate step as is generated and decides whether to accept the small model's output or to regenerate that specific step with larger model. TRIM ensures that interventions occur only when necessary at individual steps, rather than handing over the entire remaining solution to the larger model. This step-by-step generation process with targeted intervention enables TRIM to achieve efficiency by confining expensive calls to precisely those steps where intervention prevents cascading errors while allowing the smaller model to handle routine continuations. This approach reduces the primary cost bottleneck in inference: the number of tokens generated by the expensive, larger LLM.
Building this framework, we design multiple strategies for stepwise routing, each tailored to different computational and informational constraints: a simple thresholding policy that uses step- level scores to identify erroneous steps, two RL-trained policies that reason about long-horizon trade-offs between accuracy and cost (one using full sequential features, another using aggregated statistics), and a Partially Observable Markov Decision Process (POMDP) based approach that accounts for the inherent uncertainty in step-level correctness estimates while enabling efficient policy recomputation across different cost budgets. Our cost metric focuses on the number of tokens generated by the expensive model, as prefill costs can be amortized through paraliel decoding strategies (Leviathan et al., 2023; Cai et al., 2024), whereas generation tokens impose unavoidable sequential costs. We evaluate our approach on diverse mathematical reasoning benchmarks, including MATH (Hendrycks et al., 2021), Olympiad-Bertt Bench (He et al., 2024), and AIME (Di Zhang, 2025). On MATH-500, our basic thresholding policy is 5x more cost-effective than existing methods, while our trained RL. and POMDP policies achieve the performance of the expensive L.MM using only 20% of the expensive tokens. On the more challenging AIME benchmark, these policies achieve 3. and 6.33x higher cost efficiency, respectively. Moreover, routing policies trained on one dataset generalize strongly others: the RL-trained policy attains up to 11.68x higher cost efficiency when evaluated on OlympiadBench, despite being trained only on AIME. Our main contributions are: (1) We provide systematic evidence showing that small number of targeted step-level interventions can substantially enhance the efficiency of multi-step reasoning. (2) We demonstrate that this insight can be used to instantiate practical frameworks that are amenable to several routing policies, from simple thresholding policies that surpass existing query-level routing methods to advanced RI.-trained and POMDP-based approaches that achieve competitive performance with oracle routers having perfect task knowledge, all while using significantly fewer expensive model tokens. Our routing policies can be trained effectively with limited supervision data while achieving robust performance across diverse cost budgets, making the approach practical for weal-world deployment scenarios. (3) Finally, we demonstrate robust generalization of the routing strategy across datasets, with methods trained on AIME delivering strong efficiency gains on OlympiadBench and Minerva Math, suggesting that step-level difficulty patterns capture transferable structure across related reasoning benchmarks for a given base model, rather than being tightly coupled specific dataset.
2 RELATED WORK
The trade-off between model performance and computational cost has become increasingly salient as large language models grow in size and capability. Our work on targeted stepwise routing builds upon
2

---

Published as a conference paper at ICLR 2026
THUM-TRR va Oratto Query-Lovel Router on GSMIK
TRIAS Orucle QuerryLLvvel Rloter on Moth-SOt
THUM-TTe va Oracle Query-Levsl Bouter on Nitserva Math
0.572
0.03
3.822
Accumaas
1962 41 Agg Maoee
0.90
1121 Naaae)
Wong
0.546
0.31
6.70
0.33
0350
0.53
Seringe Teloss Crod per Queedion
Aereage Saken Cood per Queeslaen
Aerringe Tikee Godd per Quandiion
Figure 2: Comparison of task performance-cost trade-offs for Brocn2.5-3B-Intrruue (M.r) and Claude 3.7 Sonnet (M.), under the myopic thresholding policy (with Dween2.-MMat-PRMM- versus the Idealized Oracle Query-level Router, across multiple math benchmarks. The Oracle Router evaluated by incrementally varying the number of queries routed to Ms. selecting those solvable by M, but not by Mu. 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.
prior work from several areas: LLM routing strategies, process supervision and step-level verification, and multi-step reasoning optimizations.
LLM routing and model selection. Traditional routing approaches operate at the query level, assigning entire queries to single model based on estimated difficulty or uncertainty. Under this query-level paradigm, Hybrid-LLM (Ding et al., 2024) frames routing as a classification task using BERT-style encoders, while Zooter (Lu et al., 2023) employs swarrd-guide training for normalized reward prediction. RouteLLM (Ong et al., 2024) learns routing directly from preference data, and AutoMix (Aggarwal et al., 2023) formulates query-level routing as a POMDP with self- verification for difficulty estimation. Recent advances extend these approaches in complementary directions. BEST-Route (Ding al., 2025) augments routing with adaptive test-time compute, jointly selecting a model and number of samples to generate per query (best-of-1z) based on query difficulty and quality thresholds However, all these methods assume uniform difficulty across generation steps, leading to inefficient resource allocation in multi-step reasoning tasks where difficulty varies significantly across steps. In such settings, errors are often localized, and correcting a few critical steps often sufficient to steer the solution back onto successful path. In contrast, TRIM departs from query-level routing by operating at the granularity of individual reasoning steps, enabling targeted interventions precisely where stronger models most valuable while avoiding unnecessary escalation elsewhere.
Process supervision and step-level verification. Process reward models (PRMs) have emerged as a powerful technique for evaluating intermediate reasoning steps. Human-in-the-lowwe step verification method (Lightman et al., 2023) showed that process supervision significantly outperforms outcome supervision math reasoning tasks, motivating scalable automated alternatives. Subsequent work developed automated process supervision (Luo et al., 2024) and stepwise verification without human annotations (Wang et al., 2023). Recent studies further demonstrate that strategic computation allocation (Hwang et al., 2024; Snell et al.. 2024; Setlur et al., 2024b) can significantly improve mathematical reasoning efficiency. While these works primarily use PRMs for candidate selection or exploration shaping, our work instead uses PRM scores to inform routing decisions during generation.
Multi-step reasoning and test-time compute. Recent work has highlighted the importance of strategic computation allocation in multi-step reasoning. Research on "forking paths" (Bigelow et al., 2025) and (Wang et al., 2025) demonstrates that certain tokens represent critical decision points that dramatically alter solution trajectories. Similarly, work on reinforcement learning for mathematical reasoning (Setlur et al., 2024a; DeepSeek-Al et al., 2025; Yang et al., 2025; Team et al., 2025; Yu et al., 2025; Lin et al., 2024) and optimizing test-time compute (Snell et al., 2024; Qu et al., 2025) shows that targeted interventions at crucial steps can be more effective than uniform computation increases. The observation that high-entropy minority tokens drive effective learning (Wang et al., 2025) further supports the notion that not all generation steps are equally important.
Speculative and parallel decoding. Our approach shares some similarities with speculative decod- ing (Leviathan et al.. 2023) and parallel inference strategies (Cai et al., 2024), which also involve multiple models collaborating during generation. However, these methods focus on acceleration
3

---

Published as a conference paper at ICLR 2026
through draft-and-verinn paradigms, while our work targets the quality-cost trade-off in multi-step reasoning through selective model escalation. Neward-Guider SD (Liao et al., 2025) and SpecRea- son (Pan et al., 2025) are closer in spirit in that they use step-level signals to accept regenerate intermediate generations (RSD via PRMs, SpecReason via expensive model- -ene--geee scores), but they differ fundamentally in objective. Both operate in a fixed, high-budget regime and are optimized primarily for latency reduction with marginal accuracy gains, whereas TRIM explicitly optimizes accuracy under constrained cost budget by limiting strong- moodel token usage. Although RSD and SpecReason can be adapted to routing by varying acceptance thresholds, such adaptations remain inherently myopic: decisions depend only on the current step's s score and do not account for noise in conrectness estimates or long-horizon budget allocation. By contrast, TRIM is more general and com- plementary: it subsumes threshold-based rules while enabling more expressive policies that reason about trajectory-level progress, uncertainty in step-level correctness, and expected future benefit of intervention, leading to more robust and efficient compute allocation in multi-step reasoning.
3 PRELIMINARIES AND PROBLEM STATEMENT
We define stepwise routing within multi-step reasoning task as sequential decision process. The objective is to derive a routing policy that, at each step of generation, decides whether to (i) accept the output of a cheap language model, or (ii) re-generate the step using more capable but expensive LLM, thereby incurring an additional per-token cost. This policy must balance the trade-off between maximizing the task reward of the final solution and minimizing the serving cost, defined as the number of tokens generated from the expensive, stronger LLM. We formalize this problem below.
Problem Formulation. A multi-step reasoning task is specified by query q Q and a sequence of reasoning steps
YI.N (g. V1 1/2 .... UN) E VN.
where Y44 denotes the i-th reasoning step. For our purposes, we assume these steps are delimited by double newlines in the generated text. At time-step the current prefix is denoted yia (g, ,1.....9tt E Vt, where 3% denotes the set of all prefixes of length t. Within Vv, we distinguish two disjoint subsets: (1) Pa S Vr, the set of incomplete prefixes (partial answers) that can be extended further; (2) C, S Vv, the set of completed (terminated) answers at step t.
Let M denote set of language models, where each model M E .M maps a partial reasoning trace Y1:t E PP to the next step y1:t-1 E Vt+1: extending the prefix or producing completed solution. We consider a two-model setup consisting of two classes of models M: (1) strong expensive L.L M, that produce high-quality responses but incurs per-token cost (2) cheap LLM M,. which offer relatively lower-quality responses at negligible cost. This abstraction captures the trade-off between generation quality and inference cost that motivates routing across heterogeneous language models.
We aim to learn a routing policy ग that, at each reasoning step t, chooses an action a e regeeeraaeeee to either accept the weak model's continuation or replace it with the strong model's generation. Formally, for Ma E Mn and M, E M,
If o; continue, the next prefix updates as V1.++1 (y1.,, (x.:)) If a, = regenerate, the policy instead sets y'., = (y1n-1, (,(Yra1))) and then continues with V1.t- = (y'., M (y
For Y1a Ce. choosing a regenerate yields yí.-- (y1:-1, M, (y while choosing a-- continue leaves the prefix unchanged; in either case, the reasoning process terminates.
The quality of intermediate steps in a reasoning trace can be estimated by assigning probabilistic scores. This can be achieved through methods such as seff-verification, process reward models (PRMs), or other step-level evaluation techniques. For our experiments, we adopt PRM that evaluates partial traces and produces a sequence of step-level scores. Formally, given a reasoning trace y1.t, the PRM assigns step-level rewards as riæ = PRM(Y = (r1, rt), where r, = (xI.t) denotes the reward for step i.
We use these scores as proxies for step-level correctness, providing a signal for whether escalation to a larger model is likely to offer additional benefit given the current prefix. In practice, the PRM score for a solution can also be aggregated across steps using either the product of all step scores or the minimum score across steps (Wang al., 2023; Lightmnn et al., 2023). These aggregated scores widely used in practice for ranking and comparing multiple candidate solutions. Although PRMs have been used in prior work primarily to improve candidate selection in beam search to shape
4

---

Published as a conference paper at ICLR 2026
Step-Wing Routers
Problem Statement
de lulton
Observations
2
Step N-1
Process Revord
Reward Tekkee
20H Renarce
Amang hit Stupa
Min PRM Reward
Segenerste Nit Step
With
Nith Stop
Of
Contirus With
Reward Tekens
Index & Toleens Used for han Step
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.
exploration in search and RL (Snell et al., 2024; Setlur et al., 2024b), our use is distinct: we leverage PRM outputs inform routing decisions during generation.
TRIM AND ROUTING STRATEGY DESIGNS
We introduce TRIM, framework for step-level routing within each query, illustrated in Figure I. Instead of routing an entire query to strong model, TRIM operates the level of individual reasoning steps. At each step of the reasoning process, the cheap model Mw proposes candidate continuation yt"" (Y1:t11) The router policy then evaluates the partial reasoning trace together with y/" and either accepts this step (a, = continue) or escalates to the strong model M, (a, = regonerate), which regenerates the step as yf = M, (y ). Based on this decision, the appended step is y/t = yi" if contiiiiinnnn conntinueei and = V/ otherwise. Thus, TRIM incrementally constructs the solution trace by appending at each position either the M -generrat step or the MM-regenerated step, depending on the router's ' action (see Appendix D for a detailed discussion of the action design). This contrasts with query-level routing, which which commits to either Mi, or M, for the entire generation.
We now describe different strategies for learning routing policies within TRIM. These strategies differ in how much information they incorporate about the reasoning trajectory and whether the decisions myopic (local given step) non-myopic (awareness of the full trajectory). On one end of the spectrum, thresholding policies rely only on current step correctness, while other approached the utilize RL training and POMDP-based formulations account for long-horizon accuracy/cost tradeoffs.
4.1 TRIM-THR: MYOPIC THRESHOLDING POLICY
We first introduce simple yet effective routing policy that solely relies the PRM score of the current generated step of the cheap model My to make routing decisions. If this probability falls below predefined threshold kt the router regenerates the step with the strong model M,; otherwise, it accepts the cheap model's output and continues generation. Adjusting the threshold parameter provides principled way vary the task performance-cost trade-off. Our TRIM-Thr can be viewed as an adaptation of the fixed-threshold mechanism in Liao et al. (2025) to the routing setting, where thresholds vary with the cost budget. Formally, the policy
regenerate, if PRM(y1a)) < k. continue, otherwise.
The (V1.)) =
4.2 RL-TRAINED POLICIES
While TRIM-Th Thr demonstrates strong performance (Figure 2), it is inherently myopic in the sense that it makes routing decisions solely based on the correctness estimate of the most recent step, without considering past context or future consequences. A richer set of signals can be exploited for more effective decision-making
For instance, even if the most recent step predicted to be incorrect, regenerating it with a strong model may not be beneficial if the overall trajectory is already far from the correct solution, or if the additional cost of intervention outweighs the potential gain. Conversely, when prior steps are largely consistent and the trace remains plausibly aligned with correct solution, targeted intervention can be highly impactful.
Prior work stepwise verification (Lightman al., 2023; Wang et al., 2023) and beam search reranking highlights the importance leveraging the sequence of correctness scores accumulated over the reasoning trace, rather than focusing exclusively on the most recent one. Additionally, token
5

---

Published as a conference paper at ICLR 2026
counts at each step provide information about the cost of regenerating with Mg. Incorporating these richer signals allows the router to reason jointly about (i) whether the trace remains plausibly on track toward correct solution and (ii) whether the cost of intervention is justified, enabling more principled stepwise routing policies beyond TRIM-Thr.
TRIM-Seq: Learning to Route from Sequential Features. Formally, let C1:t = (C1y.... Ct) denote the token counts for each step y1.....9/ V/. in the reasoning trace Y1:2 and Fi:t = (fi,.... .r) be the stepwise correctness scores, then joint feature sequence is f1 = ((11. eI),.... (T,, These features capture the two quantities that are fundamental to routing decisions in multi-step reasoning: (i) semantic fidelity, via step-level correctness estimates that indicate whether the current trajectory remains on track, and (ii) marginal intervention cost, via token counts that quantify the expense of regenerating a step with the strong model. Together, they are sufficient to represent the relevant trade-off faced by the router-whether an intervention is likely improve final correctness enough to justify its cost-without requiring access to raw token sequences or model internals.
We parameterize the routing policy with a transformer-based network that processes the feature sequence f1x and outputs distribution over actions E (continnue regenerate). The policy is optimized via RL: each regeneration action (regenerate) on prefix Y1: incurs cost proportional to the number of tokens generated by the strong model M.. À M, (y 1)), where À >> 0 is a cost- performance trade-off parameter, and (y denotes the token length of the strong model's regenerated output. We measure cost solely in terms of large-model decode tokens, as these constitute the dominant and irreducible component of inference latency and compute. While prefill (K V-cchhe construction) can be amortized via chunked prefilling, executed in parallel with small-mode decoding- similar to the parallelization used in speculative and draft-and -ve-iffy decoding --r---rr---ooo tokens impose unavoidable sequential cost and cannot be parallelized across tokens. The episodic return further includes the (binary) terminal task reward R. which reflects the correctness of the final solution. The policy is thus optimized to maximize the expected return
decode
T
J(T) = E, R(yI.) A
Σ 1aa regenerate]- |M
which balances solution correctness against the cumulative generation cost of invoking M,.
Learning to Route from Aggregated Features. TRIM-Seq leverages the full sequence of stepwise correctness scores (r1, (2.,/-....p ."2..--,,tt-1) together with token lengths to model routing deci- sions. While this provides rich sequential context, simpler feature representations can capture the most salient statistics of the reasoning trace. In multi-step reasoning, errors exhibit a compounding structure: a single incorrect step can invalidate the entire solution (making the minimum score a useful "weakest-link" indicator), while sequence of mildly suspect steps can collectively cause the trajectory to drift off course (making the product of scores useful proxy for cumulative validity across the trace). Prior work (Lightman et al., 2023; Wang et al., 2023) formalizes this intuition, showing that solution-leve correctness can be effectively estimated from stepwise PRM scores via reductions such as the minimum or product. Motivated by this, we construct a reduced feature set fie = (rt, min(ri:e-1), Ct, t), where min(r = min(r)...., Tt-1). ct denotes the token length of the current step, and indexes the position in the trace. This reduced representation discards the full sequential history while retaining key aggregated indicators of correctness and cost. Using fis. we train a policy network with the same RL objective as TRIM-Seq. Empirically, this design yields substantially faster training, with negligible performance loss across all tradeoff parameters X.
TRIM-Agg:
4.3 PRIM-POMDP: POMDP- -BASED ROUTER POLICY
final remaining challenge stepwise routing methods discussed above arises from the imperfect nature of process reward model (PRM) estimates. While PRMs provide informative signals about the correctness of intermediate steps, their predictions are often noisy and can misclassify correct steps as incorrect (or vice versa). When trained on large amounts of data, routing policies like TRIM- Agg can implicitly learn to discard errors in the PRM estimates. However, RL under long- horizoo sparse rewards makes training such policies both cample-inefficien and expensive. This raises key question: how can routing decisions be improved when only noisy correctness signals are available?
Our solution is explicitly treat PRM scores as imperfect observations of an unobserved latent state reflecting true correctness of the reasoning trajectory. and attempt to infer the true latent space first when learning routing policy (akin to control in a partially-observel MDP). As shown in Figure 4,

---

Published as a conference paper at ICLR 2026
Conte-Conddtoeedd PRM Scores (Observation Function TRIM-POMDP)
Observations
State
State
Min PRN Reward Among Nirs Steps
PRM Reward for Nih Step
Step Index & Takens Used for With Stop
62-
o5
Cseeusithe Cormetens Noary
Concedatic Canectano
Price
paTrereect correy)
(oilese Decomeen deweevtooocc
Sp All Prior Stapa Ne-1) Cornect
One or More of First N-1 Staps Are incornect
All Prlor Stops Are Correct
Nith Step Correct
Nth Stop Inrolevant
Nth Stop is Incorrect
Step Indie N- Todaan Uieed for Nit Stop
Step Index N&& Takens Step
Step Index N& Tokera Used for Nith Step
130
True Underlying Stotes Representing Ground Truth Correctness
Figure 4: POMDP state space S classes and obser- vation space Q in TRIM-POMDP. The latent state space S consists of three correctness classes: So (tra- jectory correct so far). S1 (trajectory irrecoverably in- correct due to carlier error), and S₂ (most recent step incorrect prior steps correct and thus poten- tially recoverable), augmented with current step index and token cost. The observation space Q comprises the PRM-based cumulative correctness score of the prior steps, PRM score of the current step, and auxiliary fea- tures. which together serve as noisy observations of the underlying latent correctness state.
Figure Learned observation function illustrat- ing PRM noise. Hcatmaps show the empirical proba- bility density of PRM-based observations conditioned on each latent correctness state, estimated from Pro- cessBench (Omni-MATH), While density concentrates around state-consistent regions (e.g., near (1, 0) for S2. 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.
TRIM-POMDI defines the latent state using three correctness classes (augmented with the current step index and token cost): (i) Sp. where the trajectory remains correct so far, (ii) Si, where the trajectory has already diverged irrecoverably, and (iii) S2, where the most recent step is incorrect but prior steps are correct, leaving the trajectory still potentially recoverable. If this latent state were perfectly observed, the routing problem would reduce to solving fully observable MDP. In practice, however, the latent correctness state is hidden, and we only observe noisy proxies provided by the PRM. To bridge this gap, we learn an observation function that helps us map the entire history of observations (f)::) to probability distribution over the latent states. Concretely, this amounts to modeling the distribution of PRM outputs conditioned on state states. Figure 5 visualizes these conditional distributions, showing that while PRM scores tend to concentrate around state-conan consistent regions, they exhibit substantial spread, reflecting noise in step-level correctness estimates. This observation function be fit offline using process supervision datasets with ground-truth step-level annotations, such as ProcessBench (Zheng et al., 2024). Moreover, because the observation function only requires aligning PRM scores with annotated correctness labels, it can be trained once and reused across different performance-cost trade-off parameters A. Once learned, we can invoke a POMDP solver on-the-fly to compute routing policies that optimally balance accuracy and cost. A detailed analysis of TRIM's 's robustness to PRM noise and miscalibration is provided in Appendix E.
This compact POMDP formulation of the sequential routing problem enable efficient policy computa- tion using standard POMDP solvers. Moreover, policy computation with modern POMDP solvers both efficient and flexible, with offline solvers typically requiring less than minute runtime. As a result, policies can be recomputed easily for different performance cost trade-off parameters X. further advantage of TRIM-POMDP is that the resulting routing policy is largely agnostic to the specific choice of LLMs (Ma, M.), depending only their next-step accuracies provided inputs to the transition function. See Appendix for further details and the complete POMDP formulation. Figure 3 illustrates the end-to-end TRIM pipeline, showing how step-level correctness signals and token costs from partially generated solutions are transformed into routing observations and used by different router instantiations (TRIM-Seq, TRIM-Agg, and TRIM-POMDP) to decide whether to regenerate a specific step with the strong model or continue with the cheap model; we next evaluate how these design choices translate into performance- cost trade-offs.
7

---

Published as a conference paper at ICLR 2026
Performance-Cost Trade-off MATH-300
Performance-Cost Trade-off AIME
Strong (809.26 Tokens)
LLM
0.00
(462.11 Ang Tokcess)
0.25
Accurncy
0.75
Accurary
0.20
0.35
ITUSS-US
AHIN-ADM
0.65
THUMS
50
1011
100
200
3000
Aserage: Token Cosk per Qunstion
Average Tokmn Cood per Question
Figure 6: Performance-con trade-off's of TRIM routing approaches MATH-500 and AIME. TRIM-Thr provides strong low-complexity baseline, while TRiM-Age and LLLI-POMDP offer best cost pest-performace trade-offs, achicving near-M, M, accuracy with a small fraction of expensive tokens. Notably, TRIM-POMDP performs particularly well in low-budget regimes, whereas TRIM-Agg slightly dominates at higher budgets.
5 EXPERIMENTS
We evaluate TRIM by benchmarking its stepwise routing strategies against established query-level routing approaches, primarily RouteLLMM (Ong et al., 2024), Smoothie (Guha et al., 2024), and AutoMix (Aggarwal et al., 2023). RouteLLM is state-of-the-art approach for query-level routing, which uses preference data for making these decisions. Smoothie instead adopts label-free, un- supervised routing approach, fitting latent-variable Gaussian model over output embeddings to estimate per-sample model quality and select the most suitable LLM. AutoMix, in contrast, adopts non- preddiiiee routing paradigm: it first utilizes smaller, cost-efficient models and then few-shot self-verificatto signals as observations in POMDP policy that decides whether escalation to a larger model is necessary. We begin by introducing the evaluation metrics used throughout our analysis.
Metrics. We evaluate router policies by quantifying the trade-off between task performance and the cost of invoking the strong model M.. In our setting, the cost of query is measured by the number of tokens generated by the expensive (strong) model M,. Importantly, TRIM does not introduce significant wall-clock overhead: when implemented using the same system-level optimizations employed by low-latency speculative decoding, TRIM is often faster than running M, alone (see Appendix A). We adapt evaluation metrics from prior work to the setting of stepwise routing in multi-step reasoning. For query q in the set of queries Q, let C(z; T) denote the number of tokens generated by the strong model M, under router policy स, and let C.(g) be the number tokens generated when using M, alone. We define ((rr) the average number of M, tokens per query and c(T)) as the normalized fraction of tokens from M,;
Σ C(giT), d(x) = Enco C(T:T) DE
((T) = Q
(2)
Eyee Cz(g)
Following Ong et al. (2024), if s(g; E {0, 1} denotes the correctness of query q under x. the average performance and the performance gap recovered (PGR) are defined as
r()) r(Mw)
(T)
s(g:T), PGR(T) -(M,) T(M 9
where r r(M,) and r(Mw) denote the accuracies of the M, and M.. respectively. PGR(T) quantifies how much of the performance gap between Ma, and M, is recovered by ा.
To capture the cost achieve desired level of performance, utilize cost-performnnce threshold (CPT). Specifically, CPT(~ denotes the minimum token cost (in terms of C or c) required by policy и achieve PGR of æ%, providing measure of efficiency different target performance levels. Finally, following Aggarwal et al. (2023), we report the incremental benefit per cost (IBC) as: r(M,) r(Mw) 1
IBC(T)) r(m) r(Mw)
AIBC(T) IBC(TT) IBC
IBC
((T)
2 Ege Co(g:T)'
IBCRase Base
(4)
8

---

Published as a conference paper at ICLR 2026
MATH-500
AIME
Method
(PT(
(90%))
(95;
Aree
CPT(50))
CPT(80%)
(P6())))
Arac
BERT
(92.022)
131.45 (71.6%%)
685,49
0.08
331.85 (38)18%)
616.53 (70.93%))
701,58 (80.715))
0.44
MF
160.83 (34.78%)
324.13 (70.10%
(93.55)))
0.49
358.81 (11.28%%
602.62 (69.32%)
813.28 (93.56%)
0.65
SW Ranking
(40.1444
330.89 (71.66))
297.08
496.17 (57.15%)
715.72
0.79
Smoothii
220.69 (87.73%)
345.19 (74.65%)
433.77 (93.81%)
0.30
396. (45.65%)
704.50 (81.05%)
822.09 (94.57%)
0.03
Autoemix
175.67 (37.99%)
314.11 (67 93%%)
120.87 (91 (91.02%))
0.12
450.3 (51.8%)
715.31 (82.29%)
857.89 (98.69%)
0.0001
AutoMix-PRMM
110.42 (23 8
198.73 (42 98%)
249, (53 9668)
0.95
360 (43.7778)
605, 83 (9%)
703.77 (80) 96%)
0.07
TRIM- The
43.68 (9.45%)
(55.995)
(23.4798)
218)
1.81
TRIM-Age
(7.3%)
56.49
(17.21%%
(12.35%%
241.55 (27.795)
(38.015%)
2,50
TRIM-POMDP
29.27 (6.33%)
66.63 (14.411%)
83.12.(17.9))
5.866
139.21 (16.01%)
206.06 (23.71 %))
244.86 (28.17%)
5,00
Table 1: Comparison of A'RII across AIME & MATH-500 benchmarks.
OlympiadBench
Minerva Math
Method
CPT(30)))
(S0()))))
CPT(
Agac
CPT(50%)
CPT(8G)
(P6())
Arre
BERT
367.25 (55.03%)
584.14 14 (87.41%)
642.50 (96. 14%)
-0.04
209.99 (48, 90%)
578.24 (88.25%)
421 241144 (98.33%)
-0.1
MF
369.15 (55.24%)
522.68 (78.21%)
601.77 (90.05%)
-0.07
166.99 (38.96%)
249.82 (58.28%)
326.05 (76.07%)
0.42
SW Ranking
351.34 (52.57%))
511.08 (76.48)
635.05 (95.0235))
0.07
212.73
(29.63%))
342.71 (79.96/9)
421.17
0.04
Smoothie
348.59 (52.16%)
511.64 (76.56%)
615.49 (92. 10%)
-0:08
(54.75%)
(80.539%)
402.19 (93.87%)
40.09
Automix
300.4 (44.95%)
475.37 (71.13%)
644.46 (96.445)
0.02
85.35 (19.91%)
193.04 (45.04%%
231 231.50 (54.03%)
0.72
AutoMix-PRM
265 95 (39.8%)
411.47 47 (61.57%)
481.49 (72.055))
0.22
72.08 (16.82%))
1401.24 (32.725%
196, 12 (45. (45.76%))
1.35
TRIM-The
136.64 (20.4598)
220.70 (33.03%%
$13.89 (46.979))
1.31
(15.2%)
92.78 (21.6898)
44.555
2.23
TRIM-Agg
(14.13%)
190.11 (28.45%)
287.17 (42.97%)
2.57
47.37 (11.05%)
(20.89%)
138.67 (32.35%)
3.12
Table 2: Cross-Benchmark Generalization of Routers Trained on AIME
Here, IBC(TT) measures performance improvement per unit expensive-model M, token usage, IBC is the baseline corresponding to always using M,, and (T)(TT quantifies relative gain. A positive ABC indicates more effective gains than standalone LLM. For evaluation, we compute AIRC across 100 equally sized performance regions between M.. and M, and report the average.
Experimental setup. We use two-model setup with Frren2.5-3BIIntruuce as the cheap LLM (Mw) and Claude 3.7 Sonnet as the expensive model (Ma). guided by level correctness estimation. For AIME (Di Zhang, 2025), we use an approximately 50-50 train-test split across alternate years and problem sets; for MATH (Hendrycks et al., 2021), we train on the 7.5k official training set and evaluate on MATH-500 We benchmark three TRIM routing strategies (TRIM- Thr, TRIM-Agg, and TRIM-POMDP) against Auto Smoothie, and RouteLLM (Ong et al., 2024) approaches (BERT classifier, matrix factorization, and SW ranking); see Appendix G for the detailed experimental setup. To enable fair comparison with TRIM and strengthen AutoMix, we introduce our own adapted variant of AutoMix, in which we replace the original verifiiciiitton component with cumulative PRM scores during both training and evaluation. This modification strengthens AutoMix, as PRM-based verification desiineed for stepwise mathematical reasoning- -yelds more informative correctness signals than generic self-verification leading AutoMix-PRM to consistently outperform vanilla AutoMix across benchmarks (Tables 1 and 2).
for step-
Results. Table 1 reports the evaluation results across benchmarks, and Figure 6 illustrates the performance-cost trade-off curves achieved by different TRIM strategies. To further assess the generalization capability of TRIM-Agg, we evaluate routers trained on AIME in cross-dataset setting, testing on other math benchmarks of comparable difficulty, namely OlympiadBench and Minerva Math, and report results in Table 2 and the corresponding performance curves in Figure 7. Additional comparisons based on budgeted accuracy are provided in Appendix C.
Our experiments reveal distinct strengths across regimes. As shown in Figure 6, in the low-budget setting (large A), TRIM-POMDS achieves superior performance via principled long- horizon planning under uncertainty. Unlike RL-trained policies, which struggle in this regime due to sparse rewards, modern POMDP solvers efficcintly compute policies without being hindered by sparse-reward learning dynamics. In the high- high-budget regime (small A), however, RL-trained TRIM-Agg policies performs strongly, achieving 95% of the performance gap for MATH-500 while using 80% fewer expensive tokens, as policy optimization becomes significantly casier. Even our simplest approach, TRIM-Thr, achieves 5x (Aic = 4.75 vs ABSC 0.95) better cost efficiency than baselines.
Beyond budget regimes, our cross- datssett evaluations reveal a key distinction: predictive query- level routers (e.g.. RouteLLM) often fit to the intrinsic characteristies of specific datasets, while TRIM captures transferable routing behaviors that generalize across benchmarks of comparable difficulty. Query-level routers rely on coarse input-level signals and may exploit superficial correlates
9

---

Published as a conference paper at ICLR 2026
Performance-Coo Trade-off for OlympiadWeecc
Performaanc--Cot Minerra
Expensive Streng LLM Avg
(428.62 62 Aug Tolkoss)
Pxpensdive Streng 11M
Accurac
6.40
Aucuracy
1200
0.35
1130
0.30
1H0M-A30
100
2000
300
50
100
Average: Token Cost per Questton
Average: Token Cost per Question
Figure 7: Performance-cost trade-offss under dataset generalization. TRIM-Agg routers trained on AIME demonstrate strong performance across benchmarks of similar difficulty
of difficulty (e.g., dataset style, formatting, or typical problem length) that do not transfer, whereas TRIM conditions decisions on step-level correctness signals within the evolving solution trace, which more directly reflect universal failure modes in multi-step reasoning (e.g.. divergence at critical steps) and therefore generalize better across datasets. For instance, BERT achieves AIBC of 0.44 on AIME but drops dramatically to -0.04 on OlympiaBBench and -0.1 on Minerva Math, while SW Ranking similarly degrades from 0.79 0.07 and 0.04, respectively. In contrast, TRIM-Age achieves ÂJвC of 2.5 on AIME and also maintains strong performance with ABc values of 2.57 on OlympiadBench and 3.12 on Minerva Math when trained solely on AIME, demonstrating superior generalization.
Although AutoMix (both vanilla and PRM variants) generalizes better than RouteLLM, its in- distribution performance on AIME notably weaker; nonetheless, our PRM variant consistently outperforms vanilla AutoMix. We attribute this gap less reliable cumulative correctness score estimates on harder problems, limiting routing accuracy despite explicit noise modeling in the correctness estimates. Overall, these findings suggest that step-level difficulty patterns captured by TRIM reflect fundamental properties of multi-step reasoning rather than dateset-specific artifacts.
Despite being trained on fewer than 500 AIME samples, TRIM-Ags achieves strong performance on the held-out test set with 38.01% % expensive token usage at CPT(95%) and exhibits robust generalization to datasets of comparable difficulty, consistently surpassing both TRIM-Thr and query- level baselines. In parallel, TRIM-POMDP performs strongly across all cost pere-ormaraaee offee-off regimes on both MATH-500 and AIME, despite its observation function being trained on different (but comparably difficult) math datasets (detailed in Appendix G) and the solver requiring only the estimated next-step accuracies of (M M,) as input. Additional evidence of TRIM's robustness across different model pairs and its extension beyond math reasoning is presented in Appendix F.
6 DISCUSSION AND CONCLUSION
In this work, we present TRIM, an approach for targeted stepwise routing that escalates only critical steps to stronger, more expensive LLMs, intervening precisely where the partial reasoning trace risks diverging from correct solution. Our key insight is that even small number of well-placed interventions can dramatically boost task accuracy, enabling significant efficiency gains compared to conventional query-level routing. Building on this insight, we designed multiple routing strategies for TRIM that differ how much trajectory-level information they exploit when making routing decisions. While TRIM already surpasses contemporary routing methods and performs competitively with oracle query-level routers, further improvements may be possible by moving beyond stee- granularity to token-level routing. Since certain tokens disproportionately influence downstream generation (Wang et al., 2025), token-level routing offers a promising direction for achieving even finer-grained and cost-efficient interventions.
10

---

Published as a conference paper at ICLR 2026
7 REPRODUCIBILITY STATEMENT
In order to foster reproducibility of our work, outline implementation details of our approach in Appendix G and Section 5.
ETHICS STATEMENTS
All authors have read and agree to abide by the ICLR Code of Ethics. We acknowledge the usage of LLMs for editing (e.g.. grammar, spelling, word choice) purposes for polishing writing. However, the use of LLMs was limited to editing and formatting text only.
REFERENCES
Pranjal Aggarwal, Aman Madaan, Ankit Anand, Srividya Pranavi Potharaju, Swaroop Mishra, Pei Zhou, Aditya Gupta, Dheeraj Rajagopal, Karthik Kappaganthu, Yiming Yang, et al. Automix: Automatically mixing language models. arXiv preprint arXiv:2310.1296. 2023.
Eric Bigelow, Ari Holtzman, Hidenori Tanaka, and Tomer Ullman. Forking paths in neural text generation. In The Thirleenth International Conference on Learning Representations, 2025. URL
wew..wwwwe.....................w net/ / forum? weteetddddddiiiee 3=8RCmNLeeexx
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple IIm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401. 107744 2024.
DeepSeek-AI, Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang, Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongjie Ji. Erhang Li, Fangyun Lin, Fucong Dai. Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Qu, Hui Li, Jianzhong Guo, Jiashi Li, Jiawei Wang. Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, J. I.. Cai, Jiaqi Ni, Jian Liang, Jin Chen, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang. Kuai Yu, Lean Wang. Lecong Zhang, Liang Zhao, Litong Wang. Liyue Zhang, Lei Xu, Leyi Xia, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Meng Li, Miaojun Wang, Mingming Li. Ning Tian, Panpan Huang, Peng Zhang, Qiancheng Wang. Qinyu Chen, Qiushi Du, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, R. J. Chen, R. L. Jin, Ruyi Chen, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shengfeng Ye, Shiyu Wang, Shuiping Yu, Shunfeng Zhou, Shuting Pan, S. S. Li, Shuang Zhou, Shaoqing Wu, Shengfeng Ye, Tao Yun, Tian Pei, Tianyu Sun, T. Wang, Wangding Zeng, Wanjia Zhao, Wen Liu, Wenfeng Liang, Wenjun Gao, Wengin Yu, Wentao Zhang, W. L. Xiao, Wei An, Xiaodong Liu, Xiaohan Wang, Xiaokang Chen, Xiaotao Nie, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, X. Q. Li, Xiangyue Jin, Xiaojin Shen, Xiaosha Chen, Xiaowen Sun, Xiaoxiang Wang, Xinnan Song, Xinyi Zhou, Xianzu Wang. Xinxia Shan, Y. K. Li, Y. Q. Wang. Y. X. Wei, Yang Zhang, Yanhong Xu, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Wang, Yi Yu, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He. Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma. Yiyuan Liu, Yongqiang Guo, Yuan Ou, Yuduan Wang. Yue Gong, Yuheng Zou, Yujia He, Yunfan Xiong. Yuxiang Luo. Yuxiang You, Yuxuan Liu, Yuyang Zhou. Y. X. Zhu, Yanhong Xu, Yanping Huang, Yaohui Li. Yi Zheng, Yuchen Zhu, Yunxian Ma. Ying Tang, Yukun Zha, Yuting Yan, Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhicheng Ma, Zhigang Yan, Zhiyu Wu, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Zizheng Pan. Zhen Huang, Zhipeng Xu, Zhongyu Zhang, and Zhen Zhang. Deepscek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning, 2025.
URL wwttth:ss//aiiiiee org/ abs 12948.
Di Zhang. Aime 1983 2024 (revision 6283828), 2025. URL https: https::/huugiiggaacc.oo. datasets/di -iaa/ddi-----g--ae-- zwazgaaaaaaaee fdu/AIME, 288__024.

---

Published as a conference paper at ICLR 2026
Dujian Ding, Ankur Mallick, Chi Wang, Robert Sim, Subhabrata Mukherjee, Victor Ruhle, Laks VS Lakshmanan, and Ahmed Hassan Awadallah. Hybrid Ilm: Cost-efficient and quality- aware query routing, arXiv preprint arXiv:2404.14618, 2024.
Dujian Ding, Ankur Mallick, Shaokun Zhang, Chi Wang, Daniel Madrigal, Mirian Del Car- men Hipolito Garcia, Menglin Xia, Laks V. S. Lakshmanan, Qingyun Wu, and Victor Rühle. Best-route: Adaptive Ilm routing with test-time optimal compute, 2025. URL https:///axxtt srg/abs/2500 22716.
Maxim Egorov, Zachary N. Sunberg, Edward Balaban, Tim A. Wheeler, Jayesh K. Gupta, and Mykel Kochenderfer. POMDPs,jU jl: A framework for sequential decision making under uncertainty, Journal of Machine Learning Research, 18(26):1-5, 2017. URL http: /mmrr./mr...r///// org argg/papess/v18/ 16-300. htmi.
Neel Guha, Mayee Chen, Trevor Chow, Ishan Khare, and Christopher Re. Smoothie: Label free language model routing, Advances in Neural Information Processing Systems, 37:127645-127672, 2024.
Chaoqun He, Renjie Luo. Yuzhuo Bai, Shengding Hu. Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, al. Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems. arXiv preprint arXiv:2402. /4008, 2024.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv: 103. 2021.
Hyeonbin Hwang, Doyoung Kim, Seungone Kim, Seonghyeon Ye, and Minjoon Seo. Self-explore: Enhancing mathematical reasoning in language models with fine-grained rewards. In Yaser Al- Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 1444-1 1466, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 00.18633vv120222 /2t../////......iinnnna.oaao findings-emmlp.765 URL https:// aclanthology 202/ findings emnlp. 78/.
Hanna Kurniawati, David Hsu, Wee Sun Lee, et al. Sarsop: Efficient point- pomdp planning by approximating optimally reachable belief spaces. In Robotics: Science and systems, volume 2008. Zurich, Switzerland, 2008.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In International Conference on Machine Learning, pp. 19274-19286. PMLR, 2023.
Baohao Liao, Yuhui Xu, Hanze Dong, Junnan Li, Christof Monz. Silvio Savarese, Doyen Sahoo, and Caiming Xiong. howard-guided speculative decoding for efficient Ilm reasoning. arXiv preprint N9NV::2501.19324, 2025.
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lec, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. In The Twelfth International Conference on Learning Representations, 2023.
Zicheng Lin, Tian Liang, Jiahao Xu, Qiuzhi Lin, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, and Zhaopeng Tu. Critical tokens matter: Token-level contrastive estimation enhances Ilm's reasoning capability, arXiv preprint arXiv:241..199. 2024.
Keming Lu, Hongyi Yuan, Runji Lin, Junyang Lin, Zheng Yuan, Chang Zhou, and Jingren Zhou. Routing to the expert: Efficient howard-guided ensemble of large language models. arXiv preprint arXiv:2311. 1.0692, 2023.
Liangchen Luo, Yinxiao Liu, Rosanne Liu, Samrat Phatale, Meiqi Guo, Harsh Lara, Yunxuan Li, Lei Shu, Yun Zhu, Lei Meng, Jiao Sun, and Abhinav Rastogi. Improve mathematical reasoning in language models by automated process supervision, 2024. URL https : / /arxiv...v.//////////a/////avav/ www.wii..rrrr/aaoe /bs// 2406.00 06592.
12

---

Published as a conference paper at ICLR 2026
Isaac Ong, Amjad Almahairi, Vincent Wu, Wei-Lin Chiang, Tianhao Wu, Joseph E Gonzalez, M Waleed Kadous, and Ion Stoica. Routellm: Learning to route Ilms with preference data. arXiv preprint arXiv:2406. 18665, 2024.
Rui Pan, Yinwei Dai, Zhihao Zhang, Gabriele Oliaro, Zhihao Jia, and Ravi Netravali. Specreason: Fast and accurate inferencetimm compute via speculative reasoning. arXiv preprint arXiv:2504.07891, 2025.
Yuxiao Qu, Matthew YR Yang, Amrith Setlur, Lewis Tunstall, Edward Emanuel Beeching, Ruslan Salakhutdinov, and Aviral Kumar. Optimizing test-time compute via meta reinforcement fine- tuning. arXiv preprint arXiv:2503.07577 2025.
Amrith Setlur, Saurabh Garg, Xinyang Geng, Naman Garg, Virginia Smith, and Aviral Kumar. RI on incorrect synthetic data scales the efficiency of Ilm math reasoning by eight-fold. Advances in 2024a.
Neural Information Processing Systems, 37:43000
Amrith Setlur, Chirag Nagpal, Adam Fisch. Xinyang Geng, Jacob Eisenstein, Rishabh Agarwal, Alekh Agarwal, Jonathan Berant, and Aviral Kumar. Rewarding progress: Scaling automated process verifiers for IIm reasoning. arXiv preprint wxXiv:2410.08. 46, 2024b.
Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling Ilm test test-time compute optimally can be more effective than scaling model parameters. arXiv preprint 230. 2024.
Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, Zhuofu Chen, Jialei Cui, Hao Ding, Mengman Dong, Angang Du, Chenzhuang Du, Dikang Du, Yulun Du, Yu Fan, Yichen Feng. Kelin Fu, Bofei Gao, Hongcheng Gao, Peizhong Gao, Tong Gao, Xinran Gu, Longyu Guan, Haiqing Guo, Jianhang Guo, Hao Hu, Xiaoru Hao, Tianhong He, Weiran He, Wenyang He, Chao Hong. Yangyang Hu, Zhenxing Hu, Weixiao Huang, Zhiqi Huang, Zihao Huang, Tao Jiang, Zhejun Jiang, Xinyi Jin, Yongsheng Kang, Guokun Lai, Cheng Li, Fang Li, Haoyang Li, Ming Li, Wentao Li, Yanhao Li, Yiwei Li, Zhaowei Li, Zheming Li, Hongzhan Lin, Xiaohan Lin, Zongyu Lin, Chengyin Liu, Chenyu Liu, Hongzhang Liu, Jingyuan Liu, Junqi Liu, Liang Liu, Shaowei Liu, T. Y. Liu, Tianwei Liu, Weizhou Liu, Yangyang Liu, Yibo Liu, Yiping Liu, Yue Liu, Zhengying Liu, Enzhe Lu, Lijun Lu, Shengling Ma, Xinyu Ma. Yingwei Ma, Shaoguang Mao, Jie Mei, Xin Men, Yibo Miao, Siyuan Pan, Yebo Peng, Ruoyu Qin. Bowen Qu Zeyu Shang, Lidong Shi, Shengyuan Shi, Feifan Song, Jianlin Su, Zhengyuan Su, Xinjie Sun, Flood Sung, Heyi Tang, Jiawen Tao, Qifeng Teng. Chensi Wang, Dinglu Wang, Feng Wang. Haiming Wang, Jianzhou Wang, Jiaxing Wang, Jinhong Wang. Shengjie Wang, Shuyi Wang, Yao Wang. Yejie Wang. Yiqin Wang. Yuxin Wang. Yuzhi Wang, Zhaoji Wang, Zhengtao Wang. Zhexu Wang, Chu Wei, Qianqian Wei, Wenhao Wu, Xingzhe Wu, Yuxin Wu, Chenjun Xiao, Xiaotong Xie, Weimin Xiong, Boyu Xu, Jing Xu, Jinjing Xu, L. H. Xu, Lin Xu, Suting Xu, Weixin Xu, Xinran Xu, Yangchuan Xu, Ziyao Xu, Junjie Yan, Yuzi Yan, Xiaofei Yang. Ying Yang, Zhen Yang, Zhilin Yang. Zonghan Yang, Haotian Yao, Xingeheng Yao, Wenjie Ye, Zhuorui Ye, Bohong Yin, Longhui Yu, Enming Yuan, Hongbang Yuan, Mengjie Yuan, Haobing Zhan, Dehao Zhang, Hao Zhang, Wanlu Zhang, Xiaobin Zhang, Yangkun Zhang, Yizhi Zhang. Yongting Zhang, Yu Zhang, Yutao Zhang, Yutong Zhang, Zheng Zhang, Haotian Zhao, Yikai Zhao, Huabin Zheng, Shaojie Zheng. Jianren Zhou, Xinyu Zhou, Zaida Zhou, Zhen Zhu, Weiyu Zhuang, and Xinxing Zu. Kimi k2: Open agentic intelligence. 2025. URL https: / arxiv. org//ogg//....ooooooo.......o..o
Peiyi Wang, Lei Li, Zhihong Shao, RX Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. Math-shepherd: Verify and reinforce llms step-by-step without human annotations. arXiv preprint arXiv:23 2023.
Shenzhi Wang, Le Yu, Chang Gao, Chujie Zheng, Shixuan Liu, Rui Lu, Kai Dang. Xionghui Chen, Jianxin Yang, Zhenru Zhang, et al. Beyond the 80/20 rule: High-entropy minority tokens drive effective reinforcement learning for Ilm reasoning. arXiv preprint arXiv:2506.0199 2025.
An Yang. Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, Jian Yang. Jianhong Tu, Jianwei Zhang, Jianxin Yang, Jiaxi Yang. Jing Zhou, Jingren Zhou, Junyang Lin, Kai Dang. Keqin Bao, Kexin Yang, Le Yu, Lianghao Deng, Mei Li, Mingfeng Xue, Mingze Li, Pei Zhang, Peng Wang, Qin Zhu, Rui
13

---

Published as a conference paper at ICLR 2026
Men, Ruize Gao, Shixuan Liu, Shuang Luo, Tianhao Li, Tianyi Tang, Wenbiao Yin, Xingzhang Ren, Xinyu Wang. Xinyu Zhang. Xuancheng Ren, Yang Fan, Yang Su, Yichang Zhang, Yinger Zhang, Yu Wan, Yugiong Liu, Zekun Wang, Zeyu Cui, Zhenru Zhang, Zhipeng Zhou, and Zihan Qiu. Qwen3 technical report, 2025. URL https: / / /aaxiv...v.. org/ abs abs/2505 0 09388.
Jiahao Yu, Zelei Cheng, Xian Wu, and Xinyu Xing. Gpo: Learning from critical steps to improve I1m reasoning, arXiv preprint arXiv:2509. 16456, 2025.
Muru Zhang, Ofrr Press, William Merrill, Alisa Liu, and Noah A. Smith. How language model hallucinations can snowball. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. Jog, 2024.
Chujie Zheng, Zhenru Zhang, Beichen Zhang, Runji Lin, Keming Lu, Bowen Yu, Dayiheng Liu, Jin- gren Zhou, and Junyang Lin. Processbench: Identifying process errors in mathematical reasoning. arxiv preprint arXiv:2412.06552 2024.
14

---

Published as a conference paper at ICLR 2026
Appendix Contents
A. Latency Analysis of TRIM
15
B. Formalizing TRIM-POMDP
16
C. Budgeted Accuracy
17
D. Analysis of regenerate Action Design Choice
17
E. Robustness of TRIM to PRM Noise and Miscalibration
Generalizability of TRIM Across Model Pairs and Domains
G. TRIM Implementation Details
H. Effectiveness of Qwen2.55-Mtt--PRM-- for AIME
A LATENCY ANALYSIS OF TRIM
Targeted intervention enables TRIM to reduce cost, which is defined as the number of tokens generated from the expensive, stronger LLM. In practice, TRIM can be implemented using the same system-level optimization techniques that enable low-latency speculative decoding (Leviathan et al., 2023; Liao et al., 2025; Cai et al., 2024; Pan et al., 2025), thereby avoiding frequent context re-encoding ("prefill"" when switching between models and effectively ignoring the associated latency cost overhead. We now provide direct empirical evidence to this claim, demonstrating that TRIM does not introduce significant wall-clock overhead and faster than running the expensive model M, alone.
Experimental Setup. We conducted end-to-end latency and throughput measurements on fixed 2xH100 setup using vLLM with prefix caching enabled. For TRIM-Thr, GPU memory was allocated follows: the large model M, was sharded across both GPUs using tensor parallelism (55% memory utilization per GPU), while the small model Mw and the PRM each ran on separate GPUs with 40% memory utilization. For single-large-ment model baselines, both GPUs were fully dedicated serving the large model. All measurements were conducted on the MATH-500 test set, and we evaluated TRIM-Th under two model-pair configurations:
TRIM-Thr (1,5B 32B): (wwn2.5-115B-Insste Instructt as Mw. and Qwen2.5-32B Instructt as M, TRIM-Thr (1.5B 7B): Qwen2.5-1.5B Instruct M,, and Qwen2.5-7B. 7---n.---B--sttrrett M,
These results show that TRIM-Thr (1.5B 32B) achieves 1.xx2.77x speedup over running the 32B model alone and stepwise routing. Moreover, we observe that latency gains increase with larger strong models and lower routing thresholds.
12 s 1.41, 17.10/6.21 xe 2.75), despite incorporating PRM evaluation
Theoretical Justification. In TRIM, the small, inexpensive model Myr acts the draft generator, while the strong model M, performs chunked prefilling of the ongoing M,, generation to maintain a synchronized KV cache. Consequently, routing decisions do introduce sequential stalls: when escalation occurs, M, can immediately resume decoding from its pre-computed prefix.
Furthermore, for all three TRIM variants (TRIM-Thr, TRIM-Agg. and TRIM-POMDP), the routing decision overhead can be made negligible. TRIM-AgE employs lightweight MLP with two hidden layers of 128 units, whose inference cost is effectively zero relative LLM decoding. For TRIM- POMDP, we use SARSOP, an offline POMDP solver. To achieve the best performance out of the POMDP solver, we recompute the policy at every routing decision by initializing the initial state distribution with the current belief distribution (one can skip recomputation when belief changes small without any loss in performance). However, because the true state space small and each policy computation is fast (about seconds average), this enables mapping belief distributions to actions, resulting in zero latency from the router during inference.
precompute lookup table
With this implementation, overall wall-clock time is dominated by strong strong-model decoding. Because TRIM substantially reduces the number of tokens generated by M,, end-to-end inference is faster than running the expensive model alone.
15

---

Published as a conference paper at ICLR 2026
Table 3: End-to-ent latency and throughput analysis.
Model Configuration
Threshold (k)
Latency (sec/query)
Throughput (tok/sec)
Baselines
Qwen2.5-32B
17.10
31.77
Qwen2.5-7B
10.27
65,20
3*TRIM-Thr (1.5B 32B)
0.1
6.21
64.30
0.4
9.02
52.30
0.7
12.10
47.86
3*TRIM-Thr (1.5B 7B)
0.1
6.35
66.66
0.4
8.17
65,25
0.7
9.51
64.61
B FORMALIZING TRIM-POMDP
Foundations of POMDPs. A Partially Observable Markov Decision Process (POMDP) provides principled framework for sequential decision-making under uncertainty when the underlying system state is not directly observable. A POMDP is defined by the tuple (S, A, R, S2, O), where S is the state space, A the set actions, and Ω the observation space. The transition function T(s'|e ss a) specifies the probability of transitioning from S to s' S given function O maps 8 E S and E to probability distribution over the observation space Ω. The reward function R(s, a) assigns scalar reward state-action pairs.
A, while the observation
Since the true state is hidden, the agent maintains belief state b E ASS), a probability distribution over latent states, updated recursively via Bayes' rule after each obstrvobbeevvttion pair. policy is a mapping π + a, and the objective is to maximize expected discounted return by optimizing over policies.
B. FORMAL DEFINITION OF TRIM-POMDP
We define the components for TRIM-POMDP as follows:
State space (S): As shown in Figure 4, we categorize the state space into three correctness classes: So (all prior steps correct and current step correct), S₁ (at least prior step incorrect), and S₂ (prior steps correct but current step incorrect), along with a terminal absorbing state Ster- Each state is augmented with the step index and the token count o of the current step Yr within the trace y Iir-
Space (Ω): The observation space (noisy state observations) is defined as the set of aggregated features fi.r = (rz. min(ri:t-1). Ce, t), identical to the state features used in TRIM-Age These observations help obtain probability distribution over the state space S.
Observation
Action Space (A): Similar to prior router policies, the action set is A = {continue, regenerate).
Transition Function (T)): The transition dynamics capture the accuracy of M, and My as well as transitions into the terminal state Ster. Formally:
T(s' E So s So U S2 a regenerate) =
(next step accuracy of Ma)
T(s' E So | S Su, a = continue) = Pw (next step accuracy of M.))
T(s' E S s S2, a = continue) =
T(s S a E A, s' S1) =
(irrecoverability assumption)
The observation function a mapping from s E S to the probability distribution over the observation space Ω. The probabilities P(o)s), give us the likelihood of observing Ω (i.e., fi.)) given state E S. This corresponds modeling the distribution of PRM scores conditioned on each state class, which can be learned from process supervision datasets with step-level annotations (e.g., PRM800K).
Observation Function (0):
Reward Function (R): Invoking the strong model incurs a cost proportional to the number of tokens generated, i.e., R( ((s,e a regenerate, 8') - IM, Any transition into the terminal
16

---

Published as a conference paper at ICLR 2026
Table 4: Budgeted-accuracy comparison TRIM across the AIME & MATH-500 benchmarks.
MATHH-901
AIM
Misted
10%
17%
50%
13%4
.1%
15.4
'1/4
1075
31/4
MF
66 22.0
(10 596
33.30% (31,4%)
71.mm6 (3) 6M6:
11, 11.5% (40.599)
72. 3.30% 643. (63.79)
14049 (28 3443
17.01% (33.99)
17.42% (35 4500
15.33 18.35% (60.29)
Rusking
73:nt: 15-
21 11.34%
12:0K5 (14889
(31.75)
(28055)
17:495
18.8%%
Stawahle
11.3%)
53-77% (15.7%)
(14 166
(25
(3--95)
15.39%
(30.55)
Ass0tis PMM
60 083899-00000
(22.44
35:44 (22.99)
12.15% (42 99:
13.28l (52-299)
75 18.38% 161.900 (1199)
12.34% (7.099
88866 (16.99)
14.5% (20.99)
13,899
138 586-999
14.99% (51,79))
TRINS-Agc
15-4200
492
73.00 (20.9))
92 154
66 93.999 1318 (h63
5336495 (1.38 146
15,1490 (43 4h6
29 82 (55.999
22.87% (67.49))
23,229
(59.5%) 599
23.55% (55.99)
THIN POMMe
10.399-440 (64 24g
09.49)
11.3155 1G.
41.80% 053
62.415-(1055) 96
18 110055/03--99)
+92099 (96.15)
909 (65.55)
25.366 (81-5)
78467% (101 25g)
Table 5: Cross-benchmark generalization performance of AlME-trained routers under budgeted-accuracy
evaluation.
Olenpia/Boxd
Minara Mals
Meduol
13%
20%
20%
39/5
DF%
1354
21%
21%
20/5
BEST
28,00)
20.548
MI
2002 (1/19)
10 1000% (
11.00% (119%)
11.50% (17-4%)
13.995 (25.1%)
36 40.%% (16.7%
70.06% (22.3%)
33.10% (31.5%
13.35% (31.1%
13.75% (55.7%
(14.093
(22-7%)
13.04% (25.9%)
13475 (28-9%)
3147% 4%)
18:155 (8.75)
(1% 75)
31/40%
11.45% (26-15)
11.91% 176.9%
23,005:
AsatMic PRM
16.2962
11:284
11.24% (16 1%
12.41% (2% 19)
13.47% (30.9)
11.29% (36.7%
15.6% (39-30)
14.45% (16 7%)
16-09% (56 19
17.08% (66.9%
18.51% (76.5%
TRIM-TTs
31.31% (80.4g)
11:2218
40/205 (08.49)
(76222)
35.90%
15:47%
35:34 (73.2%)
41:31/8 (88.99)
(92.4R)
state Ster yields the task reward R., if and only the final state corresponds to a correct solution (i.e.,
E So).
TRIM-POMDP can be viewed as an extension of Aggarwal et al. (2023) to the setting of multi- step reasoning. In their official implementation, Automix x employs a greedy approximation to the POMDP, which is sufficient for task-level routing since it is a single-step decision problem (horizon of one). In contrast, multi-step reasoning requires planning over long horizons, making full POMDP formulation more appropriate and motivating the use of sophisticated solvers. Furthermore, our formulation introduces key differences that provide additional flexibility. In Aut omix, x, self- verification probabilities (observations) are used to obtain estimates of model performance metrics (state features), and thus retraining of the observation function O is required whenever the model pair (M,, Mz changes. However, TRIM-POMDP incorporates model accuracies into the transition function, eliminating the need for retraining the observation model when switching model pairs.
C BUDGETED ACCURACY
In addition to the Perfooaanee Threshold (CPT) and ABC metrics, we report budgeted accuracy in Table 4 and Table 5. Along with performance accuracy at each token budget, we also report the performance gap recovered (PGR) for improved comparability, with PGR percentages shown in parentheses. The columns indicate the normalized percentage of tokens generated by Mg. expressed relative to the total number of tokens that would be produced when running M, alone (c(TT).. Budgeted accuracy is a standard metric that evaluates each routing method under a fixed compute or token budget. These results enable direct comparison under matched compute budgets and complement our CPT and IBC analyses.
One could adopt a more conservative cost model by explicitly accounting for prefill overhead. Empirically, for large-scale models, the prefill cost per token is approximately 10% to 15% of the sequential decode cost. Under a conservative 20% overhead model, the reported Cost-Performance Thresholds (CPT) for TRIM would increase by constant factor of 1.2x (or 6/5). Importantly, this adjustment does not alter the qualitative conclusions: TRIM continues to dominate query-level baselines by a wide margin, and the relative ordering between methods remains unchanged.
D ANALYSIS OF regenerate ACTION DESIGN CHOICE
central design choice in TRIM is the definition of the regenerate action: after escalating a single reasoning step the strong model M,, control is returned to the weak model M,, rather than allowing M, to take over all remaining steps. We now provide both empirical motivation and ablation results supporting this choice.
Our choice to return control to Ma after a single M, intervention is motivated by an empirical observation (also supported by Yu et al. (2025); Lin et al. (2024)): when an incorrect step causes the
17

---

Published as a conference paper at ICLR 2026
Table 6: Distribution of accuracy gains by number strong-model interventions on MATH 550 and AIME.
Dataset
Threshold (A)
M. Rounds:
M. Rounds: 2
M, Rounds: 3
M. Rounds: >> 3
3*MATH-500
0.10
19
8
2
1
0.35
30
17
4
7
0.75
28
24
20
8
3*AIMB-Tes Test
0.10
14
0
0.35
16
12
0.75
20
13
Table 7: Comparison of standard TRIM-Thr against One-Step Thr, alternative regenerate action in which M. generates all remaining steps after intervention.
Dataset
Method
C5T((50%
CPT((0%)
CPT(95%)
IBC
2*MATH500
TRIM-Thr
43.68 (9.45%)
73.74 (15.95%)
115.99 (2.088)
4.75
One-Step Thr
114.60 (24.78%)
153.59 (33.22%)
179.31 (38.78%)
1.10
A*MMME
TRIM-Thr
204.01 (23.47%)
314.70 (36.20%)
372.79 (42.89%)
1.81
One-Step Thr
217.70 (25.04%)
482.48 (55.50%)
590.83 (67.97%)
0.84
OlympiadBennc
Thr
136.64 (20.45%)
220.70 (33.03%)
313.89 (46.97%)
1.31
One-Step Thr
199.10 (29.79%)
348. 12 (52.09%)
442. 12 (66. 16%)
0.63
2 MinervaMath
TRIM-Thr
65.15 (15.20%)
92.78 (21.65%)
148.55 (34.66%)
2.23
One-Step Thr
144.86 (33.80%)
189.97 (44.32%)
259.03 (60.43%)
0.50
reasoning trajectory to diverge, correcting just that critical step is often sufficient to steer solution back onto successful path. In our experiments, we observed that when M produces incorrect step, the primary issue not that all subsequent steps require M,, but rather that the trajectory has diverged at critical decision point. When that specific step is regenerated by M,, the corrected token sequence frequently realigns the reasoning process, after which My is able to continue successfully without additional intervention.
This behavior aligns with prior observations multi-step reasoning (Wang et al., 2025; Bigelow et al., 2025), where small number of pivotal steps disproportionately determine downstream correctness, and correcting these critical tokens can dramatically change the outcome. This phenomenon is empirically supported by intervention analysis in Table 6, which lists the number of MATH-500 and AIME problems solved exclusively due to TRIM- grouped by the number of intervention rounds, and shows that the vast majority of accuracy gains arise from just 1-3 targeted interventions rather than sustained takeover by M,-
intervention (i.c., not solved by M..).
D. ALTERNATIVE regenerate ACTION: FULL TAKEOVER ABLATION
We further evaluate an alternative regenerate action definition in which M, takes over generation for all remaining steps after the first escalation. We denote this variant as One-Step Thr and compare it directly with the standard TRIM-Thr used throughout the paper, reporting the results in Table 7.
This ablation shows that allowing M, to take over the remaining steps is substantially less cost- efficient, requiring significantly more expensive-model tokens to reach comparable performance. Overall, these results confirm that targeted, single-step interventions are both sufficient and optimal for correcting reasoning trajectories in multi-step reasoning tasks, validating design choice.
ROBUSTNESS OF TRIM TO PRM NOISE AND MISCALIBRATION
key strength of stepwise routing is its ability to leverage Process Reward Models (PRMs) to make fine- iineggrineed decisions during generation. TRIM is specifically designed to remain robust to the inevitable inaccuracies and miscalibration in PRM-based step-level correctness estimates, addressing this challenge through both its POMDP-hased formulation and its RI-trained routing variants.
18

---

Published as a conference paper at ICLR 2026
Table Performance comparison of TRIM methods under noisy PRM scores.
Method
(5%)
CPT(80%)
CPT(95%)
ABc
TRIM-Thr 187.37 (27.63%)
450.77 (66.46%)
564.92 (83.29%)
0.38
TRIM-Age 156.96 (23. 14%)
262.82 (38. 75%)
305.01 (44.97%)
1.16
Table 9: Performance of TRIM-Thr for the model pair (0wen2.5-1.5B-Instrrrr Qwen3-32B) on AIME & MATH-500.
MATH-500
AIME
Method (50%)
(BT()
(91())5)
Ame
CPT(5005)
(SP())))
CP1(95%)
Amc
TRIM-Thr 30.02 (4.77%)
117.32 (18.65%) 169.9 (27.016)
7.46
541.35 (32.61%)
870.86 (52.47%)
1035.62 (62.39%)
1.25
In TRIM-POMDP, the routing problem is formulated as a Partially Observable Markov Decision Pro- cess (POMDP), explicitly designed account for inaccuracies in PR.M-based step-level correctness estimates. Rather than treating PRM scores as ground-trun indicators of correctness, TRIM-POMDE interprets them as noisy observations of an underlying latent correctness state. An observation function is learned to map these imperfect PRM scores to a belief distribution over these hidden correctness states, allowing routing decisions to be made based on posterior beliefs rather than raw scores. This explicit modeling of uncertainty is the primary motivation for the POMDP formulation and makes TRIM-POMDP robust to PRM miscalibration, in contrast to threshold-based routing methods such TRIM-The Thr, which operate directly
PRM outputs.
On the other hand, TRIM-Agg provides a complementary form of robustness and implicitly learns to model the unreliability in the PRM estimates during RL training. Thus, while the performance of TRIM-Thr can degrade using noisier PRMs, TRIM- Agg maintains robust performance. To validate this, trained TRIM-Agg against TRIM- The on OlympiadBench, with both methods relying on the noisy PRM evaluations. Noise introduced by applying left-padding instead of the recommended right-padding during PRM evaluation, which introduces significant noise into step-level correctness estimates.
MATH-500 using noisy PRM scores and compared its performance
The results in Table show that TRIM-Agg maintains strong cross-dataset generalization performance despite the noisy PRM: it achieves CPTT(95%) on OlympiadBench using only 45% of expensive tokens, whereas TRIM-Thr suffers significant degradation, requiring over 80% of expensive tokens to reach C97(9 This empirically confirms that RL-trained routing policies can also learn to be robust PRM noise, while simple threshold- based methods cannot.
F GENERALIZABILITY OF TRIM ACROSS MODEL PAIRS AND DOMAINS
TRIM is model-agnostic and applies to arbitrary model pairs (Mw. M,), provided they are sufficiently compatible for stepwise continuation. We evaluate generalization along two axes: (i) robustness across model pairs, and (ii) applicability beyond mathematical reasoning.
Generalization Across Model Pairs. In addition to the cross-dataset generalization experiments reported in Table 2 and Figure 7, we evaluate TRIM-Thr using substantially different model pair: 1..n2.5-..B-Instruue as the weak model M,,, and Qwen3-32B as the strong model M,. Table reports the resulting performance on MATH-500 and AIME.
These results demonstrate that Thr continues to achieve strong efficiency gains under a different model family and scale, confirming that its routing behavior is not tied to specific model pair.
Beyond Mathematical Reasoning, While our main experiments focus on mathematical reasoning due to the availability of reliable Process Reward Models (PRMs), TRIM itself is not math-specific. The framework only requires a mechanism for estimating step-level correctness, not a domain-specinn PRM. In domains where PRMs unavailable, TRIM can be instantiated using model-gene generated self-evaluations by prompting the strong model to score intermediate steps steps-similar similar in spirit to SpecReason (Pan et al., 2025).
19

---

Published as a conference paper at ICLR 2026
Table 10: Budgeted-accuracy comparison of TRIM-Thr the model pair (Qwen2.551..B-Instrucc Qwen3-32B) on AIME & MATH-500.
MATH:500
ATME
TRIS--Er 53.8% (1625)
67.15 (5AN)
7463% (9079)
70.1R (ML2%)
33.2%% 34.2%% (1GG.9) 10.705 (258%)
11.3% (37,99)
14J00% (44.5))
18.775 (3-5%)
18.42% (4675)
Table 11: Performance of TRIM-Thr on GPOA-DIAMOND for the model pair (Qwee3-1.7B, Qwen3-32B).
Method
CPT(50%)
CPT()0%)
CPT(995))
ABC
TRIM-Thr 537.05 (2%) 1012.97 (54. 17%) 1165.85 (62.34%)
0.26
To demonstrate this, extend TRIM-Thr beyond mathematics and evaluate it on GPQA-DIAMOND, a challenging benchmark spanning biology, physics, and chemistry. We use Qwen3-1.7B as M... Qwen3-32B as Mg. and obtain step-level correctness scores by prompting M, to assign scores on 0-9 scale. The resulting performance-com cost trade-offs are shown in Table 11. TRIM-Thr achieves (PT(95%) using only 62% of expensive-moder tokens. Moreover, at high threshold (9), TRIM-Thr improves the accuracy of Qwen3-32B from 0.419 to 0,450 while using only 73% of the expensive tokens. These results indicate that TRIM remains effective even in scientific reasoning domains without dedicated PRM.
G TRIM IMPLEMENTATION DETAILS
Targeted intervention enables TRIM to reduce cost, which is defined as the number of tokens generated from the expensive, stronger LLM. While this sort of "step-level"" routing may appear to require frequent context re-encoding ("prefill") when switching models, this overhead can be largely amortized by chunked prefilling of the large model's 's KV cache in parallel with the small model's decoding and PRM evaluation. In practice, the large model can maintain a synchronized KV cache (aggregate prefill) by incrementally prefilling the tokens generated by the small model. This ensures that upon escalation, the large model can almost immediately transition to decoding the next token rather than re- encoding the history. This design mirrors speculative decoding strategies (Leviathan al., 2023; Cai et al., 2024), which similarly overlap draft and verification phases minimize latency.
In contrast, the generation (decode) phase of the large model imposes an unavoidable sequential cost and cannot be parallelized across tokens, while the prefill can be parallelized, cached, or reused across steps. Consequently, the number of tokens generated by the large model is the dominant and irreducible component of inference cost, making it both theoretically clean and operationally grounded efficiency metric. In summary, large-mode! model decode tokens capture the true marginal expense, while prefilling and routing overheads are effectively hidden by overlapping computation with small-model decoding and PRM evaluation. While we did not implement these system-level optimizations in our primary experiments-an public API endpoints for proprietary models (e.g.. Claude) do not currently expose the necessary KV cache state management- established standards in open-source serving frameworks (e.g., vLLM, SGLang) and are widely utilized by production providers. Consequently, the absence of these optimizations in our API-based evaluation does not undermine the conceptual validity of our cost model.
techniques are
TRIM-Agg Implementation. The router policy is parameterized by a simple MLP policy with two hidden layers (128 units each, Tanh activations), followed by separate actor and critic heads, and trained using PPO. The selected hyperparameters are summarized in Table 12. Training is conducted across performance-comm trade-off parameters X, ranging from 10-4 to 8 x 10-5 for AIME and from 8 x 10 to 3 10-4 for MATH, at regular intervals.
TRIM-POMDP Implementation. For TRIM-POMDE we learn the observation function using a reflected KDE estimator applied to the ProcessBench (Zheng et al., 2024) dataset. Specifically, we evaluate our PRM on step-by-step solutions in the dataset and align its outputs with step-level labels. The observation function is trained on Omni-MATH problems, while evaluation is conducted on AIME and GSM8k for MO0H-500. Importantly, the only model-specific information required by TRIM- is the next-step accuracies of the models, which estimated from the
20

---

Published as a conference paper at ICLR 2026
Table 12: Hyperparameters used for Agg.
Hyperparameter
Value
Learning rate
1.0c-4
Clipping coefficient
0.2
Entropy coefficient
0.01
Advantages
Unnormalized
Rewards
Undiscounted
GAE
0.95
corresponding training sets. For solving the POMDP, we use SARSOP (Kurniawati et al., 2008), which implement using the POMDPs. j1 framework (Egorov et al., 2017). We solve the POMDP using SARSOP (Kurniawati et al., 2008), implemented using the POMDPs j1 framework (Egorov al., 2017), with hyperparameters set the default values provided in POMDPs ji. While SARSOP can in principle handle large observation spaces, it is sensitive to the choice of the initial state distribution. To achieve the best performance, we therefore recompute the policy every step using the updated belief distribution as the initial state distribution. To improve efficiency, we employ a simple heuristic: the policy is recomputed only if the belief mass on state S2 (the case where the most recent step is incorrect but prior steps are correct) lies within 0.35-0.40 of the maximum belief state class; otherwise, we default to continuing with M (i.e., Or = continue). For all TRIM routing policies, the reasoning trace is truncated to maximum of 30 steps during both training and inference, and the solution at this cutoff is returned.
Mux Implementation. In its official implementation, AutoMix employs greedy approximation to the proposed POMDP formulation. For two-model setup, queries are categorized into three classes: (I) solvable by M,.. (2) unsolvable by both Mi and M,, and (3) solvable only by M,. Using few-shot self-verification or correctness scores, AutoMix discretizes the verifier probability space into uniformly spaced bins. For each bin, it estimates the empirical conditional distribution over these three outcome classes-essennially learning a binned classifier that maps verifier probabilities to class likelihoods. During inference, routing decisions are made greedily: given a new verifier score, the corresponding bin provides the estimated probability distribution over outcome classes, and the router selects an action that maximizes the expected return under a specified trade-off parameter balancing accuracy and cost.
Smoothie Implementation. For Smoothie, we include results for Emnothie-Trriin following the official implementation. As noted in its limitations (Guha et al., 2024), Smoothie is optimized for performance rather than explicit cost-performance trade-offs and is designed to escalate to the strongest available model within given mixture of LLMs. To populate performance-cost curve in our setting, we therefore route based on the estimated difference in LLM quality scores for each sample. Since Smoothie requires outputs from least three models to estimate these scores during training, we use Mathstral- 7B- as the third model in our Smoothie experiments.
H EFFECTIVENESS OF QWEN2.5-MATH-
FOR AIME
natural question how effective Qwen2.5-M Math-PRM-7BB is as a PRM for AIME problems, which substantially more challenging than other mathematical benchmarks considered in this work, such as MATH-500, OlympiadBench, and Minerva Math. Since step-level ground-truth annotations not available for the AIME, we assess PRM effectiveness by measuring how well aggregated stepwise PRM scores correlate with the final correctness of full solutions. Our evaluation procedure as follows:
1. We generate full reasoning traces for AIME, MATH-500, and OlympiadBench using the same weak model, Q
2. For each trace, compute PRM step-level scores and aggregate them using the minimum score across steps. This aggregation strategy is used throughout our work and is standard in prior PRM literature.
21

---

Published as a conference paper at ICLR 2026
3. We compute the AUC-ROC between these aggregated scores and the final correctness labels, quantifying how well PRM scores separate correct from incorrect solutions.
The resulting AUC-ROC values are:
AIME: 0.8397
OlynpiadBench: 0.8954
MATH500: 0.9363
These scores indicate that while PRM reliability strongest on MATH-500 and OlympiadBennh for the generated traces, it still provides meaningful discriminatory signal AIME-fat far above random (0.5) and sufficiently informative for decision-man mecision-making. The somewhat lower AUC on AIME reflects that these problems more difficult and produce longer, noisier reasoning traces, but the PRM remains useful rather than failing outright.
Crucially, the reduced PRM reliability AIME highlights the need for robustness to noisy correct- ness estimates. TRIM-POMDP is explicitly designed for this: by modeling PRM scores as noisy observations of underlying correctness state, it remains stable even when the PRM is imperfect. Consistent with this robustness role, TRIM-POMDP achieves a 2.76x improvement over TRIM- Thr in ÅBC on AIME AIME-muci much larger than the x.23x improvement on 100d--00odmmmttraiinn that uncertainty-ave -ware routing is especially valuable on harder datasets where PRM noise is more pronounced, and estimates less reliable.
22


Reviews:
[{"review_id": "koHqmWJxvl", "reviewer": "Reviewer_X7GU", "summary": "This paper proposes TRIM (Targeted Routing in Multi-step reasoning tasks), a novel hybrid inference strategy designed to improve the cost-efficiency of multi-step reasoning. The central problem it addresses is that conventional \"query-level\" routing (assigning an entire task to either a cheap or an expensive LLM) is inefficient. Relying only on cheap models leads to cascading errors, while relying only on expensive models is cost-prohibitive.\n\nTRIM's core idea is to perform **step-wise routing**. A cheap model generates each reasoning step by default. After each step, a router, guided by a Process Reward Model (PRM), decides whether to accept the cheap step or to regenerate *only that specific step* with an expensive model before handing control back to the cheap model. This \"targeted intervention\" aims to prevent cascading failures at critical steps while saving costs on routine ones.\n\nThe authors propose and evaluate four distinct routing strategies:\n1.  **TRIM-Thr:** A simple threshold policy (regenerate if PRM score < *k*).\n2.  **TRIM-Seq:** An RL-trained policy (using full step history).\n3.  **TRIM-Agg:** A simpler RL-trained policy (using aggregated step statistics).\n4.  **TRIM-POMDP:** A sophisticated policy that models the true step correctness as a hidden state and the PRM score as a noisy observation.\n\nOn mathematical reasoning benchmarks, the simple TRIM-Thr policy is shown to be 6.51x more cost-efficient than contemporary query-level routing. The more advanced RL and POMDP policies can match the performance of the expensive-only model while reducing the number of expensive tokens generated by 80%.", "questions": "1.  **Latency Cost:** Your cost model ignores the significant per-step latency of running the PRM and the router policy. How does the total wall-clock time of a TRIM-generated solution (with its N sequential stop-and-think steps) compare to a single, uninterrupted generation from the expensive model $M_s$?\n2.  **\"Regenerate\" Action:** Why did you choose to return control to the weak model $M_w$ *immediately* after an expensive $M_s$ intervention? Have you experimented with an alternative \"regenerate\" action where $M_s$ takes over for all *remaining* steps, and how does that compare on the performance-cost curve?\n3.  **PRM Sensitivity:** How sensitive are the simpler (and more practical) TRIM-Thr and TRIM-Agg policies to the quality of the Process Reward Model? How much does their performance drop if a smaller, faster, but less accurate PRM is used?\n4.  **Justifying POMDP Complexity:** The results for TRIM-Agg and TRIM-POMDP in Figure 6 appear nearly identical. Given the extreme complexity of the POMDP (requiring annotated data to train the observation function, plus an online solver), what is the practical justification for using it over the much simpler RL-based TRIM-Agg policy?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* **Novelty:** The core concept of step-wise routing, as opposed to query-level routing, is a significant and novel contribution. It directly targets the \"cascading failure\" weakness of multi-step reasoning.\n* **Impressive Results:** The efficiency gains are substantial. Achieving an 80% reduction in expensive token usage while matching the strong model's accuracy is a highly compelling result. The 6.51x cost-efficiency improvement from the *simplest* policy (TRIM-Thr) over SOTA query-level routing is also very strong.\n* **Methodological Thoroughness:** The paper explores four different router implementations, from a simple heuristic to a complex POMDP. This shows a deep understanding of the problem, from a practical, easy-to-implement solution to a more theoretically-grounded one that handles observational uncertainty (PRM noise).\n* **Generalizability:** The cross-dataset generalization experiments are a key strength, validating that the model is learning fundamental properties of reasoning difficulty rather than just overfitting to a specific dataset's quirks.", "weaknesses": "* **Incomplete Cost Model (Latency):** The paper's cost model is defined *only* as \"the number of tokens generated by the expensive model.\" This model completely ignores the very significant latency and compute costs incurred at *every single step* of the generation. The loop is: (1) $M_w$ generates a step, (2) generation pauses, (3) the 7B-parameter PRM runs, (4) the router policy runs (for TRIM-POMDP, this involves *re-solving* the policy), (5) $M_s$ *maybe* runs. This sequential, stop-and-go process introduces a massive amount of latency per step, which is not accounted for. A simple, streamed generation from the expensive model $M_s$ might be much faster in wall-clock time, even if it uses more \"expensive tokens.\"\n* **The \"Regenerate\" Action Definition:** The \"regenerate\" action is defined as $M_s$ regenerating the current step, after which control is *immediately* returned to the cheap $M_w$ for the *next* step. This seems suboptimal. If a step was so critical that $M_w$ failed and $M_s$ had to intervene, why would you trust $M_w$ with the very next step? This could lead to an immediate re-failure. A more intuitive action would be to let $M_s$ take over for the remainder of the generation. This design choice feels overly optimized for the paper's cost metric, not for maximum accuracy.\n* **Complexity vs. Performance:** The TRIM-POMDP approach is by far the most complex, requiring ground-truth step-level annotations (from ProcessBench) to learn the observation function and a POMDP solver. However, in the results (e.g., Figure 6), its performance-cost curve is almost identical to that of the much simpler `TRIM-Agg` policy (a simple MLP trained with PPO on aggregated stats). The paper does not adequately justify why this massive increase in complexity is worthwhile for such a marginal gain.\n* **Dependency on the PRM:** The entire system is critically dependent on the quality of the Qwen2.5-Math-PRM-7B model. While the POMDP policy *models* this uncertainty, the simpler and highly effective TRIM-Thr and TRIM-Agg policies *trust* the PRM scores. The paper doesn't analyze how the performance of these simpler policies degrades if a less-accurate, noisier, or miscalibrated PRM is used.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iTd7YpebIJ", "reviewer": "Reviewer_hcjw", "summary": "This paper proposes a novel method that improves the inference efficiency of large language models by performing routing decisions at the step level. The authors formalize this process as a sequential decision problem and design several routing strategies, including a simple threshold-based policy (TRIM-Thr), two reinforcement learning–based policies (TRIM-Seq and TRIM-Agg) that balance accuracy and cost, and a POMDP-based policy (TRIM-POMDP) to handle uncertainty in step correctness estimation. Experiments on multiple mathematical reasoning benchmarks demonstrate that this approach can achieve performance comparable to strong models while significantly reducing computational cost.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is clearly structured and well written, with intuitive figures and tables that effectively illustrate the key ideas and experimental results.\n2. The proposed method demonstrates strong innovation by moving beyond traditional query-level routing, offering a substantial conceptual advance in improving multi-step reasoning efficiency, and showing significant empirical effectiveness.\n3. The authors formalize the routing process as a sequential decision problem and propose multiple complementary strategies that balance methodological simplicity with theoretical depth.", "weaknesses": "1. The approach heavily relies on the Process Reward Model; inaccuracies or poor calibration in PRM estimation may negatively affect routing stability and overall performance.\n2. Experiments focus primarily on mathematical reasoning tasks, and the generalization of TRIM to other multi-step reasoning domains (e.g., code generation or scientific reasoning) remains untested.\n3. Stepwise generation introduces additional inference latency; although token cost is reduced, the paper does not provide a detailed analysis of runtime efficiency or latency trade-offs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GQMpjWtShJ", "reviewer": "Reviewer_5mDK", "summary": "The paper presents TRIM, a hybrid inference framework that uses a process reward model (PRM) to decide when to route reasoning steps from a smaller model to a larger one. The approach aims to balance inference cost and reasoning accuracy. It is primarily evaluated on mathematical reasoning tasks such as AIME using Qwen2.5-Math-PRM-7B and Claude 3.7 Sonnet. The authors introduce a Cost–Performance Threshold (CPT) metric and report improved cost–performance trade-offs compared with several baseline methods.", "questions": "* There is already a considerable body of work on stepwise routing with PRMs [1–3]. What is the fundamental difference between TRIM and these existing approaches?\n* How effective is Qwen2.5-Math-PRM-7B as a PRM for AIME problems? Compared with the datasets reported in the official documentation [4], AIME problems are substantially more difficult, and the reasoning traces generated by LLMs are usually much longer. It is therefore more challenging to evaluate the correctness of intermediate steps.\n* The Cost–Performance Threshold (CPT) metric is not commonly used. In contrast, budgeted accuracy, which compares methods under the same compute or token budget, is a standard and decision-useful metric for evaluating efficient reasoning or routing. Could the authors also report results under this metric for better comparability?\n* The baseline coverage should be expanded. In addition to stepwise routing methods, please include non-routing anchors such as Always-Small-Model and Always-Large-Model, as well as query-level routing [5,6] and step-/token-level adaptive compute approaches [1–3].\n* Line 420: The paper uses models from different families as the base (e.g., Qwen vs. Claude), which may have distinct formatting styles and output distributions. How do the authors handle these inconsistencies? Moreover, since Claude 3.7 Sonnet relies on a chat-completion API schema, continuing reasoning from a partial trace can differ from a true completion mode. Could the authors explain how this issue is addressed or justify this model choice?\n* Figures 6 and 7: Should the x-axis represent the number of tokens from the expensive model?\n* How generalizable is the proposed approach? Can the authors test TRIM with different combinations of cheap and expensive models, and on additional domains such as scientific reasoning or code to demonstrate broader applicability?\n* Can the authors provide the code and data for review?\n\n## Reference:\n[1] Liao, Baohao, et al. \"Reward-Guided Speculative Decoding for Efficient LLM Reasoning.\" Forty-second International Conference on Machine Learning.\n\n[2] Pan, Rui, et al. \"Specreason: Fast and accurate inference-time compute via speculative reasoning.\" arXiv preprint arXiv:2504.07891 (2025).\n\n[3] Liu, Yuliang, et al. \"AdaptiveStep: Automatically Dividing Reasoning Step through Model Confidence.\" Forty-second International Conference on Machine Learning.\n\n[4] Zhang, Zhenru, et al. \"The lessons of developing process reward models in mathematical reasoning.\" arXiv preprint arXiv:2501.07301 (2025).\n\n[5] Ding, Dujian, et al. \"Hybrid LLM: Cost-Efficient and Quality-Aware Query Routing.\" The Twelfth International Conference on Learning Representations.\n\n[6] Guha, Neel, et al. \"Smoothie: Label free language model routing.\" Advances in Neural Information Processing Systems 37 (2024): 127645-127672.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* **Clear motivation and relevance:** The paper tackles a practical problem of reducing inference cost while maintaining reasoning accuracy.\n* **Reasonably structured framework**: The framework is generally well organized, though its design largely follows existing PRM-based routing paradigms.\n* **Clarity:** The paper is well organized and easy to follow, with clear figures and explanations.", "weaknesses": "* **Limited novelty over existing PRM-based routing methods:** The paper does not clearly articulate how TRIM differs conceptually from prior stepwise routing approaches using process reward models (PRMs).\n* **Unclear reliability of PRM supervision for AIME problems:** The use of Qwen2.5-Math-PRM-7B raises concerns about label reliability on AIME, where reasoning traces are substantially longer and harder to evaluate.\n* **Model family mismatch and API issues:** Mixing models from different families (e.g., Qwen and Claude) may introduce inconsistencies in formatting and reasoning styles. In addition, Claude 3.7 Sonnet’s chat-completion API makes partial-trace continuation non-trivial.\n* **Insufficient baseline coverage:** Several key comparisons are missing, such as non-routing anchors (Always-Small, Always-Large), query-level routing, and adaptive compute methods.\n* **Code and data unavailable for review:** The absence of released code or data prevents verification and reproducibility of the reported results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wgH6ICRebq", "reviewer": "Reviewer_oN83", "summary": "This paper proposes TRIM, a hybrid inference method that routes reasoning steps between a small and a large LLM. The key idea is to escalate only critical reasoning steps to the stronger model, improving efficiency while maintaining accuracy. The paper includes several routing strategies, from a simple threshold-based approach to more advanced RL and POMDP-based ones and demonstrates improvement over query level baseline results on multiple math reasoning benchmarks.", "questions": "1- Section 4 feels ambiguous. While Figure 3 suggests that routing and decisions happen step-by-step, the notation y_s = Ms(y1:t−1) makes it look like the strong model regenerates the whole prefix. Which interpretation is correct? The notation and explanation should be clarified.\n\n2- What exactly is a “step”? Is it defined as a sentence, a reasoning line, or a fixed number of tokens? How do you prompt or constrain the model to generate only one step at a time? Is this step size consistent across different datasets or tasks?\n\n3- In TRIM-Seq, how is the step-level correctness score computed? Is it purely from the PRM output or based on an external reference?\n\n4- Since cost is measured only by the number of tokens from the expensive model, does this ignore switching overhead between models? Routing back and forth should add latency and communication cost, which could be nontrivial in real systems.\n\n5- The reported baseline performance for Claude 3.7 seems off, while this model reportedly has ~96.2 % accuracy on MATH-500,  Figure 6 shows it closer to 80%. Can you clarify this discrepancy or explain what experimental setup causes this gap?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The overall presentation is good, and the paper is mostly easy to read and understand. The concept of step-level routing appears practical, bridging the gap between coarse query-level routing and fine-grained token-level control. The experimental setup seems reasonable, and the reward modeling and reinforcement learning training are well-motivated.", "weaknesses": "Some parts of the paper are confusing or underspecified, especially the step definition and regeneration process. The cost metric ignores potential latency from frequent model switching, which could matter in deployment. A few experimental details (like baseline performance numbers) seem inconsistent or not clearly justified. Clarifying these points would make the technical contribution much easier to trust and replicate. See questions section for more details. I am open to increase my rating if the issues raised are adequately addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)