Process Reward Agents for Steering Knowledge-Intensive Reasoning

🎯 Pitch

Giving frozen language models step-by-step feedback from retrieved medical evidence during generation—rather than after—boosts MedQA accuracy by up to 25.7% and sets a new state of the art at the 4B scale, all without touching the underlying model.

1. Executive Summary

This paper introduces Process Reward Agents (PRA), a test-time framework that provides online, step-wise, retrieval-grounded rewards to guide a frozen reasoning model through knowledge-intensive tasks — decoupling evidence search and verification from the policy model and enabling search-based decoding to rank and prune candidate trajectories at every generation step (beam search with branching and cumulative reward scoring). Experiments on medical reasoning benchmarks using Qwen3-4B-Instruct as the frozen policy demonstrate that PRA consistently outperforms strong baselines, achieving 80.8% accuracy on MedQA — a new state of the art at the 4B scale — and generalizes across unseen policy models ranging from 0.5B to 8B parameters, improving MedQA accuracy by up to 25.7% without any policy model updates, establishing that substantial reasoning gains can be unlocked from smaller frozen models by applying process-level rewards online during generation rather than through post hoc scoring alone.

2. Context and Motivation

The Core Problem: Knowledge-Intensive Domains Lack Mechanisms for Validating Intermediate Reasoning Steps

The paper addresses a fundamental tension in deploying large language models for high-stakes, knowledge-intensive reasoning: while chain-of-thought and related techniques produce step-by-step traces, there is no mechanism to verify whether each intermediate step is actually correct as it is generated. This matters because errors in early reasoning steps can compound into incorrect final answers, and without step-level verification, there is no way to detect or correct these errors before they propagate through the entire trace.

The authors frame this problem with unusual clarity by contrasting medical reasoning with domains where verification is straightforward:

"In contrast to formal proofs or software programs, where each step can be mechanically checked against axioms, syntactic rules, or compiler constraints, medical reasoning rarely admits rigorous verification."

This contrast is not merely rhetorical — it identifies the precise structural difference that makes knowledge-intensive reasoning harder to verify than math or code. In mathematics, you can check whether a simplification follows algebraically from the previous line. In programming, you can run a compiler or unit tests. In medicine, evaluating whether "the patient likely has rheumatoid arthritis given the symmetric joint involvement" requires synthesizing clinical knowledge, epidemiological data, and diagnostic criteria that are not encoded in any simple formal system.

The stakes amplify the problem. The paper explicitly positions healthcare as a domain where "clinically correct decisions must be defensible throughout the entire reasoning trace, not only in the final answer." A model might produce the correct diagnosis through faulty reasoning — attributing the right conclusion to wrong evidence — which is medically dangerous and erodes trust in AI-assisted decision-making. The paper's impact statement underscores this: the goal is to "increase the reliability and verifiability of reasoning traces" so that "individual steps need to meet a high bar to enable trust and appropriate reliance on AI systems."

Why Prior Approaches Fall Short

The paper identifies four categories of prior work, each with specific limitations that PRA is designed to address:

1. Post-Training Approaches (SFT and RLVR)

Several recent efforts improve medical reasoning by fine-tuning language models through supervised fine-tuning or reinforcement learning from verifiable rewards (Chen et al., 2024; Zhang et al., 2025a; Liu et al., 2025a; Huang et al., 2025; Thapa et al., 2025). The paper acknowledges these as legitimate advances but identifies a structural limitation: they embed domain knowledge into the model parameters through retraining. This creates a coupling between the reasoning model and the medical knowledge it relies on, which becomes problematic because:

"it becomes prohibitively expensive to repeatedly fine-tune each new LRM backbone to remain aligned with updated medical knowledge."

Medical knowledge evolves continuously — new clinical trials, updated guidelines, emerging disease understanding, and institution-specific protocols all change over time. If every update requires retraining the entire reasoning model, deployment becomes infeasible at scale. The paper is essentially pointing to a knowledge staleness problem that is inherent to parameter-embedded approaches.

Additionally, these methods typically optimize only for final answer correctness (outcome-level rewards) even when they produce step-by-step reasoning. There is no guarantee that intermediate steps are individually correct — the model might learn to produce plausible-sounding reasoning traces that lead to the right answer without actually reasoning correctly, a form of Clever Hans behavior that is well-documented in the literature.

2. Retrieval-Augmented Generation (RAG) with Context Injection

The most direct prior approach for grounding reasoning in external evidence is retrieval-augmented generation: retrieve relevant documents, insert them into the policy model's context window, and let the model reason over them (Lewis et al., 2020; Zakka et al., 2024; Kim et al., 2025; Gao et al., 2026). The paper identifies two specific failure modes of this approach:

First, attention dilution. Simply injecting documents into context does not guarantee the model will attend to the right evidence at the right reasoning step. The model's attention mechanism must navigate a much longer context that mixes the question, retrieved documents, and progressively generated reasoning steps. There is no mechanism to ensure that step 3 of the reasoning actually uses the relevant evidence rather than overriding it with parametric knowledge or ignoring it entirely. The paper phrases this as:

"simply injecting retrieved documents into a thereby bloating policy context does not guarantee that the model will attend to the right evidence at the right time, nor does it provide a mechanism to detect and correct mistakes as they emerge."

Second, no error detection. RAG provides evidence but does not verify whether the model uses it correctly. If the model misinterprets a retrieved document, draws an incorrect inference, or ignores contradictory evidence, RAG offers no corrective signal. The model can produce errors that are fully consistent with the (misinterpreted) retrieved text, and there is no mechanism to catch this.

Some work couples grounding with training by constructing reasoning traces from structured knowledge graphs (Wu et al., 2025), but the paper notes this still relies on training or context injection — it does not provide online, step-level verification during inference.

3. Existing Process Reward Models (Including Retrieval-Augmented Variants)

This is the most directly relevant prior work, and the paper engages with it in detail. Process reward models assign scores to intermediate reasoning steps rather than just final answers, enabling more granular evaluation of reasoning quality. The paper identifies two specific prior systems in medical reasoning:

Med-PRM (Yun et al., 2025): Trains a process reward model for evaluating policy-generated reasoning traces, incorporating external medical evidence via retrieval. The key limitation: it operates post hoc, meaning "reward signals are applied only after a complete reasoning trajectory has been generated." The paper characterizes this as a fundamental limitation because it "limits intervention during reasoning, allowing errors to accumulate before any corrective signal is applied" and "precludes fine-grained control over the generation process, restricting the model's ability to explore alternative reasoning paths or prioritize evidence."
Med-S3 (Jiang et al., 2025): Jointly trains a policy model and a reward model through a self-evolving framework. The paper notes it "does not incorporate search" — meaning there is no mechanism for exploring alternative reasoning branches at inference time. The reward model exists but is not used to guide generation dynamically.

Both approaches share a critical architectural constraint: the reward signal is temporally decoupled from generation. The model generates first, and evaluation happens later. This means that when the model is at step 3 of a 10-step trace, making a decision about what to say next, it receives no guidance from the reward model about whether it is on the right track. Errors made at step 3 affect every subsequent step, and the reward model only gets to point them out after the entire trace is complete — at which point the compute has already been spent.

4. General-Purpose Process Reward Models and the Distribution Shift Problem

Beyond medical reasoning specifically, the paper engages with the broader PRM literature. The central challenge the paper highlights is generalization across policies. PRMs are typically trained on reasoning traces from a specific policy model (often the same model they will evaluate). When a PRM trained on traces from one model is applied to traces from a different model, performance degrades due to distributional mismatch — the new model produces reasoning traces with different stylistic patterns, error modes, and reasoning structures that fall outside the PRM's training distribution.

The paper cites Liu et al. (2025b) and Snell et al. (2024) on this point:

"Applying a PRM off-policy, that is, scoring reasoning traces generated by a policy different from the one used during training, often degrades performance due to distributional mismatch, particularly in settings where PRMs are used to guide inference-time search."

This is a critical practical problem because it means each new reasoning backbone requires its own separately trained PRM — exactly the kind of costly retraining cycle the paper aims to break.

In mathematical reasoning, Zhu et al. (2025) showed that retrieval could help PRMs generalize across models by providing similar questions and steps as warm-up context. But the paper notes that in medicine, retrieval-augmented PRMs (like Med-PRM) "typically retrieve evidence only after a complete reasoning trace has been generated and apply rewards post hoc." Online, retrieval-grounded step-wise evaluation that remains robust under policy shift had not been explored.

The Gap: Online, Retrieval-Grounded, Policy-Agnostic Step Verification

Synthesizing these limitations, the paper identifies a clear architectural gap. No prior system simultaneously provides:

Online verification — rewards applied during generation, at each step, so that errors can be caught and corrected before they propagate
Retrieval-grounded evaluation — verification that uses external evidence rather than just the PRM's parametric knowledge to assess step correctness
Decoupling from the policy model — a reward agent that works across different frozen reasoning backbones without requiring retraining for each one, and without modifying the policy's parameters or input context

The paper frames PRA as filling this gap by introducing "a retrieval-augmented process reward framework in which a Process Reward Agent (PRA) interacts with a frozen reasoning model" where "at each reasoning step, the PRA observes the current reasoning trace, optionally decides whether to search for external medical evidence, and assigns a local reward signal to guide generation in real time."

How the Paper Positions Itself

The paper makes three explicit contributions that position it relative to prior work:

Contribution (i): Formulating retrieval-grounded, step-wise evaluation as an online control problem. Rather than treating verification as a post hoc scoring task (as Med-PRM and Med-S3 do), the paper reframes it as a control problem where the reward agent must decide when to retrieve evidence, what score to assign, and how those scores should steer ongoing generation. This control-theoretic framing is novel — it treats the reasoning process as a dynamical system that can be guided in real time rather than audited after the fact.

Contribution (ii): Decoupling evidence search and verification from a frozen policy. The paper explicitly separates three concerns that prior work conflated: the policy model (which proposes reasoning steps), the evidence retrieval (which finds relevant external knowledge), and the verification (which scores step quality). By moving retrieval entirely to the PRA's side, the policy model never sees retrieved documents and never receives gradient signals from the reward agent. This decoupling is the mechanism that enables policy-agnostic operation — the PRA can work with any actor because it doesn't depend on the actor's training distribution. The paper states:

"the policy is never conditioned on retrieved documents and receives no gradient signal from the PRA, [so] different reasoning backbones can be substituted at deployment time with no retraining."

Contribution (iii): Enabling inference-time branching and pruning that generalizes across tasks and backbone models. The paper positions the PRA as a meta-controller that enables search-based decoding (specifically beam search with step-wise cumulative reward scoring and pruning) that would not be possible with post hoc reward models. This is a direct response to the limitation that Med-S3 "does not incorporate search" and Med-PRM only evaluates post hoc — PRA enables fine-grained search by providing step-level rewards during generation, when search decisions must be made.

The paper also positions its three architectural advantages explicitly:

Domain adaptation reduces to retraining a single reward module rather than fine-tuning every policy model, because the policy itself is never updated
Policy models are interchangeable because PRA operates purely through inference-time selection — it chooses among candidate steps the policy proposes without modifying input context or model parameters
Selective retrieval introduces a new axis of inference-time scaling where search can be invoked conditionally per step, trading off retrieval cost against reward signal quality

This last point — the conditional retrieval mechanism via the search/reward action controller — is a novel capability not present in any prior PRM system. Standard retrievers always retrieve when invoked. The PRA can decide at each step whether retrieval is necessary, which the paper explores through the margin-shift criterion and the search-accuracy trade-off analysis in Section 5.

The Broader Vision: Frozen Reasoners + Domain-Specific Reward Modules

The paper concludes its positioning with a broader architectural vision that extends beyond the immediate experimental results:

"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."

This vision inverts the prevailing paradigm where domain adaptation requires retraining the model that does the reasoning. Instead, the paper proposes keeping the reasoner general-purpose and frozen, while channeling all domain-specific work (evidence retrieval, step verification, knowledge updates) into a separate, lightweight reward agent that guides the reasoner at inference time. This modularization has practical advantages for knowledge-intensive domains where the knowledge base changes frequently (new medical literature, updated guidelines) while the reasoning capability (how to think step-by-step about a diagnosis) is relatively stable.

The experimental design — testing PRA across six different policy models ranging from 0.5B to 8B parameters, including models from different families (Qwen and Llama) — is explicitly structured to validate this vision by demonstrating that the same trained PRA can improve reasoning regardless of which frozen backbone it is paired with.

3. Technical Approach

3.1 Reader Orientation (Approachable Technical Breakdown)

What the system is: PRA is an inference-time steering mechanism where a trained reward agent watches a frozen language model generate reasoning step-by-step, retrieves medical evidence at its own discretion, and assigns a score to each step that drives a beam search — ranking and pruning candidate reasoning paths online to keep the most promising trajectories alive. What problem it solves and the "shape" of the solution: Rather than auditing reasoning traces after they are complete (post hoc) or injecting evidence into the policy model's context (RAG), PRA interposes an external controller between the policy's step proposals and the final answer selection, decoupling three concerns — evidence retrieval, step verification, and generation — so that the reward agent can guide reasoning in real time using up-to-date external knowledge without ever modifying the policy model's parameters or input.

3.2 Big-Picture Architecture (Diagram in Words)

The system has three physically separate components and one coordination procedure:

