PICO가 전부 실패한 롤아웃에서 정책 호환 궤적을 복원한다
When All Rollouts Fail: Policy-Compatible Joint Decoding for Hard-Problem Reasoning
무엇인가
이 논문은 검증 가능한 보상 기반 강화학습(RLVR)이 모델 능력 밖의 난문제에서 겪는 학습 절벽을 다룬다. RLVR은 온폴리시 탐색으로 얻은 궤적의 이진 보상에서 학습하지만, 한 문제의 모든 롤아웃이 실패하면 정답에 대한 정보가 전혀 남지 않는다. 특히 GRPO나 RLOO처럼 그룹 상대 이점을 쓰는 방법은 보상이 모두 0인 그룹에서 이점이 0으로 붕괴해 업데이트 신호가 사라진다. 전문가 궤적을 넣는 우회로가 있지만, 이 궤적은 현재 학생 정책의 저확률 영역에 놓이기 쉬워 직접 학습하면 과도한 정책 이동과 표면적 모방을 낳고, 오프정책 샘플로 섞으면 분포 불일치로 최적화가 불안정해진다. 접두사나 힌트로 유도하면 학생이 스스로 도달하지 못한 중간 상태에 의존하는 특권 문맥 문제도 생긴다. 저자들은 난문제에서 유용한 학습 신호가 검증 가능하게 정답이면서 동시에 현재 정책과 충분히 호환되어야 한다고 주장한다.
어떻게 동작하나
PICO의 핵심 구성 요소는 분포 제약 공동 디코딩이다. 교사 모델은 전문가 해답이라는 특권 정보 g에 조건화된 다음 토큰 분포 q_t를 내고, 학생은 g 없이 같은 생성 접두사 y_<t에 조건화된 p_t를 낸다. 각 디코딩 단계에서 두 모델의 상위 k 후보 집합을 교집합으로 묶어 F_t = TopK(q_t, k_T) ∩ TopK(p_t, k_S)를 만든다. 이후 하드 공동 디코딩 정책 μ_t^H(v) = p_t(v)1{v∈F_t}/Z_t를 사용해, 교사가 고른 후보 안에서 학생의 상대 토큰 확률을 그대로 보존한다. 교사는 후보 선택으로만 개입하고 소프트 재가중은 하지 않는다. 논문은 먼저 역 KL과 교사 신뢰도 항을 가진 목적함수를 두고 폐형해 μ ∝ p_t exp(λq_t)를 유도한 뒤, KL이 지원 집합 제한에 의한 −log Z_t와 소프트 재가중에 의한 추가 KL로 분해됨을 보인다. 실험적으로 지원 제한만으로 충분한 롤아웃 정확도가 나와 λ=0으로 두고, 학생의 상대 확률을 유지하는 하드 공동 디코딩을 최종 정책으로 쓴다.
무엇과 다른가
동적 top-k는 이 공동 디코딩의 분포 이동을 유한한 예산 안에 묶는 장치다. 하드 공동 디코딩에서 단계별 KL은 정확히 −log Z_t이고, Z_t는 교사 후보와 학생 상위 K_S 후보의 교집합이 남기는 학생 확률 질량이다. 따라서 Z_t ≥ Z_min을 보장하면 KL ≤ −log Z_min이 된다. 구현은 학생 컷오프 K_S를 고정하고 교사 컷오프 k를 1부터 K_Tmax까지 늘려 가며 F_t(k)와 Z_t(k)를 계산한 뒤, Z_min을 만족하는 가장 작은 k를 고른다. 즉 k_t* = min{k : Z_t(k) ≥ Z_min}이다. 만약 그런 컷오프가 없으면 해당 디코딩 단계에서는 학생 정책 p_t로 폴백한다. 이 절차가 교사 지도와 학생 정책 근접성 사이의 균형을 토큰 단위로 조정한다.
어떻게 쓰나
공동 디코딩으로 만든 궤적은 혼합 행동 정책 최적화로 학습한다. 공동 디코딩은 롤아웃 그룹의 n개 궤적이 모두 틀렸을 때만 발동한다. 이때 하드 공동 디코딩 정책으로 궤적 y^J를 하나 추가 샘플링하고, 검증을 통과하면 원래 그룹의 틀린 궤적 하나를 대체한다. 결과 그룹은 n−1개의 온폴리시 궤적과 1개의 공동 디코딩 궤적으로 구성되며, 이점은 재구성된 보상 집합의 평균을 빼서 계산한다. 학습 목적함수는 학생 궤적과 공동 궤적을 함께 클리핑된 대리 목적함수로 최적화한다. 공동 궤적에는 중요도 비율 r̂_J,t(θ) = π_θ(y_t^J|x,y_<t^J)/μ_t(y_t^J) = Z_t π_θ(y_t^J|x,y_<t^J)/π_old(y_t^J|x,y_<t^J)를 적용해, 공동 디코딩 행동 정책과 학생 정책 사이의 토큰 수준 불일치를 보정한다. 폴백 단계에서는 μ_t=p_t, Z_t=1이므로 표준 정책 비율로 줄어든다.
전제와 한계
실험은 Qwen3-4B-Base, Qwen3-8B-Base, Qwen3-4B-Instruct를 학생 정책으로 사용하고, 모든 설정에서 Qwen3-4B-Instruct가 전문가 해답에 조건화된 교사로 쓰였다. 수학·일반 추론 학습에는 DAPO-Math-17K를 쓰고, 평가는 AIME 2025, AIME 2026, HMMT 2026, MATH-500과 분포 밖 일반화용 GPQA-Diamond, HLE에서 수행했다. 베이스라인은 GRPO, LUFFY, Prefix-RFT, ReLIFT, RLRT다. PICO는 Qwen3-4B-Base에서 GRPO 대비 평균 4.1%포인트, Qwen3-8B-Base에서 3.6%포인트 향상했고, 가장 강한 경쟁 방법보다 각각 2.2%포인트와 1.8%포인트 앞섰다. AIME26과 HMMT26 같은 난문제에서 이득이 두드러졌고, GPQA-Diamond와 HLE에서도 향상이 나타나 수학 중심 학습 분포 밖으로 일반화됨을 보였다. 학습 중에는 더 높은 보상과 더 적은 solve-none 문제, 더 긴 응답이 관찰됐으며, 공동 디코딩이 전멸 그룹의 최대 80%에서 검증된 양의 보상 궤적을 복원했다고 보고한다. 에이전트 실험은 ALFWorld와 ScienceWorld에서 Qwen3-4B-Instruct를 학생과 토큰 수준 교사로 사용했고, 평균 성공률 74.3%로 GRPO보다 6.8%포인트, LUFFY보다 4.9%포인트, Prefix-RFT보다 5.2%포인트 높았다.
절제 실험은 정책 호환성 통제가 실제로 중요함을 보여준다. 공동 디코딩 없이 교사에서 직접 샘플링하면 행동 불일치가 가장 커져 평균 점수가 37.2%에서 33.8%로 떨어졌다. 고정 k 변형에서는 k를 5에서 10으로 늘릴 때 학습이 좋아졌지만, 20으로 늘리면 교사 지도가 약해져 성능이 조금 감소했고 고정 설정 중에서는 k=10이 가장 좋았다. 동적 top-k는 단계별로 남는 학생 확률 질량에 따라 후보 집합을 조정해 37.2%로 가장 높은 평균 점수를 냈다. 공동 궤적의 오프정책 중요도 비율을 제거하면 35.3%로 떨어져, 남은 행동 정책 불일치를 명시적으로 보정하는 것이 필요함을 확인했다.
개발자 관점에서 PICO는 검증기와 전문가 해답 또는 교사 모델을 학습 중에 확보할 수 있고, 난문제에서 롤아웃 그룹이 자주 전멸하는 RLVR 파이프라인에 특히 관련이 있다. 구현 시에는 학생과 교사의 토큰 분포를 모두 꺼내 상위 k 교집합을 만들고, 남는 학생 확률 질량 Z_t로 동적 top-k와 폴백을 결정하며, 전멸 그룹에만 공동 궤적을 추가하고 중요도 비율을 정확히 넣는 부분을 확인해야 한다. 추론 시에는 교사 지도가 필요 없다는 점이 실무 배포에 유리하다. 다만 원문이 명시한 전제상 학습 중에는 전문가 해답과 교사 모델이 필요하고, 공동 디코딩은 모든 롤아웃이 실패한 그룹에서만 작동한다. 후보 교집합이 비거나 Z_min을 만족하는 교사 top-k가 없으면 학생 정책으로 폴백하므로, 실제 적용에서는 폴백 비율과 분포 이동 예산, 교사·학생 토크나이저와 어휘 공유 조건을 함께 점검해야 한다.