Process Reward Models › PRM Architecture and Training

A process reward model is a neural network that takes as input a problem and a partial or complete reasoning chain and produces a scalar score for the most recent step. The architecture is typically a transformer-based language model fine-tuned to perform step-level evaluation. Understanding both the architectural choices and the training procedure is essential for appreciating both the strengths and the limitations of current PRMs.

The input to a PRM consists of the original problem statement, all preceding steps in the reasoning chain, and the step being evaluated. The model must use context from earlier steps to assess whether the current step is valid. For example, if a step claims "therefore $x = 5$," the PRM needs to know what was established in prior steps to judge whether that claim follows.

A common input format represents each step as a block of text, with steps separated by a special delimiter token. The PRM produces a prediction at each delimiter position, yielding one score per step. This framing turns the PRM into a token-level classification problem: at each step boundary, predict whether the step is correct.

The use of a full language model as the backbone is important for two reasons. First, step correctness is inherently a language understanding problem: you need to understand the semantic content of each step and compare it to the problem context and prior steps. Second, the scoring model must generalize across many problem types, mathematical domains, and notation conventions. A transformer trained on mathematical text develops the representational capacity to handle this diversity.

In the PRM800K setup from Lightman et al., the backbone was a GPT-4 class model fine-tuned specifically for step evaluation. The input was formatted as a sequence of problem text followed by steps, with a special token inserted between steps. The model predicted a score at each step-boundary token, effectively running a per-step binary classifier on top of the language model's contextualized representations.

One design choice is whether to score each step in isolation (re-running the model from scratch for each step) or to score all steps in a single forward pass. The single-pass approach is more efficient: one inference call processes the entire chain and outputs one score per step boundary. The isolation approach might seem to provide cleaner credit assignment, but in practice the contextual dependencies between steps matter for correct evaluation, so single-pass scoring is both more efficient and more accurate.

Training a PRM requires step-level labels: for each step in each solution, a binary or continuous label indicating whether the step is correct. Obtaining these labels is the core challenge of PRM training.

There are two main approaches to generating step-level labels:

Human annotation. Human raters read each step in a reasoning chain and judge whether it is mathematically or logically valid. This produces high-quality labels but is expensive and difficult to scale. The Process Reward Model dataset from Lightman et al. (OpenAI, 2023) used this approach, collecting step-level labels from human mathematicians for solutions to MATH benchmark problems. Human annotators were asked to label each step as positive (correct), negative (incorrect), or neutral (not determinative of the final answer's correctness). The resulting dataset, PRM800K, contains approximately 800,000 step-level annotations across 75,000 solutions. The quality of human annotation depends heavily on annotator expertise. Mathematical step correctness often requires domain knowledge that cannot be crowd-sourced to general-purpose annotators.

Monte Carlo estimation. For each prefix of a reasoning chain (steps 1 through $k$), generate multiple continuations using the base model and check whether any continuation reaches the correct final answer. The proportion of continuations that succeed is used as the label for step $k$. A step receives a high score if the model can frequently recover from that prefix and reach the correct answer; it receives a low score if the model rarely succeeds from that point onward.

More precisely, given a problem $p$ and a partial solution $(s_1, \ldots, s_k)$, we generate $M$ completions $(s_{k+1}^{(i)}, \ldots, s_T^{(i)})$ for $i = 1, \ldots, M$ and compute:

$$\hat{y}_k = \frac{1}{M} \sum_{i=1}^{M} \mathbf{1}[a^{(i)} = a^*]$$

where $a^{(i)}$ is the final answer reached by the $i$-th completion and $a^*$ is the correct answer. This estimated success rate $\hat{y}_k$ becomes the training label for step $k$. The intuition is that a step is "good" if the model can build on it to reach the right answer, and "bad" if the model consistently fails from that prefix.

The Monte Carlo approach allows automated label generation without human annotation, at the cost of computational overhead (many continuations must be generated per prefix) and potential noise (the success rate of a prefix is not the same as the correctness of the most recent step, but it correlates strongly with it). A subtle issue is that Monte Carlo labels measure the capability of the current base model to recover from a prefix, not the intrinsic correctness of the steps in that prefix. If the base model is unusually good at recovering from certain types of errors, those errors receive artificially high scores.

Given step-level labels, the PRM is trained with a classification or regression loss. For binary step correctness labels, a cross-entropy loss is standard:

$$\mathcal{L} = -\sum_{t=1}^{T} \left[ y_t \log p_t + (1 - y_t) \log (1 - p_t) \right]$$

where:

• $T$ is the number of steps in the reasoning chain
• $y_t \in \{0, 1\}$ is the binary label for step $t$ (1 = correct, 0 = incorrect)
• $p_t \in (0, 1)$ is the PRM's predicted probability that step $t$ is correct

The loss sums over all steps in the chain. Each step contributes independently to the gradient, which means the PRM learns to evaluate each step based on its local context without being dominated by the outcome at the end.

For continuous labels (such as Monte Carlo success rates), a mean squared error loss or binary cross-entropy with soft targets is used instead. In the soft-target case, the label $y_t$ is replaced by the Monte Carlo estimate $\hat{y}_k$, and the same cross-entropy formula applies. This treats the Monte Carlo estimate as a probability rather than a hard binary label, which is appropriate given the stochastic nature of the estimation procedure.

One important consideration is class weighting. In most reasoning datasets, correct steps are far more common than incorrect steps. A naive training procedure that minimizes unweighted cross-entropy on a 90-10 correct-incorrect split will converge to a model that predicts "correct" for every step, achieving 90% accuracy while being entirely useless as an evaluator. Proper class weighting or focal loss formulations are necessary to ensure the model learns to identify incorrect steps.

The distinction between outcome supervision and process supervision in training can be formalized. Let $s = (s_1, s_2, \ldots, s_T)$ be a solution with $T$ steps, and let $a$ be the final answer.

With outcome supervision, the reward is:

$$R_{\text{outcome}}(s, a) = \mathbf{1}[a = a^*]$$

where $a^*$ is the correct answer. The entire solution receives a single binary reward.

With process supervision, the reward at step $t$ is:

$$R_{\text{process}}(s, t) = \text{PRM}(p, s_1, \ldots, s_t)$$

where $p$ is the problem and $\text{PRM}(\cdot)$ is the process reward model's score for step $t$ given the problem and all preceding steps.

The key difference is that process supervision provides a vector of rewards $(R_1, R_2, \ldots, R_T)$ rather than a single scalar. This vector can be used in multiple ways: to directly train a policy via reinforcement learning, to rank candidate solutions step by step, or to select the best reasoning paths in a search algorithm.

When used to train a policy via reinforcement learning, process supervision typically applies a per-step policy gradient update. The policy gradient theorem gives:

$$\nabla_\theta J(\theta) = \mathbb{E}_\tau \left[ \sum_{t=1}^{T} \nabla_\theta \log \pi_\theta(s_t \mid p, s_1, \ldots, s_{t-1}) \cdot R_t \right]$$

where $\pi_\theta$ is the policy (the language model being trained), $s_t$ is the $t$-th reasoning step, and $R_t$ is the process reward at step $t$. Compared to the outcome reward case where $R_t$ would be constant across all steps and equal to the final reward, the process reward provides a differentiated signal: steps with high process reward receive positive gradient updates, while steps with low process reward receive negative ones.

One important consideration in PRM training is the relationship between the PRM and the policy model whose outputs it evaluates. In most setups, these are separate models: the policy generates candidate solutions, and the PRM evaluates their steps. But the quality of the PRM depends critically on what distribution of solutions it was trained on.

A PRM trained on solutions generated by a weak base model will see primarily simple errors and simple recoveries. When applied to solutions generated by a stronger model, it encounters different error patterns and may not generalize well. Conversely, a PRM trained on solutions from a very strong model may see too few errors to learn what incorrect steps look like.

This suggests that PRM training should be iterative: train a PRM on solutions from the current policy, use the PRM to improve the policy, then retrain the PRM on solutions from the improved policy, and so on. This bootstrap approach mirrors the iterative nature of RLHF fine-tuning more generally. In practice, maintaining this cycle is computationally expensive, so most published work uses a fixed PRM trained on a fixed dataset.