Process Reward Agents for Steering Knowledge-Intensive Reasoning

Jiwoong Sohn; Tomasz Sternal; Kenneth Styppa; Torsten Hoefler; Michael Moor

cs.AIApr 10, 2026

Process Reward Agents for Steering Knowledge-Intensive Reasoning

Authors: Jiwoong Sohn, Tomasz Sternal, Kenneth Styppa, Torsten Hoefler, Michael Moor

Organizations: ETH Zürich, Switzerland · Heidelberg University, Germany

Abstract

Reasoning in knowledge-intensive domains remains challenging as intermediate steps are often not locally verifiable: unlike math or code, evaluating step correctness may require synthesizing clues across large external knowledge sources. As a result, subtle errors can propagate through reasoning traces, potentially never to be detected. Prior work has proposed process reward models (PRMs), including retrieval-augmented variants, but these methods operate post hoc, scoring completed trajectories, which prevents their integration into dynamic inference procedures. Here, we introduce Process Reward Agents (PRA), an inference-time method for providing domain-grounded, online, step-wise rewards to a frozen policy. In contrast to prior retrieval-augmented PRMs, PRA enables search-based decoding to rank and prune candidate trajectories at every generation step. Experiments on multiple medical reasoning benchmarks demonstrate that PRA consistently outperforms strong baselines, achieving 81.9% accuracy on MedQA with Qwen3-4B, a new state of the art at the 4B scale. Importantly, PRA generalizes to unseen frozen policy models ranging from 0.5B to 8B parameters, improving their accuracy by up to 25.7% without any policy model updates. More broadly, PRA suggests a paradigm in which frozen reasoners are decoupled from domain-specific reward modules, allowing the deployment of new backbones in complex domains without retraining.

Figures & tables

Figure 1: Overview of our approach. A Process Reward Agent (PRA) observes the reasoning trace generated by the frozen policy model (reasoner), decides when to search for external evidence, and assigns step-level rewards. This interaction can steer the policy at inference time, enabling more robust and controllable reasoning, particularly in knowledge-intensive domains like medicine.

Policy	Method	ID	OOD						Average
Policy	Method	MedQA	Medbullets	MedMCQA	MMLU	GPQA	Lancet	NEJM	Average
Qwen3-4B-Instruct	Direct	61.6	48.8	55.6	77.4	51.1	60.4	45.3	57.2
	Direct + SC	61.3	48.7	55.8	77.3	50.8	60.2	46.4	57.2
	CoT	72.7	56.5	61.1	83.7	60.8	62.4	62.7	65.7
	CoT + SC	74.8	58.7	62.7	84.9	51.8	63.5	63.2	65.7
	RAG	72.2	55.7	63.3	85.4	59.2	62.1	65.7	66.2

Table 1: Main results on medical reasoning benchmarks. PRA outperforms direct answering, chain-of-thought (CoT), and retrieval-augmented generation (RAG) baselines on the in-distribution MedQA benchmark and six out-of-distribution benchmarks, achieving the best average score overall. Using Qwen3-4B-Instruct as the base policy model, PRA improves over the strongest baseline, RAG + SC, by 5.2 points on average.

Figure 2: Performance on MedQA under inference time scaling. PRA continues to benefit from additional compute, while Self-Consistency saturates quickly. For SC, we estimate per-question expected accuracy via Monte Carlo sampling (1,000 trials); shaded regions show ±1 SE computed via bootstrap resampling over questions.

Policy	Method	Acc.	$\Delta$
Llama-3.1-8B-Instruct	CoT	67.0	–
	+ Self-Consistency	75.1	+8.1
	+ PRA	82.3	+15.3
Qwen3-4B-Instruct †	CoT	72.7	–
	+ Self-Consistency	74.8	+2.1
	+ PRA	81.9	+9.2

Table 2: Cross-model generalization on MedQA. PRA trained with Qwen3-4B-Instruct † generalizes effectively to both larger and smaller policy models, with larger gains observed for smaller models. All non-daggered policies are unseen during PRA training.

Reward Agent	Trained?	Search?	Method	# Sample	Acc.
$\times$	$\times$	$\times$	CoT	1	72.7
$\times$	$\times$	$\times$	CoT + SC	64	74.8
$\times$	$\times$	✓	RAG	1	72.2
$\times$	$\times$	✓	RAG + SC	64	76.7
Qwen3-4B	$\times$	$\times$	PRA	64( $B\times b$ )	74.4
Qwen3-4B	$\times$	✓	PRA	64( $B\times b$ )	76.7

Table 3: Ablation study of PRA training and search on MedQA. All methods use the same frozen Qwen3-4B-Instruct policy model and differ only in whether a reward agent is used, whether it is trained, and whether search(retrieval) is enabled. Training the reward agent accounts for the majority of the performance gain, and combining it with search yields further improvements, achieving the highest accuracy with PRA.

Method	Reward Level	Reward Time	Search	Acc.
CoT + SC	$\times$	Post hoc	$\times$	74.8
PRA (Last)	Outcome	Post hoc	✓	75.7
PRA (Min)	Process	Post hoc	✓	74.3
PRA (Max)	Process	Post hoc	✓	77.5
PRA (Average)	Process	Post hoc	✓	77.6
PRA (Ours)	Process	Online	✓	81.9

