기술 요약: 로컬 브랜치 라우팅(Local Branch Routing)을 통한 효율적이고 학습 가능한 언어 모델 테스트 타임 스케일링

문제 정의

최근 언어 모델(LLM) 추론의 발전은 추론 시점에 추가적인 계산량을 할당하는 것(테스트 타임 스케일링)이 성능을 크게 향상시킬 수 있음을 보여줍니다. 그러나 기존 방식들은 다음과 같은 결정적인 트레이드오프에 직면해 있습니다:

• 사고 사슬(Chain-of-Thought, CoT) 샘플링: 효과적이긴 하지만, 표준 CoT는 대안들의 결과를 관찰하기 전에 하나의 토큰을 확정 짓는 단일 스레드 방식에 머물러 있습니다.
• 탐색 기반 방법(Search-Based Methods): Tree of Thoughts나 Monte Carlo Tree Search(MCTS)와 같은 방식은 중간 사고의 트리를 탐색함으로써 너비(width)를 도입합니다. 그러나 이러한 방식은 종종 입자가 거칠고(coarse-grained), 계산 비용이 많이 들며, 단일한 학습 가능한 디코딩 정책으로서 엔드 투 엔드로 최적화하기 어렵습니다.
• 소프트 토큰(Soft-Token) 방법: Soft Thinking이나 Multiplex Thinking과 같은 기술은 후보 임베딩을 연속적인 혼합체로 병합하여 토큰 수준의 너비를 제공합니다. 학습은 가능하지만, 이 방식은 후보 브랜치들의 이산적인 정체성을 흐리게 하여 개념 특화된 정보를 가릴 위험이 있습니다.

본 논문은 언어 모델이 전체 솔루션 수준의 탐색 비용을 들이거나 소프트 혼합체의 모호함을 초래하지 않으면서도, 로컬 미래 후보들의 은닉 상태(hidden states)로부터 정보를 얻을 수 있는 경량화된 토큰 수준의 테스트 타임 스케일링 메커니즘으로부터 이득을 얻을 수 있다고 주장합니다.

방법론: 로컬 브랜치 라우팅 (Local Branch Routing, LBR)

저자들은 학습 가능한 토큰 수준의 브랜치 앤 루트(branch-and-route) 디코딩 프레임워크인 LBR을 제안합니다. LBR은 표준 자기회귀 생성에 롤링 로컬 룩어헤드 트리(rolling local lookahead tree)를 결一组합니다.

디코딩 파이프라인

각 디코딩 단계 $t$에서, 확정된 접두사 $x_{<t}$가 주어지면 LBR은 4단계의 "가지치기-이동-성장(prune–shift–grow)" 과정을 실행합니다.

1. 성장 (Grow): 단일 다음 토큰을 샘플링하는 대신, LBR은 $x_{<t}$에 뿌리를 둔 너비 $K$, 깊이 $L$의 로컬 트리 $\mathcal{A}_t$를 확장합니다. 샘플링된 모든 노드는 언어 모델을 통해 **포워딩(forwarded)**되어 은닉 상태 $h_t(v)$를 계산합니다. 이는 단순히 후보 문자열의 집합이 아닌, 로컬 은닉 상태의 트리를 생성합니다.
2. 라우팅 (Route): 경량 라우터 $\rho_\phi$는 포워딩된 트리를 관찰하고 확정할 하나의 깊이-1 서브트리를 선택합니다. 라우터는 집합 주의 집중(set-attention) 메커니즘을 사용하여 각 후보 서브트리(예: $L=1$인 경우 단일 토큰, 또는 $L>1$인 경우 토큰과 그 연속어들)를 벡터로 인코딩합니다. 이 벡터들은 교차 서브트리 주의 집중(cross-subtree attention)을 통해 비교되며, 이를 통해 라우터는 형제 미래(sibling futures) 사이에서 대조적인 결정을 내릴 수 있습니다.
3. 가지치기 (Prune): 라우터는 최적의 서브트리(인덱스 $k^*_t$)를 선택합니다. 나머지 모든 서브트리는 폐기됩니다. 선택된 서브트리의 루트 토큰이 다음 토큰 $x_t$로 확정됩니다.
4. 이동 및 재성장 (Shift and Regrow): 선택된 서브트리는 앞으로 이동합니다(루트가 제거되고 그 자손들이 다음 단계의 새로운 접두사가 됨). 새로운 레이어가 트리 깊이 $L$을 복구하기 위해 성장하며, 사이클이 반복됩니다.

학습 목적 함수

LBR은 검증 가능한 보상(RLVR)을 통한 엔드 투 엔드 강화 학습을 가능하게 하는 다루기 쉬운 **트리 궤적 가능도(tree-trajectory likelihood)**를 정의합니다.

• 확률성 (Stochasticity): 두 가지 작업만이 확률적입니다: 베이스 모델에서 새로 성장하는 노드 $G_t$를 샘플링하는 것과 라우터 결정 $k^*_t$를 샘플링하는 것입니다. 가지치기, 이동, 노드 재사용은 결정론적입니다.
• 가능도 (Likelihood): 궤적 $F$의 로그 확률은 새로 성장한 토큰들의 로그 확률과 라우터의 선택 확률의 합입니다:
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• 최적화: 프레임워크는 베이스 모델의 토큰 생성과 라우터의 선택 결정을 모두 보상하는 어드밴티지 $A(F)$를 사용하는 가능도 비율 목적 함수(예: GRPO)를 사용하여 학습됩니다.

