Process Reward Models › Code Implementation
Language AI HandbookAll chapters Process Reward Models › Code ImplementationOne section · Sections in this chapter7 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
In this section, we implement a simplified process reward model. We use a small feedforward network as a proxy for the transformer backbone that would power a full-scale PRM, to illustrate the key components of the training pipeline.
Setup and Synthetic Data
We begin by generating synthetic step-level training data. In practice, you would use human annotations or Monte Carlo rollouts, but for this walkthrough we create a toy dataset to illustrate the mechanics.
import numpy as np
import torch
np.random.seed(42)
torch.manual_seed(42)
n_problems = 500
n_steps_per_problem = 5
embedding_dim = 64
# Synthetic embeddings for problems and steps
problem_embeddings = np.random.randn(n_problems, embedding_dim).astype(
np.float32
)
step_embeddings = np.random.randn(
n_problems, n_steps_per_problem, embedding_dim
).astype(np.float32)
# Correctness labels: positive activation in the first 10 dims raises error risk
correctness_labels = np.ones(
(n_problems, n_steps_per_problem), dtype=np.float32
)
for i in range(n_problems):
for j in range(n_steps_per_problem):
error_signal = step_embeddings[i, j, :10].mean()
p_correct = 1.0 / (1.0 + np.exp(5.0 * error_signal - 2.0))
correctness_labels[i, j] = float(np.random.random() < p_correct)Dataset: 500 problems, 5 steps each Total steps: 2500 Correct steps: 2006 (80.2%) Incorrect steps: 494 (19.8%)
The synthetic dataset reflects a realistic class imbalance: most reasoning steps are correct, and errors are relatively rare. The error signal is based on the mean activation in the first ten embedding dimensions, creating a learnable structure: steps with strongly positive activation in this subspace tend to be incorrect. This is a simplified proxy for the kinds of features a real PRM might learn, such as patterns in the language that correlate with mathematical errors.
Building the PRM Dataset
We structure the data into a PyTorch Dataset that concatenates the problem embedding with a running mean of all preceding steps. This reflects the autoregressive nature of step evaluation: when scoring step , the model sees the problem and all steps 1 through .
import torch
from torch.utils.data import DataLoader, Dataset
class StepRewardDataset(Dataset):
def __init__(self, problem_embs, step_embs, labels):
self.samples = []
n_probs, n_steps, emb_dim = step_embs.shape
for i in range(n_probs):
for j in range(n_steps):
if j == 0:
context = problem_embs[i].copy()
else:
prev_mean = step_embs[i, :j].mean(axis=0)
context = problem_embs[i] + prev_mean
current_step = step_embs[i, j]
label = labels[i, j]
self.samples.append(
(
torch.tensor(context),
torch.tensor(current_step),
torch.tensor(label),
)
)
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
return self.samples[idx]
train_size = int(0.8 * n_problems)
train_dataset = StepRewardDataset(
problem_embeddings[:train_size],
step_embeddings[:train_size],
correctness_labels[:train_size],
)
val_dataset = StepRewardDataset(
problem_embeddings[train_size:],
step_embeddings[train_size:],
correctness_labels[train_size:],
)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)Training samples: 2000 Validation samples: 500 Batches per epoch: 63
Each sample in the dataset consists of three tensors: the context vector encoding the problem and all prior steps, the current step vector, and the binary correctness label. By organizing data this way, we make the dependency structure explicit: step is never scored without knowledge of the problem and steps 1 through .
PRM Model Architecture
The PRM takes a concatenated context-plus-step vector and produces a single probability that the step is correct. In a full-scale implementation this would be a transformer computing attention over token sequences; here we use a feedforward network to demonstrate the training mechanics clearly.
import torch.nn as nn
class ProcessRewardModel(nn.Module):
def __init__(self, input_dim, hidden_dim=32):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim * 2, hidden_dim),
nn.LayerNorm(hidden_dim),
nn.GELU(),
nn.Dropout(0.1),
nn.Linear(hidden_dim, hidden_dim // 2),
nn.GELU(),
nn.Linear(hidden_dim // 2, 1),
nn.Sigmoid(),
)
def forward(self, context, step):
x = torch.cat([context, step], dim=-1)
return self.net(x).squeeze(-1)
model = ProcessRewardModel(input_dim=embedding_dim)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-3)
criterion = nn.BCELoss()
total_params = sum(p.numel() for p in model.parameters())Model parameters: 4,737 Input dimension: 128 (context + step concatenated)
The architecture uses Layer Normalization after the first linear layer, which stabilizes training by normalizing the hidden representations before the nonlinearity. The GELU activation function is used throughout, consistent with modern language model practice. The final sigmoid ensures the output is a probability in , suitable for binary cross-entropy training.
In a production PRM, this feedforward network would be replaced by a transformer encoder that attends to the full sequence of problem tokens and step tokens, building contextual representations that capture the semantic content of each step in relation to the full reasoning chain. The feedforward network here is a stand-in that demonstrates the training mechanics without requiring the computational resources of a full transformer training run.
Training Loop
def train_epoch(model, loader, optimizer, criterion):
model.train()
total_loss = 0.0
correct = 0
total = 0
for context, step, label in loader:
optimizer.zero_grad()
pred = model(context, step)
loss = criterion(pred, label)
loss.backward()
optimizer.step()
total_loss += loss.item() * len(label)
predicted = (pred >= 0.5).float()
correct += (predicted == label).sum().item()
total += len(label)
return total_loss / total, correct / total
def eval_epoch(model, loader, criterion):
model.eval()
total_loss = 0.0
correct = 0
total = 0
with torch.no_grad():
for context, step, label in loader:
pred = model(context, step)
loss = criterion(pred, label)
total_loss += loss.item() * len(label)
predicted = (pred >= 0.5).float()
correct += (predicted == label).sum().item()
total += len(label)
return total_loss / total, correct / total
n_epochs = 9
train_losses = []
val_losses = []
train_accs = []
val_accs = []
for epoch in range(n_epochs):
tr_loss, tr_acc = train_epoch(model, train_loader, optimizer, criterion)
va_loss, va_acc = eval_epoch(model, val_loader, criterion)
train_losses.append(tr_loss)
val_losses.append(va_loss)
train_accs.append(tr_acc)
val_accs.append(va_acc)Final training loss: 0.2575 | accuracy: 0.896 Final validation loss: 0.3954 | accuracy: 0.850
The training and validation accuracy converge closely, indicating the model has learned a generalizable mapping from step context to correctness probability. The small gap between training and validation loss shows that the model does not overfit the synthetic training set, which is expected given the relatively simple structure of the data: correctness correlates strongly with the mean activation of the first embedding dimensions, and the network learns this relationship quickly.
Scoring Candidate Solutions
We now show how to use the trained PRM to score candidate solutions using two aggregation strategies.
def score_solution_min(model, problem_emb, step_embs):
"""Score by minimum step score: one bad step sinks the whole solution."""
model.eval()
scores = []
with torch.no_grad():
for j in range(len(step_embs)):
if j == 0:
context = torch.tensor(problem_emb).unsqueeze(0)
else:
prev_mean = step_embs[:j].mean(axis=0)
context = torch.tensor(problem_emb + prev_mean).unsqueeze(0)
step = torch.tensor(step_embs[j]).unsqueeze(0)
scores.append(model(context, step).item())
return min(scores), scores
def score_solution_geomean(model, problem_emb, step_embs):
"""Score by geometric mean of all step scores."""
model.eval()
scores = []
with torch.no_grad():
for j in range(len(step_embs)):
if j == 0:
context = torch.tensor(problem_emb).unsqueeze(0)
else:
prev_mean = step_embs[:j].mean(axis=0)
context = torch.tensor(problem_emb + prev_mean).unsqueeze(0)
step = torch.tensor(step_embs[j]).unsqueeze(0)
scores.append(model(context, step).item())
log_scores = [np.log(max(s, 1e-9)) for s in scores]
geo_mean = np.exp(np.mean(log_scores))
return geo_mean, scores
test_problem = np.random.randn(embedding_dim).astype(np.float32)
n_candidates = 5
candidates = [
np.random.randn(n_steps_per_problem, embedding_dim).astype(np.float32)
for _ in range(n_candidates)
]
min_scores = []
geo_scores = []
step_score_trajectories = []
for cand in candidates:
ms, step_s = score_solution_min(model, test_problem, cand)
gs, _ = score_solution_geomean(model, test_problem, cand)
min_scores.append(ms)
geo_scores.append(gs)
step_score_trajectories.append(step_s)Candidate solution scores: Candidate Min-step Geo-mean ---------------------------------------- Solution 1 0.3449 0.6706 Solution 2 0.8354 0.9548 Solution 3 0.3070 0.7317 Solution 4 0.7847 0.8900 Solution 5 0.7099 0.8738 Best by min-step: Solution 2 Best by geo-mean: Solution 2
The two scoring strategies can select different "best" candidates. The minimum-step strategy is more conservative, penalizing any solution with a single weak step heavily. The geometric mean strategy is more lenient, allowing strong steps to compensate for weaker ones. In practice, the minimum-step strategy has shown stronger empirical performance on mathematical benchmarks, because in mathematics a single incorrect step typically invalidates the entire solution regardless of the quality of other steps.
Key Parameters
The key design choices for a process reward model are:
- Aggregation strategy: Min-step is conservative and typically performs better for math; geometric mean is more forgiving and may suit exploratory reasoning
- Label generation method: Human annotation produces high-quality labels; Monte Carlo estimation scales automatically but introduces noise
- Step granularity: Coarser steps are easier to annotate; finer steps provide more precise credit assignment
- Loss function: Binary cross-entropy for human labels; soft targets or MSE for Monte Carlo success rates. Class weighting is important when incorrect steps are rare.
- Backbone model: A larger, more capable backbone generalizes better across problem types but increases the computational cost of both training and inference.