Token-Supervised Value Models for Enhancing Mathematical Reasoning Capabilities of Large Language Models

The mathematical reasoning capabilities of LLMs can be enhanced by employing reward models as external verifiers to assess the generated reasoning paths (Cobbe et al., 2021; Uesato et al., 2022; Lightman et al., 2023; Yu et al., 2024; Wang et al., 2024). The verifier is generally trained via supervised learning on a dataset obtained by sampling multiple reasoning paths per training problem using an LLM. Specifically, given a training problem $q_{tr}$ as an input, the LLM first generates $N_{tr}$ reasoning paths, where $n$-th reasoning path is comprised of reasoning steps $\{s_{n,j}\}_{j=1}^{S_{n}}$ and a final answer $a_{n}$ for $n=1,\cdots,N_{tr}$. In token-level notation, the $n$-th reasoning path can also be expressed as a sequence of tokens $\{t_{n,k}\}^{T_{n}}_{k=1}$. Hereafter, $\{s_{n,\cdot}\}_{1}^{j}$ and $\{t_{n,\cdot}\}_{1}^{k}$ means $\{s_{n,1},\cdots,s_{n,j}\}$ and $\{t_{n,1},\cdots,t_{n,k}\}$, respectively. The final answer $a_{n}$ is correct if it is equal to the ground truth answer $\hat{a}$, and incorrect otherwise. Based on the correctness of the sampled reasoning paths, supervision is traditionally given in two ways: (i) outcome supervision (Cobbe et al., 2021; Uesato et al., 2022; Yu et al., 2024) and (ii) process supervision (Uesato et al., 2022; Lightman et al., 2023; Wang et al., 2024).

Prior works (Cobbe et al., 2021; Uesato et al., 2022; Yu et al., 2024) employ outcome supervision to label an entire reasoning path as correct if its final answer is correct (Fig. 1). The outcome reward function $r_{o}(\cdot)$ is the correctness of the final answer:

for $n=1,\cdots,N_{tr}$. An outcome-supervised reward model (ORM) $f_{ORM}$ is trained with every token in a reasoning path labeled as the outcome reward (Eq. 3). The ORM loss $\mathcal{L}_{ORM}$ is defined as

The mean squared error is typically used as a loss function $\ell(\cdot)$ in Eq. 4. Cobbe et al. (2021) demonstrated that a token-level verifier trained to judge the correctness after every token performs better than a solution-level verifier trained to determine the correctness only after the final token.

Interestingly, Yu et al. (2024) showed that ORMs can be alternatively described as modeling the cumulative reward for each token, where all intermediate rewards are zero (i.e., $r(t_{n,k})=0$ for every $n$ and $k$) and the discount factor $\gamma$ is set to 1. The cumulative reward following an intermediate token $t_{n,k}$, $R(t_{n,k})$ is calculated as

which is equivalent to $r_{o}(a_{n})$ in Eq. 3. This entails that an intermediate reasoning path is labeled as correct if the final answer is correct, and vice versa. In this sense, ORMs can indirectly and implicitly learn the potential correctness of an intermediate reasoning path (Yu et al., 2024).

Process supervision enables a more accurate assessment of a reasoning path by explicitly training a verifier on the correctness of each step with step-level supervision Lightman et al. (2023). The correctness of each reasoning step is either labeled via human annotation (Uesato et al., 2022; Lightman et al., 2023) or automation (Wang et al., 2024). Since acquiring human annotations is labor-intensive and costly, we mainly focus on process supervision without human annotations.

Nach Wang et al. (2024), an intermediate reasoning step $s_{n,j}$ can be labeled as correct if at least one of the reasoning paths starting from $s_{n,j}$ reaches the correct final answer $\hat{a}$ (Fig. 1). In practice, $s_{n,j}$ is annotated by sampling a fixed number of reasoning paths conditioned on a sequence of intermediate reasoning steps $\{s_{n,\cdot}\}_{1}^{j}=\{s_{n,1},\cdots,s_{n,j}\}$. If at least one of the sampled reasoning paths reaches the correct final answer, $s_{n,j}$ is labeled as correct with the process reward $r_{p}(s_{n,j})=1$. Otherwise, $s_{n,j}$ is labeled as incorrect and $r_{p}(s_{n,j})=0$. Using the per-step labels obtained through automation, a Process-supervised Reward Model (PRM) is trained to provide a step-level assessment by minimizing the following loss:

where $\ell$ denotes the binary cross entropy loss.

As mentioned in Sec. 2, both outcome supervision (ORMs) and process supervision (PRMs) utilize homogeneous labels determined by the correctness of either the entire reasoning path or step (Fig. 1). Consequently, we hypothesize that they are likely to be neither explicitly nor directly trained to (i) learn the discriminative details of tokens within a reasoning path or (ii) evaluate whether an intermediate reasoning path is on a promising track toward the correct final answer.

