PS-PPO: Prefix-Sampling PPO for Critic-Free RLHF

2026-06-29

AI Summary

This work addresses the high computational cost of existing critic-free RLHF methods on long reasoning trajectories, which arises from updating the policy for all tokens uniformly across the entire trajectory. To overcome this limitation, the authors propose PS-PPO, a novel approach that leverages temporal redundancy in trajectory prefixes within the PPO framework. PS-PPO introduces a prompt-conditioned truncated distribution and performs gradient updates only on randomly sampled prefixes, while employing importance weighting to ensure unbiased policy gradients. This method substantially reduces both training compute requirements and peak GPU memory usage, achieving accuracy on par with strong baselines across mathematical reasoning and RLHF benchmarks, thereby enabling efficient and unbiased policy optimization.

Abstract

Reinforcement Learning from Human Feedback (RLHF) for Large Language Models increasingly relies on critic-free methods as a practical alternative to actor--critic training. Despite their simplicity, existing critic-free approaches propagate a trajectory-level learning signal uniformly across all tokens in a trajectory. This requires full-trajectory policy updates for every rollout, leading to substantial optimization cost for long reasoning traces, even though intermediate prefixes often contain enough information to largely determine the final outcome. We propose Prefix-Sampling Proximal Policy Optimization (PS-PPO), a compute-efficient critic-free method for RLHF that exploits this temporal redundancy. PS-PPO introduces a prompt-conditioned cutoff distribution and samples a cutoff timestep for each trajectory. During the update pass, PS-PPO backpropagates only through the sampled prefix of each trajectory and applies an importance-weighting correction so that the resulting truncated gradient estimator remains unbiased with respect to the full-trajectory objective. Experiments on mathematical reasoning and RLHF benchmarks show that PS-PPO achieves large reductions in training compute and peak GPU memory, while maintaining accuracy comparable to strong critic-free baselines.

Problem

• Reinforcement Learning from Human Feedback
• critic-free RLHF
• trajectory-level learning signal
• optimization cost
• long reasoning traces

Innovation

• Prefix-Sampling
• Critic-Free RLHF
• Temporal Redundancy
• Importance Weighting
• Compute Efficiency