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

2026-01-12

AI Summary

In sparse-reward reinforcement learning, conventional Generalized Advantage Estimation (GAE) suffers from large advantage bias due to inaccurate intermediate value estimates, which degrades the performance of Proximal Policy Optimization (PPO). To address this issue, this work proposes Segment-wise Advantage Estimation (SAE), which heuristically partitions generated sequences using low-probability tokens and computes n-step variance-reduced advantages only at informative segment boundaries. By avoiding per-token accumulation of noisy and biased estimates, SAE substantially reduces advantage estimation error. Empirical results demonstrate that SAE enhances training stability, sample efficiency, and final performance of long-context language models on RLVR tasks, with consistent improvements across multiple model scales. Furthermore, the estimated advantages under SAE exhibit higher correlation with approximate ground-truth advantages, validating its effectiveness in producing more accurate credit assignment signals.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Multiagent Systems: Adversarial Agents

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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.

Problem

• Reinforcement Learning with Verifiable Rewards
• Proximal Policy Optimization
• Generalized Advantage Estimation
• Sparse Rewards
• Advantage Estimation

Innovation

• Segmental Advantage Estimation
• PPO
• RLVR
• sparse rewards
• advantage estimation