Consider a two-model setup: a weak, cheap model M_w and a strong, expensive model M_s. At each reasoning step, M_w proposes a candidate continuation, and the TRIM router decides whether to keep it or regenerate it using M_s.

This transforms routing from a one-shot decision into a sequential control problem over the reasoning trajectory. Decisions are guided by step-level correctness signals (e.g., PRM/self-verification scores) and token cost, enabling explicit optimization of the accuracy-cost trade-off.

We develop four routing strategies within TRIM, ranging from simple to highly expressive:

We instantiate TRIM with routing policies of increasing sophistication. The simplest policy, TRIM-Thr, uses only the PRM score of the current step, while the more expressive policies incorporate trajectory-level information to reason about whether intervention is still likely to improve the final solution enough to justify its cost.

We first introduce TRIM-Thr, a simple routing policy that relies solely on the PRM score of the current step generated by the weak model M_w. If this score falls below a predefined threshold k, the router regenerates the step with the strong model M_s; otherwise, it accepts the weak model’s output and continues generation.

Despite its simplicity, TRIM-Thr is highly effective. The threshold parameter k provides a principled mechanism for controlling the performance-cost trade-off: lower thresholds reduce strong-model usage, while higher thresholds allow more aggressive intervention. Empirically, this simple step-level policy already outperforms prior query-level routing methods and is competitive with an idealized oracle query-level router that selects the best model per query.

While TRIM-Thr performs remarkably well, it is inherently myopic: it makes routing decisions using only the correctness estimate of the most recent step, without accounting for past context or future consequences. In many cases, the current step alone is insufficient for deciding whether intervention is worthwhile.

For example, even if the current step appears incorrect, regeneration may not be beneficial if the overall trajectory has already diverged too far from a correct solution, or if the cost of invoking M_s outweighs the potential gain. Conversely, when earlier steps remain largely consistent, a targeted intervention can recover the trajectory and substantially improve final correctness.

This motivates richer routing policies that reason jointly about trajectory correctness and marginal intervention cost.

TRIM-Seq learns routing decisions from the full sequence of step-level signals accumulated along the reasoning trace. At each step, the router observes the sequence of correctness estimates and token counts from prior steps, allowing it to reason over how the trajectory has evolved rather than focusing only on the latest PRM score.

Concretely, TRIM-Seq uses the feature sequence (r₁, c₁), ..., (r_t, c_t), where r_i denotes the PRM score of step i and c_i denotes its token length. These features capture the two quantities most fundamental to routing in multi-step reasoning: semantic fidelity, via step-level correctness estimates, and marginal intervention cost, via token counts that quantify the expense of regeneration with the strong model.

We parameterize the routing policy with a transformer over this feature sequence and train it using RL. The objective maximizes expected return by balancing final solution correctness against the cumulative cost of invoking the strong model. Each regeneration incurs a cost proportional to the number of strong-model tokens, weighted by a trade-off parameter λ, while the final reward reflects task correctness. This enables the router to learn when an intervention is likely to improve solution quality enough to justify its cost.

While TRIM-Seq exploits the full sequential history, much of the relevant structure can be captured using a compact set of aggregate statistics. In multi-step reasoning, errors often exhibit a compounding structure: a single incorrect step can invalidate everything that follows, while a sequence of mildly unreliable steps can gradually push the trajectory off course.

Motivated by this, TRIM-Agg uses a reduced feature representation consisting of the current PRM score, the minimum of prior scores, the token count of the current step, and the step index. These aggregated features preserve key signals about whether the trajectory remains plausibly on track and whether intervention is cost-justified, while discarding the full sequential history.

We train TRIM-Agg with the same RL objective as TRIM-Seq. Empirically, this compact representation enables substantially faster training with little to no loss in performance, making it a particularly attractive practical instantiation of TRIM framework.

A final challenge in stepwise routing arises from the imperfect nature of PRM estimates. Although PRMs provide useful signals about the correctness of intermediate steps, their predictions are often noisy and can misclassify correct steps as incorrect, or vice versa. In principle, RL-trained routers can learn to compensate for this noise, but training such policies under long-horizon sparse rewards is often sample-inefficient and expensive.

TRIM-POMDP addresses this by explicitly modeling PRM scores as noisy observations of an unobserved latent state that reflects the true correctness status of the reasoning trajectory. Rather than acting directly on raw PRM outputs, the router first infers the latent state and then plans accordingly, casting routing as a partially observable Markov decision process (POMDP).

