Rewarding Progress: Scaling Automated Process Verifiers for LLM Reasoning

Thank you for the review and for a positive assessment of our paper! To address your concerns we add a new experiment that defines step-level process rewards using discounting, i.e. it takes into account the number of steps taken by the prover policy to complete the solution from the last step generated by the base policy. We also provide a pointwise responses to each of your questions below. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score.

What does it mean for prover policy to be too misaligned with the base policy?

In the right hand side of Equation 6 (Theorem 3.1), we define the alignment between prover policy $\mu$ and current base policy $\pi$ as the expected inner product between $A^\mu(s, a)$ and $A^\pi(s, a)$, where the expectation is taken over actions $a \sim \pi(\cdot \mid s )$ sampled from the base policy $\pi$, and the distribution over states $\rho$. Mathematically, the alignment is denoted as $E_{s\sim \rho} E_{a\sim \pi(a \mid s)} A^\mu(s, a) A^\pi(s, a)$. In our theoretical result in Section 3.4, we show that process rewards defined as advantages of prover policies with higher variance (under actions sampled from the base policy) guarantee a stronger improvement in the base policy (this is the distinguishability term in Equation 6). At the same time, we cannot have the prover to be simply a high variance (entropy) policy which takes high values of $A^\mu$, without being correlated with the outcome rewards at all. Thus, we want prover policies such that the process rewards $A^\mu$ are correlated with the base policy advantages $A^\pi$, where $A^\pi$ prefers actions that achieve a high outcome reward under the base policy. This is exactly what is measured by our alignment term in Equation 6.

Accounting for steps taken by the strong prover policy to arrive at the correct answer

This is a great point and we add a new experiment that defines process rewards in a way that takes into account the number of steps taken by the prover policy to complete the solution.

We add a new result (Figure 16) and discussion on using discounted rewards from provers to Appendix E of the submission. Here, we train PAVs to predict the advantages of discounted rewards from strong prover policies. Here, for the problem $\mathbf{x}$, and the state, action pair $s, a$, the process rewards are given by the effective reward from Equation 5: $Q^\pi + \alpha A^\mu$, except that the advantage $A^\mu$ is the difference in discounted rewards, i.e.:

$A^\mu(s, a) = E_{y \sim \mu(\cdot \mid s, a)} \left[ \lambda^{\mathrm{len}(y)-\mathrm{len}(s)-1} \mathrm{Rex}(y, y^\star_{x})\right] - E_{y \sim \mu(\cdot \mid s)} \left[ \lambda^{\mathrm{len}(y)-\mathrm{len}(s)} \mathrm{Rex}(y, y^\star_{x})\right],$

where the prover policy samples solution $y$ with $\mathrm{len}(y)$ steps to complete the solution, from a state $s$, which already has $\mathrm{len}(s)$ steps in it.

For this setting, we train a verifier to predict discounted rewards for the Gemma 9B prover policy. We find that the discounted process rewards from the stronger 9B prover policy performs worse than undiscounted rewards from the weaker Gemma 2B prover policy, when using either to train the 2B base policy with online RL.

The main reason for the discounted rewards to not enable the use of strong provers is because strong prover policies tend to disregard states generated by the base policy (as illustrated in Figure 2b). This means, that irrespective of whether the weak prover policy generates a partially correct or incorrect solution, when we rollout the strong prover policy from this state generated by the base policy, the strong prover directly attempts to answer the math problem with its own solution trace. Thus, from any state the strong prover is expected to complete the solution with roughly the same number of steps. This means that $A^\mu \approx 0$ even in the discounted case, which reduces the ability of the strong prover policy to distinguish steps taken by the base policy.

Thank you very much for addressing my questions. Since I am not very familiar with this area, I will wait to see if other reviewers raise their score above an 8 before considering changing my own, as I cannot speak confidently about its significance or contribution.

Thank you for the review and for a positive assessment of our paper! To address your concerns, we have made the comparison with other PRM models from prior works more explicit, and have added a new experiment where we compare RL training with PAVs and RL training with PRMs from prior works. We also respond to your other question on advantages vs. value functions. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score.

Comparing with baselines that use PRMs proposed in prior works

PRM baselines for test-time search For test-time search, our results in Fig 5a of our submission already include both baselines: 1) beam search; and 2) best-of-N search using PRMs from prior works. We make this more explicit with a new plot (Figure 13) we add in Appendix C. We find that beam search with PAVs is 1.5x-5x more compute efficient than both beam search with PRMs from prior works, and best-of-N search with PRMs from prior works. We do not attempt to compare beam search with PAVs against weaker search algorithms used in conjunction with prior PRMs. When evaluating prior work (Snell et. al. (2024)) that uses $Q^\pi$ as their proposed PRM, we run beam search where the states in the beam are ranked only using $Q^\pi$, as opposed to the effective reward $Q^\pi + \alpha A^\mu$ in our case. Thus, the search procedure in the prior work and ours is identical, and we only change the re-ranking mechanism.

Other works on PRMs (Wang et. al. (2024), Luo et. al. (2024)) that also propose to use $Q^\pi$, use PRMs for best-of-N search, where they only rank the full sequences using the trained PRM (by taking the minimum over the $Q^\pi$ values at each step in the generation sampled from the base policy $\pi$). For completeness, we also compare PAVs with these approaches (PRMs-as-ORMs), which performs similarly to using PAVs-as-ORMs in Figure 5a.

