Probepo improves reinforcement learning credit assignment in large language models
ProbePO: Sparse Probes, Redistributed Credit for Critic-Free LLM Reinforcement Learning
Machine LearningComputation and Language
Summary
Rewarding an entire answer all at once hides which parts actually helped. The authors introduce ProbePO, a way to smartly check parts of a reasoning process without needing many extra tries, improving feedback to the model. This helps train models to give better answers faster and more accurately than previous methods. They tested ProbePO on several large language models and math benchmarks, consistently getting better results with less training time.
What this means in practice
- •For machine learning engineers: Train large language models more efficiently by improving how step-by-step credit is assigned during reinforcement learning.
- •For product teams building ai assistants: Improve AI reasoning accuracy and reduce training time for complex problem-solving models.
Authors
Diyi Hu, Ruizhong Qiu, Lin Qiu, Hanghang Tong, Hanqing Zeng
Abstract
Reinforcement learning from verifiable rewards assigns one terminal reward to an entire reasoning trajectory, obscuring how individual reasoning steps contributed to the outcome. GRPO broadcasts a group-normalized advantage to every response token, while finer-grained approaches often rely on dense auxiliary rollouts or branching schemes that can also alter the data used for policy updates. We cast the missing middle as budgeted measurement: with a limited rollout budget, which partial solutions should be evaluated, and how should their value estimates refine credit? ProbePO selects partial-solution states with a reward-free uncertainty score and estimates their expected verifier rewards under a capped continuation budget. ProbePO redistributes, rather than replaces, GRPO's trajectory-level advantage across tokens while preserving each trajectory's mean token-level advantage; probe continuations serve only as measurements. Across three DeepSeekMath and Qwen-family models spanning 1.5B-7B parameters, ProbePO consistently outperforms GRPO and achieves the highest accuracy among the four evaluated methods in 16 of 18 model-benchmark comparisons. Against dense-probing VinePPO on DeepSeekMath-7B, ProbePO achieves higher accuracy on all six benchmarks with 76% less end-to-end training time.