세그먼트 기반 이점 추정: 장문 컨텍스트 LLM 학습을 위한 PPO 성능 향상

Segmental Advantage Estimation: Enhancing PPO for Long-Context LLM Training

추론 작업을 위한 대규모 언어 모델(LLM)의 학습은 점점 더 강화 학습과 검증 가능한 보상(RLVR)에 의존하고 있으며, 이 과정에서 Proximal Policy Optimization (PPO)은 안정적인 정책 업데이트를 위한 체계적인 프레임워크를 제공합니다. 그러나 PPO의 실제 적용은 희소 보상이 존재하는 RLVR 환경에서 신뢰할 수 없는 이점 추정으로 인해 어려움을 겪습니다. 이러한 문제는 희소 보상이 중간 값 예측의 부정확성을 초래하며, Generalized Advantage Estimation (GAE)을 통해 모든 토큰에서 이러한 예측을 집계할 때 상당한 편향을 발생시키기 때문입니다. 이러한 문제를 해결하기 위해, 우리는 GAE가 RLVR에서 발생시킬 수 있는 편향을 완화하는 세그먼트 기반 이점 추정(Segmental Advantage Estimation, SAE)을 제안합니다. 우리의 핵심 아이디어는 GAE와 같이 각 토큰에서 n-단계 이점을 집계하는 것이 불필요하며, 종종 과도한 편향을 초래한다는 것입니다. 왜냐하면 개별 토큰은 제한적인 정보를 담고 있기 때문입니다. 대신, SAE는 생성된 시퀀스를 일관성 있는 하위 세그먼트로 먼저 분할하며, 이를 위해 확률이 낮은 토큰을 휴리스틱 경계로 사용합니다. 그런 다음, SAE는 정보가 풍부한 세그먼트 전환에서만 분산을 줄인 이점 추정값을 선택적으로 계산하여 중간 토큰에서 발생하는 노이즈를 효과적으로 제거합니다. 우리의 실험 결과는 SAE가 우수한 성능을 달성하며, 최종 점수, 학습 안정성 및 샘플 효율성 측면에서 뚜렷한 개선을 보인다는 것을 보여줍니다. 이러한 이점은 다양한 모델 크기에 걸쳐 일관되게 나타나며, 상관관계 분석을 통해 제안된 이점 추정기가 근사적인 실제 이점과 더 높은 상관관계를 가지는 것으로 확인되었으며, 이는 제안된 방식의 우수한 성능을 뒷받침합니다.

Training Large Language Models (LLMs) for reasoning tasks is increasingly driven by Reinforcement Learning with Verifiable Rewards (RLVR), where Proximal Policy Optimization (PPO) provides a principled framework for stable policy updates. However, the practical application of PPO is hindered by unreliable advantage estimation in the sparse-reward RLVR regime. This issue arises because the sparse rewards in RLVR lead to inaccurate intermediate value predictions, which in turn introduce significant bias when aggregated at every token by Generalized Advantage Estimation (GAE). To address this, we introduce Segmental Advantage Estimation (SAE), which mitigates the bias that GAE can incur in RLVR. Our key insight is that aggregating $n$-step advantages at every token(as in GAE) is unnecessary and often introduces excessive bias, since individual tokens carry minimal information. Instead, SAE first partitions the generated sequence into coherent sub-segments using low-probability tokens as heuristic boundaries. It then selectively computes variance-reduced advantage estimates only from these information-rich segment transitions, effectively filtering out noise from intermediate tokens. Our experiments demonstrate that SAE achieves superior performance, with marked improvements in final scores, training stability, and sample efficiency. These gains are shown to be consistent across multiple model sizes, and a correlation analysis confirms that our proposed advantage estimator achieves a higher correlation with an approximate ground-truth advantage, justifying its superior performance.