Frozen Policy Model ( $\pi$ ) — an off-the-shelf language model (e.g., Qwen3-4B-Instruct) that autoregressively proposes candidate next reasoning steps given the partial trace so far. It never sees retrieved documents and receives no gradient updates. Its role is purely generative: it produces the raw material (candidate steps) that the search procedure evaluates.
Process Reward Agent ( $\mu_\phi$ ) — a fine-tuned model (Qwen3-4B-Instruct fine-tuned on teacher-labeled data) with two token-level readout heads sharing the same backbone parameters $\phi$ :
- Action controller ( $\mu_\phi^{\text{act}}$ ): inspects a partial reasoning trace and outputs a binary action $\hat{a}_t \in \{\text{search}, \text{reward}\}$ . When search, the retriever is invoked; when reward, retrieval is skipped for this step.
- Reward scoring function ( $\mu_\phi^{\text{rwd}}$ ): inspects the partial trace and (optionally) retrieved documents, then outputs a scalar reward $\hat{r}_t \in [0, 1]$ for the most recent reasoning step.
Retriever ( $\rho$ ) — a fixed dense retrieval and reranking pipeline (MedCPT; Jin et al., 2023) that, when triggered by the action controller, queries the external knowledge base $\mathcal{D}$ (medical textbooks, StatPearls, clinical guidelines, rare disease corpus) and returns the top 64 documents relevant to the current question and the last two reasoning steps.
Beam Search Coordinator — the algorithmic glue that orchestrates the three components. It maintains $B$ partial reasoning traces, uses $\pi$ to expand each with $b$ candidate next steps, uses $\mu_\phi$ to score every candidate, and prunes to the top- $B$ traces by cumulative reward. This cycle repeats until all traces are complete, at which point the trace with the highest cumulative reward supplies the final answer.

Information flow: A question $q$ enters the system → the coordinator initializes $B$ empty traces → each trace is iteratively expanded by $\pi$ → $\mu_\phi^{\text{act}}$ decides whether to retrieve → if yes, $\rho$ fetches documents → $\mu_\phi^{\text{rwd}}$ scores each candidate step → the coordinator aggregates cumulative rewards and prunes to $B$ survivors → when all traces terminate, the highest-scoring complete trace yields the answer $\hat{y}_q$ .

3.3 Roadmap for the Deep Dive

First, the formal problem formulation (Equation 2 and surrounding notation) because it defines the optimization objective and the interface boundaries — what PRA is for and what it can assume.
Second, the PRA instantiation (Equations 3–4, the two readout heads, and the training label generation pipeline), because the PRA is the novel component and everything else depends on its design.
Third, the beam search procedure (Equation 5 and the stage-level batching infrastructure), because this is how PRA's step-wise rewards actually control generation.
Fourth, the training data construction pipeline (teacher model, reasoning labels, search labels via margin shift), because understanding how PRA is trained explains why it works off-policy and what its signals mean.
Fifth, the PRA training procedure itself (model architecture, hyperparameters, prompt templates), because it closes the loop on how the reward agent acquires its capability.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and methods paper whose core idea is that an external, retrieval-augmented reward agent can guide a frozen policy model through knowledge-intensive reasoning by providing online, step-level rewards that drive beam search — and that this guidance generalizes across different policy backbones without retraining.

Formal Problem Formulation

The paper begins by casting the inference-time reasoning problem as an optimization over the parameters of an inference procedure, setting up the interface boundaries that PRA must satisfy.

Let $q \in \mathcal{Q}$ be a question drawn from some distribution over medical queries, and let $y_q \in \mathcal{Y}$ be the ground-truth answer. The paper assumes answers are verifiable — there exists a binary correctness function $\mathcal{C}$ such that:

$\mathcal{C}(\hat{y}_q, y_q) = \begin{cases} 1 & \text{if } \hat{y}_q \text{ correctly matches } y_q \\ 0 & \text{otherwise} \end{cases}$

where $\hat{y}_q$ is the predicted answer. This verifiability assumption is structural: the system optimizes for producing answers that pass this check, but the correctness function itself is only available at evaluation time (to compute accuracy), not at inference time (where PRA must operate without ground truth).

The frozen policy model $\pi$ autoregressively generates reasoning steps, producing a reasoning trace $\tau = (s_1, s_2, \ldots, s_K)$ . The paper uses cumulative indexing — $\tau_t = (s_1, \ldots, s_t)$ refers to the partial trace up to and including step $t$ , and the final step $s_K$ is defined to be the model's final answer. This cumulative notation is important because it emphasizes that the PRA always sees the full history when evaluating step $t$ , not just the step in isolation.

The knowledge base $\mathcal{D}$ is a fixed collection of domain-specific documents (medical textbooks, StatPearls, clinical guidelines, rare disease corpus) that is available at inference time but is not part of the policy model's training or parameters.

The paper then defines the inference procedure $G_\phi$ as a parameterized function that takes three inputs — the question, the frozen policy, and the knowledge base — and outputs a final answer:

$G_\phi : (q, \pi, \mathcal{D}) \mapsto \hat{y}_q$

The optimization objective is to find parameters $\phi$ that maximize the expected correctness of the produced answer over the question distribution:

$\max_\phi \mathbb{E}_{q \sim P(\mathcal{Q}), \hat{y}_q \sim G_\phi(q, \pi, \mathcal{D})} \left[ \mathcal{C}(\hat{y}_q, y_q) \right]$

where $P(\mathcal{Q})$ is the distribution over questions, $\hat{y}_q$ is sampled from the inference procedure $G_\phi$ , and $y_q$ is the ground-truth answer for question $q$ .

What this equation computes operationally: given a parameterized inference strategy $G_\phi$ , run it on a question $q$ using the frozen policy $\pi$ and knowledge base $\mathcal{D}$ to produce an answer, check whether that answer matches the ground truth, and average this correctness indicator over the question distribution. The optimization searches over $\phi$ to find the inference strategy that maximizes this expected correctness.

Why this formulation matters: it makes three architectural commitments explicit. First, $\pi$ is frozen — the policy model's parameters are not part of $\phi$ and are never updated, which means the inference procedure must work with whatever step proposals $\pi$ generates, without being able to change $\pi$ 's behavior. Second, $\mathcal{D}$ is fixed and external — the knowledge base is a static resource that $G_\phi$ can query but cannot modify, which creates a clean separation between knowledge and the inference mechanism. Third, the optimization is over the inference procedure's parameters $\phi$ , not over the model's training recipe, which reframes test-time compute allocation as a learning problem in its own right.

The paper notes that this formulation covers "an online control problem for medical reasoning" — PRA must make decisions (when to retrieve, what score to assign) that affect subsequent generation, rather than simply evaluating a completed product.

PRA Instantiation: The Two-Head Architecture

The paper instantiates $G_\phi$ as a Process Reward Agent (PRA), which is a single language model with shared parameters $\phi$ that produces two distinct outputs at each step through separate token-level readout heads.

Core design principle: The PRA separates reasoning from evidence acquisition by delegating retrieval and evaluation to a dedicated model that operates alongside the policy rather than inside it. The paper states this explicitly:

"The PRA consists of two components: an action controller $\mu_\phi^{\text{act}}$ and a reward scoring function $\mu_\phi^{\text{rwd}}$ , both implemented as separate token-level readouts from a single model with shared parameters $\phi$ ."

This means the PRA is not two separate models — it is one model that produces two tokens in sequence, each corresponding to one of the two readout tasks. At inference time, the PRA processes the partial trace and (if retrieval was invoked) the retrieved documents, then decodes two tokens: the first token encodes the reasoning reward, and the second encodes the search action. This is parameter-efficient and ensures that the representations learned for step evaluation also inform the retrieval decision (and vice versa), since both heads share the same backbone.

Action Controller (Search Decision)

The action controller observes a partial reasoning trace $\tau_t$ and selects an action:

$\hat{a}_t \sim \mu_\phi^{\text{act}}(\tau_t), \quad \hat{a}_t \in \{\text{search}, \text{reward}\}$

where $\hat{a}_t$ is a binary decision sampled from the controller's output distribution at step $t$ . The controller sees only the partial trace $\tau_t$ — it does not yet have access to retrieved documents for this step, because the retrieval decision must be made before retrieval happens. This is a chicken-and-egg problem that the paper resolves by training the controller to predict whether retrieval would change the evaluation, using the margin-shift criterion described in the training section below.

What "search" vs. "reward" means operationally:

When $\hat{a}_t = \text{search}$ , the system invokes the retriever $\rho$ with a query formed from the question $q$ and the last two reasoning steps of $\tau_t$ . The retriever returns a set of documents $\mathcal{D}_t$ , and these documents are then provided as context to the reward scoring function $\mu_\phi^{\text{rwd}}$ .
When $\hat{a}_t = \text{reward}$ , no retrieval is performed — $\mathcal{D}_t = \emptyset$ — and the reward scoring function evaluates the step using only the information already in the partial trace.

The paper's main experiments use an "always-search" configuration where $\hat{a}_t$ is fixed to search for every step, which simplifies the system and maximizes evidence grounding. The selective retrieval experiments (Figure 3, search–accuracy trade-off) use the trained action controller with a threshold $\theta_{\text{dep}}$ to decide when retrieval is worth the cost.

Reward Scoring Function (Step Evaluation)

Once the retrieval decision is made, the reward scoring function evaluates the most recent reasoning step $s_t$ :

$\hat{r}_t = \mu_\phi^{\text{rwd}}(\tau_t, \mathcal{D}_t)$

