트리 기반 오프폴리시 강화학습으로 다단계 추론 최적화

초록

본 논문은 강력한 교사 모델이 생성한 MCTS 트리에서 얻은 부분 경로들을 활용해, 그룹 상대 정책 최적화(GRPO)를 단계별 커리큘럼 형태로 재구성한다. 각 프리픽스마다 기대 보상이 다르다는 점을 고려해 제안된 Staged Advantage Estimation(SAE)으로 저분산·프리픽스 인식 어드밴티지를 계산하고, 이를 Tree‑OPO 프레임워크에 적용해 수학 문제 해결 정확도를 향상시킨다.

상세 분석

Tree‑OPO는 기존 GRPO가 전제하는 “단일 프롬프트‑완전 트레이스” 방식을 탈피한다. 교사 모델이 오프라인으로 수행한 MCTS는 하나의 문제에 대해 다수의 부분 솔루션(프리픽스)과 그에 대한 성공/실패 정보를 트리 형태로 제공한다. 이 트리는 깊이에 따라 난이도가 자연스럽게 역커리큘럼을 형성하는데, 깊은 노드일수록 더 많은 컨텍스트를 포함해 성공 확률이 높다. 이러한 구조적 차이를 무시하고 전체 그룹에 대해 평균 보상만을 기준으로 어드밴티지를 정의하면, 얕은 프리픽스와 깊은 프리픽스가 동일한 기준선에 놓여 편향된 신호가 발생한다.

SAE는 이 문제를 수학적으로 정식화한다. 먼저 각 샘플 ( (p_k, r_k) ) 에 대해 원시 어드밴티지 ( a’_k = r_k - \alpha V(p_k) ) 를 계산한다. 여기서 ( V(p_k) ) 는 해당 프리픽스의 서브트리 성공률을 추정하는 베이스라인이며, ( \alpha ) 는 베이스라인 활용 정도를 조절한다. 이후 전체 어드밴티지 벡터 ( a ) 를 제약조건을 갖는 최소제곱 문제로 최적화한다. 핵심 제약은 (1) 평균 제로 ( \mathbf{1}^\top a = 0 ) , (2) L2 노름 제한 ( |a|2 \le N ) , (3) 트리 구조에 기반한 순서 제약 ( a_i + \delta{ij} \le a_j ) 이다. 순서 제약은 “부모‑자식 관계에서 실패‑성공 전이”와 “형제 관계에서 잠재 성공 경로 관찰”을 반영한다. 즉, 동일 프리픽스에서 실패했지만 더 깊은 자식이 성공을 보인 경우, 자식의 어드밴티지가 부모보다 크게 설정돼야 한다는 의미다.

제약식은 두 가지 모드로 구현된다. 하드 제약은 ( |a|2 = N ) 과 양의 마진 ( \delta{ij}>0 ) 을 강제해 어드밴티지 차이를 크게 만들고, 소프트 제약은 완화된 L2 제한과 ( \delta_{ij}=0 ) 을 사용해 해의 유일성을 보장한다. 이론적으로는 소프트 제약이 최적의 분산 감소와 기대 보상 정렬을 제공함이 증명된다. 실험에서는 계산 효율성을 위해 하드 제약을 채택했으며, 근사 히어스틱(서브트리 성공률, 낙관적/비관적 베이스라인)도 비교했다.

그 결과, SAE를 적용한 Tree‑OPO는 동일 설정의 GRPO 대비 샘플당 그래디언트 분산이 현저히 낮아졌으며, 특히 얕은 프리픽스에서의 학습 효율이 크게 개선되었다. 이는 트리 구조를 활용한 “프리픽스‑조건부 베이스라인”이 정책 업데이트 시 불필요한 잡음을 제거하고, 어려운 단계에 대한 신호를 보존함을 의미한다. 또한, 오프라인 교사 트리를 재활용함으로써 온라인 MCTS 탐색 비용을 크게 절감하면서도, 교사의 고품질 경로를 효과적으로 학생 정책에 전이시킬 수 있었다.