기술 요약: CATPO (Critique-Augmented Tree Policy Optimization)

문제 정의

검증 가능한 보상(Verifiable Rewards, RLVR)을 활용한 강화 학습은 대규모 언어 모델(LLM)의 추론 능력을 향상시키기 위한 지배적인 패러다임이 되었습니다. TREERPO와 같은 트리 기반 방식은 트리 구조의 롤아웃을 통해 조밀한 단계별 보상 신호를 생성함으로써 평탄한 궤적 샘플링(flat trajectory sampling)보다 개선된 성능을 보여주지만, 결정적인 비효전성 문제에 직면해 있습니다: 바로 정보가 없는 트리로 인한 계산 낭비입니다.

현재의 트리 기반 접근 방식은 훈련 중에 모든 샘플링된 트리를 동일하게 취급합니다. 그러나 많은 트리는 유용한 그래디언트(gradient) 신호를 거의 또는 전혀 제공하지 못합니다:

• Dead-correct 트리: 모든 리프 노드가 성공함. 모델이 이미 문제를 마스터했으므로, 개선을 위한 대조적 신호(contrastive signal)를 제공하지 못함.
• Dead-wrong 트리: 모든 리프 노드가 실패함. 양(+)의 브랜치가 없으면 그룹 상대적 이점(group-relative advantages)이 0으로 수렴하여, 롤아웃이 그래디언트 업데이트에 무용지물이 됨.
• Stale 트리: 정책의 확신도가 보상 분포를 완벽하게 추적하고 있어, 한계 효용이 체감되는 상태.

DAPO 및 GRESO와 같은 기존 방식들은 평탄한 롤아웃에서 유사한 문제를 해결하기 위해 프롬프트를 필터링하지만, 이는 프롬프트 수준에서만 작동합니다. 이들은 트리 롤아웃의 내부 구조, 즉 특정 브랜치는 놀랍거나 오류가 있는 반면 다른 브랜치는 그렇지 않은 상황을 진단하거나 수리할 수 없습니다.

방법론

CATPO(Critique-Augmented Tree Policy Optimization)는 트리 수준에서 트리 롤아웃을 진단, 수리 및 재가중치 부여하는 4단계 파이프라인을 통해 이러한 낭비를 해결합니다.

1. 트리 정보성 점수 ($F(T)$)

CATPO는 정책 업데이트를 위한 트리의 유용성을 정량화하는 지표를 도입합니다. 이 점수는 추가 비용 없이 샘플링 과정에서 이미 생성된 데이터를 사용하여 계산됩니다. 이는 두 가지 신호를 결합합니다:

• 리프 결과 다양성 (Leaf Outcome Diversity): 리프 통과율($\hat{p}$)의 베르누이 분산 $\hat{p}(1-\hat{p})$로 측정됩니다. 이 항은 결과가 혼재되어 있을 때(일부는 정답, 일부는 오답) 최대화되며, 모든 리프가 일치할 때는 0이 됩니다.
• 정책-보상 비상관성 (Policy-Reward Decorrelation): 노드 로그 확률과 전파된 보상 사이의 피어슨 상관계수 $\rho_{\pi,r}$를 이용한 $1 - \rho^2_{\pi,r}$로 측정됩니다. 이는 정책의 확신도가 보상 구조를 선형적으로 설명하지 못하는 잔여 불확실성을 포착합니다.

최종 점수는 $F(T) = \hat{p}(1-\hat{p}) \cdot (1 - \rho^2_{\pi,r})$입니다. 트리는 임계값에 따라 네 가지 범주로 분류됩니다:

• Dead-correct/Dead-wrong: 낮은 다양성 (거의 만장일치인 결과).
• Stale: 중간 정도의 신호.
• Informative: 높은 다양성 및 비상관성.

2. 비평 가이드 치유 (Critique-Guided Healing)

Dead-wrong 트리(모든 브랜치가 실패한 경우)에 대해, CATPO는 롤아웃을 버리지 않습니다. 대신, 트리를 "치유"하려고 시도합니다:

1. 실패 식별: 너비 우선 탐색(BFS)을 통해 모든 자식 노드가 실패(보상 < $\epsilon$)한 가장 얕은 노드 $v^*$를 찾습니다.
2. 비평 생성: 정책 모델 자체가 $v^*$에서의 오류에 대한 자연어 비평을 생성하도록 프롬프트되어, 구체적인 논리적 또는 수학적 실수를 식별합니다.
3. 그래프팅 (Grafting): 모델에 문제, $v^*$까지의 부분 솔루션, 그리고 생성된 비평을 함께 제공하여 $k$개의 정제된 연속된 결과(continuations)를 생성하도록 다시 프롬프트합니다.
4. 재전파: 이 새로운 브랜치들을 $v^*$ 지점에 접목(graft)하고, 다시 점수를 매긴 후 보상을 하단에서 상단으로 전파합니다. 이를 통해 "죽은" 트리를 양(+)의 보상 다양성을 가진 트리로 전환합니다.

