TARo: LLM 테스트 시간 정렬을 위한 토큰 레벨 적응형 라우팅

TARo: Token-level Adaptive Routing for LLM Test-time Alignment

대규모 언어 모델(LLM)은 뛰어난 추론 능력을 보이지만, 일반적으로 높은 성능을 달성하기 위해서는 비용이 많이 드는 추가 훈련이 필요합니다. 최근의 테스트 시간 정렬 방법은 경량화된 대안을 제공하지만, 주로 선호도 정렬에 초점을 맞추고 있으며, 추론 능력 향상을 위한 연구는 부족했습니다. 이러한 격차를 해소하기 위해, 본 논문에서는 추론 과정에서 동결된 LLM을 구조화된 추론으로 유도하는 토큰 레벨 적응형 라우팅(TARo) 기법을 제안합니다. 구체적으로, 우리는 먼저 단계별 수학적 추적 데이터를 기반으로 보상 모델을 훈련하여 미세한 논리적 일관성 신호를 포착하고, 학습 가능한 토큰 레벨 라우터를 도입하여 보상 모델의 지침이 기본 모델에 미치는 영향을 자동으로 제어합니다. 광범위한 실험 결과, TARo는 기본 모델 대비 최대 +22.4%, 기존의 토큰 레벨 테스트 시간 정렬 방법 대비 +8.4%의 상당한 추론 성능 향상을 보여줍니다. 또한, TARo는 MedXpertQA 데이터셋에서의 일반화된 임상 추론 능력과 AlpacaEval 데이터셋에서의 지시사항 준수 능력을 향상시킵니다. 더욱이, TARo는 재훈련 없이 작은 모델부터 큰 모델까지 적용 가능하며, 테스트 시간 정렬을 단순한 선호도 최적화에서 벗어나 강력하고 다양한 분야에 적용 가능한 추론 능력 향상으로 확장합니다.

Large language models (LLMs) exhibit strong reasoning capabilities but typically require expensive post-training to reach high performance. Recent test-time alignment methods offer a lightweight alternative, but have been explored mainly for preference alignment rather than reasoning. To bridge this gap, we propose, Token-level Adaptive Routing (TARo), which steers frozen LLMs toward structured reasoning entirely at inference time. Specifically, we first train reward models on step-wise mathematical traces to capture fine-grained logical consistency signals, then introduce a learnable token-level router that automatically controls the guidance of the reward model to the base model. Extensive experiments show that TARo significantly improves reasoning performance by up to +22.4% over base model and +8.4% over existing token-level test-time alignment methods, while also boosting out-of-distribution clinical reasoning (MedXpertQA) and instruction following (AlpacaEval). Furthermore, TARo also generalizes from small to large backbones without retraining, extending test-time alignment from preference optimization to robust, cross-domain reasoning.