where $\tau_t$ is the partial trace up to step $t$ (including all previous steps and their rewards, which are provided in the PRA's prompt — see Appendix Figure 9), and $\mathcal{D}_t$ is the set of retrieved documents (possibly empty). The output $\hat{r}_t \in [0, 1]$ is a scalar score representing the PRA's estimate of whether step $s_t$ is correct given the reasoning so far and the available evidence.

How the scalar reward is computed from logits: The paper specifies a concrete readout mechanism. Let $\ell^{(1)} \in \mathbb{R}^{|\mathcal{V}|}$ be the logit vector at the first output token position of the PRA, where $|\mathcal{V}|$ is the vocabulary size. The reward is obtained by:

"applying a two-way softmax to the logits of tokens 0 and 1 and taking the normalized score assigned to token 1"

In notation, this is:

$\hat{r}_t = \frac{\exp(\ell^{(1)}_1)}{\exp(\ell^{(1)}_0) + \exp(\ell^{(1)}_1)}$

where $\ell^{(1)}_0$ is the logit for the token "0" (incorrect) and $\ell^{(1)}_1$ is the logit for the token "1" (correct). This two-way softmax constrains the output to $[0, 1]$ and gives it a probabilistic interpretation: $\hat{r}_t$ is the model's confidence that the step is correct, normalized against the alternatives.

Why a two-way softmax over binary tokens rather than a scalar regression head? The paper does not explicitly justify this choice, but the design follows naturally from the training procedure: the teacher model produces binary labels (correct/incorrect), and the PRA is trained to predict these labels with a cross-entropy objective. Using a two-way softmax over the token logits means the PRA's output is directly comparable to the teacher's binary labels during training, and the probabilistic interpretation (confidence) is well-calibrated for beam search — higher scores unambiguously mean "more likely to be correct."

Action Readout (Parallel Mechanism)

The action controller uses a parallel readout mechanism. Let $\ell^{(2)} \in \mathbb{R}^{|\mathcal{V}|}$ be the logit vector at the second output token position. The controller distribution is:

$\mu_\phi^{\text{act}}(\tau_t) = \text{softmax}(\ell^{(2)}_0, \ell^{(2)}_1)$

where the softmax is again computed only over the logits of tokens "0" and "1" — but here, "1" corresponds to search and "0" corresponds to reward. The action is sampled from this distribution: $\hat{a}_t \sim \mu_\phi^{\text{act}}(\tau_t)$ .

Why sample rather than take the argmax? The paper uses sampling for the action but not for the reward (where the argmax would be the token with the higher logit). Sampling introduces stochasticity into the retrieval decision, which could be useful for exploration during inference — occasionally retrieving even when the model is uncertain, which might catch cases where retrieval would be beneficial but the controller's confidence is miscalibrated. The paper does not ablate this choice, so it remains an implicit design decision.

Three Stated Advantages of the Architecture

The paper enumerates three advantages that follow directly from this two-head, policy-decoupled design:

Domain adaptation reduces to retraining a single reward module. When the knowledge base $\mathcal{D}$ is updated (new medical guidelines, new literature), only the PRA needs to be retrained or fine-tuned on the new evidence. The frozen policy $\pi$ remains unchanged. This is because the policy never sees retrieved documents — its behavior depends only on the question and the reasoning trace so far — so changes in $\mathcal{D}$ do not affect what steps $\pi$ proposes. The PRA may assign different scores to the same steps given new evidence, but the policy's generation distribution is invariant.
Policy models are interchangeable with no retraining. Because "the policy is never conditioned on retrieved documents and receives no gradient signal from the PRA," any autoregressive model that can generate step-by-step reasoning can serve as the policy backbone. The PRA trained on Qwen3-4B traces can guide Llama, smaller Qwen variants, or any other model — it only needs to score the steps those models produce. The paper demonstrates this in Table 2, where the same PRA improves accuracy across six different policy models ranging from 0.5B to 8B parameters without any retraining.
Selective retrieval introduces a new axis of inference-time scaling. The conditional activation of retrieval through the action controller means that search is not an all-or-nothing proposition. Some steps may require evidence grounding (e.g., checking a drug interaction), while others may be purely deductive (e.g., "given the symptom cluster, the most likely diagnosis is X"). The PRA can learn to distinguish these cases and invoke retrieval only when it matters, trading off computational cost (retrieval latency and processing) against reward signal quality. This is explored in the search–accuracy trade-off experiments (Figure 3).

PRA-Guided Beam Search: How Rewards Control Generation

The beam search procedure is the mechanism through which PRA's step-wise rewards actually influence which reasoning paths survive and which are pruned. The paper uses a standard beam search formulation (citing Boulanger-Lewandowski et al., 2012; Graves, 2012) adapted to the multi-turn, retrieval-augmented setting.

Beam state and cumulative reward.

At any step $t$ , the system maintains a beam of width $B$ — that is, $B$ partial reasoning traces, each at potentially different numbers of completed steps. Let $\tau_t^{(j)}$ denote the $j$ -th trace in the beam after $t$ steps (where different traces may have different $t$ values, but the notation abstracts over this). Each trace is scored by its cumulative reward:

$R(\tau_t^{(j)}) = \sum_{i=1}^{t} \hat{r}_i^{(j)} = \sum_{i=1}^{t} \mu_\phi^{\text{rwd}}(\tau_i^{(j)}, \mathcal{D}_i^{(j)})$

where $\hat{r}_i^{(j)}$ is the PRA's reward for the $i$ -th step of the $j$ -th trace, and $\mathcal{D}_i^{(j)}$ is the set of documents retrieved at that step (possibly empty).

What this equation computes operationally: for each trace in the beam, walk through its reasoning steps from the beginning, summing the PRA's per-step scores. The result is a single scalar $R(\tau_t^{(j)})$ that represents the accumulated "correctness signal" for that trace up to step $t$ . Higher cumulative reward means the trace has been consistently scored as correct by the PRA.

Why sum rather than average or product? The paper uses a simple cumulative sum, which has the property that longer traces with consistently high per-step scores will accumulate more total reward than shorter traces. This creates a natural preference for traces that sustain correctness over many steps — a reasonable inductive bias in medical reasoning, where a correct final answer should be supported by a chain of correct intermediate deductions. A product of probabilities (as used in some PRM work; Lightman et al., 2023) would penalize length and might favor short, insufficiently justified answers. An average would make a 2-step trace with scores [0.9, 0.9] indistinguishable from a 10-step trace with scores [0.9, 0.9, ..., 0.9], which loses information about how much verified reasoning supports the answer.

One expansion-and-pruning cycle.

At each step $t$ , the frozen policy $\pi$ extends each of the $B$ traces with $b$ candidate next steps (the branching factor), producing $B \times b$ candidate partial traces. The PRA scores every candidate step $\hat{r}_t$ as described above. The cumulative reward for each candidate is computed as $R(\tau_{t-1}^{(j)}) + \hat{r}_t$ . Then, the $B \times b$ candidates are ranked by cumulative reward, and only the top $B$ are retained — the rest are pruned (discarded).

Termination and final answer selection.

Generation terminates when all $B$ traces in the beam are complete — meaning each trace has produced a final answer step. The trace with the highest cumulative reward $R(\tau_K)$ yields the final answer $\hat{y}_q$ , which is extracted from the last step $s_K$ of that trace (using the standardized format "the answer is (option alphabet)" specified in the policy prompt; Appendix Figure 7).

Why beam search rather than best-of-N? The paper does not explicitly justify beam search over alternative search strategies, but the rationale is implied by the online nature of PRA rewards. Best-of-N (generate $N$ complete traces, score them post hoc, pick the best) would waste computation on traces that are clearly going off-track early — by step 3, a trace might already contain a fatal error, but best-of-N would still complete it. Beam search prunes such traces early, reallocating the generation budget ( $B \times b$ candidates per step) toward more promising paths. The online, step-wise nature of PRA rewards is what enables beam search — post hoc PRMs like Med-PRM cannot provide the intermediate scores needed to decide which traces to prune at step 3.

Computational budget equivalence. The paper's main experiments use $B = 4$ (beam width) and $b = 16$ (branching factor), yielding $B \times b = 64$ candidate expansions per step. For fair comparison, the Self-Consistency baselines also sample 64 complete traces. This means the total number of step proposals generated by $\pi$ is matched between PRA and baselines, controlling for the policy model's inference cost.

Stage-level batching for efficient execution.

A significant engineering contribution that the paper details only in Appendix B is the stage-level batching system. Naïvely implementing beam search with retrieval would be extremely inefficient: different traces are at different reasoning depths, some need retrieval and some don't, and processing them one at a time would leave GPUs idle.

The paper's solution is to treat each question as a state machine with four stages — REASON, REWARD, SEARCH, DONE — and to maintain a single global queue of all active traces across all questions. At each iteration, traces are partitioned by their current stage, and each partition is dispatched as a single batched operation:

All traces in the REASON stage are sent to $\pi$ in one batch for next-step generation.
All traces in the REWARD stage are sent to $\mu_\phi$ in one batch for step scoring.
All traces in the SEARCH stage are sent to $\rho$ in one batch for document retrieval.

After each batch operation completes, traces advance to their next stage and re-enter the queue. This "synchronized stage-level batching sustains high GPU utilization even as traces become desynchronized due to variable-length reasoning, early termination, and conditional retrieval" (Appendix B).

Why this matters practically: without stage-level batching, the overhead of coordinating three models ( $\pi$ , $\mu_\phi$ , $\rho$ ) across many concurrent questions with variable-length reasoning would make beam search prohibitively slow. The paper's batching infrastructure makes the approach computationally viable for benchmark-scale evaluation, which is essential for demonstrating the method's effectiveness.

Concrete configuration from the paper: the main experiments use beam width $B = 4$ and branching factor $b = 16$ , matching the 64-trace Self-Consistency budget. The retriever fetches 200 candidates from each corpus, reranks the combined pool, and retains the top 64 documents. The retrieval query consists of the question $q$ and the last two reasoning steps from the partial trace.

Training Data Construction: Teacher Model and Label Generation

The PRA needs training data that provides, for every reasoning step in a set of training traces, (1) a binary label indicating whether the step is correct, and (2) a binary label indicating whether retrieval was necessary to evaluate that step. The paper generates both using a large teacher model — Qwen3-235B-Instruct — applied to reasoning traces generated by the (much smaller) frozen policy model Qwen3-4B-Instruct.

Why a 235B teacher for a 4B student? This is a distillation setup: the large model has stronger reasoning and medical knowledge, so its judgments of step correctness serve as high-quality supervision for the smaller PRA. The PRA learns to approximate what the teacher would say about a step, but at a fraction of the inference cost — the PRA is itself a 4B model, so running beam search with PRA scoring is far cheaper than running beam search with a 235B judge at every step.

Step 1: Generate policy traces.

For each question in the MedQA training split (10,178 questions), the paper generates a single reasoning trace using Qwen3-4B-Instruct as the frozen policy, following the policy prompt in Appendix Figure 7. This prompt enforces a structured format: each step must begin with "Step number:" and the final answer must use the phrase "the answer is (option alphabet)." This structured format makes it easy to parse steps for partial trace evaluation during PRA training and to extract the final answer at inference time.

Why only one trace per question? The paper does not discuss this choice, but generating only one trace per training question keeps the training data manageable (10,178 traces rather than, say, 64× that many). The trade-off is that the PRA only sees one reasoning path per question during training, which could limit its exposure to diverse error modes. However, since the PRA generalizes across unseen policies (Table 2), this single-trace-per-question setup apparently provides sufficient coverage of the kinds of errors that different policies make.

Step 2: Retrieve documents for every partial trace.

For every partial reasoning trace $\tau_t$ (every prefix of every generated trace, at every step $t$ ), the system retrieves a corresponding set of relevant documents. The retrieval query is the question $q$ concatenated with the last two reasoning steps. This means the documents $\mathcal{D}_t$ at step $t$ are sensitive to the reasoning progress — if the trace has gone down a particular diagnostic path, the retrieved documents will reflect that path.

The retriever uses MedCPT (Jin et al., 2023) for dense retrieval and reranking. The paper retrieves 200 candidates from each corpus in the knowledge base, reranks the combined pool, and retains the top 64 documents. This retrieval configuration is fixed throughout all experiments — training, inference, ablations, and baselines all use the same retrieval parameters.

Step 3: Teacher evaluation of each step (reasoning labels).

For each step $s_t$ in each training trace, the teacher model Qwen3-235B-Instruct is conditioned on:

The partial reasoning trace $\tau_t$ (all steps up to and including $s_t$ )
The corresponding set of retrieved documents $\mathcal{D}_t$
The question $q$ with its answer options
The ground-truth correct answer (used as context so the teacher knows what the right answer is)

The teacher is then prompted (Appendix Figure 8) to classify the reasoning step as either correct (1) or incorrect (0) by emitting a single token. The prompt instructs the teacher:

"You must only return one score, and output nothing else: Reasoning Score: Score 1 if the last step is logically coherent, medically sound, and aligns with the provided evidence; otherwise, score 0. Output only a single digit of your reasoning score in the following format: 1 or 0 (1: correct, 0: incorrect)"

Why include the ground-truth answer in the teacher's context? This is a critical design choice. The teacher knows the correct answer when evaluating each step. This means the reasoning label is not just "does this step look plausible?" but "given that the correct answer is X, is this step a logically valid part of reaching X?" This makes the labels more informative — a step might be medically accurate on its own but lead away from the correct answer, and the teacher would flag it as incorrect. This is appropriate because PRA's goal is to guide the policy toward the correct final answer, not just toward plausible-sounding reasoning.

The paper also extracts the log-probabilities $\log p(0)$ and $\log p(1)$ that the teacher assigns to the two reasoning labels (0 and 1 tokens). These are used for the search label computation (margin shift), not for the reasoning label itself. The reasoning label is the binary argmax — the token with higher probability — not the probability distribution.

Step 4: Teacher evaluation without retrieval (for search labels).

To determine whether retrieval was necessary for evaluating a step, the paper runs the teacher model a second time on the same partial trace, but without providing any retrieved documents. The prompt structure is otherwise identical. This yields a second set of log-probabilities for the correct/incorrect tokens — call them $\log p_{\text{no-ret}}(0)$ and $\log p_{\text{no-ret}}(1)$ .

Why two evaluations per step? The intuition is that if the teacher's assessment of step correctness is the same with and without documents, then retrieval was unnecessary — the step's correctness is evident from the reasoning alone. If the teacher's assessment changes substantially when documents are provided, then retrieval mattered — the external evidence shifted the teacher's confidence.

Step 5: Computing search labels via margin shift.

The paper formalizes this intuition using a margin shift criterion. Define the margin $m$ as the difference between the log-probabilities of the correct and incorrect labels:

$m = \log p(1) - \log p(0)$

When the teacher is more confident that the step is correct, $\log p(1) \gg \log p(0)$ , so $m$ is large and positive. When the teacher leans incorrect, $m$ is negative. When the teacher is uncertain, $m \approx 0$ .

The paper computes two margins: $m$ (without retrieval) and $m_d$ (with retrieved documents). The margin shift measures how much retrieval changes the teacher's assessment:

$\Delta m = m - m_d$

What this equation computes operationally: take the teacher's confidence gap between correct and incorrect when it sees only the reasoning trace, subtract the same gap when it also sees retrieved documents. If $|\Delta m|$ is large, the documents substantially changed the teacher's evaluation — either making it more confident that the step is correct ( $\Delta m < 0$ , meaning $m_d > m$ ) or more confident that it's incorrect ( $\Delta m > 0$ , meaning $m_d < m$ ). If $|\Delta m|$ is small, documents didn't matter.

Why use margin shift rather than, say, the difference in predicted labels? A step might change from "correct with low confidence" to "correct with high confidence" when documents are added — the binary label doesn't change, but the retrieval was clearly valuable for increasing certainty. Margin shift captures this gradation: even if the argmax label stays the same, a large change in the confidence gap indicates that retrieval provided useful signal. Using only label flips would miss these cases and underestimate the value of retrieval.

From margin shift to binary search labels.

The paper converts the continuous margin shift to a binary search label by thresholding:

$a_t = \begin{cases} \text{search} & \text{if } |\Delta m| > \epsilon_{\text{global}} \\ \text{reward} & \text{if } |\Delta m| \leq \epsilon_{\text{global}} \end{cases}$

where $\epsilon_{\text{global}}$ is set to the median of $|\Delta m|$ across all training dataset steps. This choice means exactly 50% of reasoning steps are labeled as requiring search and 50% as not — a balanced training set for the search action prediction task.

Why the median? The paper does not justify this specific threshold, but using the median ensures the training data is balanced between search and no-search examples, which prevents the PRA from learning a degenerate "always predict the majority class" strategy. A higher threshold would label fewer steps as requiring search (potentially missing cases where retrieval helps), while a lower threshold would label more steps as requiring search (wasting computation). The median is a reasonable default that the paper later explores through the threshold sweep in the search–accuracy trade-off experiments (Figure 3).

What the PRA is actually learning from these labels. The reasoning labels teach the PRA to approximate the large teacher model's judgment: "given this partial trace and these documents, is the last step correct?" The search labels teach the PRA to predict whether retrieval would substantially change that judgment — essentially, to recognize when external evidence is likely to matter. At inference time, the main experiments bypass the search label and always retrieve (maximizing evidence grounding), while the trade-off experiments use the learned action controller to selectively trigger retrieval when it's likely to matter.

A Bayesian interpretation. The paper frames the margin shift in Bayesian terms:

"From a Bayesian perspective, the margin shift between two information sets measures how much the additional evidence changes the evaluator's posterior belief."

The teacher's log-probabilities represent a posterior over step correctness. Without documents, the posterior is $p(\text{correct} \mid \tau_t)$ . With documents, it's $p(\text{correct} \mid \tau_t, \mathcal{D}_t)$ . The margin shift measures how much the posterior changes when $\mathcal{D}_t$ is added. A large shift means the evidence was informative; a small shift means it was redundant with what the trace already contained.

PRA Training Procedure

With the training data (reasoning labels and search labels for every step of every MedQA training trace) constructed via the teacher model pipeline, the PRA is trained via supervised fine-tuning.

Model initialization. The PRA is initialized from Qwen3-4B-Instruct — the same base model as the frozen policy used in the main experiments. This is an important detail: the PRA and the policy share the same pretrained backbone, which means the PRA starts with strong medical knowledge and reasoning capability. However, after fine-tuning on the teacher labels, the PRA's behavior diverges — it learns to evaluate steps rather than generate them, and it learns to predict when retrieval is necessary.

Why Qwen3-4B-Instruct as the PRA backbone? The paper does not ablate this choice, but using the same base model as the policy has practical advantages: it ensures the PRA understands the same medical terminology and reasoning patterns that appear in the policy's traces, and it means the PRA can be loaded alongside the policy without requiring a separate model family's infrastructure. The fact that the PRA generalizes to non-Qwen policies (Llama models, Table 2) despite being trained only on Qwen3-4B traces suggests that the adaptation through fine-tuning is sufficient to create a general-purpose evaluator.

Training objective and token-level readouts.

The PRA is trained to predict two binary outputs at each reasoning step: the reasoning label and the search label. The paper specifies that the model is fine-tuned using supervised learning on these labels, with the prompt template shown in Appendix Figure 9.

The PRA prompt provides:

The system instruction describing the evaluation task
The retrieved documents (only when search was triggered during training data generation; otherwise, the documents section is omitted)
The question with answer options
The partial reasoning trace including previous steps and their rewards (if any)

The PRA is trained to output two tokens in sequence: the first token is "1" or "0" for the reasoning reward, and the second is "1" or "0" for the search action (where 1 = search, 0 = reward for this head). The training objective is standard next-token prediction (cross-entropy loss) on these two output tokens.

Training hyperparameters (from Appendix C).

The paper fine-tunes with the following configuration:

Learning rate: $3 \times 10^{-5}$
Schedule: cosine decay with 100 warmup steps
Weight decay: 0.01
Effective batch size: 16
Training epochs: 3
Precision: bfloat16

These are standard fine-tuning hyperparameters for the Qwen model family. The small batch size (16) reflects the fact that each training example is a single reasoning step with its associated context (partial trace, documents, question), and individual steps can be treated as independent training samples — the model does not need to see contiguous reasoning traces during training, only individual steps with their context.

Two training variants. The paper trains two versions of PRA, corresponding to the two experimental settings:

Always-search PRA (used in main experiments, Table 1): The search label is fixed to 1 for every reasoning step during training, regardless of the margin-shift computation. This teaches the PRA that retrieval should always be invoked. At inference time, the action controller is bypassed — retrieval is always triggered before reward scoring, ensuring maximal access to external evidence.
Selective-search PRA (used in search–accuracy trade-off, Figure 3): The search labels are derived from the margin-shift criterion as described above, yielding a mix of search=1 and search=0 labels. This teaches the PRA to predict when retrieval is necessary. At inference time, the action controller outputs a score $\hat{a}_t \in [0, 1]$ (the softmax probability of the "search" token), and retrieval is triggered only when this score exceeds a threshold $\theta_{\text{dep}}$ .

Why train two separate PRAs rather than one unified model that can do both? The paper's approach allows clean ablation: the always-search variant isolates the effect of retrieval-grounded reward scoring (since retrieval always happens), while the selective-search variant adds the action controller's learned behavior. If a single model were trained with mixed labels and then evaluated in both modes, it would be harder to attribute performance differences to retrieval frequency versus reward quality.

The PRA prompt template structure (Appendix Figure 9) includes:

A system message instructing the PRA to evaluate the last reasoning step for logical coherence, medical validity, and consistency with evidence
Instructions to output two scores (reasoning reward and search reward) in the format "0,0", "1,0", "0,1", or "1,1"
The retrieved documents (if search was performed)
The question with options
The reasoning trace including prior steps

The prompt also explicitly instructs the PRA to consider "prior steps and their rewards, if they exist," which means the PRA sees not just the raw reasoning text but also the scores it previously assigned — a form of autoregressive feedback that could help the PRA maintain consistency in its evaluation across steps.

Summary of Design Choices and Their Justifications

Two-head architecture (reward + action) sharing a single backbone: parameter-efficient, ensures the representations learned for step evaluation inform retrieval decisions (and vice versa), and enables a single model to perform both tasks rather than requiring separate models for scoring and retrieval gating.
Two-way softmax over binary tokens for reward readout: produces calibrated probability scores in $[0, 1]$ , directly compatible with the binary teacher labels used during training, and interpretable as "confidence that step is correct" for beam search ranking.
Cumulative sum for beam scoring rather than product or average: naturally prefers traces that sustain correctness over many steps without penalizing length, aligning with the medical reasoning desideratum that correct answers should be supported by chains of verified intermediate deductions.
Teacher model (235B) for label generation: provides high-quality step-level supervision without human annotation, using a model with substantially stronger medical knowledge than the 4B policy; enables training at scale (10,178 questions × multiple steps each) without expert labeling costs.
Margin-shift criterion for search labels: formalizes the notion of "retrieval necessity" as the degree to which external evidence changes the evaluator's posterior belief, enabling the PRA to learn when retrieval is worth the cost rather than always retrieving.
Median threshold for binary search labels: ensures balanced training data (50% search, 50% no-search) and provides a natural default for the selective retrieval experiments; the threshold sweep in Figure 3 validates that this choice is on the Pareto frontier.
Qwen3-4B-Instruct as PRA initialization: leverages the base model's existing medical knowledge and reasoning capability, and uses the same model family as the policy for representational compatibility, while fine-tuning adapts the model from generation to evaluation.
Always-search variant for main experiments: simplifies the system and maximizes evidence grounding by ensuring every step is evaluated with external knowledge; represents an upper bound on accuracy at the cost of maximum retrieval overhead.
Stage-level batching for inference efficiency: essential engineering contribution that makes PRA-guided beam search computationally viable by batching operations across questions and stages, maintaining GPU utilization despite variable-length reasoning and conditional retrieval.

4. Key Insights and Innovations

Innovation 1: Reframing Inference-Time Verification from Post Hoc Audit to Online Control

The paper's most fundamental conceptual move is redefining what it means to verify reasoning. Prior work in medical process reward modeling — specifically Med-PRM (Yun et al., 2025) and Med-S3 (Jiang et al., 2025) — treats verification as an audit: generate a complete reasoning trace, then score its steps retrospectively. This is the natural extension of outcome reward models to step-level granularity, and it inherits the same temporal structure — evaluation follows generation, never intervenes during it.

PRA breaks this temporal coupling. By assigning rewards during generation rather than after, verification shifts from passive evaluation to active steering. The conceptual move is from "is this trace correct?" to "given where we are now, which next step should we take?" This is a control-theoretic reframing — Section 3.1 explicitly formalizes it as optimizing parameters of an inference procedure $G_\phi$ that maps questions to answers by dynamically guiding a frozen policy — rather than a scoring problem.

Why this matters beyond performance. The online framing enables capabilities that are structurally impossible under post hoc evaluation. Beam search requires knowing at step 3 which of several candidate continuations looks most promising, so it can prune the others before spending compute on their completion. Post hoc PRMs cannot provide this signal — they only score complete traces, so pruning decisions at intermediate steps have no basis. The paper's finding that outcome-level PRA (post hoc) achieves only 75.7% while online process-level PRA reaches 80.8% (Table 4) isolates the value of the temporal reframing: even with the same trained reward model, applying rewards during rather than after generation yields most of the gain.

This is a fundamental shift, not an incremental refinement. It redefines the role of a reward model from evaluator to controller, and it implies that future work on test-time compute should optimize for when rewards are applied, not just how accurate they are. The paper's contribution is not a better scoring function — the scoring function is a fine-tuned 4B model trained with standard distillation — but the architectural decision to deploy it online within a search loop.

The evidence is in Table 4: all PRA variants share identical reward model parameters; varying only timing (post hoc vs. online) and granularity (outcome vs. process) produces a spread from 74.3% (process, post hoc, min aggregation) to 80.8% (process, online). The gap between post hoc and online process-level evaluation is larger than the gap between any two post hoc aggregation strategies, confirming that timing dominates granularity as the critical variable.

Innovation 2: Decoupling Evidence Retrieval and Verification from the Policy Model as a Generalization Strategy

Prior work couples retrieval to reasoning in one of two ways. RAG-based approaches (Zakka et al., 2024; Kim et al., 2025; Gao et al., 2026) inject retrieved documents into the policy model's context, making retrieval a pre-processing step that the policy must then integrate during generation. Process reward models like Med-PRM retrieve evidence for post hoc evaluation but remain tied to a specific policy's trace distribution — they are typically trained on traces from the same model they evaluate, and applying them off-policy degrades performance due to distribution shift (Liu et al., 2025b; Snell et al., 2024).

PRA decouples entirely: retrieval and verification happen in a separate model that never modifies the policy's input context or parameters. The policy proposes steps; the PRA scores them using whatever external evidence it retrieves. The policy never sees retrieved documents and receives no gradient from the PRA. This is not a minor engineering choice — it is the architectural basis for cross-model generalization.

Why cross-model generalization is the diagnostic result. The paper's most striking finding is not that PRA improves its training-compatible policy (Qwen3-4B, +8.1% over CoT), but that the same PRA — trained exclusively on Qwen3-4B traces — improves a Llama-3.1-8B by +13.1%, a Qwen2.5-0.5B by +25.7%, and every other tested policy model without retraining (Table 2). The paper explicitly notes this is "particularly notable because PRA does not modify the policy model itself. Instead, it operates purely at inference time through step-wise selection within beam search: at each step, the policy model autoregressively proposes candidate continuations, and PRA selects among them without altering the generation procedure, injecting additional context, or updating model parameters."

This result is significant beyond the raw accuracy numbers because it inverts the prevailing assumption that reward models must be policy-specific. The paper shows that a process reward model trained on one model's traces can guide fundamentally different models — different architectures (Qwen vs. Llama), different scales (0.5B to 8B), different training distributions — because it evaluates step quality rather than mimicking a particular policy's style. The PRA's training signal comes from the 235B teacher's medical knowledge, not from matching the 4B policy's distribution, so it learns to recognize medically sound reasoning rather than Qwen3-4B-like reasoning.

The broader implication for deployment. This finding suggests an architectural paradigm where domain-specific verification is a reusable module that can be paired with any general-purpose reasoner. When medical knowledge updates (new guidelines, new literature), only the PRA needs retraining — the policy models it guides remain frozen. When a new, more capable reasoning backbone is released, it can immediately benefit from the existing PRA without any adaptation period. This modularization separates two concerns that have historically been coupled: how to reason (the policy's job) and what counts as good reasoning (the PRA's job).

The evidence is in Table 2: the PRA's benefit is largest for the smallest models (+25.7% for Qwen2.5-0.5B), suggesting that smaller models possess underutilized reasoning capacity that only becomes accessible when an external verifier helps them navigate away from dead ends. The paper characterizes this as revealing "substantial underutilized reasoning capacity in smaller models" — a finding with practical implications for resource-constrained deployment.

Innovation 3: Margin Shift as a Principled Criterion for Selective Retrieval

Retrieval-augmented systems typically face a binary choice: retrieve always or retrieve never. Always-retrieve (as in standard RAG) is computationally expensive and may add noise when retrieval is unnecessary. Never-retrieve (as in closed-book models) misses cases where external evidence is critical. Prior work lacks a principled framework for deciding when retrieval is worth its cost.

PRA introduces a new axis: selective retrieval gated by the expected informativeness of external evidence. The paper operationalizes this through the margin-shift criterion, which quantifies how much retrieved documents change a strong evaluator's (the 235B teacher's) posterior belief about step correctness. The key insight is that retrieval should be invoked only when it would change the evaluation — when external evidence provides information not already contained in the reasoning trace itself.