3. 정보성 가중치 손실 (Informativeness-Weighted Loss)

최종 훈련 목적 함수는 정규화된 정보성 점수 $w(T) = F(T) / \bar{F}$에 따라 각 트리의 정책 그래디언트 기여도를 조절합니다.

• Informative 트리는 전체 가중치를 받아 파라미터 업데이트를 주도합니다.
• Dead-correct 및 stale 트리는 가중치가 낮아집니다.
• Healed 트리는 치유 과정이 성공하면 Informative한 것으로 간데합니다.

목적 함수는 클리핑된 서로게이트 손실(PPO/TREERPO와 유사)을 사용하며, 여기서 어드밴티지는 $w(T)$에 의해 변조되어, 전체 그래디언트 크기를 유지하면서도 가장 유용한 데이터에 업데이트를 집중시킵니다.

주요 기여

1. 트리 수준 진단: CATPO는 프롬프트 수준의 필터링을 넘어, 트리 내부의 정보성을 명시적으로 진단하고 훈련 루프 내에서 uninformative한 트리에 개입하는 최초의 방법입니다.
2. 제로 코스트 스코어링: 정보성 점수 $F(T)$는 별도의 모델이나 추가적인 순전파(forward pass) 없이 기존 트리 샘플링 통계(리프 결과 및 로그 확률)를 활용합니다.
3. 비평 가이드 치유: 정책 모델이 자신의 오류를 비평하고 수정된 브랜치를 접목하도록 하여 "dead-wrong" 트리로부터 훈련 신호를 회복하는 혁신적인 메커즘을 통해, 낭비되는 계산량을 유용한 훈련 데이터로 전환합니다.
4. 가중 최적화: 고가치 트리를 동적으로 우선순위화하여 샘플 효율성을 높이는 정보성 가중 손실을 제공합니다.

실험 결과

Qwen2.5-Math-1.5B를 MATH 데이터셋으로 학습시킨 후 AIME24, MATH-500, OlympiadBench, MinervaMath의 네 가지 벤치마크에서 실험을 수행했습니다.

• 성능: CATPO는 37.5%의 Macro Accuracy를 달성하여, 베이스라인인 TREERPO(35.6%)보다 1.9%, 평탄한 궤적 방식인 GRPO(32.7%)보다 4.8% 높은 성능을 보였습니다.
• 어려운 과제: 개선 효과는 베이스라인이 더 많은 dead-wrong 트리를 생성하는 어려운 벤치마크(AIME24 및 OlympiadBench)에서 가장 두드러졌으며, 이는 CATPO가 신호를 회복할 기회가 더 많음을 의미합니다.
• 일반화: 학습 분포 외부인 MinervaMath에서도 유의미한 이득이 관찰되었으며, 이는 이 방법이 단순한 오버피팅이 아닌 전이 가능한 추론 전략을 학습하는 데 도움을 준다는 것을 시사합니다.
• 가설 검증:
  • H1: 정보성 점수 $F(T)$는 트리당 그래디언트 노름(gradient norm)과 강한 양의 상관관계(Spearian $\rho = 0.62$)를 보였으며, 이는 $F(T)$가 훈련 유용성의 신뢰할 수 있는 대리 지표임을 입증합니다.
  • H2: 비평 가이드 치유는 적어도 하나의 정답 브랜치를 도입함으로써 **88.1%**의 dead-wrong 트리의 정보성을 성공적으로 개선했습니다.

의의 및 주장

본 논문은 CATPO가 트리 기반 샘플링의 고유한 "낭비" 문제를 직접 해결함으로써 LLM 추론을 위한 RLVR의 더 효율적이고 효과적인 접근법을 제공한다고 주장합니다. CATPO는 uninformative한 트리를 구분하고 이를 능동적으로 수리함으로써 계산 자원을 가장 가치 있는 데이터에 집중시킵니다. 저자들은 이 방식이 별도의 보상 모델이나 가치 네트워크를 필요로 하지 않으며, GRPO/TREERPO의 단순성을 유지하면서도 복잡한 추론 작업에서 성능을 크게 향고시킨다는 점을 강조합니다. 이 방법은 기존의 트리 기반 RL을 보완하는 발전된 형태로, 이전 방식들이 간과했던 롤아웃의 내부 구조를 타겟팅합니다.