Process Reward Agents — Online Step-Wise Steering for Knowledge-Intensive Reasoning

Paper. Jiwoong Sohn, Tomasz Sternal, Kenneth Styppa, Torsten Hoefler, Michael Moor. Process Reward Agents for Steering Knowledge-Intensive Reasoning. ICML 2026. ETH Zürich · Heidelberg University. [arXiv]


0. The Picture in One Paragraph

In math and code, intermediate reasoning steps are locally verifiable — you can check each algebraic step or run the code. In knowledge-intensive domains like medicine, step verification requires cross-referencing vast external knowledge bases, and a subtly wrong clinical inference may propagate through an entire reasoning trace undetected. Process Reward Agents (PRA) address this by training a dedicated reward agent — distinct from the reasoning policy — to observe the unfolding CoT trace, selectively query external knowledge via retrieval, and assign scalar reward scores to candidate continuations during beam search. The policy model is never fine-tuned; only the reward agent is trained. The result is a modular, backbone-agnostic plug-in: a frozen Qwen3-4B-Instruct reaches 80.8% on MedQA (new SOTA at the 4B scale), with up to +25.7% improvement across frozen policy models ranging from 0.5B to 8B parameters. Unlike Self-Consistency, which degrades on hard benchmarks as more samples are added, PRA continues improving monotonically as inference compute scales.


1. The Problem — Non-Local Verifiability

Process reward models (PRMs) have been applied successfully to mathematics: you train a reward model to score each step of a proof or calculation, then use those scores to guide beam search toward correct solutions. The recipe works because math steps are locally verifiable — you can determine whether an intermediate algebraic transformation is correct without knowing the final answer.

Medical reasoning is structurally different:

Property Math / Code Medicine
Local verifiability Yes — each step checkable in isolation No — verifying a clinical inference requires synthesizing clues across large knowledge bases
Error propagation Detectable early Subtle errors compound undetected through the trace
External knowledge needed Rarely Constantly (drug interactions, symptom prevalence, guidelines)

Prior retrieval-augmented PRMs score completed trajectories after the fact. They cannot steer what is generated — by the time a bad reasoning path is scored, the trajectory is already done. PRA’s contribution is making the process reward signal online and integrated into the generation loop.


2. The Architecture — Frozen Policy + Reward Agent

PRA decouples two roles that prior work combined:

Policy model (frozen)   →   generates candidate CoT continuations
         ↓
Reward agent (trained)  →   reads partial trace
                        →   decides whether to retrieve
                        →   synthesizes retrieved docs with trace
                        →   assigns reward score per candidate
         ↓
Beam search             →   prunes low-reward candidates
                        →   propagates high-reward candidates

2.1 Frozen policy model

The reasoning policy generates text. It is never updated during PRA deployment. This separation means: when a better backbone model becomes available, you swap the frozen policy with no retraining of the reward agent. Domain expertise transfers.

2.2 Reward agent

A separate LLM trained to:

  1. Read the partial CoT trace up to the current step.
  2. Decide whether to issue a retrieval query.
  3. If retrieval is useful: select a query, retrieve relevant documents (medical literature, PubMed, clinical guidelines), and synthesize them with the trace.
  4. Assign a scalar reward score to each candidate continuation produced by the frozen policy.

The reward agent’s training labels are generated by a teacher model: given the partial trace and retrieved documents at each step, the teacher assesses whether the partial step is on track. This produces “reasoning labels” that train the reward agent to assign reliable intermediate rewards at inference time.

At each beam search step:

  1. The frozen policy autoregressively proposes K candidate continuations.
  2. The reward agent scores each candidate.
  3. Beam search uses scores to rank and prune: low-scoring candidates are discarded, high-scoring ones propagate.
  4. Repeat until the final answer is generated.

This is online (rewards computed during generation, not after) and process-level (rewards at intermediate steps, not just the final answer).


3. Results

3.1 MedQA (USMLE-style)

Model Accuracy
PRA + Qwen3-4B-Instruct 80.8% (new SOTA at 4B scale)
Improvement across 0.5B–8B frozen policies up to +25.7%

All improvements are achieved without modifying the policy model.

3.2 Generalization benchmarks

PRA is evaluated out-of-distribution on:

  • MedBullets Op4 / Op5 (harder USMLE-style)
  • MedMCQA (Indian medical licensing exam)
  • MMLU-Med (medical subset of MMLU)
  • GPQA (graduate-level “Google-proof” expert questions)
  • Lancet clinical cases (real-world Lancet journal cases)
  • NEJM clinical cases (New England Journal of Medicine cases)

