Retrieval-Augmented Process Reward Model for Generalizable Mathematical Reasoning
Authors: Jiachen Zhu, Congmin Zheng, Jianghao Lin, Kounianhua Du, Ying Wen, Yong Yu, Jun Wang, Weinan Zhang
Accepted at: acl25

Retrieval-Augmented Process Reward Model for Generalizable Mathematical Reasoning

Jiachen Zhu1, Congmin Zheng1, Jianghao Lin1, Kounianhua Du1
Ying Wen1*, Yong Yu1, Jun Wang2, Weinan Zhang1**
1Shanghai Jiao Tong University, 2University College London
{gebro13, desp.zcm, chiangel, kounianhuadu, ying.wen, wnzhang}@sjtu.edu.cn,
yyu@apex.sjtu.edu.cn
jun.wang@cs.ucl.ac.uk

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 key OOD issues, including step OOD—caused by differences in reasoning patterns across model types and sizes—and question OOD, which arises from dataset shifts between training data and real-world problems. To address these issues, we introduce Retrieval-Augmented Process Reward Model (Retrieval-PRM), a novel framework designed to tackle these OOD issues. By utilizing a two-stage retrieval-enhanced mechanism, RetrievalPRM retrieves semantically similar questions and steps as a warmup, enhancing PRM’s ability to evaluate target steps and 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 RetrievalPRM model, establishing a new standard for PRM performance.

1 Introduction

While large language models (LLMs) have advanced mathematical reasoning (OpenAI, 2023; Dubey et al., 2024; Zhu et al., 2024; Shao et al., 2024; Yang et al., 2024b), 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 (Lightman et al., 2023). To address this, Process

![The distribution differences across three datasets: GSM8K, MATH and Olympiad. We use sentence-bert to encode these questions and perform t-sne visualization.](page_1012_370_367_367.png)

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) (Lightman et al., 2023; Wang et al., 2024b) have been developed to rigorously evaluate the logical validity of intermediate steps (Cobbe et al., 2021), mirroring human pedagogical practices that prioritize reasoning quality over answer correctness.