PRM baselines for online RL training We add a new experiment where we use the PRMs proposed in prior works on test-time search, (Wang et. al. (2024), Luo et. al. (2024), Snell et. al. (2024)) as step-level rewards for training a base policy with RL. Here, the PRMs are trained offline and fixed, and then used to assign scores to intermediate steps in trajectories sampled during online RL. For this, we add a new experimental result where we use $Q^\pi_{\mathrm{base}}$ as the step-level score during RL training initialized with the base policy $\pi_{\mathrm{base}}$ since $Q^\pi_{\mathrm{base}}$ is the PRM proposed in prior works. This step-level reward is used to reward trajectories sampled during online RL training, in addition to the outcome reward, similar to our effective reward in Equation 5. We find that the test accuracy drops quite quickly, even though the train rewards keep improving, since the policy just learns to hack the step-level Q-values, for the reasons we explain in L240-245. We see a similar collapse when we use $Q^\mu$ instead of $A^\mu$ (see Appendix G for qualitative examples). On the other hand, if we were to use the Q-values or advantages from the current policy iterate $\pi_t$ as the step-level reward, then that is equivalent to only optimizing the outcome reward, and the only benefit of using the step rewards would be to reduce variance during online RL iterations. Thus, for online RL as well, our proposed step-level rewards (PAVs) which uses advantages of the prover policy $A^\mu$ outperforms baselines that plugin PRMs proposed in prior works on test-time search, We have added this new experiment on online RL where step-level rewards are given by PRM $Q^\pi_{\mathrm{base}}$ proposed in prior works to Appendix E, Figure 17.

Thank you for the review! To address your concerns, we clarify that for experiments in Section 4 (test-time search), we already compare with several prior works that use PRMs for beam search and best-of-N search. If there is a comparison that we are missing, please let us know and we will be happy to add it. We also include a new experiment for RL training with PRM baselines that we used for test-time search. We also respond to your other questions and concerns below. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score. We are happy to discuss further!

Comparing beam search with PAVs and search with PRMs from prior works

For test-time search, our results in Fig 5a of our submission already include both baselines: 1) beam search; and 2) best-of-N search using PRMs from prior works. We make this more explicit with a new plot (Figure 13) we add in Appendix C. We find that beam search with PAVs is 1.5x-5x more compute efficient than both beam search with PRMs from prior works, and best-of-N search with PRMs from prior works. We do not attempt to compare beam search with PAVs against weaker search algorithms used in conjunction with prior PRMs.

When evaluating prior work (Snell et. al. (2024)) that uses $Q^\pi$ as their proposed PRM, we run beam search where the states in the beam are ranked only using $Q^\pi$, as opposed to the effective reward $Q^\pi + \alpha A^\mu$ in our case. Thus, the search procedure in the prior work and ours is identical, and we only change the re-ranking mechanism.

Other works on PRMs (Wang et. al. (2024), Luo et. al. (2024)) that also propose to use $Q^\pi$, use PRMs for best-of-N search, where they only rank the full sequences using the trained PRM (by taking the minimum over the $Q^\pi$ values at each step in the generation sampled from the base policy $\pi$). For completeness, we also compare PAVs with these approaches (PRMs-as-ORMs), which performs similarly to using PAVs-as-ORMs in Figure 5a.

PRM baselines for online RL training

We add a new experiment where we use the PRMs proposed in prior works on test-time search, (Wang et. al. (2024), Luo et. al. (2024), Snell et. al. (2024)) as step-level rewards for training a base policy with RL. Here, the PRMs are trained offline and fixed, and then used to assign scores to intermediate steps in trajectories sampled during online RL. For this, we add a new experimental result where we use $Q^\pi_{\mathrm{base}}$ as the step-level score during RL training initialized with the base policy $\pi_{\mathrm{base}}$ since $Q^\pi_{\mathrm{base}}$ is the PRM proposed in prior works. This step-level reward is used to reward trajectories sampled during online RL training, in addition to the outcome reward, similar to our effective reward in Equation 5. We find that the test accuracy drops quite quickly, even though the train rewards keep improving, since the policy just learns to hack the step-level Q-values, for the reasons we explain in L240-245. We see a similar collapse when we use $Q^\mu$ instead of $A^\mu$ (see Appendix G for qualitative examples). On the other hand, if we were to use the Q-values or advantages from the current policy iterate $\pi_t$ as the step-level reward, then that is equivalent to only optimizing the outcome reward, and the only benefit of using the step rewards would be to reduce variance during online RL iterations. Thus, for online RL as well, our proposed step-level rewards (PAVs) which uses advantages of the prover policy $A^\mu$ outperforms baselines that plugin PRMs proposed in prior works on test-time search, We have added this new experiment on online RL where step-level rewards are given by PRM $Q^\pi_{\mathrm{base}}$ proposed in prior works to Appendix E, Figure 17.

Expanding on our explanation in L369-374 on how to use PAVs during test-time.

