PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yiran Guo, Lijie Xu, Jie Liu, Dan Ye, Shuang Qiu

arXiv:2505.23564v2cs.LGcs.AIcs.CL

TL;DR

RL for reasoning LLMs must balance fine-grained credit assignment against reliable and efficient advantage estimation. SPO addresses this with modular segment-level estimation and specialized chain- and tree-based variants, which outperform PPO and GRPO on the reported mathematical-reasoning evaluations. Its long-CoT experiments remain bounded by context and sampling-efficiency constraints.

  • Problem

    Reasoning LLMs need accurate credit assignment, but token-level methods rely on difficult critic estimation while trajectory-level methods provide coarse final-reward feedback.

  • Method

    SPO uses flexible sequence segmentation, Monte Carlo segment advantage estimation without an additional critic, and policy optimization with segment advantages, including probability masks.

  • Results

    SPO-chain outperforms PPO and GRPO by 6–12 percentage points on GSM8K, while SPO-tree outperforms GRPO across the reported 2K and 4K context evaluations.

  • Takeaways & Limitations

    Segment-level advantages can outperform coarse trajectory-level advantages while using fewer estimation points than token-level approaches.

  • Takeaways & Limitations

    Long-CoT experiments are constrained by on-policy sampling, which limits how far replayed trajectories can be distributed across future iterations, and by the reported 4K training context.

Abstract

from arXiv · show

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 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.

1 Introduction

RL trains reasoning LLMs but faces difficult credit assignment because rewards are sparse and delayed. Existing token-level and trajectory-level methods trade estimation precision against reliability and efficiency, motivating SPO’s segment-level approach.

  • Credit assignment is difficult in LLM reinforcement learning because rewards are typically sparse and arrive only at response end.
  • Token-level PPO uses a critic for fine-grained advantages, but unreliable value prediction and added memory and computation can impair practical credit assignment.
  • SPO estimates advantages at segment granularity to provide more localized feedback than trajectory-level methods while using fewer estimation points than token-level methods.
  • SPO combines flexible partitioning, Monte Carlo segment advantage estimation, and policy optimization with segment advantages in a modular framework.
  • SPO-chain targets short CoT with cutpoint-based partitioning and chain-based estimation, while SPO-tree targets long CoT with tree-based estimation.

2 Background

Language generation can be formulated as an MDP with deterministic token transitions and a sparse terminal reward. PPO uses critic-based token advantages, whereas GRPO removes the critic but assigns normalized trajectory rewards uniformly across tokens.

  • In the language-generation MDP, states contain the prompt and previous tokens, actions select the next token, and transitions append that token deterministically.
  • The reward is zero during generation and binary at termination, indicating whether the complete response matches the ground truth.
  • PPO stabilizes policy updates with a clipped surrogate objective and estimates token-level advantages using GAE and a critic model.
  • GRPO avoids a critic by normalizing sampled responses’ rewards within groups and assigning each trajectory-level advantage uniformly to its tokens.
  • SPO is positioned between these extremes by using segment-level feedback to balance credit-assignment detail and estimation reliability.

3 Segment Policy Optimization

SPO partitions generated sequences into contiguous segments, estimates each segment’s incremental value with Monte Carlo methods, and updates the policy using those segment advantages. Its modular design supports flexible granularity and distinct chain- and tree-based strategies.

  • Framework overview: SPO’s modular architecture addresses sequence partitioning, segment advantage estimation, and policy updates with interchangeable strategies.
  • SPO-Chain: SPO-chain uses cutpoint-based partitioning, chain-based estimation, and probability-mask policy gradients for short CoT scenarios.
  • Flexible Segment Partition: A segment is a contiguous token sequence, and SPO permits arbitrary boundaries between token-level and trajectory-level granularity without requiring semantic completeness.
  • Segment Advantage Estimation via Monte Carlo: SPO defines a segment advantage as the difference between the expected rewards after and before generating that segment.
  • Segment Advantage Estimation via Monte Carlo: Chain-based estimation independently rolls out N trajectories from each segment-boundary state, while tree-based estimation reuses samples through bottom-up aggregation.
  • Policy Optimization Using Segment Advantages: Policy optimization may assign a segment advantage to every token, only low-probability tokens through probability masks, or other compatible optimization procedures.

