Trajectory-aware PRMs for Long CoT Reasoning
Jiaru Zou, Ling Yang, Jingwen Gu, Jiahao Qiu, Ke Shen, Jingrui He, Mengdi Wang
Process Reward Models have recently emerged as a powerful framework for supervising intermediate reasoning steps in large language models. Existing PRMs are primarily trained on model final responses and struggle to robustly evaluate intermediate thinking traces, especially in the emerging setting of trajectory–response data generated by frontier models like OpenAI-o1 and DeepSeek-R1. These silver-standard traces are increasingly used in reasoning distillation and test-time inference, yet remain noisy and difficult to supervise. In this work, we introduce TAP (Trajectory-Aware PRM), a novel reward modeling framework explicitly designed to evaluate multi-step reasoning processes. TAP incorporates both step-level and trajectory-level supervision, enabling fine-grained reward assignment aligned with structured chain-of-thought data. We further adapt TAP into both offline and online reward supervision: select high-quality data for model distillation, provide dense process-level rewards for reinforcement learning, and support reward-guided test-time scaling via a best-of-N strategy. Empirical results on challenging benchmarks such as AIME and GPQA show that TAP outperforms other strong reward model baselines, leading to over 22.7% performance improvement.