PaperlayerView exact revision on arXiv

Source-linked AI summary

Med-PRM: Medical Reasoning Models with Stepwise, Guideline-verified Process Rewards

Jaehoon Yun, Jiwoong Sohn, Jungwoo Park, Hyunjae Kim, Xiangru Tang, Yanjun Shao, Yonghoe Koo, Minhyeok Ko, Qingyu Chen, Mark Gerstein, Michael Moor, Jaewoo Kang

arXiv:2506.11474v2cs.CL

TL;DR

Medical process reward models struggle to localize reasoning errors because step-level supervision is costly and outcome-based labels can misjudge valid intermediate steps. Med-PRM retrieves clinical evidence and uses it to verify each reasoning step, improving benchmark performance and supporting expert-aligned, plug-and-play verification. The method improves base models by up to 13.50% and reaches 80.35% on MedQA with 8B-parameter models.

  • Problem

    Medical PRMs need scalable, evidence-grounded step-level supervision because human labels are costly and outcome-based auto-labeling can misclassify valid intermediate reasoning.

    Source: §1 · p. 1

  • Method

    Med-PRM uses retrieval-augmented generation and an LLM-as-a-Judge to evaluate each reasoning step against retrieved medical documents.

    Source: §1 · p. 2

  • Results

    Med-PRM achieves state-of-the-art performance across six of seven medical QA benchmarks, improves base models by up to 13.50%, and reaches 80.35% on MedQA with 8B-parameter models.

    Source: §1 · p. 2, §6.1 · p. 5

  • Takeaways & Limitations

    Retrieval-based step verification provides scalable, expert-aligned labels for policy training and test-time scaling across multiple medical reasoning tasks.

    Source: §1 · p. 2, §8 · p. 9

    Evaluation is confined to medicine, experiments focus on small models such as Llama 3.1 8B and Meerkat 8B, and diverse reinforcement-learning methods were not extensively explored.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Large language models have shown promise in clinical decision making, but current approaches struggle to localize and correct errors at specific steps of the reasoning process. This limitation is critical in medicine, where identifying and addressing reasoning errors is essential for accurate diagnosis and effective patient care. We introduce Med-PRM, a process reward modeling framework that leverages retrieval-augmented generation to verify each reasoning step against established medical knowledge bases. By verifying intermediate reasoning steps with evidence retrieved from clinical guidelines and literature, our model can precisely assess the reasoning quality in a fine-grained manner. Evaluations on five medical QA benchmarks and two open-ended diagnostic tasks demonstrate that Med-PRM achieves state-of-the-art performance, with improving the performance of base models by up to 13.50% using Med-PRM. Moreover, we demonstrate the generality of Med-PRM by integrating it in a plug-and-play fashion with strong policy models such as Meerkat, achieving over 80\% accuracy on MedQA for the first time using small-scale models of 8 billion parameters. Our code and data are available at: https://med-prm.github.io/

1 Introduction

Clinical decision making requires stepwise reasoning grounded in evolving clinical evidence, but existing medical PRMs face costly supervision and unreliable outcome-based labeling. Med-PRM addresses these challenges with retrieval-augmented step verification and reports strong benchmark, expert-alignment, and cost-efficiency results.

1 Introduction

  • Outcome-based auto-labeling can undervalue logically sound early steps when sampled continuations fail, distorting intermediate reasoning supervision.Source: §1 · pp. 1–2
  • Clinical reasoning requires domain knowledge beyond model parameters, making retrieved guidelines and medical evidence important for factual, non-hallucinatory stepwise rewards.Source: §1 · pp. 1–2
  • The approach emphasizes cost-efficient scalability, using a curated reward-model dataset costing less than $20 compared with approximately $20,000 for UltraMedical training data.Source: §1 · p. 2
  • Med-PRM evaluates each reasoning step using the clinical question and retrieved evidence, enabling fine-grained, evidence-grounded assessment.Source: §1 · p. 2
  • Med-PRM improves base-model performance by up to 13.50% at test time and reaches 80.35% on MedQA using only 8B-parameter models.Source: §1 · p. 2
  • Its verifier aligns closely with clinical experts in logical consistency and factual accuracy, while plug-and-play use with policy models yields state-of-the-art performance.Source: §1 · p. 2

2 Related Work

Related work advances medical reasoning through large language models, prompting, agents, reinforcement learning, and process rewards. Med-PRM differs from prior MCTS-based and similarity-retrieval approaches by retrieving medical evidence for LLM-judged stepwise supervision.

