LLM 추론을 위한 탐색 재정의: 확률 기반 정책 최적화와 이점 재가중 메커니즘

초록

본 논문은 RL with Verifiable Rewards(RLVR)에서 발생하는 정책의 엔트로피 붕괴와 모드 콜랩스를 해결하기 위해, 기존 GRPO에 확률 기반 이점 재가중(ARM) 기법을 도입한 ProGRPO를 제안한다. 프롬프트 퍼플렉시티와 답변 신뢰도를 활용해 과도하게 확신된 경로의 업데이트를 억제하고, 저확신 경로에 확률 질량을 재분배한다. 수학·코드 베엔치마크에서 Qwen2.5·DeepSeek 모델에 적용한 결과, Pass@1·Pass@32 모두에서 기존 GRPO 대비 5~14%p 향상을 달성하며 다양성과 정확성 사이의 균형을 개선한다.

상세 분석

본 연구는 대형 언어 모델(LLM)의 추론 단계에서 강화학습(RLVR) 기반 정책 최적화가 “고확률 경로”에 과도하게 편향되는 현상을 ‘샘플링 확률 동역학’ 관점에서 분석한다. 기존 GRPO는 그룹 내 평균을 베이스라인으로 사용해 보상 차이를 이점(A)으로 변환하지만, 보상이 높은 경로가 반복적으로 강화되면서 정책의 엔트로피가 급격히 감소하고, 결과적으로 동일한 정답이라도 다양한 사고 사슬을 생성하지 못한다는 모드 콜랩스 문제가 발생한다.

이를 해결하기 위해 저자들은 두 가지 핵심 설계를 제시한다. 첫째, Advantage Re-weighting Mechanism(ARM) 으로, 각 샘플의 프롬프트와 답변에 대한 모델 자체의 확신 점수 cθ(q)와 cθ(o|q)를 계산한다. 이 점수는 저확률 토큰 집합 T_low 에 대해 로그 확률을 평균한 지수형식으로 정의되며, ‘예측 불확실성이 높은 토큰’에 집중한다. ARM은 기존 이점 Ai에 α·(cθ(q)−cθ(o|q)) 를 더함으로써, 프롬프트에 대한 높은 확신과 답변에 대한 낮은 확신을 가진 샘플에 대해 이점을 상승시켜 학습이 과도하게 확신된 경로에 편향되지 않도록 조정한다.

둘째, Low‑Probability Token Length Normalization을 도입한다. 전체 시퀀스 길이 정규화는 대부분의 토큰이 고확률(>0.9)인 상황에서 보상 신호를 희석시키는 부작용을 낳는다. 논문은 전체 토큰 중 약 20% 에 해당하는 T_low 을 선정해 해당 토큰에만 길이 정규화를 적용함으로써, 실제로 불확실성이 큰 구간에서만 신호를 강화한다. 이는 정책이 ‘트리비얼’ 토큰에 의해 압도되지 않고, 핵심 추론 단계에서 충분히 탐색하도록 만든다.

ProGRPO의 전체 목표 함수는 기존 GRPO의 클리핑 구조에 ARM을 삽입한 형태이며, 클리핑 구간을