정규화를 넘어: 파티션 함수를 강화 학습 기반 변환 모델(RLVR)의 어려움 스케줄러로 재해석

Beyond Normalization: Rethinking the Partition Function as a Difficulty Scheduler for RLVR

보상을 최대화하는 강화 학습 방법은 LLM의 추론 성능을 향상시키지만, 종종 출력의 다양성을 감소시킵니다. 최근 연구에서는 GFlowNet을 사용하여 LLM을 학습시켜 특정 분포와 일치시키면서 동시에 해당 분포의 파티션 함수를 학습하는 방식으로 이 문제를 해결합니다. 기존 연구에서 파티션 함수를 단순한 정규화 요소로 취급하는 것과 달리, 본 연구에서는 이를 각 프롬프트별 예상 보상(즉, 온라인 정확도) 신호로 재해석하고, 이 활용되지 않은 정보를 활용하여 샘플 효율성을 향상시킵니다. 구체적으로, 먼저 파티션 함수와 프롬프트별 정확도 추정치 간의 이론적 관계를 확립합니다. 이 핵심적인 통찰력을 바탕으로, 본 연구는 Partition Function-Guided RL (PACED-RL)이라는 후처리 프레임워크를 제안합니다. PACED-RL은 정확도 추정치를 활용하여 학습 과정에서 유용한 질문 프롬프트를 우선적으로 선택하고, 정확도 추정치 오류 우선 순위 재학습을 통해 샘플 효율성을 더욱 향상시킵니다. 특히, 제안하는 방법의 두 가지 구성 요소 모두 GFlowNet 학습 과정에서 이미 생성된 정보를 재사용하여 계산 비용을 기존 최적화 과정에 효과적으로 분산시킵니다. 다양한 벤치마크에서의 광범위한 실험 결과는 GRPO 및 기존 GFlowNet 방법보다 성능이 크게 향상되었음을 보여주며, PACED-RL이 LLM의 보다 효율적인 분포 일치 학습을 위한 유망한 방향임을 시사합니다.

Reward-maximizing RL methods enhance the reasoning performance of LLMs, but often reduce the diversity among outputs. Recent works address this issue by adopting GFlowNets, training LLMs to match a target distribution while jointly learning its partition function. In contrast to prior works that treat this partition function solely as a normalizer, we reinterpret it as a per-prompt expected-reward (i.e., online accuracy) signal, leveraging this unused information to improve sample efficiency. Specifically, we first establish a theoretical relationship between the partition function and per-prompt accuracy estimates. Building on this key insight, we propose Partition Function-Guided RL (PACED-RL), a post-training framework that leverages accuracy estimates to prioritize informative question prompts during training, and further improves sample efficiency through an accuracy estimate error-prioritized replay. Crucially, both components reuse information already produced during GFlowNet training, effectively amortizing the compute overhead into the existing optimization process. Extensive experiments across diverse benchmarks demonstrate strong performance improvements over GRPO and prior GFlowNet approaches, highlighting PACED-RL as a promising direction for a more sample efficient distribution-matching training for LLMs.