TRIM-POMDP defines a compact latent state space with three correctness classes: S₀, where the trajectory remains correct so far; S₁, where the trajectory has already diverged irrecoverably; and S₂, where the most recent step is incorrect but prior steps are correct, so the trajectory may still be recoverable through intervention.

If this latent state were directly observed, routing would reduce to a standard fully observable control problem. In practice, however, the router only sees noisy PRM-based signals. This makes partial observability central to the problem rather than incidental, and motivates an explicit belief-state approach.

To bridge the gap between noisy PRM outputs and latent correctness states, TRIM-POMDP learns an observation function that maps the history of routing observations to a probability distribution over the latent states. Concretely, this amounts to modeling the distribution of PRM-based features conditioned on each correctness class.

This observation model can be fit offline using process supervision datasets with step-level annotations, such as ProcessBench. Once learned, it can be reused across different cost budgets, since it depends only on the alignment between PRM scores and ground-truth correctness labels rather than on a specific routing objective.

By explicitly separating state inference from policy optimization (via observation function), TRIM-POMDP provides a principled way to route under noisy correctness signals without requiring expensive RL training for every performance-cost trade-off.

Once the observation model is learned, TRIM-POMDP computes routing policies using standard POMDP solvers. An additional advantage is that the resulting router is largely agnostic to the specific choice of weak and strong LLMs, depending primarily on their step-level transition characteristics. This makes the formulation modular and reusable across model pairs.

Setup: We evaluate TRIM on challenging mathematical reasoning benchmarks including MATH-500, AIME, OlympiadBench, and Minerva Math. The primary two-model setup uses Qwen2.5-3B-Instruct as the cheap model (M_w) and Claude 3.7 Sonnet as the expensive strong model (M_s), with step-level signals obtained from Qwen2.5-Math-PRM-7B. We compare against RouteLLM (BERT, MF, SW Ranking), Smoothie, AutoMix, and our adapted variant of AutoMix, AutoMix-PRM.

We measure cost efficiency using CPT(x%): the fraction of expensive-model M_s tokens (normalized by full M_s usage) required to recover x% of the performance gap between M_w and M_s (lower is better). Δ_IBC measures the average performance gain per unit cost relative to a random router, capturing the overall quality of the cost-performance frontier (higher is better).

TRIM substantially improves over query-level routing across benchmarks. Even the simplest policy, TRIM-Thr, is already highly competitive and often dramatically more efficient than prior methods, while TRIM-Agg and TRIM-POMDP further improve the trade-off by leveraging trajectory-level information and uncertainty modeling.

The headline pattern is consistent across settings: routing a few critical steps is markedly better than routing an entire query. TRIM approaches strong-model accuracy using a small fraction of its token cost, with particularly large gains on harder benchmarks like AIME.

Table 1: Benchmarking TRIM on MATH-500 and AIME. TRIM consistently dominates prior query-level routing baselines. The simple threshold TRIM-Thr policy is already very strong, while TRIM-Agg and TRIM-POMDP deliver the best overall trade-offs in different budget regimes.

We now visualize the full performance-cost Pareto frontiers of all TRIM routing approaches on MATH-500 and AIME.

In practice, TRIM can be implemented using the same system-level optimization techniques that enable low-latency speculative decoding, thereby avoiding frequent context re-encoding ("prefill") when switching between models and effectively eliminating the associated latency overhead. We provide direct empirical evidence that TRIM does not introduce significant wall-clock overhead—and is in fact faster than running the expensive model M_s alone.

We measured end-to-end latency and throughput on a fixed 2×H100 setup using vLLM with prefix caching enabled. For TRIM-Thr, M_s was sharded across both GPUs via tensor parallelism (55% memory per GPU), while M_w and the PRM each ran on separate GPUs (40% memory). For single-model baselines, both GPUs were fully dedicated to the large model. All measurements were conducted on the MATH-500 test set under two model-pair configurations:

• TRIM-Thr (1.5B + 32B): Qwen2.5-1.5B-Instruct as M_w and Qwen2.5-32B-Instruct as M_s
• TRIM-Thr (1.5B + 7B): Qwen2.5-1.5B-Instruct as M_w and Qwen2.5-7B-Instruct as M_s

Table 3: End-to-end latency and throughput. TRIM-Thr (1.5B+32B) achieves a 1.4×–2.75× speedup over running the 32B model alone (17.10/12.10 ≈ 1.41, 17.10/6.21 ≈ 2.75), despite incorporating PRM evaluation and stepwise routing. Latency gains increase with larger strong models and lower routing thresholds.