Why margin shift is conceptually novel. Most work on retrieval necessity uses heuristics — retrieve when confidence is low, retrieve when the question contains entity mentions, retrieve every $k$ steps. These heuristics are proxies that may or may not correlate with retrieval's actual value. Margin shift directly measures retrieval's causal effect on evaluation quality: running the evaluator twice (with and without documents) and measuring the difference. This is a counterfactual quantity — it answers "would this evaluation be different if we had retrieved?" — rather than a correlational proxy.

The paper frames this in Bayesian terms (Section 3): the margin is the log-ratio of posterior probabilities $p(\text{correct}) / p(\text{incorrect})$ , and margin shift is the KL-relevant change in this posterior when conditioning on documents. A small margin shift means the documents are redundant with information already in the trace; a large shift means they provide genuinely new signal.

The search–accuracy trade-off as validation. The paper demonstrates that this criterion produces a meaningful Pareto frontier (Figure 3): as the retrieval threshold $\theta_{\text{dep}}$ increases (invoking retrieval less often), accuracy generally declines, but the frontier shows that comparable accuracy can be achieved with substantially fewer retrieval calls. This is not just a cost-saving measure — it validates that the margin-shift criterion successfully identifies which steps genuinely benefit from retrieval, since the steps where retrieval is skipped (low margin shift) are those where documents wouldn't have changed the evaluation anyway.

This is a fundamental contribution to the retrieval-augmented generation literature, not an incremental refinement. It provides a theoretically grounded, empirically validated answer to "when should we retrieve?" that generalizes beyond medical reasoning — any domain where a strong evaluator can compare judgments with and without external evidence could use the same criterion to train a selective retrieval controller.

