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:
- Read the partial CoT trace up to the current step.
- Decide whether to issue a retrieval query.
- If retrieval is useful: select a query, retrieve relevant documents (medical literature, PubMed, clinical guidelines), and synthesize them with the trace.
- 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.
2.3 Integration into beam search
At each beam search step:
- The frozen policy autoregressively proposes K candidate continuations.
- The reward agent scores each candidate.
- Beam search uses scores to rank and prune: low-scoring candidates are discarded, high-scoring ones propagate.
- 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:
- 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.
- 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.
- 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:
- 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.
- 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.
- 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
- Sohn, J., Sternal, T., Styppa, K., Hoefler, T., & Moor, M. (2026). Process Reward Agents for Steering Knowledge-Intensive Reasoning. ICML 2026. arXiv:2604.09482.
- Related on this site: RLHF → RULER trend note — the process reward model context this work extends; Ch.8 Adversarial Feedback Design — external patterns in critic/verifier agent design.
Enjoy Reading This Article?
Here are some more articles you might like to read next: