Efficient Process Reward Modeling via Contrastive Mutual Information

TL;DR

This paper introduces CPMI, a contrastive mutual information method for automatic reward labeling in process reward models, significantly reducing annotation effort and computational costs while improving reasoning accuracy.

Contribution

It proposes CPMI, a novel automatic reward labeling technique that leverages model probabilities to infer step-level supervision efficiently.

Findings

CPMI reduces dataset annotation time by 84%.

CPMI decreases token generation costs by 98%.

CPMI achieves higher accuracy on reasoning benchmarks.

Abstract

Recent research has devoted considerable effort to verifying the intermediate reasoning steps of chain-of-thought (CoT) trajectories using process reward models (PRMs) and other verifier models. However, training a PRM typically requires human annotators to assign reward scores to each reasoning step, which is both costly and time-consuming. Existing automated approaches, such as Monte Carlo (MC) estimation, also demand substantial computational resources due to repeated LLM rollouts. To overcome these limitations, we propose contrastive pointwise mutual information (CPMI), a novel automatic reward labeling method that leverages the model's internal probability to infer step-level supervision while significantly reducing the computational burden of annotating dataset. CPMI quantifies how much a reasoning step increases the mutual information between the step and the correct target answer relative to hard-negative alternatives. This contrastive signal serves as a proxy for the step's contribution to the final solution and yields a reliable reward. The experimental results show that CPMI-based labeling reduces dataset construction time by 84% and token generation by 98% compared to MC estimation, while achieving higher accuracy on process-level evaluations and mathematical reasoning benchmarks.