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

TL;DR

This work tackles unstable advantage estimation in PPO when training LLMs under RLVR with sparse final-token rewards. It proposes Segmental Advantage Estimation (SAE), which segments generated sequences at semantic boundaries identified by low-probability tokens and computes advantages only at those boundaries, yielding a recursive, adaptive weighting $A_t^{\text{SAE}} = \delta_t + \lambda_{\text{SAE}}(t) A_{t+1}^{\text{SAE}}$ and a bias bound that favors longer segments. Theoretical analysis shows SAE reduces bootstrapping bias relative to token-level GAE, with a bias upper bound that tightens as segment length $M$ grows. Empirically, SAE delivers higher final scores, better stability, and improved sample efficiency across 4B/8B/14B models on mathematical reasoning benchmarks (AIME, AMC, BeyondAIME) and in code/STEM domains, with stronger correlation to a ground-truth advantage. Overall, SAE provides a practical, scalable enhancement to PPO in long-context RLVR settings, enabling more reliable credit assignment for complex reasoning tasks.

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.

Figures (9)

• Figure 1: The sketch of GAE and SAE. Instead of bootstrapping at every token as GAE does, SAE first partitions the sequence into semantically coherent segments, and then computes advantage estimators only at the boundaries of these segments.
• Figure 2: Macro-averaged test-set scores across four test sets for all methods training on the Qwen3-8B-base. SAE maintains consistent improvements against baselines throughout the training process.
• Figure 3: Results of different methods across model sizes (4B/8B/14B). SAE consistently outperforms other baselines.
• Figure 4: Results of additional reasoning domains (CODE and STEM). SAE consistently outperforms other baselines.
• Figure 5: Correlation of different methods with approximate ground-truth advantage $A^*$, which is obtained by extensive Monte Carlo sampling. SAE and $A^*$ are the most highly correlated among all methods.

Theorems & Definitions (6)

• Theorem 4.1
• Definition 1: Segment-Aware Discount Function
• Theorem A.1: SAE Formulation Equivalence
• proof
• Theorem B.1
• proof