주요 기여

1. 이산적 정체성을 가진 토큰 수준 브랜칭: 후보들을 연속적인 혼합체로 병합하는 소프트 토큰 방식과 달리, LBR은 후보들을 이산적인 포워딩 브랜치로 유지합니다. 이를 통해 라우터는 포스트-후보 은닉 상태(post-candidate hidden states)—즉, 후보 토큰이 생성되고 모델에 의해 처리된 후에만 드러나는 정보—를 활용할 수 있습니다.
2. 다루기 쉬운 엔드 투 엔드 학습: 트리의 궤적에 대한 특정 가능도를 정의함으로써, LBR은 이산 토큰 추론 시스템에서 사용되는 것과 동일한 RLVR 원칙 하에 베이스 모델과 라우터를 공동 최적화할 수 있게 합니다.
3. 효율적인 테스트 타임 스케일링: 이 프레임워크는 얕은 룩어헤드($L$이 1 또는 2와 같이 작음)를 통해 토큰 수준에서 너비를 도입하여, 전체 깊이 탐색의 계산 오버헤드를 피하면서도 루트 다음 토큰 분포 이상의 증거를 활용합니다.
4. 집합 주의 집중 라우터 (Set-Attention Router): 저자들은 교차 서브트리 주의 집중을 사용하여 형제 후보들을 비교하는 라우터 아키텍처를 도입하였으며, 이를 통해 라우팅이 독립적인 재순위화(reranking)가 아닌 대조적 결정임을 입증했습니다.

실험 결과

합성 계층적 계획 (Synthetic Hierarchical Planning)

Radix-번역된 그래프 도달 가능성(radix-translated graph reachability)을 포함하는 합성 태스크에서:

• 성능: LBR은 목표 정확도 0.967을 달성하여, Discrete CoT(0.914)와 Soft Thinking(0.845)을 능가했습니다.
• 메커니즘 분석: 프로빙 실험 결과, LBR의 포스트-후보 은닉 상태가 브랜칭 전 상태보다 최종 타겟을 훨씬 더 잘 예측하는 것으로 나타났습니다. 이는 후보를 포워딩하는 것이 표준 자기회귀 디코딩에서는 불가능한 다운스트림 정보를 모델이 접근할 수 있게 함을 확인시켜 줍니다.
• 표현 (Representation): 연속적인 혼합체를 통해 개념 정체성을 흐리는 Soft Thinking과 달리, LBR은 이산적인 브랜치 정체성을 보존하여 올바른 경로를 더 잘 회복할 수 있게 했습니다.

수학적 추론 벤치마크

DeepSeek-R1-Distill-Qwen 모델(1.5B 및 7B)을 사용하여 6개의 벤치마크(Minerva, AIME, MATH500 등)에서 평가했습니다:

• 성능 향상: LBR은 다음 모델들에 비해 Pass@1 및 Pass@32 모두에서 성능을 개선했습니다:
  • Discrete Chain-of-Thought.
  • Vanilla discrete-token RLVR.
  • RL 호환 소프트 토큰 브랜칭 베이스라인 (예: Multiplex Thinking).
• 깊이 절제 연구 (Depth Ablation): 룩어헤드 깊이를 $L=1$에서 $L=2$로 늘리면 특히 7B 백본에서 성능이 더욱 향상되었으며, 이는 더 깊은 로컬 연속(continuation)이 더 많은 라우팅 증거를 제공함을 시사합니다.
• 라우터 절제 연구 (Router Ablation): 교차 서브트리 주의 집중을 제거(독립적인 라우터 사용)했을 때 성능이 낮아졌으며, 이는 형제 브랜치들의 대조적 비교가 중요하다는 것을 입증합니다.
• 학습 역학 (Training Dynamics): RLVR 학습 동안 라우터의 엔트로피는 라우터가 더 결정적으로 변함에 따라 감소했고, 라우팅된 위치에서의 베이스 모델의 엔트로피 또한 감소했습니다. 이는 모델이 더 구별 가능한 로컬 브랜치를 생성하도록 학습되는 공진화(co-adaptation)를 나타냅니다.

의의 및 주장

본 논문은 로컬 브랜치 라우팅이 언어 모델의 효율적이고 학습 가능하며 이산적인 형태의 테스트 타임 스케일링을 제공한다고 주장합니다. 주요 의의는 다음과 같습니다:

1. 포스트-후보 증거의 가치: 라우팅 결정은 단순히 브랜칭 전의 분포뿐만 아니라, 후보 토큰에 의해 유도된 은닉 상태로부터 이득을 얻습니다.
2. 이산 브랜칭의 우월성: 이산적인 브랜치 정체성을 보존하는 것은 개념 정체성을 흐릴 수 있는 소프트 토큰 혼합체보다 더 명확한 의사결정을 가능하게 합니다.
3. 확장성: 이 프레임워크는 단일 스레드 CoT와 비용이 많이 드는 전체 트리 탐색 사이의 중간 지점을 제공하며, 경량 로컬 브랜칭을 통해 효과적인 스케일링을 가능하게 합니다.

저자들은 **국소성(locality)**의 한계를 인정합니다. LBR은 얕은 룩어헤드 윈도우에 의해 제한되며 완전한 솔루션에 대한 전역 탐색을 수행하지 않습니다. 향-후 연구로는 적응형 깊이(adaptive depth) 또는 선택적 전역 탐색을 결합하는 방향이 제시되었습니다.