Unleashing Implicit Rewards: Prefix-Value Learning for Distribution-Level Optimization
Abstract
Process reward models (PRMs) provide fine-grained supervision for reasoning, but reliable PRMs often require step annotations or heavy verification pipelines, making them costly to scale and refresh during online RL. Implicit PRMs reduce this cost by training log-likelihood-ratio rewards from trajectory-level outcome labels. However, the log-ratio is constrained only as a sequence-level aggregate during training, while inference decomposes it into token- or step-level scores for partial prefixes. This train–inference mismatch leaves local credits weakly identified, so distribution-wide scoring can amplify misleading advantages. We propose Implicit Prefix-Value Reward Model (IPVRM), which directly learns the probability of eventual correctness for each prefix from outcome labels. Step signals are then obtained as temporal-difference (TD) differences between consecutive prefix values, aligning the training target with inference-time use. IPVRM markedly improves step-verification F1 on ProcessBench. To exploit these prefix values during policy optimization, we further introduce Distribution-Level RL (DistRL), which applies TD advantages to both sampled tokens and high-probability candidate tokens, providing dense counterfactual updates without additional rollouts. Experiments show that DistRL brings limited gains with unreliable implicit rewards, but consistently improves downstream reasoning when paired with IPVRM. The implementation of our method is available at https://github.com/gaoshiping/IPVRM .
Lay Summary
Large language models are often trained to solve reasoning problems such as mathematics by checking whether the final answer is correct. However, this final-answer signal does not clearly show which intermediate step helped or hurt the solution. Existing methods that try to provide step-by-step feedback can be expensive because they often require detailed annotations or additional verification. This paper proposes a more scalable way to give models useful feedback during reasoning. Instead of directly labeling every reasoning step, our method learns to estimate how likely a partial solution is to eventually reach a correct answer. By comparing this estimate before and after each new token or step, the model obtains a useful signal about whether the reasoning process is improving. We then use these signals to improve reinforcement learning by updating not only the tokens the model actually sampled, but also other likely next-token choices. Our experiments focus on verifiable reasoning tasks, especially mathematical reasoning, where final answers can be checked automatically. The results show that our method improves the reliability of process-level feedback and helps reinforcement learning make better use of that feedback.