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

Motivation

• Known: LLM의 test-time scaling은 parallel sampling-and-aggregation, sequential refinement, 그리고 이 둘을 결합한 hybrid 방식(주로 tree-search decoding 기반 MCTS/PUCT)으로 나뉘며, 예산이 커질수록 더 많은 대안을 탐색해 답변 품질을 높일 수 있음이 알려져 있다.
• Gap: 기존 tree-search 정책(MCTS with PUCT 등)은 대체로 budget-agnostic하여 토큰 예산을 단순 종료 조건으로만 사용하기 때문에, 후반부 과도한 branching으로 예산을 조기 소진하거나 반대로 예산을 다 쓰지 못하고 조기 종료하는 문제가 발생하며, early-stopping 변형들도 branching에서 refinement로의 전환을 예산에 조건화하지 않는다.
• Why: 실제 배포 환경에서는 쿼리별로 고정된 토큰 예산이 제품·설정에 따라 다양하게 주어지므로, 주어진 예산을 최대한 효율적으로 활용해 답변 품질을 극대화하는 budget-aware 탐색 정책이 실용적으로 중요하다.
• Approach: MCTS의 PUCT 선택 규칙과 tree widening 메커니즘을 남은 예산 비율(budget sufficiency ratio) ρ에 조건화하여, 탐색 초기에는 넓은 exploration을, 예산이 줄어들수록 깊은 노드의 정제와 답변 완성을 우선하는 BG-MCTS를 제안한다.

Achievement

1. BG-PUCT 스코어 제안: 표준 PUCT 식의 exploration term에 ρ를 곱해 예산 감소에 따라 탐색 강도를 점진적으로 줄이고, exploitation term에는 depth-biased completion bias를 추가한 ˜W(s, ρ)를 도입하여 후반부에 얕은 노드보다 답변 완성에 가까운 깊은 노드를 선호하도록 만들었다.
2. Budget-guided tree widening 메커니즘 제안: 남은 예산에 따라 새로운 자식 노드 생성 여부를 조절하여 후반부의 불필요한 branching을 억제한다.
3. 일관된 성능 우위 입증: MATH500, AIME24/25 수학 추론 벤치마크와 UGPhysics 물리 추론 벤치마크, 그리고 Llama-3.1-8B-Instruct, Qwen2.5-7B-Instruct, Qwen3-32B 등 open-weight LLM들에 대해 토큰 예산 B∈{10k,20k,30k} 전 구간에서 budget-agnostic tree-search 베이스라인 대비 일관되게 우수한 정확도를 달성하였다.

How

• Budget sufficiency ratio ρ = 1 − Cused/B를 정의해 탐색 정책의 조건 변수로 사용.
• BG-PUCT(p,s,ρ) = ˜W(s,ρ)/ms + ρc·P(s|p)·sqrt(ln(mp)/ms) 형태로, exploration term에 ρ를 곱해 예산 감소 시 탐색을 점진적으로 억제.
• ˜Q(x,ρ) = Q(x) + κ(1−ρ)^(d(x)/ˆdans) 형태의 completion bias를 도입하여, 답변 완성 예상 깊이(ˆdans)에 가까운 노드일수록 (1−ρ)가 커질수록 더 큰 보정값을 받도록 함. 단, 이미 답을 포함한 노드는 보정 없이 최종 평가값 그대로 사용.
• Budget-guided tree widening 메커니즘(가상의 생성 비용 등을 활용, Fig. 2)을 통해 남은 예산에 따라 새로운 자식 노드 도입 시점을 조절.
• MATH500, AIME24/25, UGPhysics 벤치마크에서 여러 open-weight LLM과 다양한 토큰 예산 하에 baseline들과 accuracy, answered-tree rate 등을 비교.

Originality

• 기존 tree-search decoding 연구들이 대부분 budget을 단순 종료 조건으로만 취급한 것과 달리, 남은 예산 비율(ρ)을 selection과 widening 두 메커니즘 모두에 명시적으로 조건화하는 최초의 시도.
• PUCT의 exploration term에 ρ를 곱하는 단순하지만 효과적인 annealing과, depth 기반 completion bias를 결합한 ˜W(s,ρ) 값 보정 방식이 novel함.
• 수학 추론뿐 아니라 물리 추론(UGPhysics)까지 확장해 일반성을 검증한 점.

Limitation & Further Study

• κ, ˆdans의 추정치(답변 완성 평균 깊이) 등 여러 하이퍼파라미터에 의존하며, 이들의 민감도나 다양한 태스크/모델에 대한 일반화 가능성에 대한 분석이 제한적일 수 있음.
• 평가에 사용된 verifier/reward model의 품질에 따라 Q(x) 자체의 신뢰도가 달라질 수 있는데, 이에 대한 강건성 분석이 본문 발췌만으로는 충분히 드러나지 않음.
• 실험이 주로 수학·물리 추론 벤치마크에 국한되어 있어, 코드 생성, 대화형 태스크 등 다른 도메인에서의 일반화 가능성은 추가 검증이 필요함.
• 후속 연구로 ρ 조건화 함수 형태(선형/지수 등)에 대한 이론적 정당화나, 더 큰 모델·더 다양한 예산 범위에서의 확장 실험이 필요해 보임.

Evaluation

총평: 고정 토큰 예산이라는 실용적 제약을 tree-search decoding 정책 설계에 명시적으로 통합한 점이 참신하며, 다양한 벤치마크와 모델에서 일관된 개선을 보여 실용성이 높은 연구이다.