arXiv:2601.10245cs.AIcs.CL2026-01中稿 · ICLR被引 5

只让大模型处理关键步骤,显著降低数学推理成本

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

  • 按步骤识别易错环节,仅对关键步骤调用大模型
  • 简单策略成本效率提升5倍,高级策略节省80%大模型算力
  • 适用于各类数学推理任务,尤其适合资源受限场景

多步推理任务(如数学解题)易因单步错误导致整个求解失败。现有LLM路由方法将整段查询分配给单一模型,忽视各步骤差异。本文提出TRIM(目标导向的多步推理路由),仅将可能引发错误的关键步骤路由至大模型,其余步骤由小模型完成。核心思想是:通过精准干预高风险步骤,大幅提高推理效率。TRIM在步骤层面运作,利用过程奖励模型识别错误步骤,并基于步骤级不确定性和预算约束做出路由决策。设计了多种路由策略,从简单阈值法到能权衡长期准确率与成本的复杂策略。在MATH-500上,最简策略成本效率超此前方法5倍;更先进策略以仅80%的大模型令牌数达到强模型性能。在更难的AIME基准上,成本效率最高提升6倍。所有方法均在不同数学推理任务中表现良好,表明步骤难度是推理的本质特征。

原文摘要 · Abstract (English)

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.

推理优化路由策略数学推理成本效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。