4 SPO-Chain for Short CoT

SPO-chain adapts segment-level policy optimization to short chain-of-thought by partitioning trajectories at informative cutpoints, estimating segment advantages with chain-based Monte Carlo, and focusing policy updates on critical tokens. Its adaptive partitioning avoids unnecessary sampling when token probabilities are high.

  • Overview: SPO-chain targets short CoT with cutpoint-based partitioning, chain-based segment advantage estimation, and policy-gradient optimization using probability masks.Short CoT makes Monte Carlo sampling manageable, while segments typically contain few tokens.
  • Adaptive Cutpoint-based Segment Partition: Adaptive cutpoint partitioning avoids unnecessary segments when within-segment token probabilities are close to 1 and boundary values barely differ.The strategy addresses wasted sampling budget caused by fixed token-count partitions.
  • Adaptive Cutpoint-based Segment Partition: Equal cutpoint allocation across segments is the optimal solution under the stated divisibility assumption.This corresponds to partitioning trajectories so each segment contains the same number of cutpoints.
  • Adaptive Cutpoint-based Segment Partition: Superior performance is reported for the adaptive cutpoint partition strategy compared with alternative partition strategies.The comparison is reported in the paper’s partition-strategy experiments.
  • Chain-based Segment Advantage Estimation: Chain-based Monte Carlo sampling independently samples N trajectories at each segment boundary and averages their returns to estimate the boundary value.Each segment advantage is computed from the difference between estimated values at consecutive segment boundaries.
  • Policy Gradient with Token-Probability Masks: Probability masks concentrate each segment’s advantage on critical tokens whose policy probabilities meet the mask condition, rather than applying it uniformly to every token.The normalization term Z counts tokens selected by the mask, and the strategy is reported to improve accuracy.

5 SPO-Tree for Long CoT

SPO-tree adapts SPO to long chain-of-thought by using fixed-length segments and a tree of shared Monte Carlo samples. Reusing these samples for both value estimation and policy optimization reduces sampling overhead and improves sample efficiency.

  • Overview: SPO-tree addresses long CoT, where chain-based Monte Carlo sampling becomes prohibitively expensive, with tree-based segment advantage estimation.Its configuration combines fixed token-count partitioning, tree-based estimation, and probability-masked policy gradients.
  • Fixed Token Count Segment Partition: Fixed token-count partitioning creates equal-length segments because long-CoT segments usually contain many tokens and are unlikely to have uniformly high transition probabilities.This partition supports the subsequent tree-based estimation method.
  • Tree-based Segment Advantage Estimation: Tree-based estimation reuses value-estimation samples during policy optimization, substantially improving sample efficiency.This directly addresses the sample waste of chain-based estimation.
  • Tree Construction: Each tree node represents a segment generated by extending its parent history with M newly sampled tokens.The trajectory sampling process is modeled as a tree whose nodes encode segment histories.
  • Tree Construction: Sibling nodes share prompts and sequence lengths, enabling fair advantage comparisons under the same token budget.Node values are estimated recursively from leaves toward the root.
  • Tree-based Segment Advantage Estimation: Tree values use leaf rewards and recursively aggregate child-node values for non-leaf nodes.This bottom-up value construction supports segment advantage computation within the tree.
  • Policy Optimization and Tree Design: Tree width controls sibling comparisons, depth controls partition granularity, and non-zero-advantage segments are selected for policy optimization.Larger width provides more comparison nodes, while deeper trees provide finer-grained signals.

6 Experiments on SPO-chain

