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

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.

Large language models are often improved by training them with feedback on their outputs. For long reasoning problems, however, this training can be expensive because existing methods typically update the model using the entire generated solution, even when the final outcome is largely determined by earlier parts of the reasoning. This paper introduces Prefix-Sampling PPO (PS-PPO), a method that makes such training more efficient by updating the model using only a sampled early portion of each solution. The sampling is designed so that the model still learns, on average, as if the full solution had been used. Across mathematical reasoning and alignment benchmarks, PS-PPO substantially reduces training time and memory use while maintaining performance comparable to strong existing methods.