The evidence is in Figure 3 and the surrounding analysis (Section 5.3, Figures 4-5): margin shift varies systematically with trajectory correctness and question difficulty. Correct traces show larger margin shifts at later steps (retrieval becomes more valuable as reasoning progresses toward a conclusion), while incorrect traces show decreasing margin shifts (their errors are detectable without external evidence). This diagnostic pattern would not emerge if margin shift were simply noise — it captures a real signal about when external knowledge matters.

Innovation 4: Step-Level Rewards as the Dominant Factor in Test-Time Compute Scaling, Not Just Reward Accuracy

The paper's ablation results (Table 4) reveal a finding with implications beyond medical reasoning: the timing and granularity of reward application matter more for inference-time scaling than the absolute quality of the reward signal. This challenges an implicit assumption in much of the process reward model literature — that the primary path to better test-time performance is building more accurate reward models.

The evidence is straightforward but striking. When the same trained PRA is applied:

As an outcome-level, post hoc scorer (like a standard ORM): 75.7% — only marginally better than Self-Consistency (74.8%)
As a process-level, post hoc scorer with average aggregation: 77.6% — a modest improvement from evaluating steps rather than just the final answer
As a process-level, online scorer within beam search: 80.8% — a jump that exceeds the combined gains from all other variations

The gap from outcome-level post hoc (75.7%) to process-level online (80.8%) is +5.1 percentage points using the identical reward model. The intermediate post hoc process-level variants span 74.3% to 77.6% depending on aggregation operator (min, max, average) — meaning that even the best post hoc use of step-level rewards underperforms online application by 3.2 points.

What this reveals about test-time compute. This pattern suggests that the primary bottleneck in prior PRM work was not reward model accuracy but reward model integration. Med-PRM and Med-S3 built capable process reward models but applied them post hoc, missing the majority of the potential gain. The paper's contribution is demonstrating that the same reward signal, applied during generation within a search framework, unlocks substantially more performance than improving the signal's accuracy would — because online application enables the system to act on reward information (pruning unpromising paths, investing compute in promising ones) rather than merely observing it after the fact.

This is a diagnostic finding with implications for how the field should allocate research effort. It suggests that work on inference-time search infrastructure (beam search, tree search, MCTS) and reward timing may be higher-leverage than work on incrementally better reward models, at least until reward models reach some quality threshold. The paper does not claim that reward model quality is unimportant — the untrained Qwen3-4B used as a reward agent achieves only 74.4% (Table 3), substantially below the trained PRA's 80.8% — but it shows that how rewards are used can dominate how good they are as a source of performance variation.

Why this is a fundamental insight, not an incremental finding. Prior work on test-time compute scaling (Snell et al., 2024) focused on how much compute to allocate and which strategy to use (best-of-N vs. beam search vs. revisions), treating the reward model as a fixed component. This paper shows that within a single strategy (beam search), the temporal deployment of the reward signal creates a larger performance spread than the choice of aggregation operator or even the choice between outcome and process granularity. This reframes the test-time compute problem as one of reward deployment architecture as much as reward quality or search algorithm — a dimension that had not been systematically studied.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary in-distribution evaluation uses the MedQA test split (Jin et al., 2020), with the training split (10,178 questions) used exclusively for generating PRA training data — all evaluation questions are held out. Out-of-distribution generalization is tested on six additional benchmarks: MedBullets (Chen et al., 2025), MedMCQA (Pal et al., 2022), MMLU-Med (Hendrycks et al., 2021; Singhal et al., 2023), GPQA (Rein et al., 2023), and clinical case datasets from The Lancet and The New England Journal of Medicine (Thapa et al., 2025).
Base model(s). The frozen policy model for the main experiments is Qwen3-4B-Instruct, chosen as a representative 4B-scale model with non-trivial but unsaturated medical reasoning capability (72.7% MedQA accuracy under chain-of-thought). For cross-model generalization experiments (Table 2), six additional frozen policy models are tested spanning 0.5B to 8B parameters and two model families: Llama-3.1-8B-Instruct, Llama-3.2-3B-Instruct, Qwen2.5-3B-Instruct, Llama-3.2-1B-Instruct, and Qwen2.5-0.5B-Instruct.
Metrics. Accuracy is measured as the fraction of questions for which the predicted answer matches the ground truth, where answers are multiple-choice options extracted from the final reasoning step (using the standardized format "the answer is (option alphabet)" enforced by the policy prompt). The paper reports straightforward accuracy percentages without confidence intervals for the main results tables, though Figure 2 includes ±1 standard error bands computed via bootstrap resampling over questions.
Baselines. The paper compares against five baselines, each evaluated at both single-sample and 64-sample (Self-Consistency) budgets: (1) Direct — answering without explicit reasoning; (2) Chain-of-Thought (CoT) — step-by-step reasoning via prompting; (3) Retrieval-Augmented Generation (RAG) — CoT with retrieved documents prepended to the policy input (Lewis et al., 2020; Zakka et al., 2024); (4) CoT + SC — majority voting over 64 CoT samples; (5) RAG + SC — majority voting over 64 RAG samples. For the ablation in Table 3, additional baselines include using the untrained Qwen3-4B-Instruct as a reward agent within beam search (with and without retrieval).
Generation budget / compute accounting. Fair comparison is enforced by matching the number of policy model generations. PRA with beam search uses beam width $B = 4$ and branching factor $b = 16$ , producing $B \times b = 64$ candidate step proposals per expansion step. The Self-Consistency baselines sample 64 complete reasoning traces, matching the policy model's total generation load. The paper does not account for the PRA's own inference cost or retrieval latency in the main budget comparison — only the policy model's generation budget is equalized.
Cross-validation / statistical protocol. The paper does not employ cross-validation or train/test splits for PRA training — the MedQA training split (10,178 questions) is used to generate training traces, and the MedQA test split is used for evaluation, with no overlap. For out-of-distribution evaluation, all six benchmarks are entirely unseen during training. The inference-time scaling analysis (Figure 2) reports expected accuracy with ±1 standard error computed via bootstrap resampling over questions and Monte Carlo sampling (1,000 trials) for estimating per-question expected Self-Consistency accuracy at each sample size. No statistical significance tests are reported for pairwise comparisons between methods.

Main Quantitative Results

In-Distribution and Out-of-Distribution Medical Reasoning (Table 1)

PRA with Qwen3-4B-Instruct achieves 80.8% accuracy on MedQA, outperforming the strongest baseline (RAG + SC at 76.7%) by 4.1 percentage points. On the six out-of-distribution benchmarks, PRA achieves the highest average accuracy (71.0%), exceeding RAG + SC (66.9%) by 4.1 points and CoT + SC (65.7%) by 5.3 points. The per-benchmark breakdown is reported in Table 1:

MedQA (in-distribution): PRA 80.8% vs. CoT + SC 74.8% vs. RAG + SC 76.7%
MedBullets: PRA 63.6% vs. CoT + SC 58.7% vs. RAG + SC 58.4%
MedMCQA: PRA 66.2% vs. CoT + SC 62.7% vs. RAG + SC 64.8%
MMLU-Med: PRA 86.6% vs. CoT + SC 84.9% vs. RAG + SC 86.2%
GPQA: PRA 64.4% vs. CoT + SC 51.8% vs. RAG + SC 54.4%
Lancet: PRA 67.0% vs. CoT + SC 63.5% vs. RAG + SC 61.0%
NEJM: PRA 68.3% vs. CoT + SC 63.2% vs. RAG + SC 66.9%

A notable pattern emerges on GPQA and Lancet: Self-Consistency degrades performance relative to single-sample CoT or RAG on these benchmarks. For GPQA, CoT achieves 60.8% but CoT + SC drops to 51.8%, and RAG achieves 59.2% but RAG + SC drops to 54.4%. The paper attributes this to the policy model "frequently producing incorrect or incomplete responses across repeated samples, causing Self-Consistency to amplify errors through majority voting." PRA avoids this degradation, achieving 64.4% on GPQA — the only method that substantially improves over single-sample CoT on this benchmark.

Inference-Time Scaling Behavior (Figure 2)

PRA and Self-Consistency show qualitatively different scaling curves as the number of sampled traces increases from 1 to 64. Self-Consistency saturates quickly — accuracy improves from approximately 73% at 1 sample to roughly 75% at 8 samples, then plateaus through 64 samples. PRA continues to benefit from additional compute, with accuracy rising from approximately 73% at 1 sample to roughly 80% at 64 samples, and the curve showing no clear sign of saturation at the maximum tested budget. The paper attributes this divergence to PRA's ability to "steer reasoning toward more promising trajectories and recover from early errors" through step-wise rewards applied during generation, whereas Self-Consistency "is constrained by the policy's initial sampling distribution and can only aggregate over completed samples."

Cross-Model Generalization (Table 2)

The PRA trained exclusively on Qwen3-4B-Instruct reasoning traces generalizes to all six tested policy models without any retraining, with larger absolute and relative gains for smaller models:

Llama-3.1-8B-Instruct: 67.0% (CoT) → 80.1% (PRA), +13.1 points
Qwen3-4B-Instruct (training policy): 72.7% → 80.8%, +8.1 points
Llama-3.2-3B-Instruct: 56.0% → 75.4%, +19.4 points
Qwen2.5-3B-Instruct: 49.5% → 69.9%, +20.4 points
Llama-3.2-1B-Instruct: 36.2% → 57.8%, +21.6 points
Qwen2.5-0.5B-Instruct: 28.4% → 54.1%, +25.7 points (+90.5% relative)

The gradient is consistent: the weaker the policy model's standalone CoT performance, the larger the absolute improvement from PRA guidance. The smallest model (Qwen2.5-0.5B) more than doubles its effective accuracy relative to chain-of-thought, though its absolute PRA performance (54.1%) still trails larger models with PRA (80.1% for Llama-3.1-8B). Self-Consistency provides smaller absolute gains across all models — for example, on Llama-3.2-1B-Instruct, SC adds +7.8 points while PRA adds +21.6 points — and the gap between SC and PRA widens for smaller models.

Ablation Studies and Robustness Checks

Training and retrieval contributions (Table 3): Disentangling the effects of reward agent training, retrieval, and search structure while keeping the policy model fixed to Qwen3-4B-Instruct reveals that trained reward agent accounts for the majority of the performance gain. Using the untrained Qwen3-4B-Instruct as a reward agent within beam search without retrieval achieves 74.4%, which is comparable to CoT + SC (74.8%) despite the different inference structure. Adding retrieval to the untrained reward agent yields 76.7%, matching RAG + SC exactly. Only training the reward agent (PRA, always-search) pushes performance to 80.8%, demonstrating that "training the reward agent is critical for effective inference-time scaling with beam search and retrieval of external evidence." The controlling role of training — rather than just restructuring inference or adding retrieval — is the central finding of this ablation.

Reward level and timing (Table 4): Fixing the trained PRA model and varying only how its rewards are applied isolates the contribution of online, process-level deployment. Outcome-level PRA applied post hoc (scoring complete traces with a single reward) achieves 75.7%, only marginally better than Self-Consistency (74.8%). Process-level PRA applied post hoc (scoring individual steps but aggregating after trace completion) yields results that depend on the aggregation operator: min achieves 74.3% (worse than outcome-level), max achieves 77.5%, and average achieves 77.6%. The full method — process-level rewards applied online within beam search — achieves 80.8%, representing a 3.2-point gain over the best post hoc process-level variant (average) and a 5.1-point gain over outcome-level post hoc. The paper interprets this as evidence that "the majority of the gain stems not only from stronger reward signals, but from enabling online, process-level control over the reasoning process itself."

Search–accuracy trade-off (Figure 3): Training PRA with selective search labels (derived from the margin-shift criterion) and sweeping the retrieval threshold $\theta_{\text{dep}}$ from 0 to 1 in increments of 0.1 produces a clear trade-off curve between search frequency and accuracy. At $\theta_{\text{dep}} = 0$ (always retrieve, 100% search frequency), accuracy reaches approximately 81%. As $\theta_{\text{dep}}$ increases and search frequency decreases, accuracy generally declines — at 50% search frequency, accuracy is roughly 78%; at 0% search frequency, accuracy drops to approximately 75%. The Pareto frontier (Figure 3a) shows that comparable or slightly higher accuracy can be achieved at reduced search frequencies — there exist operating points with substantially fewer retrieval calls that match or nearly match the accuracy of always-retrieve. The paper notes this validates that "selective retrieval can achieve comparable, and sometimes slightly higher, accuracy with fewer search calls." The always-search configuration used in the main experiments represents a practical upper bound on accuracy rather than an optimal cost-accuracy trade-off.

Margin shift analysis by trajectory correctness and position (Figure 4): Analyzing how $|\Delta m|$ varies across reasoning step positions, separated by whether the trace ultimately produces a correct or incorrect final answer, reveals systematically different patterns. For correct traces, margin shift increases toward later steps — retrieved evidence becomes more influential in the teacher model's evaluation as reasoning progresses toward a conclusion. For incorrect traces, margin shift decreases at later steps — the teacher model becomes increasingly confident in its (negative) assessment without needing external evidence, suggesting that "flaws in the reasoning become more apparent to the teacher model even without additional evidence." At the final step (typically an answer selection), margin shift is minimal for both correct and incorrect traces, consistent with final answer steps containing little substantive reasoning to evaluate.