We elucidate our assertion through cases observed in practice, as illustrated in Fig. 2. In the reasoning path ranked highest by ORM ((a) in Fig. 2), reasoning step 4 begins with ‘‘So the total difference’’ but ends with a summation, where a logical error occurs. However, ORM is unable to catch the error and maintains a score over $0.4$, the highest score among $256$ candidate reasoning paths. In the reasoning path ranked highest by PRM ((b) in Fig. 2), reasoning step 5 starts with ‘‘Thus, the total difference’’ but ends in subtracting a larger number (‘‘120’’) from a smaller number (‘‘95’’), which is the exact opposite of the definition of difference. In the reasoning path, ‘‘120’’ appears in reasoning step 4 after ‘‘95’’ appears in reasoning step 3. Since PRMs focus on assessing the correctness of the current reasoning step, the sequential appearance of numbers and the resulting subtraction are considered correct by the PRM even though the reasoning path is unlikely to lead to a correct answer. The observed failures inspire the proposal of a token-level value supervision strategy for training verifiers.

To overcome the issues above, we propose a new verifier based on token-level supervision with distinctive token-wise labels according to the potential of tokens in deducing the correct final answer. A natural choice to appropriately reflect the token-wise potential is prospective value modeling (Sutton and Barto, 2018), which is fine-grained and future-oriented compared to retrospective cumulative reward modeling (Eq. 5). Accordingly, we construct a supervision scheme for token $t_{n,k}$ in a reasoning path $\{t_{n,\cdot}\}_{1}^{k}=\{t_{n,1},\cdots,t_{n,k}\}$ with the expected cumulative reward (i.e., value):

where $r(\cdot)$ and $\gamma$ denote a reward function and the discount factor, respectively.

The primary challenge in training value models as verifiers is estimating the value labels of a generated reasoning path Yu et al. (2024). However, under the specific outcome reward formulation of Eq. 3 and no intermediate rewards, the expected cumulative reward (Eq. 7) reduces to the probability of reaching the correct final answer conditioned on the question $q_{tr}$ and intermediate reasoning path $\{t_{n,\cdot}\}_{1}^{k}$, which can be straightforwardly computed from generated reasoning paths and can indicate whether an intermediate reasoning path (i.e., $\{t_{n,\cdot}\}_{1}^{k}$) is on a promising track toward the correct final answer.

Following Proposition 3.1, we train the Token-supervised Value Model (TVM) by supervising each token with a value label empirically estimated as the probability of reaching the correct final answer given until that token. The objective of TVM is

for $n=1,\cdots,N_{tr}$ and $k=1,\cdots,T_{n}$, where $\mathbb{P}_{n,k}$ indicates the right-hand side of Eq. 8 and the loss function $\ell$ is the mean squared error.

Compared to existing verifiers, the resolution of assessment provided by the proposed token-level value supervision adequately matches the token-wise granularity of LLMs, thereby being able to capture the discriminative details of tokens within a reasoning path (Fig. 2). In contrast to ORMs, TVM is trained to directly estimate the probability of an intermediate reasoning path being on a promising track toward the correct final answer (Proposition 3.1). As a result, TVM can choose the reasoning path most likely to reach the correct final answer among candidate reasoning paths, whether they are partial or complete.

During inference, TVM can be employed to either search the reasoning path most likely to be correct over complete reasoning paths generated from an LLM (Lightman et al., 2023) or distinguish prospective candidates likely to reach the correct final answer among partially generated reasoning paths. For the latter, we conduct a detailed study in Sec. 4.3 in the setting of verifier-guided step-wise beam search (Yu et al., 2024).

As discussed in Sec. 3.2, Proposition 3.1 alleviates the practical challenges of value estimation (Eq. 7) by formulating the value as the ratio of correct reasoning paths to total reasoning paths. Following Eq. 7 and Eq. 8, the estimated value for each token $t_{n,k}$ can be represented as

In practice, Eq. 10 can be empirically estimated from $N_{tr}$ generated reasoning paths as the ratio of correct reasoning paths starting from $\{t_{n,\cdot}\}_{1}^{k}$ among $N_{tr}$ and total reasoning paths starting from $\{t_{n,\cdot}\}_{1}^{k}$ among $N_{tr}$, respectively. The value label of each token $V(t_{n,k})$ is assigned as

where $\mathbb{I}(\cdot)$ is the indicator function and $N_{tr}$ cancels out. The overall procedure of empirical value estimation is described in Figure 3. The overall algorithm is deferred to Appendix C.