Retrieval-Augmented Process Reward Model for Generalizable Mathematical Reasoning (2025)

Abstract: While large language models (LLMs) have significantly advanced mathematical reasoning, Process Reward Models (PRMs) have been developed to evaluate the logical validity of reasoning steps. However, PRMs still struggle with out-of-distribution (OOD) challenges. This paper identifies the OOD issues including step OOD, arising from differences in reasoning patterns across model types and sizes, and question OOD, due to dataset shifts between training and real-world problems. To address these issues, we introduce Retrieval-Augmented Process Reward Model (RetrievalPRM), a novel framework designed to tackle these OOD issues. By utilizing a two-stage retrieval-enhanced mechanism, RetrievalPRM retrieves semantically similar questions and steps for PRM as a warmup to stimulate its potential to judge target steps, improving generalization and reasoning consistency across different models and problem types. Our extensive experiments demonstrate that RetrievalPRM outperforms existing baselines across multiple real-world datasets. Our open-source contributions include a retrieval-enhanced dataset, a tuning framework for PRM training, and the RetreivalPRM model, establishing a new standard for PRM performance.

## introduction while large language models (llms) have advanced mathematical reasoning @xcite @xcite @xcite , they remain prone to critical flaws: explicit errors (e.g., miscalculations, logical inconsistencies) and implicit risks where correct answers mask flawed intermediate steps. even when final results are accurate, llms often generate plausible-but-incorrect reasoning chains, eroding trust in their problem-solving processes @xcite . to address this, process figure 1 : the distribution differences across three datasets: gsm8k, math and olympiad. we use sentence-bert to encode these questions and perform t-sne visualization. reward models (prms) @xcite have been developed to rigorously evaluate the logical validity of intermediate steps @xcite , mirroring human pedagogical practices that prioritize reasoning quality over answer correctness. existing works @xcite @xcite frame prm as a binary classification problem. they train prm on open-source base llms such as qwen @xcite or llama @xcite using humanannotated dataset @xcite or automated process supervision method @xcite @xcite . although these approaches show great performance and empirical success, they still face kinds of outof-distribution challenges. we believe the out-ofdistribution (ood) problem can be viewed from the following perspectives: firstly, step ood may occur because of different processes generated by different models. due to the high cost of manual annotation, there are very few accurately labeled prm expert datasets, such as prm800k and processbench, with processes generated by @xcite and ## preliminary in this section, we formulate the whole problem and introduce prm as a binary classification model. ## methodology in this section, we introduce our proposed re-trievalprm framework in detail. ## experiments in this section, we present the experimental settings and results. our implementation code of retrieval-prm is publicly available. ## conclusion in this paper, we have addressed the significant out-of-distribution (ood) challenges faced by process reward models (prms), particularly step ood and question ood. by introducing the retrieval augmented process reward model (re-trievalprm), we propose an effective solution that leverages a two-stage retrieval-enhanced mechanism to improve the generalization of prms across diverse models and problems. extensive experiments on multiple real-world datasets have shown that retrievalprm consistently outperforms existing methods, highlighting its effectiveness in tackling ood issues.