Margin shift analysis by question difficulty and correctness (Figure 5): Binning questions by difficulty — defined as the fraction of policy-generated reasoning samples that reach the correct answer — reveals that correct traces consistently show larger margin shifts than incorrect traces across all difficulty levels. The effect is strongest on harder questions: for questions with solve rates below approximately 30%, correct traces exhibit mean $|\Delta m|$ around 2.8–3.0, while incorrect traces show values around 1.5–1.7 — nearly a factor of two difference. As solve rate increases (easier questions), the gap narrows, with both correct and incorrect traces converging toward lower margin shift values on the easiest questions. The paper interprets this as evidence that "incorrect reasoning trajectories contain internal inconsistencies or errors that are detectable by the teacher model without strong reliance on external evidence," particularly on harder questions where the errors are more glaring.

Critical Assessment

Claim 1: PRA achieves state-of-the-art performance for 4B-scale models on MedQA (80.8%).

The evidence directly supports this claim with qualifications. Table 1 reports 80.8% on MedQA, which is the highest number reported in the paper and exceeds the strongest baseline (RAG + SC, 76.7%) by 4.1 points. The claim of "state of the art at the 4B scale" is comparative — the paper asserts this but only benchmarks against its own baselines (Direct, CoT, RAG, and their SC variants). No head-to-head comparison with Med-PRM (Yun et al., 2025), Med-S3 (Jiang et al., 2025), or other published 4B-scale medical reasoning systems is provided, making "state of the art" a claim about the broader literature that the paper's own experiments do not directly verify. The 4B specification matters — larger models may achieve higher MedQA accuracy, and the paper does not claim to surpass them.

A more specific question: does PRA's advantage over RAG + SC represent a fundamental improvement in reasoning quality, or is it partly an artifact of RAG's known limitations (context dilution, attention diffusion) that beam search with external verification naturally avoids? The paper's framing suggests the latter — PRA succeeds precisely because it does not inject documents into the policy context — but without a baseline that combines RAG with some form of post hoc verification, it is difficult to isolate how much of the gain comes from avoiding context dilution versus from online step-wise guidance.

Claim 2: PRA generalizes across unseen policy models without retraining, improving accuracy by up to 25.7%.

This claim is well-supported by Table 2, which demonstrates consistent improvements across six policy models spanning two families (Qwen, Llama) and a 16× range in parameter count (0.5B to 8B). The "up to 25.7%" figure is specifically for Qwen2.5-0.5B-Instruct — the smallest and weakest model — and the absolute PRA accuracy (54.1%) remains far below larger models with PRA, so the claim of "substantial underutilized reasoning capacity" is accurate but bounded: PRA reveals capacity that exists, but the revealed capacity is still limited by model scale. The cross-family generalization (Qwen-trained PRA working on Llama policies) is particularly notable since these models have different architectures, tokenizers, and training distributions — the PRA's robustness to this distribution shift is the strongest evidence for the paper's decoupling thesis.

A limitation: all tested policy models are instruction-tuned variants (all have the "-Instruct" suffix). The paper does not test on base (pre-instruction-tuning) models, which might produce qualitatively different reasoning traces (less structured, less step-by-step) that the PRA was not trained to evaluate. Since the PRA's training data consists of structured traces from an instruction-tuned model, it is plausible — but untested — that generalization would degrade on less structured outputs.

Claim 3: Inference-time scaling with PRA continues to improve while Self-Consistency saturates.

Figure 2 provides clear evidence for this claim at the tested budget range (1–64 traces). The Self-Consistency curve visibly flattens after 8 samples while the PRA curve continues to rise through 64 samples. However, two caveats apply. First, the Self-Consistency curve is estimated via Monte Carlo simulation (1,000 trials) rather than actually running SC at each sample size, while the PRA curve is from actual beam search runs — the paper does not explain whether the two curves are computed under identical assumptions about trace quality and independence. Second, the x-axis ("Number of Samples / Traces") aggregates differently for the two methods — for SC it is independent complete traces, while for PRA it is the total number of candidate step proposals generated across all beam expansion steps — and while the total policy generation budget is matched at 64, the intermediate points may not be 1-to-1 comparable in terms of policy compute spent.

The claim would be strengthened by pushing to higher budgets (128, 256 traces) to see whether PRA eventually saturates and, if so, at what accuracy ceiling. The paper does not explore this boundary.

Claim 4: Gains arise primarily from applying process-level rewards online during generation rather than from post hoc scoring alone (Table 4).

This claim is strongly supported by the within-model ablation. The +5.1-point gap between outcome-level post hoc (75.7%) and process-level online (80.8%) using the identical trained reward model isolates timing and granularity as the causal factors. The fact that process-level post hoc (average, 77.6%) substantially underperforms process-level online (80.8%) further isolates online application as the critical factor — step-level rewards help somewhat even post hoc, but most of their value is realized only when they can guide generation in real time.

A missing ablation: what happens if process-level rewards are applied online but without beam search — for example, using the PRA's step-wise scores to select among candidates at each step but always keeping only the single best continuation (greedy search with PRA guidance, beam width = 1)? This would isolate whether search (maintaining multiple hypotheses) or online guidance (step-wise scoring during generation) is the active ingredient. The paper conflates these in its "online" condition — beam search inherently requires both — so the attribution to "online, process-level control" could partially reflect the benefits of search diversity rather than online scoring per se.

Claim 5: Selective retrieval can maintain accuracy while reducing search cost (Figure 3).

Supported with qualifications. The Pareto frontier shows that reducing search frequency from 100% to approximately 60–70% can maintain accuracy near 80%, validating that some retrieval calls are indeed unnecessary. However, accuracy drops sharply below approximately 50% search frequency (reaching ~75% at 0% search), indicating that a substantial fraction of retrieval calls are genuinely valuable — the margin-shift criterion is identifying useful cases, not just noise. The paper does not report the actual retrieval cost in terms of latency or computational overhead, so the practical significance of reducing search frequency from 100% to 70% is qualitative ("can adaptively reduce search") rather than quantified. The always-search configuration used in the main experiments is acknowledged as an upper bound rather than an optimal operating point.

Genuine weaknesses in the experimental design:

Single benchmark for training (MedQA only). The PRA is trained exclusively on MedQA training traces, and its ability to generalize to six out-of-distribution medical benchmarks is a strength, but all benchmarks are still multiple-choice medical question answering. The paper does not test whether PRA generalizes to non-medical knowledge-intensive domains (legal reasoning, scientific QA outside medicine, financial analysis), leaving open whether the approach is specific to medicine or broadly applicable to any domain with a verifiable knowledge base.
No direct comparison to prior medical PRM systems. Med-PRM (Yun et al., 2025) and Med-S3 (Jiang et al., 2025) are the most directly comparable prior work, but the paper only describes their limitations qualitatively rather than reproducing them as baselines. A head-to-head comparison — even on a subset of MedQA — would clarify whether PRA's architectural innovations (online rewards, policy decoupling) translate to practical improvements over existing retrieval-augmented PRMs, or whether the gains come primarily from the base model's capability (Qwen3-4B might simply be stronger than the backbones used in prior work).
Always-search configuration as the main result. The 80.8% headline number uses the always-search PRA variant, which retrieves documents for every reasoning step. This configuration maximizes evidence grounding but also maximizes retrieval cost — the paper's own analysis (Figure 3) shows that comparable accuracy can be achieved with substantially fewer retrievals. Reporting the always-search result as the main finding without prominently featuring the cost-accuracy trade-off overstates the practical performance of a deployable system. A more informative headline might be the best accuracy achievable at, say, 50% search frequency.
No ablation on PRA model scale. The PRA is a 4B model (Qwen3-4B-Instruct fine-tuned), and the paper does not test whether a smaller PRA (e.g., 0.5B or 1B) could achieve comparable guidance quality. Since the PRA's inference cost is not accounted for in the generation budget comparison, using a 4B reward model to guide a 0.5B policy model (Table 2) involves an 8× asymmetry in model scale that is not discussed. For the smaller policy models, the PRA may be substantially more expensive to run than the policy it guides.
No ablation on retrieval corpus composition. The knowledge base aggregates multiple corpora (textbooks, StatPearls, clinical guidelines, rare disease corpus), but the paper does not ablate which corpora contribute most to PRA's performance. It is possible that a single high-quality corpus (e.g., StatPearls alone) provides most of the benefit, which would substantially reduce retrieval overhead and simplify deployment.
Test set sizes vary and are not always reported. MedQA has 1,273 test questions (a standard split), but the paper does not report the number of questions in each out-of-distribution benchmark. Without knowing test set sizes, it is difficult to assess whether the reported accuracies are statistically reliable — a 5-point difference on a 100-question test set carries different evidentiary weight than the same difference on a 1,000-question set.
Single teacher model for label generation. All PRA training labels come from Qwen3-235B-Instruct. If this teacher model has systematic biases in its evaluation of medical reasoning — judging certain reasoning patterns as correct when they are not, or vice versa — those biases are distilled directly into the PRA. The paper does not validate the teacher's labels against human expert judgments or against an alternative teacher model, so the quality of the supervision signal is assumed rather than verified. This is a common limitation in distillation-based approaches, but it is worth noting given the high-stakes nature of medical reasoning.
No reporting of PRA inference overhead. The main results (Table 1) match the policy model's generation budget (64 traces) between PRA and baselines, but PRA additionally requires running a 4B reward model (with retrieval) at every expansion step. The paper's stage-level batching system (Appendix B) makes this efficient, but the absolute computational cost of PRA inference is never quantified relative to the baselines. For deployment considerations, the relevant comparison is end-to-end latency and total FLOPs, not just the policy model's generation count.

6. Limitations and Trade-offs

1. Difficulty Estimation Cost Is Unaccounted For and May Dominate the Inference Budget

The assumption or constraint: The entire PRA framework depends on the ability to score reasoning steps using a 4B reward model that optionally retrieves and processes up to 64 external documents per step. The paper's main experiments use an "always-search" configuration, meaning retrieval and PRA scoring are invoked at every reasoning step of every beam expansion. The paper does not account for this cost in the headline generation budget comparison — only the policy model's generation count (64 traces) is matched between PRA and baselines. The paper acknowledges this implicitly in Appendix B, which describes the stage-level batching system designed to make this overhead "viable," but never quantifies the absolute computational cost.

The consequence: In a deployment setting, the wall-clock latency and total FLOPs of PRA-guided beam search are dominated not by the frozen policy (which is a 4B model generating relatively short reasoning steps) but by the PRA itself — another 4B model that processes the full reasoning trace plus up to 64 retrieved documents at every candidate step. For the main experiments (beam width B=4, branching factor b=16), the PRA evaluates 64 candidate steps per expansion step. Each evaluation involves encoding a prompt containing the partial trace, all prior step rewards, the question, and (when search is triggered) 64 retrieved documents. The total PRA inference cost across a complete beam search trajectory is plausibly an order of magnitude larger than the policy's generation cost, but the paper's 80.8% headline number is presented as matching a "64-trace" Self-Consistency budget — creating a misleading impression of computational parity. A practitioner choosing between PRA and simply running Self-Consistency with more samples (e.g., 128 or 256 traces) cannot make an informed decision because the true cost of PRA is unmeasured.

What evidence exists in the paper: No experiment quantifies PRA inference FLOPs, latency, or retrieval overhead. Table 1 matches "64 traces" between PRA and SC baselines but only equalizes the policy model's generation count. The stage-level batching description (Appendix B, Figure 6) confirms that PRA inference and retrieval are additional, unbudgeted operations. The search–accuracy trade-off experiment (Figure 3) explores reducing retrieval frequency but does not convert search frequency into actual latency or FLOPs savings — the x-axis is "% Search Frequency," not milliseconds or FLOPs. For a deployment-relevant comparison, one would need to know whether PRA at 64-traces-equivalent is faster or slower than, say, RAG + SC at 256 traces, and the paper provides no basis for answering this.

Mitigation status: Not addressed. The paper treats the policy generation budget as the sole unit of comparison and does not discuss the practical cost of running the reward agent or retriever. The selective retrieval experiments (Figure 3) show that search frequency can be reduced while maintaining accuracy, which is a step toward cost-awareness, but without quantifying what "100% search frequency" costs in absolute terms, the practical significance of reducing it to 70% is unknown. No future work is proposed on reducing PRA inference cost (e.g., model distillation, caching, speculative evaluation).

2. Hard Medical Questions Remain Effectively Unsolved — PRA Cannot Compensate for Fundamental Capability Gaps

The assumption or constraint: PRA operates by scoring and selecting among candidate reasoning steps proposed by the frozen policy. It can rank, prune, and guide, but it cannot generate steps that the policy model is incapable of producing. If the policy model lacks the medical knowledge or reasoning capability to produce a correct step — even as one candidate among many — PRA has no mechanism to create it. The paper's framework assumes that correct reasoning lies within the policy's generation distribution and that the challenge is identifying it among incorrect alternatives, not enabling it.

The consequence: On questions where the policy model's standalone performance is very low, PRA cannot help. This is directly visible in two places. First, the cross-model generalization results (Table 2) show that while PRA provides large relative improvements for small models (e.g., 28.4% → 54.1% for Qwen2.5-0.5B, a +90.5% relative gain), the absolute PRA accuracy (54.1%) remains far below what larger models achieve with PRA (80.1% for Llama-3.1-8B). PRA reveals underutilized capacity but cannot create capacity that doesn't exist — the 0.5B model is fundamentally limited, and PRA only partially compensates.

