Technical Summary: ProMed – Shapley Information Gain Guided Reinforcement Learning for Proactive Medical LLMs

1. Problem Statement

Interactive medical questioning is a cornerstone of clinical diagnosis, where physicians must proactively elicit incomplete patient information to reach accurate conclusions. However, existing Large Language Models (LLMs) deployed in clinical settings predominantly operate under a reactive paradigm. They generate predictions based solely on initial, often partial or biased, patient inputs. This limitation leads to cognitive biases in reasoning, resulting in diagnostic errors, misdiagnosis, and compromised patient safety.

Current attempts to shift LLMs toward a proactive paradigm (actively seeking information) rely on prompt engineering or Supervised Fine-Tuning (SFT) on static dialogues. These approaches face significant limitations:

• Prompt Engineering: Merely induces questioning behaviors without fostering genuine active reasoning, leading to limited performance gains.
• SFT: Lacks robustness in unpredictable real-world scenarios and, when relying solely on positive feedback, inadvertently amplifies incorrect or misleading inquiries.
• Reinforcement Learning (RL) Challenges: Applying RL to this domain is hindered by two core issues:
  1. Reward Modeling: Existing methods (e.g., heuristic scoring, leave-one-out evaluations) fail to capture the compositional nature of medical reasoning, where the utility of a single question often depends on its interaction with other clinical facts.
  2. Reward Attribution: Standard RL strategies (e.g., uniform reward distribution across tokens) fail to distinguish clinically salient questions from irrelevant ones and implicitly encourage verbosity over informativeness.

2. Methodology: ProMed Framework

The authors propose ProMed, a reinforcement learning framework designed to transition medical LLMs from reactive responders to proactive information seekers. The framework is built around a novel Shapley Information Gain (SIG) reward mechanism and a two-stage training pipeline.

2.1 Shapley Information Gain (SIG) Reward

To quantify the clinical utility of a question, ProMed introduces SIG, which moves beyond simple information recall to account for the interactions between clinical facts.

• Atomic Fact Foundation: Patient conditions are decomposed into a set of atomic facts ($F$).
• State Approximation: The model's internal belief state is approximated by a "doctor understanding" prompt that articulates the current grasp of the patient's condition.
• Shapley Values: Drawing from cooperative game theory, SIG calculates the Shapley value for each atomic fact. This measures the marginal contribution of a fact to the correct diagnosis across all possible subsets of facts, thereby capturing both the varying importance of individual facts and their complex interactions.
• Reward Calculation: The SIG reward for a question is the weighted sum of newly acquired facts, where weights are the softmax-normalized Shapley values. This encourages the model to prioritize questions that acquire novel, high-impact clinical information.

2.2 Two-Stage Training Pipeline

ProMed integrates SIG into a two-stage process to ensure stable initialization and fine-grained optimization:

Stage 1: SIG-Guided Model Initialization

• Goal: Construct high-quality interaction trajectories for Supervised Fine-Tuning (SFT).
• Mechanism: Uses Monte Carlo Tree Search (MCTS) guided by SIG rewards to explore optimal doctor-patient dialogue trees.
• Process: The MCTS simulates dialogues, selecting paths that maximize the cumulative SIG reward and final answer correctness. The resulting high-reward trajectories are used to initialize the policy, providing a strong foundation for proactive behavior and alleviating data scarcity.

Stage 2: SIG-Augmented Policy Optimization

• Goal: Refine the policy via Reinforcement Learning (specifically GRPO).
• Mechanism: Introduces a SIG-Guided Reward Distribution Mechanism.
• Innovation: Unlike standard GRPO, which assigns uniform rewards to all tokens in a trajectory, this mechanism decomposes the trajectory-level reward into action-level signals.
  • Follow-up questions receive rewards proportional to their SIG value.
  • The final answer receives a correctness reward.
  • These action rewards are further propagated to individual tokens.
• Outcome: This enables targeted, fine-grained optimization that reinforces the generation of informative questions while suppressing redundant or misleading inquiries. The SIG values are dynamically computed during training to remain strictly model-aware.

3. Key Contributions

1. Paradigm Shift: ProMed pioneers an RL framework that successfully transitions medical LLMs from reactive answering to proactive information seeking.
2. Novel Reward Mechanism: The development of the Shapley Information Gain (SIG) reward, which leverages Shapley values to model medical information interactions and quantify question utility, addressing the limitations of isolated fact evaluation.
3. Fine-Grained Optimization: A tailored reward distribution mechanism that allocates rewards based on question utility, enabling precise policy optimization.
4. Empirical Validation: Extensive evaluations demonstrating significant performance improvements and robust generalization to out-of-distribution (OOD) benchmarks.
5. Resource Contribution: The construction of two public benchmarks targeting interactive medical questioning with standardized splits to facilitate future research.

4. Experimental Results

Experiments were conducted on MedQA and CMB (Chinese Medical Benchmark) datasets using various LLM backbones (LLaMA3.1-8B, LLaMA3.2-3B, Qwen3-1.7B).

• Performance Gains: ProMed significantly outperforms state-of-the-art methods. On average, it achieves a 6.29% relative improvement over the second-best method. Compared to the reactive paradigm (direct answering), ProMed delivers a 54.45% gain in accuracy.
• Ablation Studies:
  • Both the SIG-Guided Initialization (Stage 1) and SIG-Augmented Optimization (Stage 2) are essential; removing either leads to substantial performance drops.
  • The Shapley weighting and the reward distribution mechanism are critical components; their removal degrades both in-domain and OOD performance.
• Generalization: ProMed demonstrates robust generalization to OOD settings, including open-ended medical QA, clinical reasoning tasks (CMB-Clin), and general-domain ambiguous question answering (Abg-CoQA).
• Question Quality: Human and automated evaluations confirm that ProMed generates questions with higher Helpfulness, Medical Relevance, and Logical Flow compared to baselines.
• Efficiency: The model exhibits a "cautious-proactive" strategy, asking fewer questions than baselines when information is sufficient and dynamically adjusting its strategy based on information scarcity.

5. Significance and Claims

The paper claims that ProMed addresses a fundamental problem in interactive medical consultation: the inability of current LLMs to systematically acquire clinically informative evidence. By shifting the paradigm from reactive to proactive, ProMed enables more unbiased and reliable clinical decision-making.

The authors emphasize that their approach is model-agnostic and task-agnostic, suggesting the methodology is generalizable across different populations and clinical contexts. They position ProMed not as a replacement for medical professionals, but as a research advancement that mitigates the risks of hallucination and unreliable responses in reactive medical LLMs operating under incomplete information. The work highlights the necessity of targeted training over simple prompting or static fine-tuning to achieve genuine interactive diagnostic capabilities.

Limitations Acknowledged:

• Training and evaluation occur in simulated environments, which may not fully capture real-world patient variability.
• Current experiments focus on multiple-choice benchmarks and specific open-ended tasks; further validation on complex, real-world scenarios (e.g., free-form treatment planning) is needed.
• Computational constraints limit the use of models up to 8B parameters; the performance ceiling of larger models remains unexplored.
• The current implementation relies on text-based facts; incorporating multimodal data (imaging, time series) is a future direction.