SPO-chain achieves strong GSM8K performance while reducing episode-generation time, and its segment design choices materially affect accuracy. The probability-mask strategy further improves policy optimization.

  • Comparison with Baseline Methods: SPO-chain (int5) achieves the highest GSM8K test accuracy, outperforming PPO and GRPO by 6-12 percentage points.It also exceeds VinePPO in accuracy while requiring less episode-generation time.
  • Comparison with Baseline Methods: SPO-chain (int5) significantly outperforms GRPO in validation performance under the same wall-clock time.The comparison includes evaluation time in the reported wall-clock budget.
  • Impact of Segmentation Granularity: A moderate segment interval, such as 5, provides the best accuracy-versus-training-time trade-off, whereas overly coarse intervals severely harm accuracy.Interval 2 slightly exceeds interval 5 in final accuracy, but interval 5 performs best under equal wall-clock time.
  • Comparison of Different Segment Partition Strategies: Despite the smallest sampling budget, SPO-chain (int5) achieves the best test accuracy among the compared partition strategies.The comparison includes fixed token-count and heuristic partitioning alongside cutpoint-based partitioning.
  • Ablation Study on Probability-Mask Optimization Strategy: Removing probability masking lowers SPO-chain (int5) accuracy from 56.7% to 55.0%, while adding it to GRPO raises accuracy from 45.7% to 53.6%.The strategy masks losses for tokens whose probabilities are close to 1.

7 Experiments on SPO-tree

SPO-tree is evaluated for long chain-of-thought reasoning against GRPO on MATH500 and across tree configurations. It consistently achieves stronger performance under the tested context sizes, while tree size affects when gains emerge.

  • Comparison with Baseline Methods: SPO-tree significantly outperforms vanilla GRPO under the same wall-clock time on MATH500 with a 2K context window.GRPO with probability masking improves over vanilla GRPO, but remains below SPO-tree.
  • Comparison with Baseline Methods: SPO-tree consistently outperforms GRPO under 2K and 4K context evaluation.The model is continued from a 2K checkpoint before expanding training context to 4K.
  • Comparison with Baseline Methods: The GRPO baseline used in these experiments reaches 45.7% test accuracy, slightly above the 44.6% reported in prior work.This difference is attributed to the implementation used here.
  • Comparison of Different Tree Structures and Advantage Computation Methods: Tree structures show similar performance under equal wall-clock time, with smaller trees initially stronger and larger trees eventually benefiting from more accurate value estimation.The compared structures are 4-4-4, 6-6-6, and 8-8-8.

8 Conclusion

The paper concludes that segment-level advantages provide a practical middle ground between token-level and trajectory-level credit assignment. Its evidence is strongest on mathematical tasks and contexts up to 4K tokens.

  • Conclusion: SPO provides more precise credit assignment than trajectory-level methods while using fewer advantage-estimation points than token-level methods.This enables unbiased Monte Carlo advantage estimates.
  • Conclusion: Experiments show that a small number of segment-level advantages can significantly outperform coarse trajectory-level advantages.The conclusion presents this as evidence supporting SPO’s effectiveness.
  • Conclusion: Long-CoT experiments are limited to a maximum context size of 4K tokens because of computational constraints.The authors plan additional experiments with larger context sizes.
  • Conclusion: The current experiments focus primarily on mathematical tasks, with broader evaluation such as code generation and RLHF left for future work.The conclusion identifies these as planned application scenarios.

B Policy Learning via Policy Iteration

The policy-iteration formulation treats RL training as KL-constrained policy optimization and updates the policy through estimated advantages. SPO combines this optimization with Monte Carlo process rewards and replay-buffer sampling to support steadier, more balanced learning.

  • Policy iteration frames LLM training as a KL-constrained policy optimization problem.
  • The maximum-entropy formulation modifies rewards with β log πref(a|s) and uses soft Bellman relationships to derive the policy objective.
  • Explicit value estimation is unnecessary when advantages are estimated as V(s_{h+1}) − V(s_h), simplifying the policy optimization objective.
  • Policy iteration produces steady improvement, but probability-mask policy gradients achieve validation accuracy slightly above 0.7 and perform better in the reported comparison.
  • Process reward: SPO integrates process rewards from current-policy Monte Carlo samples without a separate reward model, avoiding fitting errors and reward hacking risks.
  • Replay buffer: The replay buffer distributes trajectories across iterations to balance question coverage, while PPO clipping stabilizes optimization with reused samples.
  • Replay buffer: Replay-buffer off-policy learning makes SPO compatible with asynchronous RL, potentially enabling more parallel data collection and policy optimization.