For test-time search over responses from the base policy $\pi$, and fixing a prover policy $\mu$, the following is how we use PAVs for beam search at test-time. At any given time step, the beam holds partially generated responses (states) to a math problem. We score each of the states using our effective reward $Q^\pi + \alpha A^\mu$. Since $A^\mu$ is computed using $Q^\mu$ (Equation 2 in submission), we need only train verifiers to predict $Q^\pi$ and $Q^\mu$. Finally, at the end of beam search, when only we have complete responses in the beam, we pick the best response using an ORM trained only to predict the correctness of final answers. In practice, we find that using a trained ORM to rate complete responses at the end of beam search is not critical, and using the PRM that predicts $Q^\pi$ to rank the full responses also performs similarly. So, the training of an ORM model can be avoided. We have added this discussion to Appendix D. Please let us know if it is still unclear and we would be happy to expand on this further.

Is there a need to generate simpler, repetitive and incorrect steps to explore and discover the correct answer?

We agree with you that it is unclear if there is a need to generate repetitive or incorrect steps to discover the correct answer. But, these steps enable the model to spend more training time compute to discover answers to hard problems, as we describe below.

On hard questions (that were unsolved by the SFT model with pass@256), we clearly do not have coverage over the correct solution trace. Note that these are samples from the SFT model that was only trained to predict the correct final answer in least number of steps, and without any incorrect ste[s in between. So, for these hard questions, we hypothesize that the it is possible for the policy to spend more training time compute, by sampling simpler, repetitive or incorrect steps, that makes it easier to discover the answers for hard math questions. At the same time, simply generating repetitive steps is also bad, which is why we need process advantage verifiers to evaluate a step exactly based on its ability to improve the likelihood of arriving at the correct solution, as opposed to its mathematical relevance and correctness, as judged by a human. This is precisely what we set out to do in Section 3, where we choose process rewards to be potential functions that enable the optimization of outcome rewards (final answer correctness).

Confusion on fixed vs. varying prover, i.e., is $\mu = \mathrm{BoK}(\pi_{t})$?

The prover policy is always fixed. When we say that the class of ''best-of-K'' policies serves as a class that contains good, ``complementary’’ prover policies, we always mean the best-of-K over the original base policy. Thus, we use $\mathrm{BoK}(\pi_{\mathrm{base}})$ to update the current policy $\pi_t$ during online RL. Here, $\pi_t$ is the policy at the $t^\mathrm{th}$ RL training iteration, and $\pi_0$ (initialization of RL training) is set to be the base policy $\pi_{\mathrm{base}}$.

We have matching optimal policies with potential functions, but the right hand side of the bound in Theorem 3.1 can be negative?

Yes, you are correct! Using the theoretical result in Ng et. al. (1999), we can claim that the set of optimal policies that optimize the outcome reward alone, should match the set of optimal policies for our effective reward $Q^\pi + \alpha A^\mu$, since the $\alpha A^\mu$ term is a potential-based reward shaping function (we have added discussion on this in Appendix I).

The above argument does not break, even when $\mu$ and $\pi$ are highly misaligned. This is because, when $\mu$ and $\pi$ are misaligned, $E_\pi [A^\mu A^\pi]$ ends up being negative, which can make the right hand side of the lower bound in Theorem 3.1 negative, as you correctly note. But, this just makes the guarantee from policy improvement weaker. In other words, even though $E[V^\pi_{t+1}]$ is better than $E[V^\pi_t]$, we cannot guarantee the improvement using Theorem 3.1. On the other hand, whenever $\mu$ and $\pi$ are aligned, i.e., $E_\pi [A^\mu A^\pi] > 0$ and $Var_\pi [A^\mu]$ is large, we are guaranteed a stronger improvement in the policy iterates. We hope this clarifies your confusion about Theorem 3.1, and we are happy to explain further if there is still any confusion.

Clarifying our Remark 3.1 on the Best-of-K(base policy) set containing a good set of provers

We have updated the submission to explain the main motivation behind considering the class of best-of-K (over base policy). It is mainly to study a class of policies of increasing strengths (test performance), compared to the base policy. This class is conveniently parameterized (with $K$) for us to run a search over and identify a good prover policy. In Appendix F.4, (to which we add a reference in main submission) we explain what choice of $K$ can give us the best policy improvement lower bound from Theorem 3.1 .

Grid search on $\alpha$

For our RL experiments, we ran a grid search over $\alpha \in [0.5, 6.0]$, using intervals of $0.5$, and found optimal values of $3.0$ for the 9B model and $5.0$ for the 2B model. While this search was computationally expensive, we noted that all values of $5.5 > \alpha > 1.5$ significantly improved the performance of both 2B and 9B models trained with PAVs, over models trained with only ORMs. In fact tuning $\alpha$ on the smaller model (2B) is enough, and using the same for the 9B model already improves performance over ORM. This means that the choice of $\alpha$ is not a very sensitive one and in practice, we expect it to transfer across related problem instances.

Thank you for the review and for a positive assessment of our paper! To address your concerns, we have made the comparison with other PRM models from prior works more explicit, and have added a new experiment where we compare RL training with PAVs and RL training with PRMs (and not just ORMs) from prior works. We have also added another new experiment on using discounted rewards from strong provers, and provide answers to other questions below. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score.

Comparing with baselines that use PRMs proposed in prior works

PRM baselines for test-time search For test-time search, our results in Fig 5a of our submission already include both baselines: 1) beam search; and 2) best-of-N search using PRMs from prior works. We make this more explicit with a new plot (Figure 13) we add in Appendix C. We find that beam search with PAVs is 1.5x-5x more compute efficient than both beam search with PRMs from prior works, and best-of-N search with PRMs from prior works. We do not attempt to compare beam search with PAVs against weaker search algorithms used in conjunction with prior PRMs. When evaluating prior work (Snell et. al. (2024)) that uses $Q^\pi$ as their proposed PRM, we run beam search where the states in the beam are ranked only using $Q^\pi$, as opposed to the effective reward $Q^\pi + \alpha A^\mu$ in our case. Thus, the search procedure in the prior work and ours is identical, and we only change the re-ranking mechanism.

