rePIRL: 역강화 학습을 활용한 LLM 추론을 위한 프로세스 리워드 모델 학습

rePIRL: Learn PRM with Inverse RL for LLM Reasoning

프로세스 리워드는 심층 강화 학습에서 훈련 효율성을 높이고, 분산을 줄이며, 리워드 해킹을 방지하기 위해 널리 사용됩니다. LLM 추론 분야에서 기존 연구들은 전문가 정책의 도움을 받아 또는 받지 않고 효과적인 프로세스 리워드 모델(PRM)을 학습하는 다양한 방법을 탐구합니다. 그러나 기존 방법들은 전문가 정책에 대한 강한 가정을 필요로 하거나(예: 리워드 함수 필요), 내재적인 한계(예: 엔트로피 붕괴)로 인해 성능이 낮거나 일반화 능력이 부족한 PRM을 생성하는 경우가 많습니다. 본 논문에서는 전문가 정책에 대한 최소한의 가정을 통해 효과적인 PRM을 학습하는 역강화 학습 기반 프레임워크인 rePIRL을 소개합니다. 구체적으로, 정책과 PRM을 번갈아 가며 업데이트하는 이중 학습 프로세스를 설계했습니다. 우리의 학습 알고리즘은 기존 역강화 학습을 LLM에 적용할 때 발생하는 문제들을 해결하기 위한 맞춤형 기술을 포함합니다. 우리는 제안하는 학습 프레임워크가 온라인 및 오프라인 PRM 학습 방법을 통합할 수 있음을 이론적으로 증명했으며, 이를 통해 rePIRL이 최소한의 가정만으로 PRM을 학습할 수 있음을 입증했습니다. 표준화된 수학 및 코딩 추론 데이터 세트에 대한 실험 결과는 rePIRL이 기존 방법보다 효과적임을 보여줍니다. 또한, 학습된 PRM을 테스트 시간 훈련, 테스트 시간 스케일링 및 어려운 문제에 대한 초기 신호 제공에 활용하는 방법을 보여줍니다. 마지막으로, 상세한 분석을 통해 우리의 훈련 방식과 주요 설계 선택 사항을 검증합니다.

Process rewards have been widely used in deep reinforcement learning to improve training efficiency, reduce variance, and prevent reward hacking. In LLM reasoning, existing works also explore various solutions for learning effective process reward models (PRM) with or without the help of an expert policy. However, existing methods either rely on strong assumptions about the expert policies (e.g., requiring their reward functions) or suffer intrinsic limitations (e.g., entropy collapse), resulting in weak PRMs or limited generalizability. In this paper, we introduce rePIRL, an inverse RL-inspired framework that learns effective PRMs with minimal assumptions about expert policies. Specifically, we design a dual learning process that updates the policy and the PRM interchangeably. Our learning algorithm has customized techniques to address the challenges of scaling traditional inverse RL to LLMs. We theoretically show that our proposed learning framework can unify both online and offline PRM learning methods, justifying that rePIRL can learn PRMs with minimal assumptions. Empirical evaluations on standardized math and coding reasoning datasets demonstrate the effectiveness of rePIRL over existing methods. We further show the application of our trained PRM in test-time training, test-time scaling, and providing an early signal for training hard problems. Finally, we validate our training recipe and key design choices via a detailed ablation study.