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

1. 문제 정의 (Problem Statement)

• 연쇄적 실패 (Cascading Failures): 수학 문제 해결과 같은 다단계 추론 (Multi-step Reasoning) 작업에서는 단일 단계의 오류가 전체 솔루션의 붕괴로 이어지는 연쇄적 실패가 발생하기 쉽습니다.
• 기존 라우팅의 비효율성: 현재 대부분의 LLM 라우팅 방법은 전체 쿼리 (질문) 를 하나의 모델에 할당합니다. 이는 모든 추론 단계를 동일한 난이도로 간주하여, 단순한 연속 단계에도 고비용의 대형 모델을 사용하거나, 반대로 어려운 결정 단계에서 소형 모델을 사용함으로써 발생하는 비효율성을 초래합니다.
• 비용 - 성능 트레이드오프: 대형 모델은 성능은 뛰어나지만 추론 비용이 매우 높고, 소형 모델은 비용은 낮지만 품질이 떨어집니다. 기존 방법은 이 두 가지 중 하나를 선택하는 이분법적 접근을 취하여, 비용 효율성을 극대화하지 못합니다.

2. 제안 방법론: TRIM (Targeted Stepwise Routing)

TRIM 은 쿼리 단위가 아닌 추론 단계 (Step-level) 단위로 라우팅을 수행하여, 오직 임계적이고 오류가 발생할 가능성이 높은 단계에서만 대형 모델로 전환하는 하이브리드 추론 프레임워크입니다.

핵심 메커니즘

1. 단계별 생성 및 평가:
   • 소형 모델 ($M_w$) 이 다음 추론 단계를 생성합니다.
   • **프로세스 보상 모델 (PRM, Process Reward Model)**을 사용하여 생성된 단계의 정확도 점수를 실시간으로 평가합니다.
2. 라우팅 결정 (Routing Decision):
   • Continue: 단계가 정확하다고 판단되면 소형 모델의 출력을 수용하고 다음 단계를 계속 생성합니다.
   • Regenerate: 단계가 오류 가능성이 높다고 판단되면, 해당 단계만 대형 모델 ($M_s$) 에게 재생성 (Regenerate) 을 요청합니다.
   • 중요한 설계: 대형 모델이 한 번 개입하여 오류를 수정한 후, 다시 소형 모델로 제어권을 넘겨 나머지 단계를 생성하게 합니다. (대형 모델이 전체 남은 과정을 가져가는 것이 아님)
3. 비용 측정: 비용은 대형 모델이 생성한 토큰 수로 정의됩니다. (프리필 비용은 병렬 처리로 상쇄 가능하므로, 순차적 디코딩 비용인 생성 토큰 수를 주요 비용 지표로 사용)

TRIM 의 라우팅 전략 (3 가지 변형)

논문은 다양한 정보 제약 조건에 맞춰 세 가지 라우팅 정책을 제안합니다.

1. TRIM-Thr (Myopic Thresholding):
   • 현재 단계의 PRM 점수가 임계값 ($k$) 미만일 때만 대형 모델로 전환하는 단순한 임계값 기반 정책입니다.
   • 계산 비용이 매우 낮고 구현이 간단합니다.
2. TRIM-Agg (RL-Trained Policy):
   • 강화학습 (RL) 을 사용하여 학습된 정책입니다.
   • 현재 단계의 정확도뿐만 아니라, **과거 단계들의 누적 정확도 (최소값, 곱셈 등)**와 **토큰 수 (비용)**를 특징으로 사용하여 장기적인 정확도 - 비용 트레이드오프를 고려합니다.
   • PPO(Proximal Policy Optimization) 를 통해 최적화됩니다.
3. TRIM-POMDP (Partially Observable MDP):
   • PRM 점수가 불완전한 관찰 (Noisy Observation) 일 수 있다는 점을 명시적으로 고려합니다.
   • 추론 경로의 실제 상태 (정확함, 복구 불가능한 오류, 최근 단계만 오류) 를 **잠재 상태 (Latent State)**로 정의하고, PRM 점수를 이를 관찰하는 노이즈가 있는 신호로 간주합니다.
   • SARSOP 솔버를 사용하여 불확실성 하에서 최적의 라우팅 정책을 계산합니다. 이는 희소 보상 (Sparse Reward) 환경에서도 안정적인 성능을 보입니다.