2 Related Work

  • Medical language models use corpora such as PubMed, MIMIC-III/IV, and UMLS to support inference beyond surface-level recall.Source: §2 · p. 2
  • CoT prompting, agentic frameworks, reinforcement learning, verifier feedback, and PRMs have increasingly been used to refine medical reasoning traces.Source: §2 · p. 2
  • MedS3 uses MCTS-based auto-labeling, whereas Med-PRM combines process-level rewards with retrieval-augmented generation and an LLM-as-a-Judge framework.Source: §2 · p. 2
  • Med-PRM’s RAG-AS-A-JUDGE retrieves medical knowledge and evidence, including textbooks or clinical guidelines, rather than relying on automatic scoring or similar QA pairs.Source: §2 · pp. 2–3

3 Preliminaries

Reward models score either complete reasoning traces or individual steps, while auto-labeling makes process supervision scalable but can misclassify valid intermediate reasoning.

3.1 Reward Model

  • Outcome reward models score an entire reasoning trace with a sigmoid correctness value trained by cross-entropy.The gold trace label is yS = 1 for correct traces and yS = 0 otherwise.Source: §3.1 · p. 3
  • Process reward models assign confidence scores to individual reasoning steps and use the minimum step score as the trace-level reward.Each step receives a correctness confidence, and the minimum across steps defines RM(q, S).Source: §3.1 · p. 3
  • MCTS-style auto-labeling estimates step quality from sampled continuations, enabling scalable PRM training without human supervision.The method generates multiple subsequent reasoning processes from each partial trace to support step labeling.Source: §3.1 · p. 3
  • Auto-labeling can produce false negatives when no sampled continuation reaches the correct answer, mislabeling factually correct and coherent intermediate steps.This motivates retrieval-based fact-checking during labeling.Source: §3.1 · p. 3

4 Method

RAG-AS-A-JUDGE labels each reasoning step by binary classification using the question, gold answer, and retrieved medical documents.

4.1 RAG-AS-A-JUDGE

  • The method uses retrieved evidence as context for judging intermediate reasoning rather than relying only on outcome-based labels.The supplied method description explicitly frames retrieved documents as inputs to stepwise evaluation.Source: §4.1 · p. 3
  • RAG-AS-A-JUDGE evaluates each reasoning step with binary classification conditioned on the question, gold answer, and retrieved documents.It produces step-level labels yRAG for determining whether each step is correct.Source: §4.1 · p. 3

1. Trace generation

Med-PRM retrieves evidence alongside generated reasoning traces, scores traces by their weakest step, and uses those scores for inference-time selection and policy fine-tuning.

1. Trace generation

  • The overview pipeline generates traces, retrieves relevant documents, assigns stepwise labels, trains Med-PRM, and applies it for inference-time evaluation or policy training.The figure presents these four stages as the overall Med-PRM workflow.Source: Figure 1 · p. 4
  • Med-PRM evaluates reasoning traces with retrieved documents during both training and evaluation, distinguishing it from PRMs that omit external evidence.The retrieval query concatenates the question with the reasoning trace; training uses the correct trace, whereas inference uses a randomly sampled trace.Source: §1 · p. 4
  • Best-of-N selects the trace with the highest minimum step reward, while SC+RM groups traces by final answer and selects the answer with the highest total reward.These strategies support test-time scaling by ranking multiple generated traces.Source: §1 · p. 4

4.2 Policy Model Fine-tuning

  • Policy fine-tuning retains traces with the highest minimum step scores for supervised training, then reuses Med-PRM to rescore generations at inference.Questions answered correctly consistently are excluded to focus training on more challenging examples.Source: §4.2 · pp. 4–5

5 Experimental Setup

Med-PRM is trained on mixed medical question-answering data and evaluated across multiple-choice and open-ended medical benchmarks with baseline, ablation, and scaling comparisons.

5.1 Training of Med-PRM

  • The model is fully fine-tuned from Llama-3.1-8B-Instruct on one NVIDIA A100 using a 4096-token maximum sequence length.Training uses AdamW, bfloat16 precision, gradient checkpointing, Flash Attention V2, and global batch size 64.Source: §5.1 · p. 5
  • Training uses MedQA, MedMCQA, PubMedQA, and MMLU data, with filtered reasoning traces and retrieved documents forming stepwise-supervision inputs.The setup uses all 10,178 MedQA training questions and 500 sampled instances from each remaining dataset.Source: §5.1 · p. 5

5.2 Evaluation of Med-PRM

  • Evaluation covers MedQA, MedMCQA, six medical MMLU subsets, DDXPlus, and two open-ended AgentClinic variants.AgentClinic uses an open-ended format and is evaluated with Gemini-2.0-flash.Source: §5.2 · p. 5
  • Baselines include proprietary, open-source, general-purpose, reasoning, medical, and process reward models across multiple-choice and open-ended benchmarks.Internal ablations compare PRMsoft, PRMhard, Med-PRM without retrieval, and Med-PRM with retrieval under Best-of-N and SC+RM.Source: §5.2 · p. 5
  • Test-time scaling increases generated reasoning traces to N = 64, with Best-of-N and SC+RM selecting answers using reward scores.Best-of-N uses the highest rSj score, while SC+RM uses the answer group with the highest summed reward.Source: §5.2 · p. 5

