9.2.1 Med-PRM | AI Roadmap

🔗9.2.1 Med-PRM

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

TL;DR

Med-PRM is not just a reward model; it’s a full pipeline:

Step-wise, RAG-based reward modeling

Auto-labeling based on RAG

Trains policy model via SFT (data collected from rejection sampling), not RL

Usage mechanism:

Surprisingly, despite being a “reward model,” the Med-RPM is not used for reinforcement learning. Instead of encouraging exploration through RL like PPO/GRPO, Med-PRM is used along with rejection sampling for rewarding reasoning steps and collecting the SFT data, and only the best-rewarded output will be used for training the policy model using SFT.

Step-wise PRM:

Med-PRM scores reasoning steps, not individual tokens. Traditional PRM rewards at the token level and looks like this

Each token in the output will be assigned a reward score, which is usually done by adding a separate MLP header (reward header). When it comes to the Med-PRM, the reward is done in step-level: Each step is delimited with a <step> token, and the reward is assigned at the step level by applying the reward header on each <step> token.

RAG-based PRM

RAG is rooted in both training and inference. The purpose of RAG is this: Each reasoning step is verified against retrieved medical docs (e.g. PubMed, guidelines), which means the retrieved knowledge is used as part of prompts, so the PRM can give more accurate rewards.

During the training of Med-PRM, the reward model is given inputs in the form of Input = [Docs] + [Question] + [Trace]

During the inference of Med-PRM and the training of the policy model, the input for Med-PRM is still in the above form. Only care about the rewards derived from the <step> token in [Trace].

Auto-labeling w/ RAG

Instead of using MCTS, they use an LLM-as-a-Judge method (augmented with RAG as well) for labeling. Steps are labeled as correct/incorrect based on whether they align with retrieved evidence.

Med-PRM Training pipeline:

Generate reasoning traces using LLM

Retrieve evidence using RAG

Label each step (+/-) using RAG-as-a-Judge

Train Med-PRM using step-wise supervision at <step> markers

Med-PRM Inference/Policy Model Training pipeline:

Generate candidate reasoning traces using the policy model

Retrieve context using RAG

Use Med-PRM to score steps for each output/trace

Final score for an output/trace = min(step scores)

Select the best trace via Best-of-N （reward model） or SC+RM (Self-Consistency + Reward Model)

Use the collected data to train the policy model with SFT

CivitAI’s Payment Issue

Awesome-AI-Tutorials

Catalog