TRIM: 다단계 추론 과제에서 타겟 단계별 라우팅을 통한 하이브리드 추론

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

수학적 문제 해결과 같은 다단계 추론 과제는 단 하나의 잘못된 단계가 전체 풀이의 실패로 이어지는 연쇄적 실패(cascading failures)에 취약합니다. 기존의 LLM 라우팅 방법들은 모든 추론 단계를 동등하게 취급하여 전체 쿼리를 하나의 모델에 할당합니다. 이에 우리는 TRIM(다단계 추론 과제에서의 타겟 라우팅)을 제안합니다. TRIM은 풀이를 망칠 가능성이 높은 중요한 단계(critical steps)만을 더 큰 모델로 라우팅하고, 일반적인 진행 과정은 작은 모델이 처리하도록 합니다. 우리의 핵심 통찰은 더 강력한 모델이 연쇄적 오류를 방지할 수 있는 정확한 단계에만 비용이 많이 드는 호출을 제한함으로써, 타겟 단계별 개입이 추론 효율성을 근본적으로 변화시킬 수 있다는 점입니다. TRIM은 단계 수준에서 작동하며, 프로세스 보상 모델을 사용하여 오류 단계를 식별하고 단계별 불확실성과 예산 제약을 기반으로 라우팅 결정을 내립니다. 우리는 단순한 임계값 기반 정책부터 장기적인 정확도-비용 상충 관계와 단계별 정답 추정의 불확실성을 고려하는 더 고도화된 정책까지 TRIM 내에서 다양한 라우팅 전략을 개발했습니다. MATH-500에서 가장 단순한 임계값 전략조차도 5배 더 높은 비용 효율성으로 기존 라우팅 방법을 능가했으며, 더 발전된 정책들은 고비용 모델 토큰을 80% 적게 사용하면서도 강력하고 비싼 모델의 성능을 달성했습니다. AIME와 같은 더 어려운 벤치마크에서 TRIM은 최대 6배 더 높은 비용 효율성을 달성했습니다. 이 모든 방법은 수학 추론 과제 전반에 걸쳐 효과적으로 일반화되며, 이는 단계별 난이도가 추론의 근본적인 특성을 나타냄을 보여줍니다.

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.