Second, on the hardest individual questions within each benchmark, PRA's advantage over baselines likely shrinks or disappears. The paper does not report accuracy binned by question difficulty (analogous to the difficulty-quintile analysis in Snell et al., 2024, which the paper cites in Section 2.2), so this cannot be verified from the reported data. But the structural limitation is clear: if the policy's pass@1 for a question is near zero — no correct reasoning trace exists in its generation distribution even with 64 independent samples — then no amount of beam search guidance will produce a correct answer, because there is no correct trace to guide toward.

The margin-shift analysis (Figure 4) provides indirect evidence: incorrect traces show decreasing margin shift at later steps, meaning the teacher model can detect errors without external evidence — but detecting an error is not the same as proposing a correction. PRA can recognize that a trace is going off-track, but if the policy cannot generate a corrective step when the beam is pruned and re-expanded, the search will still converge to an incorrect answer.

What evidence exists in the paper: The cross-model generalization table (Table 2) shows the absolute accuracy ceiling varying dramatically with policy model scale — PRA cannot close the gap between a 0.5B and an 8B model. The paper does not analyze performance stratified by question difficulty within benchmarks, so there is no direct evidence of PRA's failure mode on the hardest questions. The margin-shift analysis (Figures 4-5) indicates that incorrect traces are detectable (margin shift is smaller for incorrect traces, especially on harder questions in Figure 5), but detectability of errors is not equivalent to correctability.

Mitigation status: Partially acknowledged. The paper frames PRA as a method for "steering" and "revealing underutilized reasoning capacity" (Section 6), which implicitly recognizes the boundary — PRA works when capacity exists but is underutilized. However, the paper never explicitly states that PRA cannot help on questions where the policy has zero probability of generating a correct trace, nor does it characterize what fraction of benchmark questions fall into this regime for different policy models. The generalisation discussion (Section 5, Table 2 text) notes that "substantial gains in reasoning performance can be achieved by more effectively exploiting the latent capabilities of existing models," but "latent capabilities" is not formally defined — the paper does not attempt to estimate the upper bound of what PRA could achieve with infinite search budget, which would quantify the "unutilized capacity" ceiling.

3. The PRA Is Trained on a Single Teacher Model's Judgments, Not on Ground-Truth Step Correctness

The assumption or constraint: All PRA training labels — both reasoning labels (correct/incorrect) and search labels (retrieval necessary/unnecessary) — come from a single teacher model, Qwen3-235B-Instruct. The teacher evaluates each reasoning step given the partial trace, retrieved documents, and the ground-truth final answer, then emits a binary judgment. The PRA is trained to approximate these judgments via distillation. The paper never validates the teacher's labels against human expert judgments or against an alternative teacher model. The quality of the entire PRA training signal rests on an untested assumption: Qwen3-235B-Instruct's step-level evaluations are sufficiently accurate to serve as ground truth.

The consequence: If the teacher model has systematic biases in its evaluation of medical reasoning, those biases are distilled directly into the PRA. For example, the teacher might:

Overrate plausible-sounding but incorrect reasoning — common in medical domains where multiple plausible diagnostic pathways exist, only one of which is correct given the specific case details
Underrate correct but non-standard reasoning — penalizing reasoning that reaches the right answer through an unconventional but valid chain of inference
Be miscalibrated on edge cases — medical questions where the correct answer depends on nuanced clinical details that even a 235B model may misinterpret
Exhibit position bias — systematically rating later steps as more or less likely to be correct independent of their content, which could affect the cumulative reward used for beam search ranking

The consequence for beam search is particularly concerning: if the teacher (and hence the PRA) systematically assigns higher scores to a specific style of reasoning that the teacher prefers, beam search will select for that style regardless of whether it correlates with actual correctness. This could produce a form of reward hacking where the search finds traces that score highly under the PRA but are not actually more likely to be correct — analogous to the verifier over-optimization documented in Snell et al. (2024) for mathematical reasoning.

The paper's result that outcome-level PRA (75.7%, Table 4) barely outperforms Self-Consistency (74.8%) is consistent with this concern: if the PRA's step-level scores genuinely reflected step correctness, even post hoc aggregation of those scores should substantially outperform majority voting over unweighted traces. The fact that it doesn't suggests that PRA's scores may contain limited information beyond what majority voting already captures — and that the online beam search benefit (+5.1 points) comes more from the search architecture (maintaining multiple hypotheses, pruning early) than from the quality of the reward signal itself.

What evidence exists in the paper: The paper does not validate the teacher's labels against any external standard. There is no human evaluation, no comparison to an alternative teacher model (e.g., GPT-4, Claude, Med-PaLM), and no analysis of teacher calibration or bias. The training data construction (Section 4.1) describes the teacher prompting procedure but does not report teacher accuracy, inter-annotator agreement (if multiple teachers or runs were used), or confidence calibration. The paper relies entirely on the assumption that "Qwen3-235B-Instruct" is a sufficiently capable medical evaluator.

The ablation in Table 3 provides indirect evidence: the untrained Qwen3-4B (74.4% with beam search, no retrieval) underperforms the trained PRA (80.8%), confirming that training on teacher labels adds signal beyond the base model's priors. But this does not tell us whether the teacher's labels are correct — only that they are different from the base model's priors and that learning them improves performance on the MedQA test set. The improvement could reflect learning the teacher's systematic biases, which happen to correlate with MedQA answer correctness without reflecting genuine step-level evaluation quality.

Mitigation status: Not addressed. The paper does not acknowledge the absence of teacher label validation as a limitation. No human evaluation, no alternative teacher comparison, and no analysis of teacher agreement with ground-truth step correctness (which would require human-annotated step labels for a subset of the data).

4. Single-Domain Training and Evaluation — Generalization Beyond Medicine Is Unsupported

The assumption or constraint: The PRA is trained exclusively on the MedQA training split (10,178 medical multiple-choice questions) and evaluated on six additional medical multiple-choice benchmarks (MedBullets, MedMCQA, MMLU-Med, GPQA, Lancet, NEJM). The knowledge base consists of medical corpora (textbooks, StatPearls, clinical guidelines, rare disease corpus). The retriever is MedCPT, a biomedical-domain dense retriever. Every component — training data, teacher model's domain expertise, knowledge base, retriever — is specialized to medicine.

The consequence: The paper presents PRA as a general framework for "knowledge-intensive reasoning" and suggests that it "suggests a paradigm in which frozen reasoners are decoupled from domain-specific reward modules" (Section 6, Abstract). But the experiments provide no evidence that this paradigm works outside of medicine. Several aspects of the approach may be medicine-specific:

The teacher model (Qwen3-235B) was presumably trained on substantial medical data, and its ability to evaluate medical reasoning steps may not transfer to other knowledge-intensive domains (law, finance, scientific research outside biomedicine) where the teacher's expertise is weaker
The margin-shift criterion for selective retrieval depends on the teacher's ability to evaluate steps both with and without external evidence — if the teacher lacks domain expertise, both evaluations may be noisy, and the margin shift may not capture meaningful retrieval necessity
The knowledge base structure (curated medical corpora with a specialized dense retriever) represents a best-case scenario for retrieval quality that may not be replicable in domains with less organized knowledge sources
Medical multiple-choice questions have a specific structure (patient case → question → 4-5 answer options) that the PRA prompt template is designed around — adapting to open-ended generation, short-answer questions, or domains without multiple-choice structure would require prompt engineering and possibly architectural changes

The claim that PRA "generalizes to unseen, frozen policy models" (Abstract) is well-supported for medical policy models specifically (Table 2), but the paper does not test whether a medically-trained PRA could guide reasoning in a non-medical domain. The broader claim of cross-domain generalization — implicit in the "paradigm" language — is entirely untested.

What evidence exists in the paper: All seven evaluation benchmarks are medical multiple-choice question answering datasets (Table 1). The out-of-distribution generalization is across medical benchmarks (different question sources, different difficulty profiles), not across domains. The paper does not evaluate on any non-medical reasoning benchmark (e.g., legal reasoning, financial QA, general science, history). The "knowledge-intensive domains" framing in the abstract and introduction is broader than the experimental evidence supports.

Mitigation status: Partially acknowledged by scope. The paper consistently describes its contributions in terms of "medical reasoning" and "knowledge-intensive domains such as medicine" (Introduction), and the experimental section is explicitly bounded to medical benchmarks. The broader "paradigm" claim in the abstract and conclusion is aspirational rather than empirically supported, which is a common but meaningful gap. The paper does not propose future work on non-medical domain adaptation or speculate on what would need to change for PRA to transfer to law, finance, or general science.

The assumption or constraint: The PRA is initialized from Qwen3-4B-Instruct and fine-tuned on reasoning traces from Qwen3-4B-Instruct. The policy model in the main experiments is also Qwen3-4B-Instruct. The cross-model generalization experiments (Table 2) test the PRA on policies from two model families: Qwen (Qwen3-4B, Qwen2.5-3B, Qwen2.5-0.5B) and Llama (Llama-3.1-8B, Llama-3.2-3B, Llama-3.2-1B). While the Llama results demonstrate cross-architecture generalization, there is an important subtlety the paper does not address: all tested models are instruction-tuned chat models (all have the "-Instruct" suffix), and the PRA was trained on similarly structured, instruction-tuned outputs. The PRA has never seen reasoning traces from base (pre-instruction-tuning) models, models fine-tuned for different tasks, or models with substantially different output formats.

The consequence: The claim that PRA "operates purely at inference time through step-wise selection within beam search... without altering the generation procedure, injecting additional context, or updating model parameters" (Section 4.2) is technically true but elides an important dependency: the PRA was trained to evaluate a specific style of reasoning — structured, step-by-step traces with explicit step numbering and a standardized answer format — enforced by the policy prompt (Appendix Figure 7). If a policy model produces reasoning in a different format — no step markers, implicit reasoning chains, different answer formatting — the PRA might fail not because the reasoning is incorrect but because the trace structure is out-of-distribution for the reward model.

This is a subtle form of distribution shift that differs from the "different model, same task" shift that Table 2 addresses. Table 2 shows that the PRA works on Llama models, but all Llama models were prompted with the same structured policy prompt that enforces "Step number:" formatting. We don't know whether the PRA would work on a Llama model prompted differently, on a base model outputting free-form reasoning, or on a model fine-tuned to produce a different reasoning style. The cross-model generalization result, while impressive, is bounded by the shared prompt template — it demonstrates robustness to model-specific distribution shift (different tokenizers, different pretraining distributions, different parameter counts) but not to format-level distribution shift.

Additionally, all tested policy models are in the 0.5B–8B range. The paper's title and framing emphasize enabling "frozen reasoners" and "decoupling from domain-specific reward modules," which suggests applicability to much larger models (e.g., 70B, 405B, or frontier models) that a practitioner might actually want to deploy without retraining. But if the PRA (a 4B model) is used to guide a 70B or 405B policy, the reward model becomes the bottleneck — it may lack the medical expertise to evaluate reasoning steps that the larger model produces, particularly on questions where the larger model's reasoning sophistication exceeds the PRA's evaluation capability. The paper's cross-model generalization stops at 8B, leaving the scaling behavior for the asymmetric case (small reward model, large policy) unexplored.

What evidence exists in the paper: Table 2 demonstrates cross-model generalization within the 0.5B–8B range and across two model families, but all models use the same structured policy prompt. The paper does not ablate the policy prompt format to test whether PRA's evaluation quality depends on the structured step-by-step formatting. No experiments test PRA with base models, differently prompted models, or models larger than 8B.

Mitigation status: Not addressed. The paper does not discuss the dependence on prompt format, the potential failure mode of format-level distribution shift, or the scaling behavior when the PRA is substantially smaller than the policy it guides. The cross-model generalization result is presented as evidence of model-agnosticism without the caveats about prompt structure and scale asymmetry.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a conceptual reorientation rather than a new model architecture or training recipe. The shift is from treating verification as post hoc audit to treating it as online control — a change in when reward signals are applied, not just how accurate they are. This is not a paradigm overthrow on the scale of "attention replaces recurrence," but it is more than an incremental refinement. It addresses a specific structural limitation that has been hiding in plain sight across the process reward model literature: PRMs have historically been temporally decoupled from generation, evaluating traces after completion rather than guiding them during construction. The paper's ablation (Table 4) demonstrates that this temporal coupling — not reward model accuracy, not aggregation strategy — accounts for the majority of observable gain, with online process-level rewards (80.8%) outperforming the best post hoc variant of the same model (77.6%) by 3.2 points and the outcome-level variant (75.7%) by 5.1 points.

This reframing has a diagnostic function for the field. It explains why prior medical PRM work (Med-PRM, Yun et al., 2025; Med-S3, Jiang et al., 2025) achieved only modest improvements over outcome-level baselines despite investing in sophisticated retrieval-augmented reward models: they applied those rewards post hoc, leaving the dominant source of gain — online guidance — untapped. The paper does not argue that Med-PRM or Med-S3 built bad reward models; it argues they deployed good reward models in a suboptimal temporal architecture. This diagnosis, if validated across domains, redirects research effort: the priority for test-time compute is not incrementally better scoring functions but infrastructure for online reward deployment — beam search, tree search, and other frameworks that let rewards steer generation rather than merely describe it.

The paper also reconciles a tension between two lines of evidence in the PRM literature. On one hand, process-level rewards are theoretically superior to outcome-level rewards because they provide denser signal and can identify where reasoning goes wrong (Lightman et al., 2023; Wang et al., 2023). On the other hand, empirical gains from process-level rewards in post hoc settings have been surprisingly modest — sometimes negligible over outcome-level baselines. The paper's Table 4 resolves this: process-level rewards do provide useful signal (post hoc process-level with average aggregation achieves 77.6%, outperforming post hoc outcome-level at 75.7%), but most of their value is only realized when they are applied online (80.8%). The modest post hoc gains observed in prior work are not evidence against process rewards — they are evidence that post hoc application wastes most of the value. This is a satisfying resolution because it explains both the theoretical appeal of process rewards and the empirical disappointment, without requiring either side to be wrong.

