Aligning Tree-Search Policies with Fixed Token Budgets in Test-Time Scaling of LLMs

기술 요약: LLM의 테스트 타임 스케일링에서 고정 토큰 예산에 맞춘 트리 탐색 정책 정렬

문제 정의

대규모 언어 모델(LLM)은 모델 파라미터를 업데이트하지 않고 추론 중에 추가 연산량을 할당하는 **테스트 타임 스케일링(Test-Time Scaling)**을 통해 응답 품질을 향상할 수 있습니다. 트리 탐색 디코딩(예: 몬테카를로 트리 탐색, MCTS)은 효과적인 테스트 타임 스케일링의 한 형태이지만, 실제 배포 환경에서는 대개 고정된 쿼리당 토큰 예산이 부과됩니다.

기존의 트리 탐색 정책들은 대체로 **예산에 무관(budget-agnostic)**합니다. 이들은 토큰 예산을 단순히 종료 조건으로 취급할 뿐, 탐색 동작을 가이드하는 동적인 제약 조건으로 다루지 않습니다. 이러한 불일치는 다음과 같은 두 가지 주요 실패 모드를 초래합니다:

1. 후기 단계의 과도한 브랜칭(Late-stage over-branching): 탐색 과정 후반부에 얕은 노드들을 계속 확장하여, 유망한 후보를 정교화하거나 완성하기도 전에 예산을 소진해 버립니다.
2. 조기 종료(Premature termination): 탐색이 너무 일찍 멈추어 할당된 예산의 일부를 사용하지 못한 채 남겨둡니다.

현재의 조기 종료 변형 방식들은 토큰을 절약하는 데는 집중하지만, 남은 예산이 감소함에 따라 탐색 전략을 광범의 탐색(exploration)에서 정교화(refinement)로 동적으로 전환하는 정책은 결여되어 있습니다.

방법론: 예산 가이드 MCTS (BG-MCTS)

저자들은 남은 토큰 예산에 따라 탐색 정책을 명시적으로 조건화하는 트리 탐색 디코딩 알고리즘인 **BG-MCTS(Budget-Guided MCTS)**를 제안합니다. 핵심 메커니즘은 다음과 같이 정의되는 예산 충분도 비율 $\rho \in [0, 1]$에 의존합니다:
$$\rho = 1 - \frac{C_{used}}{B}$$
여기서 $C_{used}$는 누적 생성된 토큰 수이고 $B$는 총 예산입니다. $\rho$가 감소함에 따라 ($\rho \to 0$), 정책은 탐색(exploration)에서 착취(exploitation) 및 완성(completion)으로 전환됩니다.

BG-MCTS는 두 가지 예산 조건부 제어 메커즘을 도입합니다:

1. 예산 가이드 선택 (BG-PUCT)

표준 PUCT 선택 방식은 탐색 보너스를 어닐링(annealing)하고 완성 편향(completion bias)을 도입하도록 수정되었습니다:
$$\text{BG-PUCT}(p, s, \rho) = \frac{\tilde{W}(s, \rho)}{m_s} + \rho \cdot c \cdot P(s|p) \sqrt{\frac{\ln(m_p)}{m_s}}$$

• 탐색 어닐링(Exploration Annealing): 탐색 항에 $\rho$를 곱합니다. 예산이 고갈됨에 ($\rho \to 0$) 새로운 브랜치를 탐색하려는 유인이 감소합니다.
• 완성 편향(Completion Bias): 누적된 가치 $\tilde{W}(s, \rho)$는 노드 깊이 $d(x)$와 $(1-\rho)$에 비례하는 항에 의해 보정됩니다. 이는 이미 완성된 답변을 포함하는 노드를 인위적으로 높이지 않으면서, 최종 답변 생성에 더 가까운 깊은 노드를 선택하도록 유도합니다.

2. 예산 가이드 트리 와이드닝 (Tree Widening)

후기 단계의 낭비적인 브랜칭을 방기하기 위해, BG-MCTS는 "와이드닝"(부모로부터 새로운 자식을 생성하는 것)을 기존 자식들과 함께 선택 가능한 액션으로 취급합니다. 가상 생성 자식(virtual generative child) $s_{gen}(p)$가 선택 집합에 추가됩니다. 그 점수는 다음과 같습니다:
$$E_{gen}(p, \rho) = \mu(p) + \lambda \cdot \rho \cdot \sigma^2(p)$$

• $\mu(p)$와 $\sigma^2(p)$는 기존 자식들의 평균과 분산을 나타냅니다.
• 이 항은 $\rho$에 의해 스케일링됩니다. 예산이 풍부할 때는 와이드닝이 권장됩니다 (특히 자식들 간의 의견이 불일치할 때). 예산이 부족해지면 와이드닝에 대한 유인이 억제되어, 기존의 유망한 궤적을 심화(deepen)시키는 데 집중하게 됩니다.