E Pseudo-code of Tree-Structured Advantage Estimation Method

The tree-structured advantage estimator recursively builds and expands a response tree, computes rewards for completed children, and aggregates child rewards into node statistics. The reported GRPO comparison highlights rapid overfitting during training.

  • Algorithm 1 constructs a rooted response tree, recursively expands nodes, and computes advantages after tree construction.
  • Completed child responses receive rewards, while length-truncated branches are expanded further before parent statistics are computed.
  • Each node records the mean and standard deviation of its children’s rewards for tree-based estimation.
  • GRPO on GSM8K rapidly overfits: training accuracy rises, unique responses decline, and validation accuracy saturates early.

F Additional Experiment Results

Additional experiments evaluate SPO across models, datasets, and reasoning settings. SPO-chain performs competitively on MATH, while broader tests report higher accuracy than GRPO across mathematical and logic benchmarks.

  • SPO-chain on DeepSeekMath 7B achieves performance comparable to VinePPO and surpasses the other reported baselines on MATH.Finer-grained advantages after reducing the interval from 10 to 5 continue improving performance.
  • Figure 10 compares the number of sampling points used by SPO and VinePPO, with MATH trajectories averaging approximately 8 reasoning steps.
  • Across Qwen2.5 instruct and math models on GSM8K and MATH, SPO consistently outperforms GRPO in test accuracy.
  • On the 3ppl Knights-and-Knaves benchmark, SPO also outperforms GRPO, extending the reported gains beyond mathematical reasoning.
  • Figure 11 compares results across Qwen2.5 models on GSM8K, MATH, and Knights-and-Knaves, separating four model–dataset settings.

G Limitations and Future Work

For long chain-of-thought, SPO-tree provides fine-grained rewards with lower sampling costs, but its replay-buffer strategy remains constrained by on-policy learning and limited future reuse.

  • SPO-tree reduces sampling costs for fine-grained reward estimation in long-CoT scenarios, but constructing a tree still requires many trajectories per problem.
  • Because on-policy algorithms limit how far trajectories can be reused, future work may investigate more off-policy methods for efficient sample utilization.

H Hyperparameters and Compute Resources

The experiments use different hyperparameter configurations and single-GPU setups for SPO-chain, SPO-tree, and comparison methods. The section also illustrates segment partitioning and model-output examples.

  • Compute resources: SPO-chain experiments use RhoMath 1.1B on a single A100 GPU with 40GB memory.The configuration includes cutpoint-based settings and Monte Carlo samples.
  • Compute resources: SPO-tree experiments use DeepSeek-R1-Distill-Qwen-1.5B on a single A100 GPU with 80GB memory for 2K and 4K contexts.
  • SPO-chain configuration: SPO-chain uses a 0.9 probability-mask threshold, nine Monte Carlo samples, and two epochs per iteration.Its listed training configuration uses batch size 64 and eight samples per prompt.
  • SPO-tree configuration: SPO-tree (6-6-6) uses 30 tokens per level, branch factors [6], and a 0.9 probability-mask threshold.The configuration lists 16 dataset samples per iteration and one epoch per iteration.
  • Segment partition example: The segment-partition example marks tokens below 0.9 probability as cutpoints and uses blue lines to show SPO segment boundaries.The example associates several low-confidence cutpoints with model mistakes, including an incorrect multiplication digit.
Loading 2505.23564v2…