예산에 맞춘 트리 탐색 디코딩: 고정 토큰 한도에서 효율적인 MCTS 설계

초록

본 논문은 테스트‑타임 스케일링을 위해 고정된 토큰 예산을 고려한 새로운 트리‑검색 디코딩 알고리즘인 Budget‑Guided MCTS(BG‑MCTS)를 제안한다. 예산 비율 ρ에 따라 탐색·확장 정책을 동적으로 조정함으로써 초기에는 넓게 탐색하고, 예산이 소진될수록 깊은 노드의 정제와 답변 완성에 집중한다. MATH500 및 AIME24/25 벤치마크에서 8‑10B 파라미터 오픈‑웨이트 LLM을 사용해 다양한 토큰 예산(10k, 20k, 30k) 하에 기존 예산‑무관 트리‑검색 기법들을 일관되게 능가한다.

상세 분석

BG‑MCTS는 기존 MCTS의 PUCT 선택 규칙을 예산 충분도 ρ = 1 − C_used/B 로 조정한다. 선택 점수는 BG‑PUCT(p,s,ρ)= \tilde W(s,ρ)·(m_s/∑_z m_z) + ρ·c·P(s|p)·ln(m_p/m_s) 로 정의되며, ρ가 감소하면 탐색 보너스가 점차 사라져 평균 가치가 높은 깊은 노드에 집중한다. 또한 \tilde W(s,ρ)는 깊이 보정값 κ·(1−ρ)·d(x)/\hat d_ans 를 더해, 예산이 부족해질수록 답변 완성에 유리한 깊은 경로를 선호한다.

폭넓은 탐색을 유지하기 위해 각 비단말 노드에 가상 생성 자식 s_gen(p)를 도입한다. 이 옵션은 기존 자식과 동일한 선택 과정에 포함되며, 선택 시 새로운 실제 자식을 하나 생성한다. s_gen(p)의 점수는 E_gen(p,ρ)= μ(p) + λ·ρ·σ²(p) 로 계산한다. 여기서 μ와 σ²는 현재 자식들의 평균·분산이며, 분산이 클수록(불확실성이 클수록) 폭넓은 확장이 장려된다. ρ가 작아지면 λ·ρ·σ²(p) 항이 감소해 예산이 거의 남지 않은 단계에서는 새로운 브랜치를 만들지 않도록 억제한다.

알고리즘 흐름은 (1) 현재 ρ 계산, (2) BG‑PUCT와 E_gen을 이용한 통합 선택, (3) 선택된 노드가 s_gen이면 폭넓은 확장, (4) 선택된 표준 자식이면 기존 확장 폭(k)만큼 새로운 자식 생성, (5) 각 새 노드에 보상 모델을 통해 Q값 부여, (6) 백프로파게이션으로 W와 m 업데이트, (7) 토큰 사용량 C_used을 누적하고 B에 도달할 때까지 반복한다.

실험에서는 Llama‑3.1‑8B‑Instruct와 Qwen‑2.5‑7B‑Instruct 두 모델을 사용했으며, 보상 모델로는 GenPRM‑7B를 적용했다. 고정 토큰 예산 10k, 20k, 30k 하에서 BG‑MCTS는 Repeated Sampling, Sequential Refinement, 기존 MCTS, AB‑MCTS‑M, LiteSearch 등과 비교해 정확도(정답 비율)와 토큰 효율성 모두에서 평균 2‑5%p 상승을 기록했다. 특히 예산이 작을수록(10k) 폭넓은 탐색을 억제하고 깊은 정제에 집중하는 BG‑MCTS의 이점이 두드러졌다.

추가 분석에서는 ρ‑스케줄링 없이 고정 탐색 폭을 사용한 변형, κ와 λ 파라미터 민감도, 그리고 가상 생성 자식 없이 순수 깊이‑우선 탐색을 수행한 경우를 비교했다. 모든 경우에서 ρ‑조건부 탐색·확장이 성능을 유지하거나 개선함을 확인했다. 제한점으로는 보상 모델의 품질에 크게 의존한다는 점과, 현재는 토큰 비용만을 고려해 예산을 정의했으므로 실제 GPU 메모리·시간 비용을 동시에 최적화하는 다중 제약 문제에 대한 확장은 필요하다.