세그먼트 정책 최적화: 대규모 언어 모델을 위한 강화 학습에서의 효과적인 세그먼트 수준 신용 할당

Segment Policy Optimization: Effective Segment-Level Credit Assignment in RL for Large Language Models

May 29, 2025

저자: Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu

cs.AI

초록

대규모 언어 모델의 추론 능력을 강화하기 위해 강화 학습(RL)을 효과적으로 활용하는 것은 여전히 중요한 과제로 남아 있다. 기존 접근 방식은 주로 두 가지 대조적인 이점 추정 세분화를 채택한다: 토큰 수준 방법(예: PPO)은 세밀한 이점 신호를 제공하는 것을 목표로 하지만, 정확한 비평 모델을 훈련하기 어려워 추정이 부정확한 문제가 있다. 반면, 궤적 수준 방법(예: GRPO)은 최종 보상에서 얻은 거친 이점 신호에만 의존하여 정확한 신용 할당이 어렵다. 이러한 한계를 해결하기 위해, 우리는 중간 세분화 수준에서 세그먼트 수준 이점 추정을 활용하는 새로운 RL 프레임워크인 세그먼트 정책 최적화(SPO)를 제안한다. SPO는 궤적 수준 방법보다 더 정확한 신용 할당을 제공하고 토큰 수준 방법보다 더 적은 추정 지점을 필요로 하여, 비평 모델 없이도 몬테카를로(MC) 기반의 정확한 이점 추정을 가능하게 한다. SPO는 세 가지 구성 요소와 새로운 전략을 특징으로 한다: (1) 유연한 세그먼트 분할; (2) 정확한 세그먼트 이점 추정; (3) 세그먼트 이점을 활용한 정책 최적화(새로운 확률 마스크 전략 포함). 우리는 SPO를 두 가지 특정 시나리오에 적용한다: (1) 짧은 사고 사슬(CoT)을 위한 SPO-chain은 새로운 컷포인트 기반 분할과 사슬 기반 이점 추정을 특징으로 하여 GSM8K에서 PPO와 GRPO보다 6-12% 포인트 정확도 향상을 달성한다. (2) 긴 CoT를 위한 SPO-tree는 새로운 트리 기반 이점 추정을 특징으로 하여 MC 추정 비용을 크게 줄이고, MATH500에서 2K 및 4K 컨텍스트 평가에서 GRPO보다 7-11% 포인트 정확도 향상을 달성한다. 우리는 코드를 https://github.com/AIFrameResearch/SPO에서 공개한다.

English

Enhancing the reasoning capabilities of large language models effectively using reinforcement learning (RL) remains a crucial challenge. Existing approaches primarily adopt two contrasting advantage estimation granularities: Token-level methods (e.g., PPO) aim to provide the fine-grained advantage signals but suffer from inaccurate estimation due to difficulties in training an accurate critic model. On the other extreme, trajectory-level methods (e.g., GRPO) solely rely on a coarse-grained advantage signal from the final reward, leading to imprecise credit assignment. To address these limitations, we propose Segment Policy Optimization (SPO), a novel RL framework that leverages segment-level advantage estimation at an intermediate granularity, achieving a better balance by offering more precise credit assignment than trajectory-level methods and requiring fewer estimation points than token-level methods, enabling accurate advantage estimation based on Monte Carlo (MC) without a critic model. SPO features three components with novel strategies: (1) flexible segment partition; (2) accurate segment advantage estimation; and (3) policy optimization using segment advantages, including a novel probability-mask strategy. We further instantiate SPO for two specific scenarios: (1) SPO-chain for short chain-of-thought (CoT), featuring novel cutpoint-based partition and chain-based advantage estimation, achieving 6-12 percentage point improvements in accuracy over PPO and GRPO on GSM8K. (2) SPO-tree for long CoT, featuring novel tree-based advantage estimation, which significantly reduces the cost of MC estimation, achieving 7-11 percentage point improvements over GRPO on MATH500 under 2K and 4K context evaluation. We make our code publicly available at https://github.com/AIFrameResearch/SPO.