6 Results

Med-PRM outperforms prior models across medical benchmarks and improves test-time reasoning when used as a verifier or training signal. Its plug-and-play use with stronger policy models reaches 80.35% on MedQA with an 8B model.

6.1 Main Results

  • 72.59% average accuracy with Best-of-N and 73.50% with SC+RM outperform all existing open-source models below 10B parameters, with state-of-the-art results on 2 and 4 of 7 benchmarks, respectively.The benchmark covers multiple-choice and open-ended medical QA tasks.Source: §6.1 · p. 5, Table 1 · p. 6
  • 12.50% and 10.75% accuracy gains on AgentClinic* under SC+RM and Best-of-N show larger improvements on complex clinical reasoning than knowledge-centric tasks.AgentClinic* is designed to resemble real-world diagnostic workflows.Source: §6.1 · pp. 5–6
  • Med-PRM improves over MedS3 by 2.44% with SC+RM and 3.08% with Best-of-N on average across benchmarks.MedS3 is identified as the previous state-of-the-art process reward model at the 8B scale.Source: §6.1 · p. 6

6.2 Reward Model as Verifier

  • Med-PRM improves diverse policy models during inference, including a 78.24% Llama-3.1-8B-Instruct score with SC+RM and a 76.98% Best-of-N score.Table 2 compares average sampling, Self-Consistency, SC+RM, and Best-of-N over 64 solutions.Source: §6.2 · p. 6, Table 2 · p. 6
  • Med-PRM π reaches 79.18% on MedQA, while pairing Med-PRM with Meerkat-8B reaches 80.35%, the first reported result above 80% for an 8B-scale model.The reward model is used both to select high-quality training traces and as a plug-and-play verifier.Source: §6.2 · pp. 6–7

7 Analysis

Ablations, open-ended diagnosis, expert alignment, and case studies show that retrieval-grounded stepwise supervision improves robustness and localizes clinically meaningful reasoning errors.

7.1 Ablation Study

  • Adding retrieval further improves the retrieval-free Med-PRM variant, while every component produces gains across both test-time strategies.The ablation compares conventional auto-labeling, LLM-based step evaluation without retrieval, and the full framework.Source: §7.1 · p. 7
  • Med-PRM consistently outperforms Self-Consistency in both SC+RM and Best-of-N, whereas conventional PRMs underperform Self-Consistency in Best-of-N.The Best-of-N setting is more discriminative for comparing reward models because it selects only the top solution.Source: §7.1 · p. 7

7.2 Open-Ended Clinical Tasks

  • 11.81% accuracy improvement on unseen AgentClinic* demonstrates effective test-time scaling on an open-ended diagnostic benchmark, exceeding Self-Consistency by 4.87% and conventional PRMs by 4.32%.Open-ended diagnosis more closely resembles practical clinical settings than multiple-choice evaluation.Source: §7.2 · pp. 7–8

7.3 Expert Alignment

  • Med-PRM correlates with expert step-level judgments at 0.74 on easy cases and 0.71 on hard cases, while soft and hard labels fall to 0.34 and 0.31 on hard examples.The expert-alignment analysis uses 180 step-level annotations.Source: §7.3 · p. 8, Table 3 · p. 8

7.4 Case Study

  • In the Graves’ ophthalmopathy case, Med-PRM labels valid early steps as correct and later erroneous steps as incorrect, unlike ORM and conventional PRM pipelines.The retrieved document distinguishes autoimmune orbitopathy from sympathetic overactivity, enabling localization of the error rather than assigning one label to the whole trace.Source: §7.4 · pp. 8–9
  • In the polyhydramnios case, Med-PRM maintains high rewards through Step 4 and sharply lowers them from Step 5 after retrieved evidence identifies the PUV error.Conventional PRMs assign relatively high rewards to the erroneous later steps without helpful external documents.Source: §7.4 · p. 9, Figure 4 · p. 9

8 Conclusion

Med-PRM verifies intermediate medical reasoning against retrieved documents rather than relying on outcome-based auto-labeling. The resulting expert-aligned labels improve performance across multiple-choice and open-ended clinical benchmarks and support both policy training and test-time verification.

8 Conclusion

  • Med-PRM uses RAG-AS-A-JUDGE to verify intermediate reasoning steps against retrieved medical documents, addressing limitations of outcome-based auto-labeling.The framework is presented as a robust and generalizable approach for improving medical AI accuracy, transparency, and trustworthiness.Source: §8 · p. 9

Limitations

The paper reports limitations in domain coverage, model scale, and reinforcement-learning exploration, while describing a retrieval pipeline and related-work positioning for Med-PRM.