Other works on PRMs (Wang et. al. (2024), Luo et. al. (2024)) that also propose to use $Q^\pi$, use PRMs for best-of-N search, where they only rank the full sequences using the trained PRM (by taking the minimum over the $Q^\pi$ values at each step in the generation sampled from the base policy $\pi$). For completeness, we also compare PAVs with these approaches (PRMs-as-ORMs), which performs similarly to using PAVs-as-ORMs in Figure 5a.

PRM baselines for online RL training We add a new experiment where we use the PRMs proposed in prior works on test-time search, (Wang et. al. (2024), Luo et. al. (2024), Snell et. al. (2024)) as step-level rewards for training a base policy with RL. Here, the PRMs are trained offline and fixed, and then used to assign scores to intermediate steps in trajectories sampled during online RL. For this, we add a new experimental result where we use $Q^\pi_{\mathrm{base}}$ as the step-level score during RL training initialized with the base policy $\pi_{\mathrm{base}}$ since $Q^\pi_{\mathrm{base}}$ is the PRM proposed in prior works. This step-level reward is used to reward trajectories sampled during online RL training, in addition to the outcome reward, similar to our effective reward in Equation 5. We find that the test accuracy drops quite quickly, even though the train rewards keep improving, since the policy just learns to hack the step-level Q-values, for the reasons we explain in L240-245. We see a similar collapse when we use $Q^\mu$ instead of $A^\mu$ (see Appendix G for qualitative examples). On the other hand, if we were to use the Q-values or advantages from the current policy iterate $\pi_t$ as the step-level reward, then that is equivalent to only optimizing the outcome reward, and the only benefit of using the step rewards would be to reduce variance during online RL iterations. Thus, for online RL as well, our proposed step-level rewards (PAVs) which uses advantages of the prover policy $A^\mu$ outperforms baselines that plugin PRMs proposed in prior works on test-time search, We have added this new experiment on online RL where step-level rewards are given by PRM $Q^\pi_{\mathrm{base}}$ proposed in prior works to Appendix E, Figure 17.

Thank you for your response. To address the questions above, we add a new experiment where we aim to explicitly evaluate the exploration capabilities of PAV. Note that our goal is not to claim that PAVs are the best approach to exploration, but instead to compare PAVs to an existing exploration method to understand if building on our approach of connecting process rewards with exploration and potential functions, and designing novel forms of process rewards could be a fruitful endeavor for future research on exploration in LLMs.

Specifically, we compare test-time beam search guided by process supervision from PAVs with the importance weighted search approach outlined in AlphaLLM [1], an approach that runs MCTS for search. We use the heuristic from AlphaLLM and our preliminary results show that PAVs are 8x more compute efficient at test-time beam search, as outlined below. We also respond to your other questions below. Please let us know if this addresses your remaining concerns, and if so, we would be grateful if you are willing to raise your score.

Are PAVs comparable with other strategies for exploration?

Background on Alpha LLM: For the importance weighted search approach of AlphaLLM, we implement beam search in the following way. At any given point, the beam consists of $N$ states $s_1, s_2, s_3, \ldots, s_N$. These are partially unrolled solutions from the base policy $\pi$, up until state $s_i$ (prefix). We then expand each node in the beam 3 times, by conditionally sampling from $\pi$ (conditioned on each of the states in the beam). We get $N \times 3$ states, which we rank with the following scoring function and then select the top $N$ state. Each of the new expanded states is of the form $s, a$, where $s$ is the previous state in the beam and $a \sim \pi(\cdot \mid s)$ is the new sampled action (step). Following, Section 4.3 from Ye Tian et. al. (2024), the score for the new state $(s, a)$ is $Q^\pi(s, a) + C \cdot U(s)$ where $U(s)$ is the uncertainty bonus for the state $s$, which is computed as $U(s) = \sqrt{ \frac{n(s)}{\sum_{i=1}^{N} n(s_i)}}$. We use $C=0.25$, which we identified by tuning performance over a hold out validation set we use for PAVs as well. This resembles UCB or UCT-style exploration. Concretely, $n(s)$ is the effective number of children for node $s$ (Section 4.3.2 in Ye Tian et. al. (2024)). The term $n(s) = C^\prime \cdot I(s)$ is computed by linearly scaling the importance $I(s)$ defined as $I(s) = \max_a |V^\pi(s) - Q^\pi(s,a)|$, where $a$ is one of the $3$ actions sampled from state $s$ when expanding the beam. We tune and set $C^\prime=2.0$. Intuitively, AlphaLLM chooses to explore states that can change the $Q$-values by a lot, when we continue to sample from them. When the $Q$-values deviate by a lot, the $I(s)$ term increases. Consequently, so does $N(s)$ (effective children count) and $U(s)$ (uncertainty bonus).

