Source-linked AI summary
CDPR: Counterfactual Advantage-based Credit Assignment for Cost-Aware Sequential Medical Diagnosis
Qi Peng, Yi Cai, Changmeng Zheng, Xin Wu, Jiayuan Xie, Qing Li
TL;DR
Medical diagnosis requires sequential, cost-aware examination decisions, but one-shot models and outcome-only rewards do not distinguish efficient from wasteful diagnostic trajectories. The paper proposes CDPR, which assigns counterfactual process rewards at uncertain states and integrates them with GRPO. Across one in-domain and two out-of-domain benchmarks, CDPR improves diagnostic accuracy while reducing examination burden and cost.
Problem
Most medical language models treat diagnosis as one-pass classification, while outcome-only rewards cannot distinguish efficient from wasteful sequential diagnostic trajectories.
Method
CDPR selects high-uncertainty states and compares chosen actions with counterfactual alternatives through short rollouts, integrating the resulting process reward with GRPO.
Results
CDPR improves diagnostic accuracy while substantially reducing examination number and cost across one in-domain and two out-of-domain benchmarks.
Takeaways & Limitations
Dense, annotation-free, critic-free process supervision redistributes sparse outcome signals onto intermediate actions in cost-aware sequential diagnosis.
Abstract
from arXiv · showhide
Clinical diagnosis is a step-by-step, cost-aware process: a physician orders examinations one at a time, observes the results, and updates the diagnosis before reaching a final conclusion. Most medical language models instead treat diagnosis as a one-pass classification task and ignore the trade-off between a test's value and its cost. We model diagnosis as a cost-aware sequential decision process and train the policy with reinforcement learning. The main difficulty is credit assignment: the only reliable signal comes once at the end of a long trajectory, so it scores a wasteful workup the same as an efficient one. We propose CDPR (Counterfactual Diagnostic Process Reward), which needs no expert labels and no learned critic. CDPR first finds the states where the policy hesitates, using the uncertainty of its action distribution, and then scores the chosen action by its advantage over the alternatives the policy itself would consider, estimated with short rollouts under a utility that balances correctness against test count, cost, and infeasible requests. A rollout cache reuses within-batch trajectories to keep the cost low. We integrate CDPR into GRPO and test it on one in-domain (MIMIC-IV) and two out-of-domain (ClinicalBench and a private hospital dataset) benchmarks. CDPR improves diagnostic accuracy while clearly reducing the number and cost of examinations.
Introduction
The paper frames diagnosis as a sequential, cost-aware process rather than one-shot classification, then addresses the resulting sparse-reward credit-assignment problem with CDPR.
- Clinical diagnosis iteratively selects examinations, observes results, updates the differential, and commits only when evidence is sufficient.
- Static diagnosis consumes the complete record in one pass, including irrelevant or costly examinations, while dynamic methods acquire evidence step by step.
- Outcome-only rewards score efficient and wasteful trajectories identically when both reach the same diagnosis, creating a central credit-assignment problem.
- CDPR estimates process reward by comparing chosen actions with plausible counterfactual alternatives at states where the policy hesitates.
- CDPR integrates state-level counterfactual advantages with GRPO and uses cached rollouts to provide dense, annotation-free supervision.
- Across MIMIC-IV, ClinicalBench, and a private hospital benchmark, CDPR improves diagnostic accuracy while reducing examination number and cost.
Related Work
Related work spans static diagnosis and newer sequential approaches that collect evidence step by step, while process-level methods seek denser supervision for intermediate decisions.
- Static diagnosis: Most existing medical diagnosis work treats the complete case description as a single-shot classification input.
- Sequential diagnosis: Sequential diagnostic methods model evidence collection as iterative examination and reasoning actions that reflect clinical practice.
- Sequential diagnosis: Illustrated sequential trajectories interleave symptoms, examinations, results, reasoning steps, and a final diagnosis.
- Process-level rewards: Agentic reinforcement-learning approaches use sparse outcome rewards, creating severe credit-assignment difficulty along long diagnostic trajectories.
- Process-level rewards: Recent methods explore step-wise or process-level rewards to assign credit to intermediate diagnostic actions.
Methodology
CDPR models diagnosis as a sequential decision process and assigns counterfactual process rewards at states where the policy hesitates. It compares plausible actions using short-horizon utility estimates while reusing cached trajectories to reduce evaluation cost.
- Sequential Diagnosis as an MDP: The policy represents diagnosis as sequential examination and diagnosis actions, with examinations extending the state and diagnosis actions terminating the episode.States contain the initial presentation plus examination history and observed results.
- Credit Assignment Challenge: Outcome-only rewards create credit-assignment difficulty because efficient and wasteful trajectories reaching the same diagnosis receive identical feedback.Wasteful trajectories may issue redundant or expensive examinations without affecting the final correctness reward.
- Action Uncertainty-guided State Selection: Counterfactual evaluation focuses on hesitant states identified through action-distribution uncertainty and disagreement among aligned cached trajectories.The two signals are unified to select states where next actions are unstable and evaluation is most informative.
- Short-horizon Continuation: For each selected state, CDPR estimates candidate action values with short-horizon continuations, using cached suffixes when available and fresh rollouts otherwise.Rollouts are bounded by horizon H, with truncated continuations forced to emit a diagnosis.
- Utility Function: The continuation utility combines diagnostic correctness with examination count, monetary cost, and unavailable-examination requests.These effort signals are weighted against the outcome reward when comparing candidate actions.
- Counterfactual Process Reward: CDPR avoids expert annotation and learned critics by assigning process reward through counterfactual comparison with alternatives the policy itself would plausibly consider.The chosen action receives an advantage over the average utility of the candidate actions.
Experiments
Experiments evaluate CDPR across in-domain and out-of-domain benchmarks, ablations, training efficiency, hyperparameter sensitivity, expert ratings, and dynamic-versus-static diagnosis. CDPR improves accuracy while reducing examination burden and cost, with process reward, cost penalties, state selection, and caching each contributing distinct effects.
- Main Results: CDPR-4B achieves the highest cross-dataset Diagnostic Accuracy (47.42) and lowest Average Exam Cost (23.52 USD) among reported methods.The cross-dataset average covers MIMIC-IV, ClinicalBench, and Private benchmarks.
- Ablation Studies: Removing process reward lowers accuracy from 53.07 to 47.33 and increases both examination count and cost.This ablation removes the dense step-level credit-assignment signal at high-uncertainty decision points.
- Training Efficiency Analysis: Full CDPR reaches approximately 53.07% accuracy while converging substantially closer to GRPO’s speed than the no-state-selection and no-cache variants.GRPO converges earliest but plateaus at approximately 47%, whereas removing the cache produces the slowest training.
- Hyperparameter Sensitivity: Accuracy peaks at M=4, K=4, H=3, β=0.5, η=0.5, and n=4, which are adopted as the default hyperparameters.Each hyperparameter is swept independently while the others remain at their defaults.
- Expert Evaluation of Intermediate Decisions: CDPR leads on all five expert-rated dimensions, including Diagnostic Gain at 4.7 versus 4.0 and Cost-effectiveness at 4.5 versus 4.0.Ratings were averaged across three physicians and 300 sampled decision states with model identity blinded.
- Dynamic vs. Static Diagnosis: CDPR slightly surpasses static diagnosis in accuracy, 53.07 versus 50.23, while reducing examination cost by 90.7%, from 395.40 to 36.65 USD.Static diagnosis receives the full case description in one shot, whereas CDPR sequentially orders examinations under cost constraints.
Conclusion
CDPR provides counterfactual, uncertainty-targeted process rewards for cost-aware sequential diagnosis and integrates them into reinforcement learning. Across one in-domain and two out-of-domain benchmarks, it improves accuracy while reducing examination cost.
- CDPR targets high-uncertainty decision states and evaluates chosen actions against counterfactual alternatives using short rollouts.Outcome differences become step-level rewards for sequential diagnosis.
- A rollout cache reuses training trajectories to reduce the cost of counterfactual evaluation under a tight rollout budget.
- Experiments on one in-domain and two out-of-domain benchmarks show higher diagnostic accuracy and substantially lower examination cost than other methods.
Dataset Description
The evaluation uses MIMIC-IV as an in-domain corpus and includes a representative MIMIC-IV case with structured clinical information and diagnosis.
- MIMIC-IV is a large-scale, publicly accessible critical-care database of de-identified electronic health records from a tertiary academic medical center.The corpus primarily covers intensive-care and emergency-department admissions.
- The dataset construction converts each discharge summary into a structured patient state.The passage identifies the chief complaint, history of present illness, and past medical history as components of that state.
- A representative MIMIC-IV case is presented in Table 4.
- The example case contains abdominal pain with clinical history, physical examination findings, laboratory results, and a final diagnosis of acute appendicitis.The case summary describes right lower quadrant pain and associated clinical records.
Prompt Templates
The framework uses separate prompts for sequential diagnosis, examination-key matching, and binary diagnosis evaluation. Together, they connect free-form agent actions to available examinations and terminal rewards.
- Sequential Diagnosis Agent Prompt: The Sequential Diagnosis Agent Prompt instructs the policy to maintain differential diagnoses, select a discriminative examination, and commit when evidence is sufficient.
- Exam Key Matching Prompt: The Exam Key Matching Prompt maps a free-form examination request to an available examination key so the environment can return the matching result without leaking values.The matcher uses abbreviations, synonyms, and alternative names, returning exactly the matching available name or “NO MATCH”.
- Diagnosis Evaluation Prompt: The Diagnosis Evaluation Prompt converts the predicted final diagnosis and ground truth into a binary “Correct”/“Wrong” verdict.This verdict serves as the outcome judge for the terminal reward.
Hyperparameters and Training Details
Training uses the verl framework with the actor and stale counterfactual actor colocated on two NVIDIA H20 GPUs. Rollouts use Qwen3-4B-Instruct-2507 with bounded multi-turn generation.
- Hardware and actors: All training runs use 2 NVIDIA H20 GPUs, with the actor and frozen stale actor colocated through separate GPU memory budgets.The stale actor provides counterfactual continuation.
- Backbone and rollout: The actor is initialized from Qwen3-4B-Instruct-2507 and trained with multi-turn trajectories capped at 8 assistant turns.
- Backbone and rollout: Training limits prompts to 2048 tokens and responses to 4096 tokens, with rollouts served by SGLang using FP8 quantization.
Expert Evaluation Protocol
The expert evaluation assessed intermediate diagnostic decisions across three benchmarks using independent physician ratings. It measured information use, evidential reasoning, disease relevance, diagnostic gain, and cost-effectiveness on five-point scales.
- Sampling and protocol: Three board-certified physicians independently evaluated 300 randomly sampled intermediate decision states across MIMIC-IV, ClinicalBench, and Private test sets.The sample included 100 states from each test set, and per-method scores averaged ratings across all states.
- Evaluation dimensions: Information completeness measures whether the model accurately exploits symptoms, history, prior results, and key positive or negative findings.Scores range from missing or misreading key information to comprehensively capturing all relevant findings.
- Evaluation dimensions: Evidence support measures whether each reasoning step is grounded in available evidence without unwarranted leaps or contradictions.The five-point scale ranges from unsupported or contradictory reasoning to inferences explicitly traceable to current evidence.
- Evaluation dimensions: Disease relevance measures whether the proposed action targets the current disease, suspected diseases, or critical differentials.Higher scores indicate tighter focus on the most likely disease or most critical differential.
- Evaluation dimensions: Diagnostic gain measures whether an action provides valuable new information by reducing uncertainty, separating differentials, or unblocking subsequent reasoning.The highest score denotes a pivotal action that markedly reduces diagnostic uncertainty.
- Evaluation dimensions: Cost-effectiveness measures whether an action is economically and clinically proportionate in monetary cost, time, invasiveness, risk, and redundancy.Higher scores reflect high diagnostic value, low burden, and clear necessity at the current diagnostic stage.