Limitations

  • The evaluation is confined to medicine, experiments use small models because of computational constraints, and diverse reinforcement-learning methods remain underexplored.The authors identify broader domains, model architectures and sizes, and more sophisticated reinforcement-learning strategies as future directions.Source: Limitations · p. 10

A Notation

B Retrieval

  • Med-PRM uses MedCPT bi-encoder and cross-encoder models for dense retrieval and reranking from four biomedical corpora: Clinical Guidelines, StatPearls, Medical Textbooks, and Rare Disease Corpus.Source: B Retrieval · p. 13
  • The retrieval implementation uses Milvus with Max Inner Product Search on AWS EC2 and reranks documents using an NVIDIA RTX 3090.The pipeline retrieves 400 documents per query before reranking.Source: B Retrieval · p. 13

D API Usage and Cost

  • Gemini-2.0-flash supplies RAG-AS-A-JUDGE labels and AgentClinic scoring through the Google Generative AI API at an approximately $20 training-set curation cost.API calls used temperature 0 and standard rate limits.Source: D API Usage and Cost · p. 13

E Related Works

  • Related work situates Med-PRM among medical reasoning, process rewards, LLM-as-a-Judge evaluation, and retrieval-augmented PRMs.Unlike mathematical RAG-PRM, Med-PRM retrieves medical evidence and knowledge rather than similar question-answer pairs.Source: E Related Works · pp. 13–14

F Benchmarks

The study evaluates Med-PRM across established medical knowledge benchmarks and more complex diagnostic settings, including reformulated and open-ended clinical tasks.

F Benchmarks

  • MedQA contains 12,730 English USMLE questions testing diagnosis, treatment, medical concepts, and clinical reasoning with four- or five-option formats.The study evaluates the English subset exclusively.Source: F Benchmarks · p. 14
  • MedMCQA contributes 500 four-option questions sampled from a dataset spanning more than 194,000 questions, 2,400-plus healthcare topics, and 21 medical subjects.The source questions come from Indian medical entrance examinations.Source: F Benchmarks · p. 14
  • The MMLU medical evaluation uses 1,089 four-option questions across anatomy, clinical knowledge, college medicine, medical genetics, and professional medicine.These subsets span basic through professional medical expertise.Source: F Benchmarks · pp. 14–15
  • DDXPlus is a large synthetic dataset of approximately 1.3 million patient cases with hierarchical symptoms, antecedents, differential diagnoses, and ground-truth pathologies.Its structure supports differential-diagnosis reasoning from richer clinical information than binary symptom datasets.Source: F Benchmarks · p. 15
  • AgentClinic is a multimodal, sequential clinical benchmark, but the study uses AgentClinic* as a single-step inference adaptation to reduce API calls and computational overhead.The adaptation preserves evidence-grounded clinical reasoning while balancing realism and tractability.Source: F Benchmarks · p. 15
  • Evaluation covers four medical QA datasets, six medical MMLU subsets, DDXPlus, and two AgentClinic variants requiring more complex clinical reasoning.DDXPlus uses differential diagnosis, while AgentClinic variants provide open-ended clinical scenarios.Source: F Benchmarks · p. 15

G Full Ablation Study of PRMs vs SC

Med-PRM outperforms MedS3 when scaling test-time computation, with evaluations combining trace selection, prompts for stepwise reasoning, and human step-level assessment.

G Full Ablation Study of PRMs vs SC

H Prompts

  • The prompts require multiple-choice solutions to use numbered reasoning steps and end with an explicit option-based answer phrase.Open-ended prompts similarly require numbered steps and a final diagnosis label.Source: H Prompts · p. 16
  • RAG-AS-A-JUDGE labels wrong medical solutions step by step using retrieved documents, the question, and the correct answer, assigning binary scores for reasoning quality.The evaluation requires a final score summary for each reasoning step.Source: H Prompts · p. 16

I Human Evaluation Details

  • Human evaluation uses one physician and two senior medical students, with disagreements reviewed and adjudicated by the physician.A step receives 1 for factual accuracy, problem-solving contribution, and coherent reasoning; critical errors propagate zero scores to dependent steps.Source: I Human Evaluation Details · p. 17
  • Human evaluators assess factual accuracy, problem-solving relevance, and logical coherence within the clinical context.Logical coherence requires reasoning grounded in appropriate medical knowledge and internally consistent clinical reasoning.Source: I Human Evaluation Details · p. 17

J Scaling Test-Time Computation with PRMs Across Multiple Models

  • Table 6 compares Llama 3.1 8B Instruct, Llama-3-8B-UltraMedical, and Med-PRM π under Self-Consistency and PRM-based Best-of-N and SC+RM strategies.The table caption states that Med-PRM achieved the highest average score across scaling methods.Source: Table 6 · p. 18

Related papers

Loading 2506.11474v2…