In Appendix K (Figure 18), we show our results for the experiment of test-time beam search with PAVs, vs. the UCB style metric for exploration in Alpha LLMs. Since we only had two days to implement and run experiments for AlphaLLM (with no code base available), our findings are preliminary. Nevertheless, we find that PAVs are 8x more compute efficient than AlphaLLM at test-time exploration for the discovery of the correct solution trace. This is likely because the exploration metric in AlphaLLM uses the absolute magnitude of the advantage under the base policy, vs. PAVs which use the signed advantage under the prover policy. Thus, our exploration metric prefers steps that increase the likelihood of a complementary prover to discover the correct solution, as opposed to preferring steps that simply change the previous state’s value function (under the base policy) by the largest magnitude (which can also be negative).

[1] Tian, Y., Peng, B., Song, L., Jin, L., Yu, D., Mi, H., & Yu, D. (2024). Toward Self-Improvement of LLMs via Imagination, Searching, and Criticizing. arXiv preprint arXiv:2404.12253.

Thank you for the review and for a positive assessment of our work! To address your concerns, we make several clarifications in the paper regarding why PAVs improve exploration; we add new experiments with discounted process rewards from strong provers; and clarify the relation to potential functions. We also respond to your other questions and concerns below, and have modified the paper to include discussion on potential functions. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score. We would be very happy to discuss further.

Are PAVs used for test-time or training time improvement?

Once we train PAVs, we use them in two different ways: 1) to rank intermediate generations during test-time beam search, where we see that PAVs improve the compute-efficiency of test-time search (Section 4); and 2) as step-level process rewards when training the base policies with online RL, where we find that PAVs improve the sample efficiency of online RL (Section 5). We have made this more clear in the section headings to avoid misunderstanding as to which section refers to train usage vs test usage.

Why does the proposed effective reward scheme improve exploration?

This is a great question! We view ''exploration'' as a scheme that enables the discovery of high-outcome reward solutions. Under this definition, exploration bonuses are not required strictly to be count-based quantities like entropy or pseudocounts, but rather in our case they correspond to advantages under some prover policy. To empirically show that PAVs improve exploration (i.e., speed up discovery of high-reward solutions), in Figure 6, we provide empirical evidence for why our process rewards (PAVs) enable exploration during online RL training, with discussion in L421-448 of our submission. Specifically, we define the set of hard questions as those that remain unsolved even after running a best-of-256 search over 256 independent samples from the SFT model, and using the most accurate outcome verifier (i.e., we check if the final answer matches). We then check how many of these are solved by the policy trained with PAVs, if we run the best-of-N search using N<256 samples. Compared to a policy trained with ORMs, the PAV trained policy is able to solve 4-5x more fraction of the hard questions.

Essentially, the additive $A^\mu$ term in the effective reward serves as a bonus to increase the likelihood of steps that induce ''progress'' under the prover policy,. i.e., improve the chances of the prover policy to discover the correct answer. Thus, even though a trajectory sampled from the base policy fails to reach the correct answer, some steps in the trajectory still end up being up-weighted, which would not be the case without the additive term $A^\mu$ (when we only use outcome rewards). This results in improved coverage over solutions at test-time and aids the policy in finding a good sequence without committing to a myopic set of prefixes early on in beam search or RL. Corroborating this intuition is the better best-of-N performance for the policy trained with our process rewards (PAVs), compared to ORMs (Figure 6). For test-time beam search, advantages under the prover policy promote the coverage of solution traces in the beam which aligns better with the canonical definition of exploration; please see the discussion in L177-187 of our submission which explains this with an illustration (Figure 2). Hence, we say that PAVs improve exploration – if you think a different terminology will be better here, please do let us know and we would be happy to make a change in the paper.

I thank the authors for their detailed response to the review. I have increased the presentation score.

Some further comments:

On the question of exploration

I don't think their definition of "exploration" is the one commonly used. This is an issue of intuition: their definition is essentially the general RL goal (finding high outcome reward solutions), so saying that "their method works well because it explores well" is essentially saying "it works well because it works well".

This distinction is important (only) in the explanation of the intuition behind the proposed method. For example, a question is: Why does having a better model act as a prover / reward helper better than simply training the prover model further? I think this also makes the discussion in Section 5, Result 3 a bit imprecise.

(The authors argue in line 346 that since a weak prover can also improve the base policy, so their method is not something similar to knowledge distillation. I can accept this argument.)

Based on the beginning of Section 3.1, I would say what the authors really do during beam-search is exploration in action space. More precisely, using advantages and not Q values results in searching for actions that are of high value, independently of the states.

This is a unique feature of this environment: the action space is really large, and states are used to a large degree as a sort of stepping stone for the next action. (This is in some ways the opposite of planning: instead of thinking ahead, try to come up with a thought that will help form the next thought.)

Minor comments:

• Takeaways in lines 452-453: The Best-of-K approach is highlighted here, but it is not mentioned with regards to exploration in the section of the main text.
• Takeaways in lines 251-252: "better explore-exploit tradeoff ... and online RL": in the preceding text, exploration is only mentioned in the context of beam search, not that of training (and I presume online RL refers to the training phase).
• Does this method improve exploration in general, or does it make the base policy explore in the directions the prover considers good?