PRA consistently outperforms baselines (greedy CoT, RAG, Self-Consistency, Best-of-N, post-hoc PRM scoring) across all benchmarks under matched sampling budgets.

3.3 Inference-time scaling behavior

This is the most practically important finding:

Self-Consistency degrades on hard benchmarks. As the number of sampled trajectories increases past ~8, majority vote begins amplifying errors — the policy frequently produces incorrect responses across samples, so the plurality answer is wrong. Self-Consistency actively hurts accuracy on GPQA and Lancet at high sample counts.

PRA continues improving. Because step-wise steering recovers from early errors (rather than averaging over them), PRA’s accuracy increases monotonically with inference compute. There is no observed saturation point in the evaluated range.

The mechanism is clear: when the policy starts down a wrong path, the reward agent’s low score causes that branch to be pruned before the error compounds. Self-Consistency has no such pruning — it samples independently and votes.

3.4 Ablations

Ablation Finding
Process-level vs. outcome-level rewards Process-level substantially better
Online vs. offline (post-hoc) timing Online substantially better
Both process + online Necessary; removing either reduces performance substantially

4. Relation to Prior Work

Post-hoc PRMs

Standard PRMs and retrieval-augmented PRMs score completed trajectories. They can identify good vs. bad completions after the fact, but they cannot intervene during generation to prevent bad paths from developing. PRA makes the reward signal generative — it shapes what is produced, not just evaluates what was produced.

RAG (Retrieval-Augmented Generation)

Standard RAG appends retrieved documents to the prompt at the beginning. PRA retrieves selectively at each step based on what has been reasoned so far. The reward agent decides whether to retrieve and what to query based on the partial trace — a more targeted and dynamic use of retrieval.

Best-of-N

Best-of-N samples N complete trajectories and picks the highest-reward one. PRA prunes during generation, preventing bad branches from completing. This is more efficient and produces better results because early pruning stops error compounds before they occur.


5. Why It Matters

Three reasons:

  1. Process-level rewards for knowledge-intensive domains. The PRM literature has largely focused on math. PRA is a concrete recipe for extending the process reward paradigm to domains where step verification requires external knowledge — a much larger class of real-world tasks.
  2. Frozen policy + trained reward agent = modular deployability. The decoupled design means domain expertise lives in the reward agent, not the policy. You can upgrade the backbone reasoning model without retraining domain knowledge. This is a practically important architectural property for medical AI deployment.
  3. Online steering vs. post-hoc scoring changes the scaling curve. Self-Consistency’s failure mode at high compute (degrading on hard tasks) is well-documented but underappreciated in practice. PRA’s demonstration that online process rewards maintain monotonic improvement where Self-Consistency fails is a direct argument for moving beyond best-of-N and majority vote as inference scaling strategies.

6. Limitations Worth Knowing

  • Not a deployment system. The paper explicitly states this is a method contribution, not a ready-to-deploy medical decision-support system.
  • Hallucination not eliminated. PRA reduces the propagation of incorrect intermediate steps but does not guarantee their absence.
  • Inference cost. Every generation step incurs retrieval overhead plus reward computation, making PRA more expensive than greedy decoding or standard RAG.
  • Hard tasks remain hard. On very challenging benchmarks (GPQA, Lancet), PRA is more stable than Self-Consistency but still operates under the inherent difficulty of the domain.

7. The Takeaway for a First Reader

If you remember three things:

  1. Knowledge-intensive reasoning has a non-local verification problem — you can’t check a clinical inference without consulting external knowledge. Prior PRMs were post-hoc; PRA trains a dedicated reward agent that reads the partial trace at each step, selectively retrieves relevant knowledge, and assigns step-level reward scores during beam search — without touching the policy model.
  2. Self-Consistency degrades on hard benchmarks as samples increase (majority vote amplifies errors). PRA continues improving monotonically because early pruning prevents bad paths from compounding into completed trajectories.
  3. PRA + Qwen3-4B-Instruct achieves 80.8% on MedQA (SOTA at 4B scale), with up to +25.7% improvement across frozen 0.5B–8B policy models and consistent gains across 7 medical reasoning benchmarks.

References




    Enjoy Reading This Article?

    Here are some more articles you might like to read next:

  • The Expressive Power of Transformers with Chain of Thought
  • Landscape of Thoughts — Visualizing Where LLM Reasoning Actually Goes
  • Magellan — Guided MCTS for Escaping the Gravity Wells of LLM Creativity
  • PriorZero — Injecting LLM Priors into MuZero-Style World Models at the MCTS Root
  • SuperThoughts — Reasoning Tokens in Superposition