불확실성 인식 트리 탐색으로 강화된 외부 추론

초록

본 논문은 외부 추론 프레임워크에서 사용되는 프로세스 보상 모델(PRM)의 에피스테믹 불확실성이 OOD(분포 외) 상황에서 성능 저하를 일으킨다는 점을 실증하고, 이를 이론적으로 분석한다. Monte Monte Dropout을 이용해 불확실성을 추정하고, 강화학습 기반 컨트롤러가 탐색 예산을 동적으로 할당하는 Uncertainty‑Aware Tree Search(UATS)를 제안한다. MA​TH‑500·AIME24 등 수학 추론 벤치마크에서 UATS는 다양한 정책·PRM 조합에서 기존 탐색 방법보다 높은 정확도와 계산 효율성을 달성한다.

상세 분석

이 논문은 외부 추론 시스템을 구성하는 세 요소—정책 모델(πθ), 탐색 알고리즘, 그리고 프로세스 보상 모델(PRM, Rϕ)—사이의 상호작용을 깊이 파고든다. 기존 연구는 PRM을 “검증자”로 활용해 후보 추론 경로를 재점수화하고 최종 답을 선택하도록 설계했지만, PRM은 제한된 학습 데이터에 기반해 훈련되기 때문에 훈련 분포와 크게 벗어나는 추론 단계에 대해 높은 에피스테믹 불확실성을 보인다. 저자들은 두 가지 핵심 실험을 수행한다. 첫째, Math‑Shepherd‑PRM‑7B와 Qwen2.5‑Math‑PRM‑7B라는 공개 PRM을 LLemma‑7B와 Qwen‑2.5‑Instruct‑7B 정책 모델이 생성한 추론 트레이스를 평가한다. 결과는 OOD 트레이스에 대해 정확도가 급격히 떨어지고, Monte Monte Dropout을 통한 점수 분산(σ²)이 크게 증가함을 보여준다. 이는 PRM이 과신(over‑confident)한 잘못된 점수를 부여할 위험을 시사한다. 둘째, 이 현상을 정량화하기 위해 베이지안 관점에서 PRM의 예측 불확실성을 총 변동성으로 분해하고, 에피스테믹 불확실성이 데이터·파라미터 부족에서 기인한다는 이론적 근거를 제시한다.

이론적 분석에서는 탐색 단계 t에서 M개의 후보 h_{t,i}가 제시되고, PRM이 제공하는 점수 Rϕ(h)와 실제 보상 R*(h) 사이의 차이를 regret으로 정의한다. PRM이 불확실성을 무시하고 greedy하게 최고 점수를 선택하면, OOD 후보가 지속적으로 선택돼 누적 regret이 O(T) 수준으로 선형 증가한다. 반면, 불확실성을 고려해 점수와 분산을 동시에 평가하는 정책은 기대 regret이 O(√T) 이하의 서브선형 수렴을 보인다. 이는 불확실성‑가중 탐색이 장기적으로 더 나은 해를 찾을 확률을 수학적으로 보증한다는 의미다.

위 이론을 실제 알고리즘에 적용하기 위해 저자들은 UATS(Uncertainty‑Aware Tree Search)를 설계한다. 핵심 구성은 (1) Monte Monte Dropout을 이용한 다중 샘플링으로 각 후보의 점수 평균과 분산을 추정하고, (2) 불확실성이 높은 후보에 대해 추가적인 재평가를 수행하는 휴리스틱 탐색 규칙, (3) 탐색 예산 C를 상태·액션 공간으로 보는 마코프 결정 과정(MDP)으로 모델링하고, 정책‑가치 네트워크를 강화학습(RL)으로 학습시켜 예산을 동적으로 배분한다. RL 컨트롤러는 탐색 깊이, 재평가 횟수, beam width 등을 상황에 맞게 조정함으로써 “탐색·예산·불확실성” 삼각형을 최적화한다.

실험에서는 MA​TH‑500(500문제)과 AIME24(수학 올림피아드) 두 데이터셋에 대해 Qwen‑2.5, Llama 3.1/3.2 등 최신 LLM을 정책으로, 다수의 공개 PRM을 검증자로 사용하였다. 다양한 컴퓨팅 예산(예: 4‑step beam, 8‑candidate final) 하에서 UATS는 기존 Beam Search, Best‑of‑N, Tree‑of‑Thoughts 등과 비교해 평균 정확도가 3~7%p 상승했으며, 특히 정책·PRM 간 분포 차이가 클 때(예: LLemma‑7B + Qwen‑PRM) 효과가 두드러졌다. 또한, 불확실성 추정 비용이 전체 예산의 10% 이하에 불과함을 보여, 실용적인 비용 효율성도 입증했다.

결과적으로 이 논문은 PRM의 에피스테믹 불확실성을 무시하면 외부 추론 시스템이 근본적인 취약점을 드러낸다는 점을 실증하고, 불확실성을 정량화·활용하는 탐색 전략이 이론적·실험적 측면 모두에서 성능 향상을 가져온다는 강력한 증거를 제공한다.