On the questions of potential functions

Let me rephrase my original question. In line 158, the authors say that in Section 3.1, they evaluate two different choices of potential functions. This leads the reader to believe that, since the two choices in Section 3.1 are $Q$ vs $A$, they mean that their reward scheme is analogous to potential functions. (See also lines 196-197: "we view process rewards as potential functions".)

My understanding, however, is that $Q$ is a potential function, and $A$ is (or can be) a reward scheme based on a potential function ($Q$).

On the question of discounting

What I understand from their reasoning is that with strong provers, they usually throw the base policy partial solution out of the window as a first step, regardless of discounting. I can accept this.

On the question of "example" in L197

Minor comment: "based on the example in Figure 2a" (in is missing)

On the question of confidence intervals

Thanks for the clarification. It seems as though that no multiple training seeds were used for the figures.

Thank you for responding to us! We have updated the submission to address the minor concerns (and are happy to adjust any wording if there are remaining concerns even after the paper update deadline passes). We clarify the latest comments below; for the remainder of the points that you accept in the response above – they sound great to us! Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score.

Clarification on the definition of exploration

You are right that there are several definitions [1, 2, 3, 4] of exploration in RL. For instance, in more theoretical work on bandits [3, 4], exploration is defined as the algorithm that helps discovery of optimal actions that maximize a given reward function, and its efficacy is measured in terms of cumulative regret or sample efficiency of learning; in unsupervised RL [1], exploration is typically defined as maximizing coverage or diversity of states since a target reward function is not known; in meta RL [2], exploration often refers to identifying the underlying MDP so that optimal policies can then be found.

In this paper, we used the term “exploration” to refer to the first definition – a procedure that helps discover the optimal action as quickly as possible. With this in mind, we justify the phrase “PAVs enable better exploration” by showing that PAVs improve sample efficiency (i.e., reduces cumulative regret) for RL, and enables finding optimal actions via test-time search. PAV advantages are akin to exploration bonuses that one could add on top of a standard greedy search that is run with respect to the policy’s Q-function. We are happy to use a different term for exploration if you would suggest, but we hope that the above explanation clarifies the meaning we referred to, for avoiding any misunderstanding.

[1] Jin, C., Krishnamurthy, A., Simchowitz, M., & Yu, T. (2020, November). Reward-free exploration for reinforcement learning. In International Conference on Machine Learning (pp. 4870-4879). PMLR.

[2] Gupta, A., Mendonca, R., Liu, Y., Abbeel, P., & Levine, S. (2018). Meta-reinforcement learning of structured exploration strategies. Advances in neural information processing systems, 31.

[3] Auer, P. "Finite-time Analysis of the Multiarmed Bandit Problem." (2002).

[4] Abbasi-Yadkori, Y., Pál, D., & Szepesvári, C. (2011). Improved algorithms for linear stochastic bandits. Advances in neural information processing systems, 24.

Why does having a better model act as a prover / reward helper better than simply training the prover model further?

While it is correct that if the prover policy is really good, we could have it serve as the base policy and improve it further via some sort of RL (either via outcome reward maximization or by using some process reward model), our conceptual insights and practical results show that improving a strong policy further would benefit from PAV rewards in turn computed using a new weak prover policy. Specifically, on one hand, we see that using Gemma 9B as the prover policy for test-time beam search over samples from the Gemma 2B base policy worked best (Figure 5c). At the same time, note that to train the stronger Gemma 9B policy with dense rewards during online RL, we used the weaker Gemma 2B prover policy (Figure 7c, Figure 19) to get good results – implying that weak policies can still serve as good provers for improving strong policies further. Does this address the question?

Based on the beginning of Section 3.1, I would say what the authors really do during beam-search is exploration in action space. More precisely, using advantages and not Q values results in searching for actions that are of high value, independently of the states. This is a unique feature of this environment.

Yes, this is correct! We do not claim that our approach is a general method for exploration, but that using advantages as exploration bonuses helps ``explore’’ actions that are of high value under the prover policy, and as you correctly note, independent of the previous states. We empirically find that this form of exploration is helpful for LLM math reasoning problems but some other forms of process rewards / exploration bonuses might work for other settings.

My understanding, however, is that Q is a potential function, and A is (or can be) a reward scheme based on a potential function (Q).

Thanks for pointing this out! Our language here was indeed a bit confusing: we have now edited the paper to indicate that the potential function is $Q^\mu$, and that $A^\mu$ function we study is the effective reward shaping term (i.e., $\Psi(s_{h+1}) - \Psi(s_h)$ in Ng et al. 1999), where $\Psi(s_h) = Q^\mu(s_{h-1}, a_h)$ is the potential function. We have also clarified that we do not evaluate two potential functions, but compare two types of process rewards / reward shaping terms. We have made this clear in Section 3 and in Appendix I where we discuss the connection to potential-based reward-shaping terms in detail.

Does this method improve exploration in general, or does it make the base policy explore in the directions the prover considers good?

Note that our goal is not to claim that PAVs are the best approach for exploration, but instead to compare PAVs to an existing exploration method and understand if building on our approach of connecting process rewards with exploration and potential functions, and designing novel forms of process rewards could be a fruitful endeavor for future research on exploration in LLMs.