3. 주요 기여 (Key Contributions)

1. 타겟팅된 개입의 효율성 입증: 전체 쿼리를 대형 모델에 맡기는 대신, 소수의 **임계 단계 (Critical Steps)**에만 대형 모델을 투입함으로써 다단계 추론의 효율성을 획기적으로 향상시킬 수 있음을 체계적으로 증명했습니다.
2. 실용적인 프레임워크 및 정책 개발: 단순한 임계값 정책부터 RL 기반, POMDP 기반의 정교한 정책까지 다양한 라우팅 전략을 제시했습니다. 이는 제한된 감독 데이터로도 훈련 가능하며, 다양한 비용 예산 (Budget) 에 대해 견고한 성능을 보입니다.
3. 강력한 일반화 능력: AIME 데이터셋으로 훈련된 라우팅 정책이 OlympiadBench, Minerva Math 등 다른 수학 벤치마크에서도 뛰어난 성능을 발휘함을 확인했습니다. 이는 단계별 난이도 패턴이 데이터셋 특유의 것이 아니라 추론 작업의 보편적 특성임을 시사합니다.

4. 실험 결과 (Results)

실험은 Qwen2.5-3B (소형) 와 Claude 3.7 Sonnet (대형) 을 모델 쌍으로 사용하며, MATH-500, AIME, OlympiadBench 등에서 평가되었습니다.

• 비용 효율성 (Cost Efficiency):
  • MATH-500: 가장 간단한 TRIM-Thr 정책만으로도 기존 라우팅 방법 (RouteLLM, AutoMix 등) 대비 5 배 높은 비용 효율성을 달성했습니다.
  • 고급 정책 (TRIM-Agg/POMDP): 대형 모델의 성능을 유지하면서 80% 적은 대형 모델 토큰을 사용했습니다.
  • AIME (더 어려운 벤치마크): TRIM-POMDP 는 기존 방법 대비 6 배 이상 높은 비용 효율성을 보였습니다.
• 교차 데이터셋 일반화:
  • AIME 에서 훈련된 TRIM-Agg 는 OlympiadBench 에서 11.68 배의 비용 효율성 향상을 보였습니다. 반면, 기존 쿼리 단위 라우팅 방법 (BERT 등) 은 데이터셋이 바뀌면 성능이 급격히 저하되었습니다.
• 지연 시간 (Latency):
  • TRIM 은 스펙큘레이티브 디코딩 (Speculative Decoding) 기법과 유사한 시스템 최적화 (KV 캐시 공유 등) 를 통해 적용 시 지연 시간 오버헤드가 거의 없으며, 대형 모델 단독 실행보다 더 빠른 경우가 많음을 입증했습니다.

5. 의의 및 결론 (Significance)

• 추론의 본질적 특성 파악: 다단계 추론에서 오류는 특정 '중요한 결정 단계'에 집중되어 있으며, 이 단계만 정확히 교정하면 전체 경로를 성공적으로 유도할 수 있음을 보여줍니다.
• 비용 최적화의 새로운 패러다임: "모든 쿼리를 가장 큰 모델로" 또는 "모든 쿼리를 작은 모델로"라는 이분법을 넘어, 단계별 (Step-level) 로 자원을 동적으로 할당하는 것이 추론 작업의 효율성을 극대화하는 핵심임을 증명했습니다.
• 실무 적용 가능성: 제한된 예산 하에서도 대형 모델의 성능을 유지할 수 있는 실용적인 솔루션을 제공하며, 추론이 필요한 다양한 분야 (수학, 과학, 코드 생성 등) 에 적용 가능한 확장성을 가집니다.

결론적으로, TRIM 은 LLM 추론의 비용 장벽을 낮추면서도 고난도 추론 작업의 정확도를 유지하기 위한 지능적이고 타겟팅된 하이브리드 라우팅 프레임워크로서, 추론 효율성 분야에서 중요한 진전을 이룩했습니다.