Table 4: Ablation of outcome-level and process-level inference-time usage of the same trained PRA on MedQA. All PRA variants share identical reward model parameters; only the timing and granularity of reward application differ.

Figure 3: Search–accuracy trade-off on MedQA. We sweep the search threshold and report accuracy versus search frequency; the Pareto frontier highlights the best operating points for a given search budget.

Figure 4: Margin shift across step positions in reasoning trajectories on MedQA, separated by traces with correct and incorrect final answers. Correct traces show larger margin shifts at later steps, whereas incorrect traces show the opposite pattern.

Figure 5: Mean absolute margin shift over reasoning steps across questions of varying difficulty, where difficulty is defined by the fraction of policy-generated reasoning samples that reach the correct answer. Correct traces consistently exhibit larger margin shifts than incorrect traces, and margin shift for correct traces is highest on harder questions and gradually decreases as solve rate increases.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Description
System
$q$	question
$\mathcal{Q}$	space of questions
$\mathcal{Y}$	space of answers
$y_{q}$	ground-truth answer for question $q$
$\hat{y}_{q}$	predicted answer for question $q$

Appendix

Table 5: Summary of notations

Panel of Figure 6 — Figure 6: Simplified pseudocode. Each Trace manages one question and cycles through four stages. The global queue collects all active traces, partitions them by pending stage, and dispatches each partition as a batched operation to the policy ( $\pi$ ), retriever ( $\rho$ ), or reward agent ( $\mu_{\phi}$ ), regardless of per-trace step index.

Panel of Figure 7 — Figure 7: Policy prompt for frozen-reasoner generation during PRA-guided beam search. The system message requires explicit step-wise reasoning (‘Step {number}:’) and a standardized final-answer phrase for parsing; the user message lists the question as ‘Question: …’ with ‘(letter): text’ options and ends with ‘Answer:’. Retrieved documents are not shown to the policy.

Panel of Figure 8 — Figure 8: Teacher prompt for step-level reasoning labels. The user message provides retrieved documents, the question with ‘(letter): text’ options, the correct answer, and the reasoning trace through the current step; the model assesses only the last step and returns a single score as [REASONING_SCORE]: 0 or 1 . The documents block is omitted in a second pass to estimate the no-retrieval margin used for search-label construction.

Panel of Figure 9 — Figure 9: PRA prompt for beam-search inference and supervised fine-tuning. The documents block is included only on the post-retrieval reward pass.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 926 papers

MedPRMBench: A Fine-grained Benchmark for Process Reward Models in Medical Reasoning
Apr 19, 2026Xiang Zheng, Lingyan Wu, Weiqi Zhai +4Healthcare Process Reward Models
Alibaba Group, China
Learning Process Rewards via Reasoning State Propagation
Sep 30, 2026Kai Gan, Zi-Hao Zhou, Bo Ye +3Process Reward Models RL for Language Model Reasoning
School of Computer Science and Engineering, Southeast University, Nanjing 210096, China · Key Laboratory of Computer Network and Information Integration (Southeast University), Ministry of Education, China · Zhongguancun Academy +1
Latent Reward Steering: An Adaptive Inference-Time Framework that Implicitly Promotes Cognitive Behaviors in Reasoning LLMs
May 30, 2026Jiakang Li, Guanyu Zhu, Can Jin +8Reward Modeling Language Model Steering
Rutgers University · South China Agricultural University · Columbia University +4

Apr 19, 2026J/K move · Enter open · S save

MedPRMBench: A Fine-grained Benchmark for Process Reward Models in Medical Reasoning

Xiang Zheng, Lingyan Wu, Weiqi Zhai, Wei Wang+3

Alibaba Group, China

Healthcare Process Reward Models Medical Reasoning Clinical Language Model Evaluation

Open paper AskPDF arXiv

Medical reasoning can reach a plausible answer while overlooking an unsafe or unsupported intermediate step. Process-Level Reward Models (PRMs) can verify such steps, yet existing PRM benchmarks largely target mathematics rather than patient-specific clinical errors. We introduce MedPRMBench, a step-level medical reasoning benchmark with 14 clinical error types, including five new and two clinically re-scoped types, and rule-defined severity grades. Built from seven medical QA sources via structured error construction and review by 86 physicians from Grade-A tertiary hospitals, it provides 6{,}500 evaluation questions, 13{,}000 reasoning chains, and 113{,}910 step labels, plus a leakage-clean 6{,}870-question training split. The best of 20 prompted critics reaches only 75.4% PRMScore, and critics differ sharply in error-type profiles and operating points. A Qwen3-8B PRM trained on the split reaches 86.6% PRMScore and 89.0% strict first-error localization; in matched ablations, physician review raises error recall by 6.55 points and cuts clean-chain false alarms from 5.6% to 2.1%. The training signal replicates across model families; without further training, the PRM transfers to ProcessBench and PRMBench, exceeding six of nine math-trained PRMs on each.

Abstract

LLM Reasoning

Explore similar work

MedPRMBench: A Fine-grained Benchmark for Process Reward Models in Medical Reasoning

Learning Process Rewards via Reasoning State Propagation

Latent Reward Steering: An Adaptive Inference-Time Framework that Implicitly Promotes Cognitive Behaviors in Reasoning LLMs