Existing works (Wang et al., 2024a; o1 Team, 2024; Zheng et al., 2024) frame PRM as a binary classification problem. They train PRM on open-source base LLMs such as Owen (Yang et al., 2024b) or Llama (Dubey et al., 2024) using human-annotated dataset (Lightman et al., 2023) or automated process supervision method (Wang et al., 2024b; Luo et al., 2024; Qin et al., 2024). Although these approaches show great performance and empirical success, they still face kinds of out-of-distribution challenges. We believe the out-of-distribution (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 GPT (OpenAI, 2023) and
Question: Five different numbers are added together in pairs, and the results are 101, 102, 103, 104, 105, 106, 107, 108, 109, 111. Which is the largest of the five numbers?

ChatGPT-4o Process:
Let’s denote the five numbers as a, b, c, d, e with a<b<c<d<e
Step 1: Sum of all pairwise sums.
The total of all ten pairwise sums is:
This is equivalent to 4(a+b+c+d+e)
Let S=a+b+c+d+e. Then, the total sum of all pairwise sums is 4S.
Step 2: Find the sum of the pairwise sums.
Now, calculate the total sum of the pairwise sums given in the problem...Thus, we have 4S=1066 so S=1066/4=266.5
...
Chain Length: 6 steps
Solution Style: Analyze, Calculate
Answer: 56 ✗

Qwen-2.5-Math-72B-instruct Process:
To solve this problem, Let’s denote the five numbers as a, b, c, d, e such that a<b<c<d<e.
Step 1: First, we note that the smallest sum a+b=101. The next smallest sum is a+c=102 and the next largest sum is c+e=108. The third smallest sum is a+d=103 and the largest sum is b+e=108.
Step 2: We can solve these equations step by step. First, subtract the first equation from the second: (a+c)-(a+b)=102-101 → c-b=1 → c=b+1
...
Chain Length: 10 steps
Solution Style: Solve Equations
Answer: 57 ✔

Qwen-2.5-Math-1.5B-instruct Process:
Step 1: To determine which number is the largest among the five given sums.
101, 102, 103, 104, 105, 106, 107, 108,
we need to identify the pair that produces the highest sum when each number is added to itself.
Step 2: First, let’s list all possible pairs (101,101), (101,102), (101,103) (101,104), (101,105), (101,106)...
...
Chain Length: 3 steps
Solution Style: Enumerate
Answer: 56 ✗

Figure 2: Processes and problem-solving ideas for the same question vary from different models with the perspectives of model types and model sizes. GPT tends to analyze and calculate, while Qwen-72B tends to solve equations. Qwen-1.5B is small and relatively weak. It can only enumerate, and its thinking chain is short, so its answers are also very wrong.

Qwen (Yang et al., 2024b), respectively. However, different model types (e.g., GPT, Qwen, Llama(Dubey et al., 2024)) approach problem-solving differently. As is shown in Figure 2, when facing the same question, GPT-4o tends to analyze and calculate, while Qwen-72B tends to solve questions directly. They have different solution styles. Therefore, using process data generated by one model to train a PRM and then applying it to guide another model leads to an OOD issue. Moreover, models of different sizes also exhibit different reasoning processes. Larger models, like exceptional students, tend to have clearer and more accurate reasoning steps, while smaller models tend to have very short reasoning chains, as shown in Figure 2.

Secondly, Question OOD emerges because of dataset shift. Current PRM datasets contain only a limited number of problems. For example, Math Shepherd and PRM800K cover problems from the GSM8K and MATH datasets, with GSM8K being at the elementary to middle school level and MATH at the high school to university level. However, real-world problems are far more diverse, such as those in the Olympic math competition dataset (He et al., 2024), leading to OOD issues in other datasets. As shown in the Figure 1, we used Sentence-BERT (Reimers, 2019) to encode all the problems from the three datasets and visualized the distribution with t-SNE. It is evident that the distributions differ, and since both Olympic and MATH problems are typically from high school-level exams, they are semantically closer to each other than to GSM8K.

To address this issue, we propose a new framework, Retrieval Augmented Process Reward Model (RetrievalPRM), which leverages a Two-stage Retrieval-enhanced Mechanism to help PRMs solve the OOD problem. we retrieve relevant questions and steps in these two stages to address the issues of question OOD and step OOD, respectively. Specifically, when predicting a step for a given question, we select semantically similar questions based on their embeddings, placing them at the beginning of the entire prompt. Additionally, we select more fine-grained, similar steps and use them as references when predicting the correctness of the step. These retrieved questions and steps serve as a kind of warm-up for PRM, acting as example problems for reference. They not only help stimulate PRM’s potential by warming up but also allow the system to handle more difficult problems by identifying similarities, thus alleviating OOD issues.

Our main contributions are summarized as follows:

• To the best of our knowledge, we are the first to highlight the key OOD problems in Process Reward Models (PRMs), particularly the question OOD and step OOD, which arise due to differences in reasoning patterns across model types (e.g., GPT, Qwen), model sizes (1.5B, 72B) and varying problem difficulties in real-world datasets.

• We introduce the Retrieval-Augmented Process Reward Model (RetrievalPRM) framework, which utilizes a Two-stage Retrieval-enhanced
Mechanism to address OOD issues by incorporating both Question-level Retrieval and Step-level Retrieval, thereby enhancing PRM’s ability to generalize across diverse problem-solving scenarios.

• We build a Retrieval-enhanced dataset for training PRM using RetrievalPRM framework. We have made our code publicly available.1 Our dataset2 and model3 are open-sourced.

• Extensive experiments on the Process-Bench (Zheng et al., 2024) on four public real-world datasets demonstrate that Retrieval-PRM outperforms strong baselines and that the Out-of-distribution issue has been alleviated due to our retrieval approach.

2 Preliminary

In this section, we formulate the whole problem and introduce PRM as a binary classification model.

2.1 Problem Formulation

We denote the Math dataset as \( \mathcal{D} = \{(q_i, s_i, y_i)\}_{i=1}^N \), where \( N \) is the number of data instances. The input \( q_i \) is the \( i^{th} \) Math question. \( s_i = \{s_i^1, s_i^2, \ldots, s_i^{n_i}\} \) are the solution steps, where \( n_i \) is the step number of solution \( s_i \). \( y_i = \{y_i^1, y_i^2, \ldots, y_i^{n_i}\} \) and the label \( y_i^j \) indicates the correctness from the \( 1^{st} \) step to the \( j^{th} \) step.

\[
y_i^j = \begin{cases}
1, & (s_i^1, \ldots, s_i^j) \text{ is correct for } q_i; \\
0, & \text{otherwise}.
\end{cases}
\] (1)

2.2 ORM vs. PRM

Outcome-supervised Reward Models are introduced (ORM) by (Cobbe et al., 2021), where verifiers are trained for judging the final correctness of generated solutions. ORM only predicts the final label \( \hat{y}_i^{n_i} \), which can be formulated as

\[
\forall i, \hat{y}_i^{n_i} = \mathrm{ORM}(q_i, s_i^1, \ldots, s_i^{n_i}).
\] (2)

Building on this, the concept of process reward models (PRM) is introduced as a more granular and transparent approach. Not only does PRM evaluate the final solutions but it also assesses intermediate processes, where \( \hat{y}_i^j \) represents the predicted label for the \( j^{th} \) step by PRM.

\[
\forall i, j, \hat{y}_i^j = \mathrm{PRM}(q_i, s_i^1, \ldots, s_i^j).
\] (3)

2.3 Large Language Model for PRM scoring

When directly adopting LLMs as the PRM for scoring, we need to convert the data \( (q_i, s_i, y_i) \) with a hard prompt template. The whole template example is illustrated in Appendix B.2.

The textual input consists of the question \( q_i \) and steps \( s_i \), followed by a binary question about the correctness of these steps.

To obtain the floating-point correctness estimation \( \hat{y}_i^j \in [0, 1] \) instead of discrete word tokens '+’ or '-’, we apply bidimensional softmax over the corresponding logits of the binary key answer tokens (ie., + & -) from LLMs to accomplish the correctness estimation during evaluation:

\[
\hat{y}_i^j = \frac{\exp(l_{i,+})}{\exp(l_{i,+}) + \exp(l_{i,-})} \in (0, 1).
\] (4)

where \( l_{i,+} \) and \( l_{i,-} \) are the logits of token + and - in the \( i^{th} \) instance, respectively.

It is important to note that the estimated PRM scoring \( \hat{y}_i^j \) is used solely for evaluation on the testing set. If training is involved, we maintain the standard instruction tuning and causal language modeling paradigm for LLMs. In this way, we don’t need to replace the language model head with binary classification head which is the last layer of LLM.

3 Methodology

In this section, we introduce our proposed RetrievalPRM framework in detail.

3.1 Overview of RetrievalPRM

The RetrievalPRM is developed to address the problem of out-of-distribution (OOD) scenarios in mathematical problem-solving, specifically focusing on both question OOD and step OOD. According to Figure 3, traditional PRM models are constrained by predefined solution steps and are unable to handle unseen questions or steps effectively, especially when the problem context shifts or the solution process deviates from previously seen examples. RetrievalPRM overcomes this challenge by incorporating a Two-stage Retrieval-enhanced Mechanism that dynamically fetches relevant questions and steps from a large pool of questions and their
Figure 3: The model structure of our proposed RetrievalPRM framework and its difference with traditional PRM. We design a Two-stage Retrieval Module to retrieve reference questions and steps in each stage.

solutions. These retrieved questions and steps serve as a kind of warm-up for PRM, acting as example problems for reference. They not only help stimulate PRM’s potential by warming up but also allow the system to handle more difficult problems by identifying similarities.

3.2 Two-stage Retrieval-enhanced Mechanism
The core of RetrievalPRM is the Two-stage Retrieval-enhanced Mechanism, which consists of two key phases: Question-level Retrieval and Step-level Retrieval.

3.2.1 Question-level Retrieval
The first stage of retrieval tackles the question OOD issue. As is shown in Figure 3, the retrieval pool is the question database \( \mathbb{D}_q = \{ q_i \}_{i=1}^N \). During retrieval process, we treat:

• Query: the target question \( q_t \).

• Key: all \( q_i \) in the retrieval pool.

• Value: all the \( (q_i, s_i) \) pair in the retrieval pool.

We calculate their similarities \( <q_i, q_t> \) to match the most similar n questions. Specifically, all questions will first pass through a Sentence-BERT model to encode questions and obtain their semantic representations.

\[
\{ e_{q_i} \}_{i=1}^N = \text{SentenceBERT}(\{ q_i \}_{i=1}^N)
\] (5)

where \( e_{q_i} \in \mathbb{R}^D \) is the embedding vector of the question \( q_i \).

And then all the embeddings undergo Principle Component Analysis (PCA) (Kurita, 2021) for dimensionality reduction to extract the most important dimensions.

\[
\{ e'_{q_i} \}_{i=1}^N = \text{PCA}(\{ e_{q_i} \}_{i=1}^N)
\] (6)

where \( e'_{q_i} \in \mathbb{R}^d \) is the embedding after dimension reduction.

Finally, we compute the cosine similarity between the target question and the entire question pool, selecting the top-k most similar questions and inputting them into the text.

\[
\langle q_t, q_i \rangle = \frac{e'_{q_t} \cdot e'_{q_i}}{|e'_{q_t}| \cdot |e'_{q_i}|}.
\] (7)

Now we sort the vector \( \{ \langle q_i, q_t \rangle \}_{i=1}^N \) of similarity and choose top-k \( (q_i, s_i) \) pairs as reference questions \( q_r \) and put them in RetrievalPRM’s input together with the target question. Furthermore, we store all the solutions \( \{ s_i \}_{i=1}^m \) of top-m (\( m > k \)) questions in a new database to conduct a further step-level retrieval.

3.2.2 Step-level Retrieval
We place step-level retrieval in the second stage of the two-stage retrieval process, rather than as a separate module, for two key reasons:

Firstly, for a solution to be meaningful, both the question and the steps must be similar. For example, two different types of questions might both use the letter "p" to represent an unknown variable, but in some problems, "p" represents a prime number, while in others, it represents probability. This results in steps that may appear similar but have entirely different meanings, rendering the retrieved steps potentially unhelpful.

Secondly, since there are many possible solutions to a question, this leads to a large number of steps. If the majority of these steps are irrelevant, the time spent calculating similarities becomes inefficient. By placing step-level retrieval in the second stage, we can save both time and computational resources.
Therefore, after retrieving the top- \( m \) most similar questions, we inject all their solutions into a new steps database \( \mathbb{D}_s \). Then, we use the target step as the query to retrieve reference steps from this new database. The similarity for retrieval is still calculated using Sentence-BERT, PCA, and cosine similarity, as mentioned in 3.2.1.

3.3 Retrieval-based System Prompt

In RetrievalPRM, The system prompt serves as the instruction set for the model, framing the problem and directing it to evaluate each step of the solution. Besides the traditional system prompt for PRM, the Retrieval-based System Prompt (RetSP) is extended with additional instructions, as shown in the red sentence in Figure 3, which encourages the model to leverage knowledge from reference questions. For example, we inform PRM that step labels "+" and "-" represent correct and incorrect steps, respectively. At the same time, to avoid noise, we specify that if the reference question or step contains no relevant or helpful information, it should not be considered. These retrieval-based system prompts give PRM a more flexible thinking process, enabling it to actively decide whether to use retrieval-based knowledge.

We define reference questions of \( q_i \) as \( q_i^r \) and reference steps as \( s_i^r \). The whole input \( x_i^j \) of predicting the \( j_{th} \) step of \( q_i \) in RetrievalPRM can be formulated as:

\[
x_i^j = (RetSP, q_i^r, q_i, s_i^1, \ldots, s_i^{j-1}, s_i^r, s_i^j, y_i^j),
\]
\[
\hat{y}_i^j = \text{PRM}(x_i^j)
\]

where \( s_i^j \) is the \( j_{th} \) step of solution \( s_i \).

According to the input template above, it is worth noting that when predicting step n, we assume that steps 1 through n-1 are correct (Luo et al., 2024; Zheng et al., 2024). At this point, the most important task for PRM is to predict step n, so PRM can only access the reference steps for step n and cannot see the reference steps for steps \( 1 \sim n - 1 \).

4 Experiments

In this section, we present the experimental settings and results. Our implementation code of Retrieval-PRM is publicly available.

4.1 Experiment Setup

4.1.1 Datasets

Datasets are categorized into two kinds: Math reasoning datasets, and prm training datasets.

Math Reasoning Datasets

We conduct experiments on four public and widely used datasets in mathematical reasoning tasks: GSM8K (Cobbe et al., 2021) which contains math problems from elementary to middle school, MATH (Hendrycks et al., 2021) which contains math problems from basic to university level, OlympiadBench (He et al., 2024) which involves questions from the Mathematical Olympiad, Omni-MATH (Gao et al., 2024b) which covers multi-domain high-difficulty problems. Further details are provided in Appendix C.

Except for GSM8K, which focuses on grade school math problems, the other three datasets feature problems of competition or Olympiad-level difficulty.

PRM training datasets

We conduct experiments on two publicly available datasets for PRM:

PRM800K (Lightman et al., 2023): Based on the MATH dataset, it contains 800,000 manually annotated step-level correctness labels for training the Process Reward Model. It relies on expensive manual annotations.

Math-Shepherd (Wang et al., 2024b): It generates 400,000 machine-annotated step-level labels (covering MATH and GSM8K datasets) by automatically building process supervision data, without manual annotation.

4.1.2 Evaluation Metrics

We evaluate our model in a public PRM benchmark ProcessBench (Zheng et al., 2024). The aim is to judge whether PRM can find the first wrong step. It divides data into two parts: samples with incorrect and correct final answers and then conducts harmonic mean on the accuracy of these two parts to get the final F1-score. Moreover, we think since the sample number of each part isn’t balanced, We add an additional metric: weighted arithmetic mean of these two parts, which is shown in Table 1 as ArithACC.

4.1.3 Baselines

Following (Zheng et al., 2024), we divide all baselines into two parts:

(1) Open-source PRM, including Skywork (o1 Team, 2024), Qwen2.5-PRM (Zheng
Table 1: The performance of different models on ProcessBench. The best result is given in bold, and the second-best value is underlined. See Table 6 in Appendix D for breakdown of evaluation results.

<table>
  <tr>
    <th rowspan="2">Model</th>
    <th colspan="2">GSM8k</th>
    <th colspan="2">MATH</th>
    <th colspan="2">OlympiadBench</th>
    <th colspan="2">OmniMATH</th>
    <th rowspan="2">Avg.F1</th>
  </tr>
  <tr>
    <th>ArithACC</th><th>F1</th>
    <th>ArithACC</th><th>F1</th>
    <th>ArithACC</th><th>F1</th>
    <th>ArithACC</th><th>F1</th>
  </tr>
  <tr>
    <td rowspan="7">Open-source PRM</td>
    <td colspan="8">RetrievalPRM-7B(Ours)</td>
    <td>76.0</td><td>74.6</td><td>70.6</td><td>71.1</td><td>59.1</td><td>60.2</td><td>55.2</td><td>57.3</td><td>65.8</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Math-7B-PRM800K</td>
    <td>73.5</td><td>68.2</td><td>65.1</td><td>62.6</td><td>53.2</td><td>50.7</td><td>43.4</td><td>44.3</td><td>56.5</td>
  </tr>
  <tr>
    <td colspan="8">Skywork-PRM-7B</td>
    <td>71.6</td><td>70.8</td><td>54.5</td><td>53.6</td><td>25.6</td><td>22.9</td><td>23.7</td><td>21.0</td><td>42.1</td>
  </tr>
  <tr>
    <td colspan="8">Skywork-PRM-1.5B</td>
    <td>59.9</td><td>59.0</td><td>49.1</td><td>48.0</td><td>20.5</td><td>19.3</td><td>19.7</td><td>19.2</td><td>36.4</td>
  </tr>
  <tr>
    <td colspan="8">Math-Shepherd-PRM-7B</td>
    <td>58.3</td><td>47.9</td><td>45.1</td><td>29.5</td><td>39.7</td><td>24.8</td><td>34.8</td><td>23.8</td><td>31.5</td>
  </tr>
  <tr>
    <td colspan="8">RLHFFlow-PRM-Mistral-8B</td>
    <td>62.3</td><td>50.4</td><td>42.1</td><td>33.4</td><td>22.3</td><td>13.8</td><td>19.1</td><td>15.8</td><td>28.4</td>
  </tr>
  <tr>
    <td colspan="8">RLHFFlow-PRM-Deepseek-8B</td>
    <td>56.9</td><td>38.8</td><td>45.1</td><td>33.8</td><td>26.5</td><td>16.9</td><td>23.2</td><td>16.9</td><td>26.6</td>
  </tr>
  <tr>
    <td rowspan="14">Language Models as Critic</td>
    <td colspan="8">QwQ-32B-Preview</td>
    <td><b>87.9</b></td><td><b>88.0</b></td><td><b>78.5</b></td><td><b>78.7</b></td><td><b>59.2</b></td><td><b>57.8</b></td><td><b>61.1</b></td><td><b>61.3</b></td><td><b>71.5</b></td>
  </tr>
  <tr>
    <td colspan="8">GPT-4o</td>
    <td>80.2</td><td>79.2</td><td>63.4</td><td>63.6</td><td>50.1</td><td>51.4</td><td>50.1</td><td>53.5</td><td>61.9</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-72B-Instruct</td>
    <td>77.9</td><td>76.2</td><td>65.4</td><td>61.8</td><td>59.8</td><td>54.6</td><td>55.1</td><td>52.2</td><td>61.2</td>
  </tr>
  <tr>
    <td colspan="8">Llama-3.3-70B-Instruct</td>
    <td>83.7</td><td>82.9</td><td>63.7</td><td>59.4</td><td>54.3</td><td>46.7</td><td>51.0</td><td>43.0</td><td>58.0</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Coder-32B-Instruct</td>
    <td>72.0</td><td>68.9</td><td>64.5</td><td>60.1</td><td>57.0</td><td>48.9</td><td>52.5</td><td>46.3</td><td>56.1</td>
  </tr>
  <tr>
    <td colspan="8">Llama-3.1-70B-Instruct</td>
    <td>75.3</td><td>74.9</td><td>52.6</td><td>48.2</td><td>50.0</td><td>46.7</td><td>43.2</td><td>41.0</td><td>52.7</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-14B-Instruct</td>
    <td>72.3</td><td>69.3</td><td>59.2</td><td>53.3</td><td>50.2</td><td>45.0</td><td>43.5</td><td>41.3</td><td>52.2</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.72B-Instruct</td>
    <td>67.8</td><td>67.6</td><td>52.3</td><td>49.2</td><td>43.3</td><td>42.1</td><td>39.3</td><td>40.2</td><td>49.8</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-32B-Instruct</td>
    <td>70.6</td><td>65.6</td><td>61.9</td><td>53.1</td><td>53.5</td><td>40.0</td><td>47.7</td><td>38.3</td><td>49.3</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Math-72B-Instruct</td>
    <td>70.3</td><td>65.8</td><td>59.6</td><td>52.1</td><td>56.1</td><td>32.5</td><td>55.1</td><td>31.7</td><td>45.5</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Coder-14B-Instruct</td>
    <td>61.9</td><td>50.1</td><td>54.2</td><td>39.9</td><td>51.4</td><td>34.0</td><td>55.6</td><td>27.3</td><td>37.8</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-7B-Instruct</td>
    <td>37.8</td><td>36.5</td><td>36.9</td><td>36.6</td><td>29.9</td><td>29.7</td><td>27.3</td><td>27.4</td><td>32.6</td>
  </tr>
  <tr>
    <td colspan="8">Meta-Llama-3-70B-Instruct</td>
    <td>62.4</td><td>52.2</td><td>48.3</td><td>22.8</td><td>46.2</td><td>21.2</td><td>44.8</td><td>20.0</td><td>29.1</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Math-7B-Instruct</td>
    <td>54.4</td><td>26.8</td><td>50.3</td><td>25.7</td><td>43.1</td><td>14.2</td><td>41.6</td><td>12.7</td><td>19.9</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.7B-Instruct</td>
    <td>25.1</td><td>8.4</td><td>20.4</td><td>19.0</td><td>16.1</td><td>14.7</td><td>13.8</td><td>12.1</td><td>13.6</td>
  </tr>
  <tr>
    <td colspan="8">Meta-Llama-3-8B-Instruct</td>
    <td>27.1</td><td>13.1</td><td>17.3</td><td>13.8</td><td>14.2</td><td>4.8</td><td>19.7</td><td>12.6</td><td>11.1</td>
  </tr>
  <tr>
    <td colspan="8">Qwen2.5-Coder-7B-Instruct</td>
    <td>49.1</td><td>14.3</td><td>46.3</td><td>6.5</td><td>47.2</td><td>4.1</td><td>48.9</td><td>1.8</td><td>6.7</td>
  </tr>
  <tr>
    <td colspan="8">Llama-3.1-8B-Instruct</td>
    <td>27.3</td><td>10.9</td><td>20.5</td><td>5.1</td><td>16.0</td><td>2.8</td><td>15.0</td><td>1.6</td><td>5.1</td>
  </tr>
</table>

Table 2: The performance of different variants of RetrievalPRM on ProcessBench. We remove different components of RetrievalPRM to evaluate the contribution of each part to the model. The best result is given in bold, and the second-best value is underlined. See Table 7 in Appendix D for breakdown of evaluation results.

<table>
  <tr>
    <th rowspan="2">Retrieval Components</th>
    <th colspan="2">GSM8k</th>
    <th colspan="2">MATH</th>
    <th colspan="2">OlympiadBench</th>
    <th colspan="2">OmniMATH</th>
    <th rowspan="2">Avg.F1</th>
  </tr>
  <tr>
    <th>Question-level</th><th>Step-level</th>
    <th>ArithACC</th><th>F1</th>
    <th>ArithACC</th><th>F1</th>
    <th>ArithACC</th><th>F1</th>
  </tr>
  <tr>
    <td>✓ ✓</td>
    <td colspan="8">76.0</td>
    <td>74.6</td><td>70.6</td><td>71.1</td><td>59.1</td><td>60.2</td><td>55.2</td><td>57.3</td><td>65.8</td>
  </tr>
  <tr>
    <td>✓ ×</td>
    <td colspan="8"><b>77.8</b></td>
    <td><b>74.9</b></td><td><b>70.7</b></td><td><b>71.2</b></td><td><b>58.4</b></td><td><b>59.8</b></td><td>50.5</td><td>54.4</td><td>65.0</td>
  </tr>
  <tr>
    <td>× ✓</td>
    <td colspan="8">73.8</td>
    <td>67.5</td><td>69.5</td><td>69.2</td><td>58.2</td><td>58.9</td><td>52.2</td><td>56.3</td><td>63.0</td>
  </tr>
  <tr>
    <td>× ×</td>
    <td colspan="8">71.0</td>
    <td>65.6</td><td>67.3</td><td>67.5</td><td>54.3</td><td>55.8</td><td>47.2</td><td>50.9</td><td>59.9</td>
  </tr>
</table>

et al., 2024), Math-Shepherd (Wang et al., 2024b) and RLHFFlow (Xiong et al., 2024). These models are binary classification PRMs.

(2) Language Models as Critic, including Llama (Dubey et al., 2024), Qwen2 (Yang et al., 2024b), Qwen2.5 (Team, 2024), Qwen2.5-MATH (Yang et al., 2024a), Qwen2.5-Coder (Hui et al., 2024), GPT-4o (OpenAI et al., 2024). These models are promoted to judge the steps with the help of majority voting.

Further details of these baselines are provided in Appendix A due to article length limitations.

4.1.4 Implementation Details

Details like base models, hyperparameters, prompts, and training sizes are provided in Appendix B due to the article length limitations.

4.2 Overall Performance

We evaluate RetrievalPRM against existing baselines on ProcessBench, and the results are presented in Table 1. The findings are as follows:

• RetrievalPRM-7B surpasses all open-source PRM baselines, achieving the highest performance. Notably, the most significant improvement is observed on OmniMATH, the most challenging dataset, with performance gains increasing as dataset difficulty rises. This phenomenon may stem from the fact that most baseline PRMs are trained on human- or machine-annotated datasets such as PRM800K or Math-Shepherd, which primarily focus on GSM8K or MATH and exhibit OOD issues when applied to more complex datasets. In contrast, our RetrievalPRM effectively mitigates the OOD problem through its
Figure 4: We show the F1 scores of Retrieval-PRM on four datasets and their average, as the number of retrieval questions varies. Specifically, Top-0 means no retrieval questions.

retrieval-based approach, demonstrating the efficacy of our Two-stage Retrieval-enhanced Mechanism.

• When comparing models of different scales, RetrievalPRM outperforms all evaluated language models, including Qwen2.5-72B-Instruct and Llama3.3-70B-Instruct, with the sole exception of QwQ-32B-Preview. Remarkably, RetrievalPRM achieves this with a model size of just 7B. This highlights that PRMs, being both lightweight and task-specific, maintain strong competitiveness and potential compared to LLMs as critics.

4.3 Ablation Study

We analyze two main components in the Two-stage Retrieval-enhanced Mechanism: Question-level Retrieval and Step-level Retrieval—through the following ablations:

RetrievalPRM (Ours): The complete version of our proposed method.
RetrievalPRM (w/o Step-level Retrieval): This variant retains only the Question-level Retrieval, removing Step-level Retrieval during both training and inference.
RetrievalPRM (w/o Question-level Retrieval): This variant retains only the Step-level Retrieval, removing Question-level Retrieval during both training and inference.
RetrievalPRM (w/o Question-level and Step-level Retrieval): In this variant, both Question-level and Step-level Retrieval are removed during training and inference.

The performance of these variants is presented in Table 2, from which we can draw the following observations:

• The performance of RetrievalPRM (w/o Step-level Retrieval) remains almost identical to that of RetrievalPRM on GSM8K and MATH but exhibits a slight decline on OlympiadBench and OmniMATH. This can be attributed to the fact that Step-level Retrieval information is partially absorbed by Question-level Retrieval. As a result, Question-level Retrieval alone may be sufficiently effective for relatively easy datasets, as the reference steps it provides contain adequate knowledge for step prediction. However, for more challenging datasets, Step-level Retrieval becomes significantly more crucial, as it offers finer-grained guidance essential for handling complex problem-solving processes.

• RetrievalPRM (w/o Question-level Retrieval) shows lower performance, as it relies solely on Step-level Retrieval. The model lacks knowledge of reference questions, which is useful to alleviate question OOD, restricting its overall performance.

• RetrievalPRM (w/o both Retrieval) performs the worst, which is expected, demonstrating the effectiveness of both question-level and Step-level Retrieval.

4.4 Hyperparameter Study

Figure 4 illustrates the impact of the number of retrieval questions on the model’s performance. The findings are as follows:

Compared to Top-0, where no retrieval questions are used, models that incorporate retrieval questions show improved performance, highlighting the importance of Question-level Retrieval. It inspires us that Reference questions are important for PRM to get warmup, no matter how many reference questions there are.

The performance of Top-3 exhibits a slight decline, potentially due to two factors: (1) An excessive number of reference questions may lead
to an overly long input prompt, making it difficult for PRMs to comprehend or extract key information effectively. (2) A limited retrieval pool might result in later reference questions being less relevant than earlier ones, increasing the likelihood of misjudgments in the model’s predictions.

Table 3: Performance across Different Step Numbers

<table>
  <tr>
    <th>step num</th>
    <th>GSM8K MATH</th>
    <th>Olympiad-Bench</th>
    <th>Omni-MATH</th>
  </tr>
  <tr>
    <td>0</td>
    <td>74.9</td>
    <td>71.2</td>
    <td>59.8</td>
    <td>54.4</td>
  </tr>
  <tr>
    <td>1</td>
    <td>74.6</td>
    <td>71.1</td>
    <td>60.2</td>
    <td>57.3</td>
  </tr>
  <tr>
    <td>2</td>
    <td>73.0</td>
    <td>70.7</td>
    <td>59.6</td>
    <td>56.7</td>
  </tr>
  <tr>
    <td>3</td>
    <td>74.7</td>
    <td>69.7</td>
    <td>58.2</td>
    <td>56.9</td>
  </tr>
</table>

We have conducted additional experiments on the hyperparameter: step number. And the results are shown in Table 3. From the experimental results, we observe that changes in the number of steps have little impact on performance. These results suggest that our retrievalPRM exhibits strong robustness with respect to the number of reasoning steps.

4.5 Retrieval Strategy

Regarding the retrieval strategy, we experimented with simple information retrieval methods such as random selection and TF-IDF, and we also explored replacing Sentence-BERT with a larger language model like Qwen-1.5B to get question or step embeddings. The corresponding experimental results are presented in Table 4. From the experimental results, we have the following observations:

• As we move from no retrieval to increasingly advanced retrieval strategies—ranging from random selection to TF-IDF to Sentence-BERT—we observe that the model is better able to retrieve similar questions and steps. This leads to improved prediction accuracy of retrieval-enhanced PRM, demonstrating that the choice of retrieval strategy significantly impacts performance.

• Once the retrieval method reaches a certain level of quality, the performance of Sentence-BERT and Qwen-1.5B becomes comparable. We believe that, given the current retrieval pool, Sentence-BERT is already capable of retrieving highly relevant items, which explains the similar performance to Qwen-1.5B.

• However, as retrieval methods become more sophisticated, their computational cost also increases. Therefore, choosing a retrieval strategy involves a trade-off between efficiency and effectiveness. Under the current setting, Sentence-BERT offers a good balance, making it the preferable choice.

4.6 Results of Other Benchmarks

In order to evaluate the performances of retrieval-PRM completely, we conducted additional experiments employing two mainstream evaluation benchmarks: PRMBench and Best-of-N. Details are shown in Appendix E

4.7 Retrieval Pool Diversity

In this paper, our retrieval pool comprises both the MathShepherd and PRM800K datasets. To empirically examine the effect of the retrieval pool, we conducted supplementary experiments in which the retrieval pool was augmented with the recently released OpenReason dataset (OpenReason, 2024). From the results in Table 5. Detailed explanations are provided in Appendix F.

5 Related Works

5.1 Process Reward Models

Process reward models (PRMs) have demonstrated significant advantages over traditional outcome reward models (ORMs) (Cobbe et al., 2021) in enhancing process-level reasoning accuracy and improving long-process reasoning abilities in model training (Lightman et al., 2023; Uesato et al., 2022). A growing number of PRMs have been proposed for application in process-level reinforcement learning with human feedback (RLHF) (Wang et al., 2024b; Qin et al., 2024; Xia et al., 2025; o1 Team, 2024). For instance, Lightman et al. (2023) made a substantial contribution by releasing a large set of human-annotated data at the process level, opening up new research opportunities for multi-step reasoning.

Additionally, Wang et al. (2024b) introduces an automatic, self-supervised pipeline for generating process-level labels and training PRMs, enabling efficient data generation. Xia et al. (2025) employs PRMs as automatic evaluators to assess the accuracy of multi-step reasoning in language models (LMs). With the surge in PRM-focused research and data curation, numerous PRMs (o1 Team, 2024; Xiong et al., 2024; Sun et al., 2024; Gao et al.,
Table 4: Performance Comparison of Different Retrieval Strategies. Time consumption refers to the time required to retrieve 400 questions along with their corresponding steps.

<table>
  <tr>
    <th>Retrieval Strategy</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>OlympiadBench</th>
    <th>OmniMATH</th>
    <th>Time Consumption</th>
  </tr>
  <tr>
    <td>Random Select</td>
    <td>73.9</td>
    <td>70.7</td>
    <td>57.7</td>
    <td>55.3</td>
    <td>3s</td>
  </tr>
  <tr>
    <td>Tf-idf</td>
    <td>74.1</td>
    <td>70.2</td>
    <td>58.3</td>
    <td>56.7</td>
    <td>70s</td>
  </tr>
  <tr>
    <td>Sentence-Bert</td>
    <td>74.6</td>
    <td>71.1</td>
    <td>60.2</td>
    <td>57.3</td>
    <td>50min</td>
  </tr>
  <tr>
    <td>Qwen-1.5B</td>
    <td>73.1</td>
    <td>71.7</td>
    <td>58.9</td>
    <td>57.3</td>
    <td>12h</td>
  </tr>
</table>

2024a; Wang et al., 2024a) have been proposed. Additionally, several studies focus on leveraging natural language feedback from large language models (LLMs) as rewards, which are termed critic models (McAleese et al., 2024; Zhang et al., 2024; Gao et al., 2024a).

However, most existing PRMs trained on math datasets such as GSM8K and MATH inevitably encounter Out-of-distribution issues, which can be divided into two categories: question OOD, where PRMs trained on simpler or medium-difficulty datasets lack understanding of questions from more challenging datasets, and step OOD, where different base models and model sizes in LLMs lead to different step distributions for the same question. This is reflected in differences in chain length, problem-solving approaches, and methods. To address these issues, we propose the RetrievalPRM framework to tackle the OOD problems encountered in the current PRM field, achieving promising results.

5.2 Retrieval-Augmented Generation

Retrieval-augmented generation (RAG) enhances language models by dynamically integrating external knowledge, pioneered by (Lewis et al., 2021) through their joint retrieval-generation architecture. Subsequent advances refined this paradigm. Guu et al. (2020) introduced REALM to co-train retrieval and generation modules via masked language modeling, while Izacard and Grave (2021) proposed Fusion-in-Decoder (FiD) to process multi-document contexts efficiently. Research further optimized retrieval precision through dense passage embeddings (Karpukhin et al., 2020) and scaled retrieval to web-level corpora (Borgeaud et al., 2022).

6 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 (RetrievalPRM), 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.

7 Limitation

RetrievalPRM has two main limitations. Firstly, the retrieval pool is only constructed from PRM800K and Math-Shepherd at present, which is relatively small and limits the diversity and breadth of the mathematical problems. Second, using Sentence-BERT to embed questions and steps struggles to capture the full complexity of mathematical problems as semantic similarity doesn’t mean knowledge similarity in Math problems. As a result, the naive cosine similarity calculated through embeddings may fail to accurately reflect the true similarity between two questions.

Acknowledgments

The Shanghai Jiao Tong University team is partially supported by National Key RD Program of China (2022ZD0114804), Shanghai Municipal Science and Technology Major Project (2021SHZDZX0102) and National Natural Science Foundation of China (624B2096, 62322603, 62177033).

References

Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin
Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. 2022. Improving language models by retrieving from trillions of tokens. Preprint, arXiv:2112.04426.

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. 2021. Training verifiers to solve math word problems. CoRR, abs/2110.14168.

Abhimanyu Dubey, Abhinav Jauhari, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. 2024. The llama 3 herd of models. arXiv preprint arXiv:2407.21783.

Bofei Gao, Zefan Cai, Runxin Xu, Peiyi Wang, Ce Zheng, Runji Lin, Keming Lu, Dayiheng Liu, Chang Zhou, Wen Xiao, Junjie Hu, Tianyu Liu, and Baobao Chang. 2024a. Llm critics help catch bugs in mathematics: Towards a better mathematical verifier with natural language feedback. Preprint, arXiv:2406.14024.

Bofei Gao, Feifan Song, Zhe Yang, Zefan Cai, Yibo Miao, Qingxiu Dong, Lei Li, Chenghao Ma, Liang Chen, Runxin Xu, Zhengyang Tang, Benyou Wang, Daoguang Zan, Shanghaoran Quan, Ge Zhang, Lei Sha, Yichang Zhang, Xuancheng Ren, Tianyu Liu, and Baobao Chang. 2024b. Omni-math: A universal olympiad level mathematic benchmark for large language models. Preprint, arXiv:2410.07985.

Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. 2020. Realm: Retrieval-augmented language model pre-training. Preprint, arXiv:2002.08909.

Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, Jie Liu, Lei Qi, Zhiyuan Liu, and Maosong Sun. 2024. Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems. Preprint, arXiv:2402.14008.

Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. 2021. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874.

Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, Keming Lu, Kai Dang, Yang Fan, Yichang Zhang, An Yang, Rui Men, Fei Huang, Bo Zheng, Yibo Miao, Shanghaoran Quan, Yunlong Feng, Xinzhang Ren, Xuancheng Ren, Jingren Zhou, and Junyang Lin. 2024. Owen2.5-coder technical report. Preprint, arXiv:2409.12186.

Gautier Izacard and Edouard Grave. 2021. Leveraging passage retrieval with generative models for open domain question answering. Preprint, arXiv:2007.01282.

Vladimir Karpukhin, Barlas Oğuz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen tau Yih. 2020. Dense passage retrieval for open-domain question answering. Preprint, arXiv:2004.04906.

Takio Kurita. 2021. Principal component analysis (pca). In Computer vision: a reference guide, pages 1013–1016. Springer.

Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. 2021. Retrieval-augmented generation for knowledge-intensive nlp tasks. Preprint, arXiv:2005.11401.

Wendi Li and Yixuan Li. 2025. Process reward model with q-value rankings. In Proceedings of the International Conference on Learning Representations (ICLR).

Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. arXiv preprint arXiv:2305.20050.

Liangchen Luo, Yinxiao Liu, Rosanne Liu, Samrat Phatale, Meiqi Guo, Harsh Lara, Yunxuan Li, Lei Shu, Yun Zhu, Lei Meng, Jiao Sun, and Abhinav Rastogi. 2024. Improve mathematical reasoning in language models by automated process supervision. Preprint, arXiv:2406.06592.

Nat McAleese, Rai Michael Pokorny, Juan Felipe Ceron Uribe, Evgenia Nitishinskaya, Maja Trebacz, and Jan Leike. 2024. Llm critics help catch llm bugs. Preprint, arXiv:2407.00215.

Skywork ol Team. 2024. Skywork-ol open series. https://huggingface.co/Skywork.

OpenAI, ;, Aaron Hurst, Adam Lerer, Adam P. Goucher, Adam Perelman, Aditya Ramesh, Aidan Clark, AJ Ostrow, and Akila Welihinda et al. 2024. Gpt-4o system card. Preprint, arXiv:2410.21276.

R OpenAI. 2023. Gpt-4 technical report. arxiv 2303.08774. View in Article, 2(5).

OpenReason. 2024. Openreason-preview. https://huggingface.co/datasets/OpenReason/OpenReason-preview. Accessed: 2025-05-29.

Yiwei Qin, Xuefeng Li, Haoyang Zou, Yixiu Liu, Shijie Xia, Zhen Huang, Yixin Ye, Weizhe Yuan, Hector Liu, Yuanzhi Li, and Pengfei Liu. 2024. O1 replication journey: A strategic progress report – part 1. Preprint, arXiv:2410.18982.

N Reimers. 2019. Sentence-bert: Sentence embeddings using siamese bert-networks. arXiv preprint arXiv:1908.10084.
Zhihong Shao, Peiying Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.

Mingyang Song, Zhaochen Su, Xiaoye Qu, Jiawei Zhou, and Yu Cheng. 2025. Prmbench: A fine-grained and challenging benchmark for process-level reward models. arXiv preprint arXiv:2501.03124.

Zhiqing Sun, Longhui Yu, Yikang Shen, Weiyang Liu, Yiming Yang, Sean Welleck, and Chuang Gan. 2024. Easy-to-hard generalization: Scalable alignment beyond human supervision. Preprint, arXiv:2403.09472.

Qwen Team. 2024. Qwen2.5: A party of foundation models.

Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. 2022. Solving math word problems with process- and outcome-based feedback. Preprint, arXiv:2211.14275.

Jun Wang, Meng Fang, Ziyu Wan, Muning Wen, Jiachen Zhu, Anjie Liu, Ziqin Gong, Yan Song, Lei Chen, Lionel M. Ni, Linyi Yang, Ying Wen, and Weinan Zhang. 2024a. Openr: An open source framework for advanced reasoning with large language models. Preprint, arXiv:2410.09671.

Peiying Wang, Lei Li, Zhihong Shao, Runxin Xu, Damai Dai, Yifei Li, Deli Chen, Yu Wu, and Zhifang Sui. 2024b. Math-shepherd: Verify and reinforce ilms step-by-step without human annotations. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 9426–9439.

Shijie Xia, Xuefeng Li, Yixin Liu, Tongshuang Wu, and Pengfei Liu. 2025. Evaluating mathematical reasoning beyond accuracy. Preprint, arXiv:2404.05692.

Wei Xiong, Hanning Zhang, Nan Jiang, and Tong Zhang. 2024. An implementation of generative prm. https://github.com/RLHFlow/RLHF-Reward-Modeling.

An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, Keming Lu, Mingfeng Xue, Runji Lin, Tianyu Liu, Xingzhang Ren, and Zhenru Zhang. 2024a. Qwen2.5-math technical report: Toward mathematical expert model via self-improvement. Preprint, arXiv:2409.12122.

An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et al. 2024b. Qwen2.5-math technical report: Toward mathematical expert model via self-improvement. arXiv preprint arXiv:2409.12122.

Lunjun Zhang, Arian Hosseini, Hritik Bansal, Mehran Kazemi, Aviral Kumar, and Rishabh Agarwal. 2024. Generative verifiers: Reward modeling as next-token prediction. Preprint, arXiv:2408.15240.

Chujie Zheng, Zhenru Zhang, Beichen Zhang, Runji Lin, Keming Lu, Bowen Yu, Dayiheng Liu, Jingren Zhou, and Junyang Lin. 2024. Processbench: Identifying process errors in mathematical reasoning. arXiv preprint arXiv:2412.06559.

Qihao Zhu, Daya Guo, Zhihong Shao, Dejian Yang, Peiying Wang, Runxin Xu, Y Wu, Yukun Li, Huazuo Gao, Shirong Ma, et al. 2024. Deepseek-coder-v2: Breaking the barrier of closed-source models in code intelligence. arXiv preprint arXiv:2406.11931.

A Baselines

A.1 Open-source PRM

• Skywork-PRM (o1 Team, 2024) is a Qwen2.5-Math-based PRM published by KunLun.

• Qwen2.5-PRM (Zheng et al., 2024) is trained by fine-tuning the Qwen2.5-Math-7B-Instruct model on the PRM800K dataset.

• Math-Shepherd (Wang et al., 2024b) generates process labels for each step by estimating the empirical probability that a given step leads to the correct final answer and trains a PRM based on their published dataset.

• RLHFlow-PRM (Xiong et al., 2024) is an 8-billion-parameter reward model trained with process supervision.

A.2 Language Models as Critic

• Llama (Dubey et al., 2024) is an open-source model developed by Meta (formerly Facebook), designed for natural language understanding and generation tasks.

• Qwen2 (Yang et al., 2024b) is a large language model developed by Alibaba Cloud, offering multilingual support and strong capabilities in language understanding and generation.

• Qwen2.5 (Team, 2024) is an advanced iteration of the Qwen series, pretrained on 18 trillion tokens, enhancing knowledge retention, programming, and mathematical reasoning.

• Qwen2.5-MATH (Yang et al., 2024a) is a specialized model for mathematical problem-solving, trained on extensive math-focused data and incorporating Chain-of-Thought (CoT) and Tool-Integrated Reasoning (TIR).
<table>
  <tr>
    <th>System:</th>
    <td>I want you to act as a math teacher. I will provide a mathematical question and several solution steps, and it will be your job to judge whether these steps are correct or not.</td>
  </tr>
  <tr>
    <th>Input:</th>
    <td>
      <b>Question:</b><br>
      How many seconds are in 5.5 minutes?<br>
      <b>Process:</b><br>
      Step 1 : 5.5 minutes is the same as 5 minutes and 0.5 minutes.<br>
      Step 2 : Since there are 60 seconds in a minute, then there are 300 seconds in 5 minutes.<br>
      Step 3 : And since there are 60 seconds in a minute, there are 30 seconds in 0.5 minutes.<br>
      Is that Step Correct? You should ONLY tell me + or -.
    </td>
  </tr>
  <tr>
    <th>Output:</th>
    <td>+</td>
  </tr>
</table>

Figure 5: The illustration of PRM input template.

• Qwen2.5-Coder (Hui et al., 2024) is a programming-oriented model trained on 5.5 trillion code-related tokens, excelling in code generation, debugging, and multilingual programming tasks.

• GPT-4o (OpenAI et al., 2024) is a multimodal AI model developed by OpenAI that processes and generates text, audio, and images in real-time, with enhanced speed and natural interaction capabilities.

B Implementation Details

B.1 Basemodel and Training hyperparameters

We selected Qwen-2.5-Math-7b-instruct(Team, 2024) as the foundational large language model (LLM) for our experiments. All computations were performed using H100 GPUs. To enhance training resource efficiency, we employed Parameter-Efficient Fine-tuning techniques LoRA. The LoRA configuration was set with a rank of 32, an alpha value of 64, and dropout set to 0.1. LoRA update matrices were specifically applied to the query and value projection matrices within the attention blocks.

We use PRM800K as our training data and both PRM800K and Math-Shepherd as our retrieval pool. The training process was carried out with batch sizes chosen from {64, 128, 256, 512} and initial learning rates selected from \( \{1 \times 10^{-3}, 1 \times 10^{-4}, 3 \times 10^{-4}, 1 \times 10^{-5}, 3 \times 10^{-5}\} \) using a linear scheduler.

B.2 Prompts

In this section, we show our training prompts for PRM in details as is shown in Figure 5 and Figure 6.

C Datasets

GSM8K (Cobbe et al., 2021): Grade School Math is a dataset for basic to intermediate math problems, covering arithmetic, algebra, geometry and other fields. Its difficulty is suitable for math problems in elementary to middle school.

MATH (Hendrycks et al., 2021): The MATH dataset contains a variety of math problems from basic to university level, covering multiple mathematical fields such as algebra, geometry, calculus, number theory, etc.

OlympiadBench (He et al., 2024): The Olympiadbench dataset contains questions from the Mathematical Olympiad. The questions are of high difficulty and involve complex combinatorial mathematics, number theory, geometry and other advanced mathematical fields.

Omni-MATH (Gao et al., 2024b): Omni-MATH is a general Olympiad-level mathematics benchmark dataset for large language models, covering multi-domain and high-difficulty mathematics problems, and is designed to evaluate the reasoning
System:
I want you to act as a math teacher. I will provide a mathematical question and several solution steps, and it will be your job to judge whether these steps are correct or not. First I will give you some similar questions and their steps for reference. For each step, if the step is correct, the step is labeled as +. If the step is wrong, the step is labeled as -. If there is no relevant or helpful information in the provided questions and steps, try to answer yourself.

Input:
Reference Question 1:
What is the equivalent number of seconds in 7.8 minutes?
Process:
Since there are 60 seconds in a minute, we can find the number of seconds by multiplying the number of minutes by 60. (+) So, 7.8 minutes is equal to 7.8 * 60 = 46 seconds. The answer is: 46 (-)
Reference Question 2:
...
Process:
...
Target Question:
How many seconds are in 5.5 minutes?
Process:
Step 1 : 5.5 minutes is the same as 5 minutes and 0.5 minutes.
Step 2 : Since there are 60 seconds in a minute, then there are 300 seconds in 5 minutes.

Reference Step1:
0.3 hours equal to 0.3 * 60 = 18 minutes. This reference step is correct.
Reference Step2:
...
Target Step 3 :
And since there are 60 seconds in a minute, there are 30 seconds in 0.5 minutes.
Is the Step Correct? You should ONLY tell me + or -.

Output:
+.
ability of models in various mathematical fields.

Except for GSM8K, which focuses on grade school math problems, the other three datasets feature problems of competition or Olympiad-level difficulty.

D Supplementary Evaluation Results

In this section, we show the breakdown of our main results in Table 6 and ablation results in Table 7

E Results of Other Benchmarks

In order to evaluate the performances of retrieval-PRM completely, we conducted additional experiments employing two mainstream evaluation benchmarks: PRMBench and Best-of-N.

PRMBench is a comprehensive benchmark designed to assess process-level reward models across multiple dimensions, including simplicity, soundness, and sensitivity (Song et al., 2025). It consists of 6,216 carefully curated problems and 83,456 step-level labels, enabling fine-grained evaluation of model robustness and generalization. Experimental results demonstrate that RetrievalPRM achieves the highest overall performance on PRMBench, outperforming significantly larger models such as OI-mini and Qwen-72B. Furthermore, RetrievalPRM excels in individual metrics, including soundness and simplicity, further underscoring the robustness and comprehensive effectiveness of our approach.

To further assess practical utility, we adopted the Best-of-N evaluation (Li and Li, 2025). In this setting, a policy model generates \( n \) candidate solutions for each problem, and the PRM assigns scores to each step within the candidate solutions. The final answer is selected by identifying the candidate with the highest trajectory score, which is computed using one of several aggregation strategies: \( min \) (using the minimal step score), \( mean \) (average score), or \( last \) (final step score). For each policy and PRM pair, we report the results using the best aggregation method. Our experiments leverage the publicly available solution steps released in the PQM paper to perform the Best-of-N assessment. RetrievalPRM consistently outperforms existing open-source PRMs across various policy models, providing further evidence that retrieval-augmented approaches can effectively mitigate the step-level out-of-distribution (Step-OOD) challenges arising from policy model variations.

These additional experiments in Table 8 and Table 9 further validate the effectiveness of Retrieval-PRM and demonstrate its practical advantages.

F Retrieval Pool Diversity

In this paper, our retrieval pool comprises both the MathShepherd and PRM800K datasets. To empirically examine the effect of the retrieval pool, we conducted supplementary experiments in which the retrieval pool was augmented with the recently released OpenReason dataset (OpenReason, 2024). From the results in Table 5. Specifically, we observe that

• Compared to the original retrieval pool, expanding the pool leads to a noticeable improvement in overall performance. This supports our initial hypothesis that the composition of the retrieval pool has a significant impact on the results.

• The OpenReason dataset includes the challenging AIME benchmark. Adding solutions from this dataset into the retrieval pool results in larger performance gains on difficult test sets. This provides further evidence that the diversity of the retrieval pool plays a crucial role in enhancing model performance.

• Additionally, we conducted further experiments where the model was trained using only MathShepherd and PRM800K as the retrieval pool. At test time, however, we expanded the pool by including OpenReason. Even without retraining, the model achieved improved performance, which indicates: (1) RetrievalPRM has learned to effectively utilize retrieved knowledge in a generalizable way - even when the retrieval pool at test time differs from that at training time. (2) Leveraging this property, we can support retrievalPRM in solving harder or unseen problems simply by expanding the retrieval pool, without needing to retrain the model. This greatly enhances the reusability and scalability of retrieval-PRM.
Table 5: Performance on Different Retrieval Pool

<table>
  <tr>
    <th>Train Ret-Pool</th>
    <th>Test Ret-Pool</th>
    <th>GSM8K</th>
    <th>MATH</th>
    <th>OlympiadBench</th>
    <th>OmniMATH</th>
  </tr>
  <tr>
    <td>MathShepherd+ PRM800K</td>
    <td>MathShepherd+ PRM800K</td>
    <td>74.6</td>
    <td>71.14</td>
    <td>60.23</td>
    <td>57.33</td>
  </tr>
  <tr>
    <td>MathShepherd+ PRM800K</td>
    <td>MathShepherd+ PRM800K+</td>
    <td>73.0</td>
    <td>72.80</td>
    <td>61.71</td>
    <td>58.78</td>
  </tr>
  <tr>
    <td>MathShepherd+ PRM800K+</td>
    <td>MathShepherd+ PRM800K+</td>
    <td>72.0</td>
    <td>72.94</td>
    <td>61.75</td>
    <td>59.53</td>
  </tr>
  <tr>
    <td>MathShepherd+ PRM800K+</td>
    <td>MathShepherd+ PRM800K+</td>
    <td></td>
    <td></td>
    <td></td>
    <td></td>
  </tr>
</table>

Table 6: Breakdown of evaluation results of different models on ProcessBench. The best result is given in bold, and the second-best value is underlined.

<table>
  <tr>
    <th rowspan="2">Model</th>
    <th colspan="3">GSM8k</th>
    <th colspan="3">MATH</th>
    <th colspan="3">OlympiadBench</th>
    <th colspan="3">OmniMATH</th>
  </tr>
  <tr>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
  </tr>
  <tr>
    <td rowspan="10">Open-source PRM</td>
    <td>RetrievalPRM-7B(Ours)</td>
    <td>64.7</td>
    <td>88.1</td>
    <td><b>74.6</b></td>
    <td>67.2</td>
    <td>75.6</td>
    <td><b>71.1</b></td>
    <td>56.0</td>
    <td>65.2</td>
    <td><b>60.2</b></td>
    <td>52.8</td>
    <td>62.65</td>
    <td><b>57.33</b></td>
  </tr>
  <tr>
    <td>Qwen2.5-Math-7B-PRM800K</td>
    <td>53.1</td>
    <td>95.3</td>
    <td>68.2</td>
    <td>48.0</td>
    <td>90.1</td>
    <td><u>62.6</u></td>
    <td>35.7</td>
    <td>87.3</td>
    <td><u>50.7</u></td>
    <td>29.8</td>
    <td>86.3</td>
    <td>44.3</td>
  </tr>
  <tr>
    <td>Skywork-PRM-7B</td>
    <td>61.8</td>
    <td>82.9</td>
    <td>70.8</td>
    <td>43.8</td>
    <td>69.2</td>
    <td>53.6</td>
    <td>17.9</td>
    <td>31.9</td>
    <td>22.9</td>
    <td>14.0</td>
    <td>41.9</td>
    <td>21.0</td>
  </tr>
  <tr>
    <td>RLHFlow-PRM-Mistral-8B</td>
    <td>33.8</td>
    <td>99.0</td>
    <td>50.4</td>
    <td>21.7</td>
    <td>72.2</td>
    <td>33.4</td>
    <td>8.2</td>
    <td>43.1</td>
    <td>13.8</td>
    <td>9.6</td>
    <td>45.2</td>
    <td>15.8</td>
  </tr>
  <tr>
    <td>RLHFlow-PRM-Deepseek-8B</td>
    <td>24.2</td>
    <td>98.4</td>
    <td>38.8</td>
    <td>21.4</td>
    <td>80.0</td>
    <td>33.8</td>
    <td>10.1</td>
    <td>51.0</td>
    <td>16.9</td>
    <td>10.1</td>
    <td>51.9</td>
    <td>16.9</td>
  </tr>
  <tr>
    <td>Skywork-PRM-1.5B</td>
    <td>50.2</td>
    <td>71.5</td>
    <td>59.0</td>
    <td>37.9</td>
    <td>65.3</td>
    <td>48.0</td>
    <td>15.4</td>
    <td>26.0</td>
    <td>19.3</td>
    <td>13.6</td>
    <td>32.8</td>
    <td>19.2</td>
  </tr>
  <tr>
    <td>Math-Shepherd-PRM-7B</td>
    <td>32.4</td>
    <td>91.7</td>
    <td>47.9</td>
    <td>18.0</td>
    <td>82.0</td>
    <td>29.5</td>
    <td>15.0</td>
    <td>71.1</td>
    <td>24.8</td>
    <td>14.2</td>
    <td>73.0</td>
    <td>23.8</td>
  </tr>
  <tr>
    <td rowspan="10">Language Models</td>
    <td>QwQ-32B-Preview</td>
    <td>81.6</td>
    <td>95.3</td>
    <td><b>88.0</b></td>
    <td>78.1</td>
    <td>79.3</td>
    <td><b>78.7</b></td>
    <td>61.4</td>
    <td>54.6</td>
    <td><b>57.8</b></td>
    <td>55.7</td>
    <td>68.0</td>
    <td><b>61.3</b></td>
  </tr>
  <tr>
    <td>GPT-4o</td>
    <td>70.0</td>
    <td>91.2</td>
    <td>79.2</td>
    <td>54.4</td>
    <td>76.6</td>
    <td><u>63.6</u></td>
    <td>45.8</td>
    <td>58.4</td>
    <td>51.4</td>
    <td>45.2</td>
    <td>53.5</td>
    <td>61.9</td>
  </tr>
  <tr>
    <td>Qwen2.5-72B-Instruct</td>
    <td>62.8</td>
    <td>96.9</td>
    <td>76.2</td>
    <td>46.3</td>
    <td>93.1</td>
    <td>61.8</td>
    <td>38.7</td>
    <td>92.6</td>
    <td><u>54.6</u></td>
    <td>36.6</td>
    <td>90.9</td>
    <td>52.2</td>
  </tr>
  <tr>
    <td>Llama-3.3-70B-Instruct</td>
    <td>72.5</td>
    <td>96.9</td>
    <td><u>82.9</u></td>
    <td>43.3</td>
    <td>94.6</td>
    <td>59.4</td>
    <td>31.0</td>
    <td>94.1</td>
    <td>46.7</td>
    <td>28.2</td>
    <td>90.5</td>
    <td>43.0</td>
  </tr>
  <tr>
    <td>Qwen2.5-32B-Instruct</td>
    <td>49.3</td>
    <td>97.9</td>
    <td>65.6</td>
    <td>36.7</td>
    <td>95.8</td>
    <td>53.1</td>
    <td>25.3</td>
    <td>95.9</td>
    <td>40.0</td>
    <td>24.1</td>
    <td>92.5</td>
    <td>38.3</td>
  </tr>
  <tr>
    <td>Qwen2.5-14B-Instruct</td>
    <td>54.6</td>
    <td>94.8</td>
    <td>69.3</td>
    <td>38.4</td>
    <td>87.4</td>
    <td>53.3</td>
    <td>31.5</td>
    <td>78.8</td>
    <td>45.0</td>
    <td>28.3</td>
    <td>76.3</td>
    <td>41.3</td>
  </tr>
  <tr>
    <td>Qwen2.5-Coder-32B-Instruct</td>
    <td>54.1</td>
    <td>94.8</td>
    <td>68.9</td>
    <td>44.9</td>
    <td>90.6</td>
    <td>60.1</td>
    <td>33.4</td>
    <td>91.2</td>
    <td>48.9</td>
    <td>31.5</td>
    <td>87.6</td>
    <td>46.3</td>
  </tr>
  <tr>
    <td>Qwen2.5-Coder-14B-Instruct</td>
    <td>33.8</td>
    <td>96.4</td>
    <td>50.1</td>
    <td>25.4</td>
    <td>92.4</td>
    <td>39.9</td>
    <td>20.7</td>
    <td>94.1</td>
    <td>34.0</td>
    <td>15.9</td>
    <td>94.2</td>
    <td>27.3</td>
  </tr>
  <tr>
    <td>Qwen2.5-Coder-7B-Instruct</td>
    <td>7.7</td>
    <td>100.0</td>
    <td>14.3</td>
    <td>3.4</td>
    <td>98.3</td>
    <td>6.5</td>
    <td>2.1</td>
    <td>99.1</td>
    <td>4.1</td>
    <td>0.9</td>
    <td>98.3</td>
    <td>1.8</td>
  </tr>
  <tr>
    <td>Qwen2.5-Math-7B-Instruct</td>
    <td>15.5</td>
    <td>100.0</td>
    <td>26.8</td>
    <td>14.8</td>
    <td>96.8</td>
    <td>25.7</td>
    <td>7.7</td>
    <td>91.7</td>
    <td>14.2</td>
    <td>6.9</td>
    <td>88.0</td>
    <td>12.7</td>
  </tr>
  <tr>
    <td>Llama-3.1-70B-Instruct</td>
    <td>64.3</td>
    <td>89.6</td>
    <td>74.9</td>
    <td>35.4</td>
    <td>75.6</td>
    <td>48.2</td>
    <td>35.1</td>
    <td>69.9</td>
    <td>46.7</td>
    <td>30.7</td>
    <td>61.8</td>
    <td>41.0</td>
  </tr>
  <tr>
    <td>Meta-Llama-3-70B-Instruct</td>
    <td>57.7</td>
    <td>96.9</td>
    <td>52.2</td>
    <td>13.0</td>
    <td>93.3</td>
    <td>22.8</td>
    <td>12.0</td>
    <td>92.0</td>
    <td>21.2</td>
    <td>11.2</td>
    <td>91.7</td>
    <td>20.0</td>
  </tr>
  <tr>
    <td>Qwen2.72B-Instruct</td>
    <td>57.0</td>
    <td>82.9</td>
    <td>67.6</td>
    <td>37.7</td>
    <td>70.9</td>
    <td>49.2</td>
    <td>34.0</td>
    <td>55.2</td>
    <td>42.1</td>
    <td>32.3</td>
    <td>53.1</td>
    <td>40.2</td>
  </tr>
  <tr>
    <td>Qwen2.5-7B-Instruct</td>
    <td>40.6</td>
    <td>33.2</td>
    <td>36.5</td>
    <td>30.8</td>
    <td>45.1</td>
    <td>36.6</td>
    <td>26.5</td>
    <td>33.9</td>
    <td>29.7</td>
    <td>26.2</td>
    <td>28.6</td>
    <td>27.4</td>
  </tr>
  <tr>
    <td>Qwen2.7B-Instruct</td>
    <td>40.6</td>
    <td>4.7</td>
    <td>8.4</td>
    <td>30.5</td>
    <td>13.8</td>
    <td>19.0</td>
    <td>22.4</td>
    <td>10.9</td>
    <td>14.7</td>
    <td>20.0</td>
    <td>8.7</td>
    <td>12.1</td>
  </tr>
  <tr>
    <td>Llama-3.1-8B-Instruct</td>
    <td>44.4</td>
    <td>6.2</td>
    <td>10.9</td>
    <td>41.9</td>
    <td>2.7</td>
    <td>5.1</td>
    <td>32.4</td>
    <td>1.5</td>
    <td>2.8</td>
    <td>32.0</td>
    <td>0.8</td>
    <td>1.6</td>
  </tr>
  <tr>
    <td>Meta-Llama-3-8B-Instruct</td>
    <td>42.5</td>
    <td>7.8</td>
    <td>13.1</td>
    <td>28.6</td>
    <td>9.1</td>
    <td>13.8</td>
    <td>27.1</td>
    <td>2.7</td>
    <td>4.8</td>
    <td>26.1</td>
    <td>8.3</td>
    <td>12.6</td>
  </tr>
</table>

Table 7: Breakdown of evaluation results of different variants of RetrievalPRM on ProcessBench. We remove different components of RetrievalPRM to evaluate the contribution of each part to the model. The best result is given in bold, and the second-best value is underlined.

<table>
  <tr>
    <th rowspan="2">Retrieval Components</th>
    <th colspan="3">GSM8k</th>
    <th colspan="3">MATH</th>
    <th colspan="3">OlympiadBench</th>
    <th colspan="3">OmniMATH</th>
    <th rowspan="2">Avg.F1</th>
  </tr>
  <tr>
    <th>Step-level</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
    <th>error</th>
    <th>correct</th>
    <th>F1</th>
  </tr>
  <tr>
    <td>✓</td>
    <td>✓</td>
    <td>64.7</td>
    <td>88.1</td>
    <td><b>74.6</b></td>
    <td>67.2</td>
    <td>75.6</td>
    <td><b>71.1</b></td>
    <td>56.0</td>
    <td>65.2</td>
    <td><b>60.2</b></td>
    <td>52.8</td>
    <td>62.65</td>
    <td><b>57.33</b></td>
    <td><b>65.8</b></td>
  </tr>
  <tr>
    <td>✓</td>
    <td>×</td>
    <td>61.8</td>
    <td>94.8</td>
    <td><b>74.9</b></td>
    <td>62.1</td>
    <td>83.3</td>
    <td><b>71.2</b></td>
    <td>48.7</td>
    <td>77.3</td>
    <td><u>59.8</u></td>
    <td>43.2</td>
    <td>73.4</td>
    <td>54.4</td>
    <td><u>65.0</u></td>
  </tr>
  <tr>
    <td>×</td>
    <td>✓</td>
    <td>51.7</td>
    <td>97.4</td>
    <td>67.5</td>
    <td>57.2</td>
    <td>87.4</td>
    <td>69.2</td>
    <td>46.0</td>
    <td>82.0</td>
    <td>58.9</td>
    <td>43.9</td>
    <td>78.4</td>
    <td><u>56.3</u></td>
    <td>63.0</td>
  </tr>
  <tr>
    <td>×</td>
    <td>×</td>
    <td>50.7</td>
    <td>92.7</td>
    <td>65.6</td>
    <td>57.9</td>
    <td>81.0</td>
    <td>67.5</td>
    <td>46.9</td>
    <td>68.7</td>
    <td>55.8</td>
    <td>39.7</td>
    <td>71.0</td>
    <td>50.9</td>
    <td>59.9</td>
  </tr>
</table>
Table 8: Results of PRMBench

<table>
  <tr>
    <th>Model</th>
    <th>Overall</th>
    <th>NR.</th>
    <th>NCL.</th>
    <th>S1</th>
    <th>ES.</th>
    <th>SC.</th>
    <th>DC.</th>
    <th>CI.</th>
    <th>S2</th>
    <th>PS.</th>
    <th>DR.</th>
    <th>MS.</th>
    <th>S3</th>
  </tr>
  <tr>
    <td>RetrievalPRM(Ours)</td>
    <td>68.9</td>
    <td>51.9</td>
    <td>58.6</td>
    <td>55.3</td>
    <td>76.2</td>
    <td>74.8</td>
    <td>69.5</td>
    <td>79.5</td>
    <td>75.0</td>
    <td>62.3</td>
    <td>73.0</td>
    <td>99.2</td>
    <td>78.2</td>
  </tr>
  <tr>
    <td>o1-mini</td>
    <td>68.8</td>
    <td>65.6</td>
    <td>63.7</td>
    <td>64.6</td>
    <td>74.5</td>
    <td>67.7</td>
    <td>73.8</td>
    <td>72.3</td>
    <td>72.1</td>
    <td>61.8</td>
    <td>64.8</td>
    <td>100.0</td>
    <td>75.5</td>
  </tr>
  <tr>
    <td>Qwen2.5-Math-PRM-72B</td>
    <td>68.2</td>
    <td>50.4</td>
    <td>58.8</td>
    <td>54.6</td>
    <td>73.7</td>
    <td>71.1</td>
    <td>72.2</td>
    <td>78.6</td>
    <td>73.9</td>
    <td>60.3</td>
    <td>71.2</td>
    <td>99.4</td>
    <td>77.0</td>
  </tr>
  <tr>
    <td>GPT-4o</td>
    <td>66.8</td>
    <td>57.0</td>
    <td>62.4</td>
    <td>59.7</td>
    <td>72.0</td>
    <td>69.7</td>
    <td>70.7</td>
    <td>71.1</td>
    <td>70.9</td>
    <td>62.5</td>
    <td>65.7</td>
    <td>99.2</td>
    <td>75.8</td>
  </tr>
  <tr>
    <td>Qwen2.5-Math-PRM-7B</td>
    <td>65.5</td>
    <td>49.0</td>
    <td>55.1</td>
    <td>52.1</td>
    <td>71.8</td>
    <td>67.3</td>
    <td>66.3</td>
    <td>78.5</td>
    <td>71.0</td>
    <td>57.6</td>
    <td>69.1</td>
    <td>99.7</td>
    <td>75.5</td>
  </tr>
  <tr>
    <td>Skywork-PRM-7B</td>
    <td>65.1</td>
    <td>56.4</td>
    <td>62.8</td>
    <td>59.6</td>
    <td>69.4</td>
    <td>67.1</td>
    <td>67.7</td>
    <td>69.9</td>
    <td>68.5</td>
    <td>60.9</td>
    <td>65.8</td>
    <td>93.2</td>
    <td>73.3</td>
  </tr>
  <tr>
    <td>QwQ-Preview-32B</td>
    <td>63.6</td>
    <td>57.2</td>
    <td>55.6</td>
    <td>56.4</td>
    <td>67.4</td>
    <td>72.3</td>
    <td>66.2</td>
    <td>66.9</td>
    <td>68.2</td>
    <td>57.8</td>
    <td>62.7</td>
    <td>100.0</td>
    <td>73.5</td>
  </tr>
  <tr>
    <td>RLHFFlow-PRM-Mistral-8B</td>
    <td>54.4</td>
    <td>46.1</td>
    <td>47.3</td>
    <td>46.7</td>
    <td>56.6</td>
    <td>55.1</td>
    <td>54.4</td>
    <td>63.8</td>
    <td>57.5</td>
    <td>51.5</td>
    <td>56.2</td>
    <td>97.9</td>
    <td>68.5</td>
  </tr>
  <tr>
    <td>Llemma-PRM800k-7B</td>
    <td>52.0</td>
    <td>49.3</td>
    <td>53.4</td>
    <td>51.4</td>
    <td>56.4</td>
    <td>47.1</td>
    <td>46.7</td>
    <td>53.3</td>
    <td>50.9</td>
    <td>51.0</td>
    <td>53.5</td>
    <td>93.6</td>
    <td>66.0</td>
  </tr>
  <tr>
    <td>MathShepherd-Mistral-7B</td>
    <td>47.0</td>
    <td>44.0</td>
    <td>50.3</td>
    <td>47.1</td>
    <td>49.4</td>
    <td>44.5</td>
    <td>41.3</td>
    <td>47.7</td>
    <td>45.7</td>
    <td>47.2</td>
    <td>48.6</td>
    <td>86.1</td>
    <td>60.7</td>
  </tr>
</table>

Table 9: Results of Best-of-n Search with different PRMs

<table>
  <tr>
    <th>Dataset</th>
    <th>Policy Model</th>
    <th>RetrievalPRM (Ours)</th>
    <th>RLHFFlow-PRM</th>
    <th>Skywork-PRM</th>
    <th>MathShepherd-Mistral</th>
  </tr>
  <tr>
    <td>MATH</td>
    <td>Llama-3-70B-Instruct</td>
    <td>48.8</td>
    <td>42.0</td>
    <td>38.0</td>
    <td>46.2</td>
  </tr>
  <tr>
    <td>MATH</td>
    <td>MetaMath-Mistral-7B</td>
    <td>37.4</td>
    <td>29.8</td>
    <td>24.0</td>
    <td>32.6</td>
  </tr>
  <tr>
    <td>MATH</td>
    <td>Muggle-Math-13B</td>
    <td>33.8</td>
    <td>23.6</td>
    <td>19.0</td>
    <td>28.4</td>
  </tr>
  <tr>
    <td>GSM-Plus</td>
    <td>Llama-3-70B-Instruct</td>
    <td>74.5</td>
    <td>70.8</td>
    <td>67.7</td>
    <td>72.6</td>
  </tr>
  <tr>
    <td>GSM-Plus</td>
    <td>MetaMath-Mistral-7B</td>
    <td>63.7</td>
    <td>53.5</td>
    <td>48.0</td>
    <td>60.3</td>
  </tr>
  <tr>
    <td>GSM-Plus</td>
    <td>Muggle-Math-13B</td>
    <td>62.4</td>
    <td>50.8</td>
    <td>43.2</td>
    <td>59.5</td>
  </tr>
</table>

Reviewer paper summaries: This paper addresses out-of-distribution (OOD) challenges in Process Reward Models (PRMs) for mathematical reasoning. The authors identify two key OOD issues: step OOD (caused by differences in reasoning patterns across model types and sizes) and question OOD (arising from dataset shifts between training and real-world problems). To address these issues, they introduce RetrievalPRM, a novel framework utilizing a two-stage retrieval-enhanced mechanism that retrieves semantically similar questions and steps as a warmup, enhancing PRM's ability to evaluate target steps and improving generalization across different models and problem types. Experiments demonstrate that RetrievalPRM outperforms existing baselines across multiple real-world datasets, including GSM8K, MATH, OlympiadBench, and OmniMATH.

The paper proposes a method to address the out-of-distribution (OOD) challenges faced by Process Reward Models (PRMs) in mathematical reasoning, specifically ​step OOD and ​question OOD. The authors propose ​RetrievalPRM, a framework that uses a two-stage retrieval mechanism: ​(1) Question-level retrieval to find semantically similar questions. (2) Step-level retrieval to identify analogous reasoning steps. These retrieved examples are integrated into the PRM’s input prompt as a demonstration and to improve generalization. Experiments on four datasets show that RetrievalPRM outperforms existing PRMs and LLM-based critics.

This paper presents Retrieval-Augmented Process Reward Model (RetrievalPRM) designed to enhance the capabilities of large language models (LLMs) by integrating external knowledge retrieval. The proposed method utilizes a retrieval module to fetch relevant information from a large corpus, which is then incorporated into the input to the LLM. Experimental results demonstrate significant improvements in performance across multiple real-world datasets, such as GSM8k, MATH, OlympiadBench, etc.
Summary of strengths: - Novel Problem Identification: The paper clearly identifies and formalizes two types of OOD problems in mathematical reasoning PRMs, supported with clear examples and visualizations.
- Technical Innovation: The two-stage retrieval mechanism is a well-designed solution to address both question-level and step-level OOD issues, with a thoughtful implementation strategy.

1. The paper studies OOD issues in PRMs, with a focus on model-type/size variability and dataset shifts.
2. The paper introduces a dual retrieval mechanism for PRMs, which is intuitive and effective.
3. The paper demonstrates significant gains on challenging datasets.

- 1) The authors have conducted a thorough evaluation of the proposed method across multiple benchmarks, such as MATH, GSM8K, OlympiadBench, and OmniMATH, demonstrating its versatility and effectiveness in different scenarios.

- 2) This work is well-written and well-organized. Readers can quickly grasp the key motivations of the paper.
Summary of weaknesses: - In lines 048-050, the authors frame PRM as a binary classification problem without rigorous justification for this choice. They don't explore alternative formulations such as multi-class classification or generation-based approaches for evaluating reasoning steps.
- The authors identify different solution styles across model types and sizes as a cause of step OOD issues, but fail to explore whether this could be addressed through context learning or prompting techniques to achieve more consistent solution styles. The paper lacks experiments comparing their retrieval approach with alternative solutions that might standardize solution styles through prompting.
- Lines 207-209 state that traditional PRM models are constrained by "predefined solution steps" and "unable to handle unseen questions," but this claim is not clearly explained or substantiated.
- As acknowledged in the limitations section, the retrieval pool is constructed only from PRM800K and Math-Shepherd, which limits the diversity of problems available for retrieval. This constraint may prevent the system from finding truly relevant questions or steps for target problems, particularly for more complex mathematical domains.
- There is minimal qualitative analysis of what kinds of questions and steps are being retrieved and how they specifically help with different types of mathematical problems.
- The paper lacks human evaluation to verify whether the improvements align with human judgment on the correctness of mathematical reasoning steps.

1. The retrieval dataset is restricted to PRM800K and Math-Shepherd, which may limit diversity and overlap heavily with the training set.
2. The paper acknowledges that Sentence-BERT may not capture mathematical reasoning similarity well. However, the paper does not explore alternative embedding strategies.
3. Figure 4 shows the study on the choice of k for question-level retrieval but a lack of study for step-level retrieval.

- 1) The retrieval pool is only constructed from PRM800K. The authors should thoroughly investigate how the composition and diversity of the retrieval pool affect the performance of LLMs on tasks such as math and GSM8K. This could provide insights into the scalability of the proposed retrieval-based approach.

- 2) Based on my understanding, leveraging retrieval to introduce external knowledge can lead to performance improvements. As a result, it is unclear how different retrieval strategies or architectures might influence the results.
Comments, suggestions and typos: 016-018: "By utilizing a two-stage retrieval-enhanced mechanism" is repeated.

1. How do irrelevant retrieved questions/steps hurt performances?
2. Figure 3 is informative but quite dense. However, the explanation is not self-contained. Examples can be simplified. Figure 2 also has a similar problem.

Please refer to Part of [Summary Of Weaknesses]
Overall assessment scores: 2

3.5

3
Metareview score: 3
Metareview: This paper focuses on improving the generalization of Process Reward Models (PRMs) for mathematical reasoning, specifically tackling out-of-distribution (OOD) challenges. The authors identify two key OOD issues: "step OOD" arising from varying reasoning patterns across model types/sizes, and "question OOD" due to dataset shifts. They propose RetrievalPRM, a framework employing a two-stage retrieval mechanism (retrieving similar questions and steps) to enhance PRMs. The core idea is to use retrieved information as a "warm-up" or context to aid the PRM in evaluating target reasoning steps. Experiments demonstrate superior performance over existing baselines on multiple mathematical reasoning benchmarks.
Summary of reasons to publish: 1. This paper explicitly identifies and distinguishes between "step OOD" and "question OOD" within mathematical reasoning PRMs, addressing an important and practical problem.  

2. The authors leverages retrieval augmentation to provide contextual information for mitigating OOD issues, which is an intuitive and plausible technical direction.  

3. Experiments demonstrate significant performance improvements across multiple datasets, particularly on challenging ones, validating the method's effectiveness.  

4. The authors provide code, an enhanced dataset, and the model, facilitating reproducibility and future work by the community.  

5. The authors were responsive in the rebuttal phase, adding experiments (e.g., expanding the retrieval pool, varying strategies, in-context learning baselines, ablations) that further strengthen the paper's claims.