At the same time, to answer your concern, we would like to point to a new experiment we run to address Reviewer Ys2B’s comments. Here, we compare test-time beam search guided by process supervision from PAVs with the importance weighted search approach outlined in AlphaLLM [1], an approach that runs MCTS for search. We use the heuristic from AlphaLLM and our preliminary results show that PAVs are 8x more compute efficient at test-time beam search.

Background on Alpha LLM: For the importance weighted search approach of AlphaLLM, we implement beam search in the following way. At any given point, the beam consists of $N$ states $s_1, s_2, s_3, \ldots, s_N$. These are partially unrolled solutions from the base policy $\pi$, up until state $s_i$ (prefix). We then expand each node in the beam 3 times, by conditionally sampling from $\pi$ (conditioned on each of the states in the beam). We get $N \times 3$ states, which we rank with the following scoring function and then select the top $N$ state. Each of the new expanded states is of the form $s, a$, where $s$ is the previous state in the beam and $a \sim \pi(\cdot \mid s)$ is the new sampled action (step). Following, Section 4.3 from Ye Tian et. al. (2024), the score for the new state $(s, a)$ is $Q^\pi(s, a) + C \cdot U(s)$ where $U(s)$ is the uncertainty bonus for the state $s$, which is computed as $U(s) = \sqrt{ \frac{n(s)}{\sum_{i=1}^{N} n(s_i)}}$. We use $C=0.25$, which we identified by tuning performance over a hold out validation set we use for PAVs as well. This resembles UCB or UCT-style exploration. Concretely, $n(s)$ is the effective number of children for node $s$ (Section 4.3.2 in Ye Tian et. al. (2024)). The term $n(s) = C^\prime \cdot I(s)$ is computed by linearly scaling the importance $I(s)$ defined as $I(s) = \max_a |V^\pi(s) - Q^\pi(s,a)|$, where $a$ is one of the $3$ actions sampled from state $s$ when expanding the beam. We tune and set $C^\prime=2.0$. Intuitively, AlphaLLM chooses to explore states that can change the $Q$-values by a lot, when we continue to sample from them. When the $Q$-values deviate by a lot, the $I(s)$ term increases. Consequently, so does $N(s)$ (effective children count) and $U(s)$ (uncertainty bonus).

In Appendix K (Figure 18), we show our results for the experiment of test-time beam search with PAVs, vs. the UCB style metric for exploration in Alpha LLMs. Since we only had two days to implement and run experiments for AlphaLLM (with no code base available), our findings are preliminary. Nevertheless, we find that PAVs are 8x more compute efficient than AlphaLLM at test-time exploration for the discovery of the correct solution trace. This is likely because the exploration metric in AlphaLLM uses the absolute magnitude of the advantage under the base policy, vs. PAVs which use the signed advantage under the prover policy. Thus, our exploration metric prefers steps that increase the likelihood of a complementary prover to discover the correct solution, as opposed to preferring steps that simply change the previous state’s value function (under the base policy) by the largest magnitude (which can also be negative).

Thank you for the review and for a positive assessment of our paper! To address your concerns, we include experiments for RL training with PRM baselines that we used for test-time search, and make the comparison with prior works on automated PRMs (Wang et. al. (2024), Luo et. al. (2024), Snell et. al. (2024)) for test-time search more explicit. We also provide qualitative examples for process rewards from good vs. bad prover policies, and respond to your other questions and concerns below. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score.

Discussion on “Why do we use a separate prover policy?”.

As you correctly note, the fact that we use a separate prover policy is indeed a key point of distinction with prior work Shao et. al. (2024). In L233-249 of our submission (Section 3.2), we explain in detail why we need a prover policy $\mu$, different from the base policy $\pi$ we optimize. In a nutshell, there are two main reasons for this choice: 1) as we note in L244 of our submission, training with process rewards where $\pi = \mu$ during RL mathematically leads to gradients that are equivalent to those observed when purely optimizing the outcome reward ($\ell_{ORM-RL}$); and 2) when the outcome rewards under a poor base policy are very sparse, $Q^\pi \approx 0$, on most states, and consequently $A^\pi \approx 0$. Thus, we need a separate prover that can distinguish steps generated by even a poor base policy. In particular, we find “complementary” provers that both distinguish actions taken by the base policy, and are not too misaligned with it, lead to largest improvements in the base policy. We show this theoretically in Theorem 3.1. We have updated the discussion in this section to more clearly signpost this argument, early on in the paper.

Introduction of separate prover policies in Section 3.2.

Thanks for the question! We have updated the paper to now clarify that the goal of Section 3.1 is primarily to motivate per-step advantages over Q-functions, with the search example in mind. The very same arguments in this section also apply to a different prover policy, but we stuck with the same policy in Section 3.1 due to the clarity of explaining one concept (Q-values vs advantages) before the next idea of separate prover policies comes through.

Flow of Section 3: In the beginning of Section 3, we define process rewards as potential-shaped step-level rewards that, when optimized during RL or test-time search, should yield better performance as measured by the outcome rewards. With this intuition, we explored a number of choices for the design of process rewards. We begin with Q-values, and explain with an illustrative example in Figure 2, why Q-value serves as a poor choice of potential functions, compared to advantages. Now, given this choice, the next question is "what policy should advantage be measured under?". To answer this question, we come back to the notion of potential functions and demonstrate a certain set of prover policies $\mu$, distinct from the base policy $\pi$ is better for this use case. In Section 3.2, to expand the set of possible potential functions that conform to the definition of potential functions in Ng et. al. (1999), we allow the advantages to be computed under a broader choice of prover policies. Following this section, in the rest of the submission, we also allow the prover policy $\mu$ to be different from the base policy $\pi$.

