Training LLMs for EHR-Based Reasoning Tasks via Reinforcement Learning

Jiacheng Lin, Zhenbang Wu, Jimeng Sun

TL;DR

Reinforcement learning with verifiable rewards (RLVR) has achieved striking success in math and coding, but applying it to messy Electronic Health Records (EHRs) introduces new failure modes. The authors introduce EHRMIND, a framework demonstrating that while RLVR can turn a small 3B model into a clinical reasoning engine that outperforms GPT-4 and DeepSeek-R1, it fails when foundational medical knowledge is missing. By using training-set sampling efficiency (Pass@k) as an early diagnostic, EHRMIND applies a lightweight supervised fine-tuning (SFT) warm-up before RLVR to prevent reasoning collapse and bootstrap clinical reasoning.


What Problem Does This Paper Address?

Adapting LLMs to clinical reasoning over Electronic Health Records is difficult due to noisy timelines, specialized medical terminology, and multi-step diagnostic logic. While RLVR (e.g., DeepSeek-R1) optimizes models using only outcome-level verification, EHR tasks suffer from two distinct failure modes:

  1. Misapplied Knowledge: The model has the necessary medical facts stored in its weights but fails to extract and apply them in context.
  2. Missing Knowledge: The base model simply lacks essential domain facts (e.g., specific drug dosage conversion ratios). Because it cannot generate correct answers by chance, rewards remain completely sparse, causing RL to stall or collapse.

The authors set out to determine whether pure RLVR can induce clinical reasoning on EHR data and when supervised warm-up is strictly required.


Methods

The authors propose EHRMIND, which trains a 3B parameter model (LLaMA-3-3B-Instruct) using Group Relative Policy Optimization (GRPO) with rule-based outcome rewards (e.g., exact-match accuracy). EHRMIND introduces two key methodological recipes:

  • Diagnostic Pass@k and Reliable Pass@k: Before expensive RL training, the authors measure whether the base model can solve instances across $k=12$ rollouts. For discrete classification tasks where random guessing artificially inflates success, they introduce Reliable Pass@k, requiring both a frequency threshold ($\tau_p = \lceil k/C \rceil + 2$) and low output prediction entropy ($H(D) < 0.8 \log C$). This diagnostic takes ~1 hour, saving ~20 hours of failed RL runs.
  • Two-Stage SFT Warm-up + RLVR: If Pass@k is low, EHRMIND first runs a lightweight SFT warm-up on ~2,000 GPT-4o-generated reasoning chains. This injects task-specific inductive biases and prevents "reasoning collapse" (where RL shortcuts into outputting raw labels without chain-of-thought rationales), followed by GRPO policy optimization.

Datasets

The authors evaluate EHRMIND across three distinct EHR reasoning setups:

  1. MEDCALC-BENCH: Medical calculation questions based on clinical notes across seven categories (lab, physical, date, dosage, risk, severity, diagnosis).
  2. TREC 2021 Clinical Trials: Patient-trial matching classifying patient note eligibility against trial criteria into Eligible, Excluded, or Irrelevant.
  3. EHRSHOT: Real longitudinal EHR data from Stanford Medicine across four 1-year disease forecasting tasks (Acute Myocardial Infarction, Hyperlipidemia, Hypertension, Pancreatic Cancer), serialized into structured text using high-utility event types (condition and procedure occurrences).

Key Insights and Ideas

  • Outcome RL unlocks latent knowledge, but cannot fabricate missing facts: In tasks where the base model had moderate prior knowledge (Physical, Date), pure RL boosted accuracy dramatically. But in knowledge-deficient tasks like Dosage Conversion, pure RL plateaued at zero gain ($R^2 = 0.91$ correlation between training Pass@12 and RL gain).
  • SFT warm-up prevents "reasoning collapse": Pure RL on classification tasks incentivized the model to omit reasoning steps and guess directly. The SFT warm-up preserved detailed, medically plausible reasoning traces throughout GRPO training.
  • RL generalizes better than pure imitation: When models trained solely on Hyperlipidemia were evaluated out-of-distribution on other diseases (e.g., Pancreatic Cancer), RLVR and SFT-RLVR transferred substantially better than SFT-only models, showing that reinforcement learning develops generalizable clinical deduction rather than superficial pattern memorization.

Key Results

  • Medical Calculations (MEDCALC): Pure RLVR elevated LLaMA-3-3B from a 9.74% zero-shot baseline to 41.26%, surpassing GPT-4 (37.92%) and Claude-3.5 Sonnet (41.18%). Adding lightweight SFT warm-up (EHRMIND-SFT-RLVR) reached 51.96%, outperforming specialized reasoning models like o3-mini (46.42%) and DeepSeek-R1 (48.13%).
  • Trial Matching (TREC): EHRMIND-SFT-RLVR achieved 63.14% Balanced Accuracy and 44.47% Macro F1, heavily outperforming proprietary baselines (GPT-4o: 38.16% BACC; o3-mini: 39.60% BACC). SFT specifically rescued performance on the rare Irrelevant and Eligible classes where pure RL failed.
  • Diagnosis Forecasting (EHRSHOT): On Pancreatic Cancer, EHRMIND-RLVR reached 81.79% BACC (compared to 70.67% for GPT-4o and 56.90% for o3-mini). In pairwise qualitative evaluations judged by GPT-4o, rationales from SFT-RLVR were preferred over pure RLVR across all four diagnostic conditions.

Limitations and Future Work

  • Sequence Length Scaling: Complete longitudinal patient histories regularly exceed 10,000 events. The current approach relies on heuristic filtering to fit context windows; integrating native hierarchical or multimodal architectures remains future work.
  • Reliance on Verifiable Outcomes: EHRMIND requires deterministic, rule-based reward verification (e.g., exact matches, structured diagnoses). Extending this paradigm to open-ended clinical workflows (e.g., discharge summary generation or clinical dialogue) remains an open challenge.

Loading related content...