통합 알고리즘

알고리즘은 표준 MCTS 루프를 따르되, 선택(selection)과 와이드닝(widening) 결정을 모두 $\rho$에 의존하게 만듭니다. 이는 wide-to-deep 스케줄을 유도합니다: 즉, 탐색 초기에는 넓은 탐색을 수행하고, 예산이 소진됨에 따라 집중적인 정교화와 답변 완성을 수행합니다.

주요 기여

1. 예산 인식 정책 설계: 본 논문은 예산을 단순히 정적 중단 신호로 취급하는 것이 아니라, 남은 토큰 예산에 따라 탐색 동작(탐색 vs 착취 vs 와이드닝)을 동적으로 정렬하는 최초의 트리 탐색 디코딩 프레임워크를 소개합니다.
2. 새로운 메커니즘: BG-PUCT(예산 조건부 탐색 및 깊이 편향 가치 보정 포함)와 예산 인식 와이드닝을 위한 **가상 생성 자식(Virtual Generative Children)**을 제안합니다.
3. 실증적 검증: 오픈 웨이트 LLM(Llama-3.1-8B, Qwen2.5-7B, Qwen3-32B)을 사용하여 수학적 추론 벤치마크(MATH500, AIME24/25)와 물리 추론 벤치마크(UGPhysics)에 대해 포괄적인 평가를 수행했습니다.

실험 결과

저자들은 10k, 20k, 30k 토큰의 고정 예산 하에서 BG-MCTS를 예산 무관 베이스라인(Standard MCTS, AB-MCTS-M, LiteSearch) 및 스케일링 베이스라인(Repeated Sampling, Sequential Refinement)과 비교 평가했습니다.

• 성능: BG-MCTS는 모든 모델과 예산에 대해 모든 베이스라인을 일관되적으로 능가했습니다. 예를 들어, Qwen2.5-7B를 사용한 MATH500에서 BG-MCTS는 30k 토큰 기준 **71.1%**의 정확도를 달이트었으며, 이는 차순위 트리 탐색 베이스라인인 MCTS의 **65.9%**보다 높은 수치입니다.
• 예산 활용도: 종종 조기에 종료되는 LiteSearch와 달리, BG-MCTS는 예산이 소진될 때까지 정확도를 지속적으로 향상시킵니다.
• 질적 행동: 트리 탐색 분석 결과, BG-MCTS는 초기에 넓은 폭을 유지하다가 탐색 후반부에 특정 브랜치를 심화시키는 방향으로 전환됨을 보여주었습니다. 반면, 표준 MCTS는 예산이 소진될 때까지 얕은 노드들을 계속 넓히는 경향이 있어, 유망한 경로를 완성하는 데 실패합니다.
• 물리 추론: 이 방법의 효과는 UGPhysics에서도 입증되었으며, 이는 'wide-to-deep'의 예산 정렬 동작이 순수 수학적 추론을 넘어 유익함을 보여줍니다.

의의 및 주장

본 논문은 트리 탐색 결정을 단순히 중단 조건으로 사용하는 것보다 남은 예산에 따라 조건화하는 것이 더 효과적임을 주장합니다.

• 예측 가능한 스케일링: BG-MCTS는 사용자가 토큰 예산을 지정하면 정책이 해당 비용 범위 내에서 탐색-완성 스케줄을 자동으로 조정하는, 예측 가능한 테스트 타임 스케일링을 위한 간단한 인터페이스를 제공합니다.
• 데이터 합성 시사점: 저자들은 BG-MCTS가 문제 인스턴스당 답변된 노드 수는 적지만, 훨씬 높은 정밀도(정확성)를 가진다는 점에 주목했습니다. 이러한 "더 적은 인스턴스, 더 나은 답변" 체계는 높은 품질의 추론 궤적이 광범위한 커버리지보다 더 중요한 데이터 합성 분야에서 특히 가치가 높습니다.
• 강건성: 이 방법은 약하거나 포화된 신호를 제공하는 불완전한 프로세스 보상 모델(PRM)을 사용할 때도 효과적이며, 이는 예산 인식 구조적 변화가 노이즈가 있는 평가 신호를 보완할 수 있음을 시사합니다.

결론적으로, 본 연구는 고정된 토큰 제약에 맞춰 트리 탐색 정책을 정렬하는 것이 자원이 제한된 실제 환경에서 LLM을 배포하는 데 있어 매우 중요한 단계임을 밝히고 있습니다.