If there is a particular part of the submission where this choice is particularly unclear, please let us know and we would be happy to edit the submission to make this more clear.

I thank the authors for their detailed response to my review, and for addressing or clarifying most of my concerns satisfactorily. I have raised my overall score as well as presentation score in response.

Thank you for the review and for a positive assessment of our work! To address your concerns, we have added discussion on the computational overhead of PAVs, compared to ORMs in Appendix H. Please let us know if your concerns are addressed, and if so, we would be grateful if you are willing to raise your score. We are very happy to discuss further!

Computational cost of training and using PAVs vs. ORMs

Thank you for the question. We have added a detailed discussion on comparison of computational cost of training and using PAVs to Appendix H of the paper. In summary, while training PAVs takes more compute than ORMs, when we use PAVs for test-time search or during online RL training they are more compute-efficient than ORMs for attaining the same level of test performance. We detail our calculations below.

Training PAVs: The training data for PAVs is collected by running rollouts to estimate the Q-values for each step. Since learning to predict expected future success at earlier steps is a statistically more challenging task than predicting the final score at the end of the complete sequence (as in ORM) , the training data for PAVs is larger and scales with the average number of steps in the LLM’s output ($\approx 10$ in our case), in order to be able to achieve the same level of prediction accuracy on all steps in a generation. With a 10x larger dataset, naturally the computational cost of training PAVs (compared to ORMs which only predict final outcomes) also scales by roughly the same factor.

Using PAVs: Despite a larger cost for training PAVs, once trained, PAVs are much more compute-efficient than ORMs on test questions. In other words, while we do incur a larger training cost, this cost is amortized over rounds of deployment. Concretely, as we show in Section 4.1, for test-time search, PAVs are 1.5-5x more compute efficient than ORMs, at achieving the same level of performance (Figure 4 in submission). This means that if we were to use a verifier at least twice upon deployment, PAV is already more compute-efficient than an ORM, accounting for both the training and test-time compute cost of PAVs.

For online RL, as we show in Section 5 (Figure 8c), PAVs achieve the same level performance in 1/6th of the RL training iterations that ORMs require. But in our implementation, for each iteration, PAVs score each step by feeding each prefix in a generation separately through the trained model, and ORMs only score the full trajectory. Thus, in this implementation PAVs consume 10x more compute per batch to score the generations in a batch (since we have 10 steps on average per generation). A more efficient implementation would score all prefixes in a generation in a single forward pass through the trained PAV. Nevertheless, the reduction in RL iteration complexity with PAVs is big enough that even with our naïve implementation, the overall computational cost is less, compared to ORMs. For example, using the formula for training and inference FLOPs in [1], during online RL, to train the Gemma 2B base policy with PAVs, we need to spend about $2.5 \times 10^{18}$ FLOPs, but to achieve the same performance with ORMs, we need about $5.9 \times 10^{18}$ FLOPs, resulting in at least half the total computational FLOPs.

[1] Hoffmann, Jordan, et al. "Training compute-optimal large language models." arXiv preprint arXiv:2203.15556 (2022).

Expanding PAVs to other tasks and model families.

We chose math reasoning domains in order to be able to perform a direct comparison with prior works (Snell et. al. (2024), Wang et. al. (2024), Lightman et. al. (2023), Shao et. al. (2024), Cobbe et. al. (2021)) that study process and outcome reward models. Within this domain, most of these prior works study the harder MATH benchmark (Hendrycks et. al. (2021)), since performance on some other reasoning datasets like GSM8K is already saturated. For example, the performance of some of the base LLMs we consider (Gemma2-9B and Gemma2-27B) is itself $>85$% on GSM8K. Please note that the scope of tasks in our paper is comparable, if not larger than several of these prior works.

At the same time, the conceptual framework we present for process reward models, and our approach PAV, is broad and not specific to math reasoning. It only requires access to a base policy and an accurate outcome reward model (this is the regular expression matcher $\mathrm{Rex}$ in our case) that can be queried without much cost on all generations sampled from the base policy, at least on a fixed set of input prompts . While extending our results to other settings like coding is definitely possible, we defer this direction of study to future work since it is unclear how to define a ''step'' for a code output. Additionally, while prior work [2] studies an initial design of steps, the codecontests [3] dataset is challenging for our models. Thus to draw meaningful conclusions on coding, we will have to choose a different base model capable enough to generate code in a format with a natural partition of steps.

On the MATH benchmark, we are trying to add results where we use PAVs to improve the performance of models from other families (like Mistral/Llama), but are unsure if it will complete during the span of the rebuttal period, since we need to train base models (which requires a different infrastructure), collect data to train PAVs, train the PAVs, and then use it for search/RL.

[2] Zheng, Qinkai, et al. "Codegeex: A pre-trained model for code generation with multilingual benchmarking on humaneval-x." Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2023.

[3] Li, Yujia, et al. "Competition-level code generation with alphacode." Science 378.6624 (2022): 1092-1097.