A less obvious but equally important shift concerns modularity as a deployment strategy. The paper demonstrates that a reward agent trained on traces from one policy model (Qwen3-4B) can guide six different models spanning two families and a 16× parameter range (Table 2), with gains of +8.1 to +25.7 points. This inverts the prevailing assumption that reward models must be policy-specific or co-trained with the policy they evaluate. The implication is that domain adaptation — keeping reasoning grounded in current medical knowledge — can be refactored from "retrain every policy model when knowledge updates" to "retrain one reward agent when knowledge updates, and use it to guide any frozen policy." This is not just a performance claim; it is an architectural claim about how to organize AI systems in knowledge-intensive domains where the knowledge base evolves independently of reasoning capability. The paper does not fully deliver on this vision — the PRA was only tested on medical benchmarks, and the 4B reward agent may not scale to guiding much larger policies — but the cross-model generalization result provides the first concrete evidence that such modularization is feasible.

Finally, the paper redirects attention toward the reward deployment architecture as the primary bottleneck in test-time compute scaling. The inference-time scaling curves (Figure 2) show PRA continuing to benefit from additional compute while Self-Consistency saturates by 8 samples. This is not because PRA has a better reward model (the ablation in Table 4 uses the same reward model in both post hoc and online settings) but because online rewards enable the system to act on information (pruning unpromising paths) rather than merely observe it. This finding suggests that work on search infrastructure — how rewards are integrated into generation, what search algorithms are used, how pruning decisions are made — is higher-leverage than work on incrementally more accurate reward models, at least at current reward model quality levels. This is directly analogous to how the reinforcement learning community recognized that reward design matters less than exploration strategy in many settings, and it opens a parallel research agenda for inference-time compute.

Follow-Up Research This Work Enables

Stress-test cross-domain generalization by training a PRA on medical data and evaluating on legal or scientific reasoning benchmarks. The paper demonstrates that a medically-trained PRA generalizes across medical benchmarks (Table 1) and across policy models (Table 2), but the broader claim that PRA enables "knowledge-intensive reasoning" without domain qualification remains untested. A strong follow-up would take the same MedQA-trained PRA and evaluate it on legal reasoning (e.g., LegalBench, CaseHOLD) or general scientific reasoning (e.g., ScienceQA, BioASQ outside medicine), measuring whether the PRA's step-level evaluation capability transfers to domains where the teacher model (Qwen3-235B) has less specialized expertise. A null result — PRA fails to improve non-medical reasoning — would clarify that the approach requires domain-specific PRA training and is not a general-purpose verification module. A positive result — even modest gains on out-of-domain benchmarks — would substantially strengthen the "paradigm" claim. The critical measurement is not just final answer accuracy but whether PRA's step-level scores correlate with actual step correctness in the new domain (requiring human-annotated step labels for a subset).

Quantify the PRA's inference overhead in FLOPs and wall-clock time relative to Self-Consistency at matched accuracy. The paper matches policy generation budgets (64 traces) between PRA and baselines but does not account for the cost of running the 4B PRA model (with retrieval) at every expansion step. A deployment-relevant comparison would measure: (1) total FLOPs for PRA-guided beam search at B=4, b=16, including PRA inference and retrieval latency; (2) total FLOPs for Self-Consistency at N=64, 128, 256, 512 traces; and (3) the crossover point where the two methods achieve equal accuracy at equal total cost. The hypothesis is that at low budgets PRA is more expensive per unit accuracy (due to PRA overhead), but at high budgets PRA's superior scaling (Figure 2) makes it more cost-effective. The paper's stage-level batching system (Appendix B) provides the infrastructure for this measurement, but the numbers are absent. A follow-up would also ablate PRA model size — fine-tune a 0.5B or 1B PRA using the same teacher labels and measure whether guidance quality degrades, establishing the minimum viable reward agent size for a given policy scale.

Validate the teacher model's step-level labels against human expert judgments. The entire PRA training pipeline depends on Qwen3-235B-Instruct's step-level evaluations, which are never validated. A follow-up would sample 200–500 reasoning steps from MedQA training traces (stratified by step position, difficulty, and whether the final answer is correct), have 2–3 board-certified physicians independently label each step as correct/incorrect, compute inter-annotator agreement, and then measure the teacher model's agreement with the human consensus. If the teacher's agreement is high (e.g., Cohen's κ > 0.7), the distillation pipeline is validated. If agreement is low or systematically biased (e.g., the teacher overrates certain reasoning styles), the paper's 80.8% result may partly reflect learning the teacher's idiosyncrasies rather than genuine step-level evaluation quality. This experiment would also reveal whether the teacher's errors correlate with question difficulty or reasoning step position, informing where PRA guidance is most reliable.

Test whether PRA-guided search exhibits verifier over-optimization at high search budgets, and whether ensembling or adversarial training mitigates it. The paper's inference-time scaling curve (Figure 2) shows PRA accuracy still rising at 64 traces, but the curve has not asymptoted — we don't know whether PRA performance eventually saturates or degrades due to over-optimization (the phenomenon where search finds traces that score highly under the reward model but are actually incorrect, documented for mathematical reasoning in Snell et al., 2024). A follow-up would push PRA beam search to B=8, b=32 (256 candidate expansions per step) or higher, measuring whether accuracy plateaus or declines. If over-optimization occurs, test two mitigations: (1) training an ensemble of PRAs (e.g., 3 independently fine-tuned Qwen3-4B models) and averaging their scores, which should be more robust than a single PRA; (2) adversarial training where the PRA is fine-tuned on traces generated by beam search itself (on-policy data) rather than on the original Qwen3-4B training traces (off-policy data), which should reduce distribution shift between PRA training and inference. The key measurement is the accuracy ceiling and whether it can be raised through ensemble or on-policy training.

Replace the teacher model distillation pipeline with reinforcement learning from ground-truth final answer correctness. The paper uses a 235B teacher model to generate step-level labels, which is expensive (requires running a 235B model on every training step), introduces potential teacher bias (as discussed in limitation 3), and does not directly optimize for the objective that matters — final answer correctness. A follow-up would train the PRA using reinforcement learning, where the reward signal is the correctness of the final answer reached by beam search guided by the PRA's step-level scores. This is analogous to how Math-Shepherd (Wang et al., 2023) used Monte Carlo rollouts rather than teacher judgments, but applied online within beam search. The PRA would be updated to maximize the expected correctness of beam search outcomes, with the search itself as the environment. This would align the PRA's training objective with its deployment objective, potentially producing a reward model that is more robust to over-optimization (since it was trained under optimization pressure) and that does not depend on a large teacher model. The experiment would compare PRA trained via teacher distillation versus PRA trained via RL from final answer correctness, measuring both MedQA accuracy and—critically—generalization to unseen policy models (does RL-trained PRA still transfer across policies?).

Develop continuous difficulty estimation and dynamic strategy allocation, extending the compute-optimal framework of Snell et al. (2024) to knowledge-intensive domains with retrieval. The paper demonstrates that PRA-guided search improves over Self-Consistency, but it does not explore adaptive allocation of the search budget — all questions receive the same beam width and branching factor. A follow-up would use the PRA's own step-level scores (or their variance, or the margin-shift estimates from the selective retrieval variant) to estimate question difficulty early in the search process, then dynamically adjust the search budget. Easy questions (where the beam quickly converges on a high-confidence answer) could terminate early, saving compute. Hard questions (where beam traces diverge or score low) could be allocated additional budget or escalated to a larger policy model. This connects the PRA framework to the compute-optimal test-time scaling paradigm (Snell et al., 2024) and would measure whether difficulty-adaptive PRA achieves higher accuracy at matched total cost than uniform-budget PRA. The paper's margin-shift analysis (Figures 4-5) already shows that difficulty-correlated signals exist in the PRA's outputs, making this a natural extension.

Practical Applications and Downstream Use Cases

Cost-efficient deployment of small models in medical question-answering systems. The paper shows that a 0.5B-parameter model guided by PRA achieves 54.1% MedQA accuracy, compared to 28.4% for the same model alone — a 90.5% relative improvement (Table 2). For organizations deploying medical QA systems on edge devices or with strict latency budgets where large models are infeasible, pairing a small policy model with a PRA (which can run on a separate server or be quantized for on-device use) could provide accuracy competitive with much larger standalone models at a fraction of the parameter cost. The PRA's cross-model generalization means the same reward agent can serve multiple policy models, amortizing the PRA's training and inference cost across deployment scenarios. The selective retrieval variant (Figure 3) further enables trading off retrieval cost against accuracy on a per-deployment basis — a clinic with limited bandwidth could use a higher retrieval threshold (lower search frequency) while a research hospital with dedicated infrastructure could use always-search for maximum accuracy.

Updatable medical AI systems that track evolving clinical knowledge without retraining the reasoning backbone. Medical guidelines, drug interactions, and diagnostic criteria change continuously. The paper's decoupling of the reward agent from the policy model means that when new clinical evidence emerges, only the PRA needs to be updated — retrained or fine-tuned on traces reflecting the new knowledge — while the reasoning model (which may be a large, expensive-to-retrain backbone) remains frozen. In practice, a hospital system could deploy a single large reasoning model (e.g., Llama-3.1-8B) and periodically update a much smaller PRA (e.g., 4B) as new guidelines are published, without ever modifying the reasoning model. The PRA's retrieval mechanism ensures that even between PRA updates, the latest evidence is incorporated at inference time through document retrieval. This addresses a genuine operational bottleneck identified in the paper's motivation: "it becomes prohibitively expensive to repeatedly fine-tune each new LRM backbone to remain aligned with updated medical knowledge."

Self-improvement pipelines for generating high-quality medical reasoning training data. The paper demonstrates that PRA-guided beam search produces higher-accuracy reasoning traces than the policy model alone (80.8% vs. 72.7% on MedQA). These PRA-guided traces could be used as training data to fine-tune the policy model itself, creating a self-improvement loop: policy generates traces → PRA guides search to select the best traces → selected traces are used to fine-tune the policy → the improved policy generates better traces → repeat. Unlike prior self-improvement approaches that rely on final-answer correctness (which requires ground-truth answers or majority voting), PRA provides step-level quality signals that could enable more granular filtering — traces with consistently high step-level PRA scores could be used for training even if they don't achieve majority consensus, increasing the volume and diversity of training data. A particularly promising variant uses the PRA's own per-step scores as a curriculum: early iterations of self-improvement use traces where all steps score highly (easy cases), later iterations include traces with some low-scoring early steps but high-scoring late steps (cases where the model recovered from initial errors — teaching error correction).

When to Prefer This Method

The paper positions PRA against three categories of alternatives, and the experimental results support specific decision rules:

Prefer PRA over standard RAG (retrieved documents in policy context) when: (1) the policy model is small relative to the complexity of the knowledge base, making attention dilution a significant concern — the paper shows RAG + SC achieving 76.7% while PRA achieves 80.8% on MedQA (Table 1); (2) you need to deploy multiple different policy backbones without maintaining separate retrieval-augmented versions of each — PRA's cross-model generalization (Table 2) avoids the N×M matrix of policy-model × retrieval-configuration variants; (3) you want step-level verifiability of reasoning rather than just final-answer accuracy — PRA's online rewards enable beam search pruning and trace selection that RAG's context injection does not provide.
Prefer PRA over post hoc process reward models (e.g., Med-PRM) when: (1) you have the infrastructure to run beam search with online rewards (the stage-level batching system in Appendix B) — the +3.2 point gap between online and post hoc process-level PRA (Table 4) only materializes if you can deploy the reward model during generation; (2) your policy model produces diverse candidate steps (branching factor b > 1) — post hoc PRMs see only one trace per question and cannot exploit step-level diversity; (3) latency is not the binding constraint — online PRA adds serial dependencies (policy generation → PRA scoring → pruning → next step) that post hoc scoring avoids.
Prefer PRA over direct fine-tuning of the policy model (SFT/RLVR) when: (1) the knowledge base updates frequently — PRA retraining is cheaper than policy retraining because the PRA can be smaller and trained on fewer examples (the paper trains on ~10K MedQA traces with a 4B model); (2) you need to deploy multiple policy models and want a single verification module — Table 2 shows one PRA serving six policies; (3) you cannot modify the policy model (e.g., API-only access to a proprietary model, or a model deployed on read-only hardware) — PRA operates entirely through inference-time selection without parameter updates.
Prefer Self-Consistency or RAG + SC over PRA when: (1) the policy model is already very strong (standalone accuracy near saturation) and the marginal gain from PRA's search guidance is small — the paper shows PRA providing +4.1 points over RAG + SC on MedQA (Table 1), but on MMLU-Med the gap narrows to +0.4 points (86.6% vs. 86.2%), suggesting diminishing returns as baseline performance rises; (2) the knowledge base is small or well-covered by the policy model's parametric knowledge, making retrieval unnecessary — the ablation shows untrained PRA without retrieval (74.4%, Table 3) is comparable to CoT + SC (74.8%), so retrieval is the source of most of PRA's advantage; (3) inference latency is absolutely critical and the PRA's scoring overhead cannot be amortized through batching — Self-Consistency can run fully parallel while PRA beam search has sequential dependencies within each trace.

These decision rules are bounded by the paper's experimental scope (medical multiple-choice QA, 0.5B–8B policy models, a single PRA backbone) and should be validated on other domains and scales before being treated as general principles.