adu2021 StepWiser: Stepwise Generative Judges for Wiser Reasoning
Core Concept
StepWiser reframes process reward modeling as a reasoning task itself. Instead of classifying intermediate steps as correct/incorrect, a generative judge model performs meta-reasoning by generating explanations before predicting verdicts. Trained via reinforcement learning on relative outcomes, StepWiser achieves better step-level accuracy than existing methods while being interpretable and generalizable to new problem distributions.
Architecture Overview
- Generative Judge Model: Produces reasoning before verdict
- Meta-Reasoning: Model explains why steps are good/bad
- RL Training: Learn from relative outcome comparisons
- Bidirectional Feedback: Training-time and inference-time improvements
- Process Supervision: Fine-grained intermediate step evaluation
Implementation Steps
Stage 1: Generative Judge Architecture
Design a model that generates step explanations before verdicts.
# Generative stepwise judge
import torch
from torch import nn
from typing import Dict, List, Tuple
class GenerativeStepwiseJudge(nn.Module):
"""Judge that generates reasoning before verdicting on steps"""
def __init__(
self,
model_dim: int = 4096,
vocab_size: int = 32000,
max_explanation_len: int = 128
):
super().__init__()
self.model_dim = model_dim
self.vocab_size = vocab_size
self.max_explanation_len = max_explanation_len
# Encoder for policy step
self.step_encoder = nn.Sequential(
nn.Linear(256, model_dim), # Embed step text
nn.ReLU(),
nn.Linear(model_dim, model_dim)
)
# Encoder for problem context
self.context_encoder = nn.Sequential(
nn.Linear(512, model_dim),
nn.ReLU(),
nn.Linear(model_dim, model_dim)
)
# Explanation generation head (autoregressive)
self.explanation_head = nn.LSTM(
input_size=model_dim,
hidden_size=model_dim,
num_layers=2,
batch_first=True
)
self.explanation_decoder = nn.Linear(model_dim, vocab_size)
# Verdict head (binary: good or bad step)
self.verdict_head = nn.Sequential(
nn.Linear(2 * model_dim, model_dim),
nn.ReLU(),
nn.Linear(model_dim, 2) # good/bad
)
def forward(
self,
problem: str,
previous_steps: List[str],
current_step: str
) -> Dict:
"""
Judge current step with explanation.
Returns:
- explanation_tokens: generated explanation
- verdict_logits: good/bad classification
- confidence: meta-confidence in verdict
"""
# Encode inputs
step_embed = self.step_encoder(self.embed_text(current_step))
context_embed = self.context_encoder(self.embed_context(problem, previous_steps))
# Generate explanation
combined = torch.cat([context_embed.unsqueeze(0), step_embed.unsqueeze(0)], dim=-1)
explanation_hidden, _ = self.explanation_head(combined)
explanation_logits = self.explanation_decoder(explanation_hidden)
explanation_tokens = explanation_logits.argmax(dim=-1)
# Verdict based on step + generated reasoning
verdict_input = torch.cat([step_embed, explanation_hidden.squeeze(0)], dim=-1)
verdict_logits = self.verdict_head(verdict_input)
# Confidence: entropy of verdict
verdict_probs = torch.nn.functional.softmax(verdict_logits, dim=-1)
confidence = 1.0 - torch.nn.functional.entropy(verdict_probs)
return {
"explanation_tokens": explanation_tokens,
"explanation_logits": explanation_logits,
"verdict_logits": verdict_logits,
"verdict_probs": verdict_probs,
"confidence": confidence
}
def embed_text(self, text: str) -> torch.Tensor:
"""Embed step text"""
return torch.randn(256) # Placeholder
def embed_context(self, problem: str, prev_steps: List[str]) -> torch.Tensor:
"""Embed problem + previous steps"""
return torch.randn(512) # Placeholder
Stage 2: RL Training via Comparative Outcomes
Train judge using RL based on which verdict predictions are correct.
# RL training for generative judges
class GenerativeJudgeRLTrainer:
"""Train judge using reinforcement learning"""
def __init__(
self,
judge: GenerativeStepwiseJudge,
policy_model,
lr: float = 1e-5
):
self.judge = judge
self.policy_model = policy_model
self.optimizer = torch.optim.Adam(judge.parameters(), lr=lr)
def generate_step_pair(
self,
problem: str,
previous_steps: List[str]
) -> Tuple[str, str, bool]:
"""
Generate two different step candidates and evaluate which is better.
Returns:
step_a, step_b, label (True if A better)
"""
# Generate two step candidates from policy
step_a = self.policy_model.generate_step(problem, previous_steps)
step_b = self.policy_model.generate_step(problem, previous_steps)
# Determine which leads to correct solution
outcome_a = self.evaluate_trajectory(problem, previous_steps + [step_a])
outcome_b = self.evaluate_trajectory(problem, previous_steps + [step_b])
# Label: is A better than B?
label = outcome_a > outcome_b
return step_a, step_b, label
def evaluate_trajectory(self, problem: str, steps: List[str]) -> float:
"""Evaluate trajectory quality (0-1, higher is better)"""
# Placeholder: in practice, check if leads to correct answer
return 0.5
def train_step(self, batch: List[Dict]) -> Dict:
"""Single RL training step"""
losses = {"verdict_loss": 0, "explanation_loss": 0}
for example in batch:
problem = example["problem"]
steps = example["steps"]
correct_verdicts = example["verdicts"] # List of correct/incorrect
# Judge each step
for step_idx, step in enumerate(steps):
output = self.judge(
problem,
steps[:step_idx],
step
)
# Ground truth: is this step correct?
is_correct = correct_verdicts[step_idx]
gt_verdict = torch.tensor([is_correct], dtype=torch.long)
# Verdict loss: classification
verdict_loss = torch.nn.functional.cross_entropy(
output["verdict_logits"],
gt_verdict
)
losses["verdict_loss"] += verdict_loss
# Explanation loss: generation (optional, for interpretability)
# In practice, could supervise with human explanations
explanation_loss = 0 # Placeholder
losses["explanation_loss"] += explanation_loss
# Normalize
num_steps = sum(len(ex["steps"]) for ex in batch)
for key in losses:
losses[key] /= max(num_steps, 1)
# Backward
total_loss = sum(losses.values())
self.optimizer.zero_grad()
total_loss.backward()
self.optimizer.step()
return losses
Stage 3: Inference-Time Step Validation
Use the judge to improve reasoning at inference time.
# Inference-time step validation
class ReasoningWithStepValidation:
"""Generate reasoning with stepwise validation"""
def __init__(self, policy_model, judge: GenerativeStepwiseJudge):
self.policy = policy_model
self.judge = judge
def generate_with_validation(
self,
problem: str,
max_steps: int = 20,
temperature: float = 0.7,
validation_threshold: float = 0.5
) -> Dict:
"""
Generate solution step-by-step with validation.
Steps with low judge confidence are regenerated or pruned.
"""
steps = []
confidences = []
all_verdicts = []
for step_idx in range(max_steps):
# Generate step candidate
step = self.policy.generate_step(
problem,
steps,
temperature=temperature
)
# Judge the step
verdict_output = self.judge(problem, steps, step)
verdict_prob = verdict_output["verdict_probs"][1] # P(good)
confidence = verdict_output["confidence"]
# Check if step passes validation
if confidence > validation_threshold:
# Accept step
steps.append(step)
confidences.append(confidence.item())
all_verdicts.append(verdict_prob.item())
else:
# Reject: try generating different step
for retry in range(3):
step = self.policy.generate_step(
problem,
steps,
temperature=temperature + 0.2 # Increase diversity
)
verdict_output = self.judge(problem, steps, step)
confidence = verdict_output["confidence"]
if confidence > validation_threshold:
steps.append(step)
confidences.append(confidence.item())
all_verdicts.append(verdict_output["verdict_probs"][1].item())
break
else:
# Give up, use step anyway
steps.append(step)
confidences.append(confidence.item())
all_verdicts.append(0.5)
# Check if done
if self.policy.is_solution_complete(problem, steps):
break
return {
"steps": steps,
"confidences": confidences,
"verdicts": all_verdicts,
"avg_confidence": sum(confidences) / len(confidences) if confidences else 0
}
Stage 4: Process Reward Training Integration
Use StepWiser judges to improve policy models.
# Process reward training with generative judges
class ProcessRewardImprovement:
"""Use judge verdicts to improve policy"""
def __init__(self, policy_model, judge: GenerativeStepwiseJudge):
self.policy = policy_model
self.judge = judge
self.policy_optimizer = torch.optim.Adam(policy_model.parameters(), lr=1e-5)
def improve_policy_with_process_rewards(
self,
rollouts: List[Dict]
) -> float:
"""
Train policy to generate steps that pass judge validation.
Reward = judge confidence in step correctness.
"""
total_loss = 0
for rollout in rollouts:
problem = rollout["problem"]
steps = rollout["steps"]
ground_truth = rollout["answer"]
# Trace through rollout with judge
for step_idx, step in enumerate(steps):
# Policy log probability for this step
log_prob = self.policy.get_log_prob(step, problem, steps[:step_idx])
# Judge verdict
verdict_output = self.judge(problem, steps[:step_idx], step)
judge_confidence = verdict_output["confidence"]
# Reward: how confident is judge in step correctness?
# Also: does path lead to solution?
path_quality = self.evaluate_path_quality(problem, steps[:step_idx+1], ground_truth)
reward = 0.5 * judge_confidence + 0.5 * path_quality
# Policy gradient: maximize log_prob * reward
loss = -(log_prob * reward)
total_loss += loss
total_loss /= sum(len(r["steps"]) for r in rollouts)
# Backward
self.policy_optimizer.zero_grad()
total_loss.backward()
self.policy_optimizer.step()
return total_loss.item()
def evaluate_path_quality(self, problem: str, steps: List[str], target: str) -> float:
"""Evaluate if current path is on track to solution"""
return 0.5 # Placeholder
Stage 5: Evaluation
Measure judge accuracy and policy improvement.
# Evaluation framework
class StepwiserEvaluator:
"""Evaluate judge quality and reasoning improvements"""
def __init__(self, judge: GenerativeStepwiseJudge, policy_model):
self.judge = judge
self.policy = policy_model
def evaluate_judge_accuracy(
self,
test_problems: List[Dict],
max_problems: int = 500
) -> Dict:
"""
Evaluate judge accuracy on intermediate steps.
Metrics:
- Accuracy: % of steps correctly classified
- F1: balance between precision and recall
"""
correct_verdicts = 0
total_verdicts = 0
tp, fp, fn = 0, 0, 0
for problem_data in test_problems[:max_problems]:
problem = problem_data["problem"]
ground_truth_steps = problem_data["steps"]
correct_step_labels = problem_data["correct_steps"]
for step_idx, step in enumerate(ground_truth_steps):
output = self.judge(
problem,
ground_truth_steps[:step_idx],
step
)
pred_is_correct = output["verdict_logits"].argmax().item() == 1
is_correct = correct_step_labels[step_idx]
if pred_is_correct == is_correct:
correct_verdicts += 1
if is_correct:
tp += 1
else:
if pred_is_correct:
fp += 1
else:
fn += 1
total_verdicts += 1
accuracy = correct_verdicts / total_verdicts
precision = tp / (tp + fp) if (tp + fp) > 0 else 0
recall = tp / (tp + fn) if (tp + fn) > 0 else 0
f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0
return {
"accuracy": accuracy,
"precision": precision,
"recall": recall,
"f1": f1
}
def evaluate_reasoning_quality(
self,
test_problems: List[Dict]
) -> Dict:
"""
Evaluate quality of generated reasoning using judge.
Does validation improve solution quality?
"""
results = {
"without_judge": 0,
"with_judge": 0,
"improvement": 0
}
for problem_data in test_problems:
problem = problem_data["problem"]
target = problem_data["answer"]
# Generate without judge
solution_no_judge = self.policy.generate_solution(problem)
correct_no_judge = self.check_solution(solution_no_judge, target)
# Generate with judge validation
reasoner = ReasoningWithStepValidation(self.policy, self.judge)
output = reasoner.generate_with_validation(problem)
solution_with_judge = " ".join(output["steps"])
correct_with_judge = self.check_solution(solution_with_judge, target)
if correct_no_judge:
results["without_judge"] += 1
if correct_with_judge:
results["with_judge"] += 1
total = len(test_problems)
results["without_judge"] /= total
results["with_judge"] /= total
results["improvement"] = results["with_judge"] - results["without_judge"]
return results
def check_solution(self, solution: str, target: str) -> bool:
"""Check if solution is correct"""
return solution.strip() == target.strip()
Practical Guidance
Training Recipe
- Phase 1: Train judge on synthetic step labels (easy supervision)
- Phase 2: RL training on relative outcome comparisons (harder data)
- Phase 3: Fine-tune on domain-specific problems
Explanation Integration
- Explanations are emergent from RL training (not explicitly supervised)
- Can optionally supervise with human-written explanations
- Improves interpretability without hurting performance
When to Use
- Reasoning tasks requiring intermediate step validation
- Multi-step math and coding problems
- Scenarios where explainability is important
- Improving reasoning through search (beam search, etc.)
When NOT to Use
- Single-step tasks without intermediate verification
- Real-time systems (judge adds latency)
- Domains without clear step correctness
Performance Expectations
- Judge accuracy: 85-92% on intermediate steps
- Reasoning improvement: +3-8% on complex problems
- Inference overhead: 1.5-2x slowdown (due to validation)
Reference
StepWiser: Stepwise Generative Judges for Wiser Reasoning. arXiv:2508.19229