PaperlayerView exact revision on arXiv

Source-linked AI summary

EHRSHOT: An EHR Benchmark for Few-Shot Evaluation of Foundation Models

Michael Wornow, Rahul Thapa, Ethan Steinberg, Jason A. Fries, Nigam H. Shah

arXiv:2307.02028v3cs.LGcs.AIcs.CL

TL;DR

Healthcare ML lacks shared EHR datasets and pretrained models for reproducible foundation-model evaluation. EHRSHOT provides longitudinal structured EHR data, releases CLMBR-T-base weights, and defines 15 few-shot tasks; its results show gains in low-label settings while identifying room for improvement.

  • Problem

    Healthcare ML lacks shared EHR datasets and pretrained model weights needed to reproducibly evaluate foundation-model benefits such as sample efficiency and task adaptability.

  • Method

    The paper releases EHRSHOT, CLMBR-T-base, reproducible preprocessing and code, and 15 few-shot clinical prediction tasks using structured longitudinal EHR data.

  • Results

    Pretrained foundation models yield significant AUROC/AUPRC gains over a traditional supervised baseline in few-shot settings, although many tasks retain substantial room for improvement.

  • Takeaways & Limitations

    EHRSHOT provides a reproducible point of comparison for technical approaches to developing and improving clinical foundation models in low-label settings.

  • Takeaways & Limitations

    The release contains only structured data, evaluates one foundation model, and covers a task-selected cohort smaller than 1% of the source EHR database.

Abstract

from arXiv · show

While the general machine learning (ML) community has benefited from public datasets, tasks, and models, the progress of ML in healthcare has been hampered by a lack of such shared assets. The success of foundation models creates new challenges for healthcare ML by requiring access to shared pretrained models to validate performance benefits. We help address these challenges through three contributions. First, we publish a new dataset, EHRSHOT, which contains deidentified structured data from the electronic health records (EHRs) of 6,739 patients from Stanford Medicine. Unlike MIMIC-III/IV and other popular EHR datasets, EHRSHOT is longitudinal and not restricted to ICU/ED patients. Second, we publish the weights of CLMBR-T-base, a 141M parameter clinical foundation model pretrained on the structured EHR data of 2.57M patients. We are one of the first to fully release such a model for coded EHR data; in contrast, most prior models released for clinical data (e.g. GatorTron, ClinicalBERT) only work with unstructured text and cannot process the rich, structured data within an EHR. We provide an end-to-end pipeline for the community to validate and build upon its performance. Third, we define 15 few-shot clinical prediction tasks, enabling evaluation of foundation models on benefits such as sample efficiency and task adaptation. Our model and dataset are available via a research data use agreement from our website: https://ehrshot.stanford.edu. Code to reproduce our results are available at our Github repo: https://github.com/som-shahlab/ehrshot-benchmark

1 Introduction

Healthcare ML lacks shared EHR datasets and pretrained models needed for reproducible evaluation of clinical foundation models. EHRSHOT addresses this gap with longitudinal structured data, released model weights, and 15 few-shot tasks.

  • 1 Introduction: Open EHR datasets and pretrained clinical models remain scarce, making reproducible healthcare ML research and foundation-model evaluation difficult.Shared models are especially valuable because most hospitals lack the resources to train foundation models, while benefits such as sample efficiency and task adaptability are difficult to evaluate.
  • 1 Introduction: EHRSHOT combines a 6,739-patient longitudinal structured EHR benchmark, released CLMBR-T-base weights, and 15 few-shot clinical classification tasks.The dataset includes full coded medical timelines, while the model was pretrained on 2.57M patient EHRs.
  • 1 Introduction: Unlike ICU-focused benchmarks, EHRSHOT captures longitudinal health trajectories across departments, including patients not seen in the ICU or emergency department.EHRSHOT contains 41.6 million clinical events across 921,499 encounters and averages 2.3x more events and 95.2x more encounters per patient than MIMIC-IV.
  • 1 Introduction: CLMBR-T-base is a 141M-parameter transformer pretrained autoregressively on structured EHR codes from 2.57M patients, with full weights released for evaluation and reuse.The model predicts the next code in a patient timeline and processes structured information rather than clinical text.
  • 1 Introduction: Across 15 few-shot tasks, the pretrained model yields significant AUROC/AUPRC gains over a traditional supervised baseline, while many tasks retain substantial room for improvement.Several tasks have naturally low prevalence, creating a realistic low-label evaluation setting.
  • 1 Introduction: The workflow uses FEMR to retain structured diagnoses, medications, and labs, pretrains on the global training split, and evaluates CLMBR-T-base against a count-based GBM.The pipeline measures AUROC and AUPRC for each model-task pair and releases reproduction code and dataset access through a research data use agreement.

2 Related Work

Existing EHR benchmarks are often ICU-centered, use bespoke schemas, and omit standardized few-shot evaluation and pretrained model weights. EHRSHOT addresses these limitations through longitudinal, interoperable data and a reproducible FM-focused benchmark.

  • 2 Related Work: Most public EHR datasets emphasize ICU data, limiting representation of patients’ full health trajectories and the range of evaluable tasks.Even multi-department datasets such as MIMIC-IV remain anchored to ICU or ED admissions.
  • 2 Related Work: Prior EHR datasets commonly require bespoke schemas and substantial downstream work to construct preprocessing pipelines, patient splits, and task definitions.Custom schemas make transferring models across datasets and sites more difficult.
  • 2 Related Work: Most existing clinical foundation models have not released their weights, hindering reproducibility, cross-model evaluation, and transfer learning.Without pretrained weights, researchers cannot readily reuse existing models to shortcut development for new tasks.
  • 2 Related Work: EHRSHOT releases full longitudinal EHRs across departments, targeting broader patient trajectories than department-specific benchmarks.Its benchmark is designed specifically for few-shot evaluation of pretrained foundation models.
  • 2 Related Work: EHRSHOT uses the interoperable OMOP-CDM and an open preprocessing pipeline, while adding a benchmark centered on few-shot evaluation of pretrained foundation models.OMOP-CDM is described as an open standard used by more than 100 health systems.

3 Dataset

EHRSHOT is a longitudinal Stanford Medicine EHR benchmark designed for few-shot evaluation, covering 6,739 patients and 15 classification tasks. It provides structured patient timelines, canonical splits and k-shot samples, and reproducible preprocessing and dataset-generation code.

  • 3.2 Tasks: EHRSHOT provides canonical train/validation/test splits and k-shot samples for 15 classification tasks across broad clinical settings.The tasks were selected using clinician input and prior benchmarks, and span operational outcomes, lab values, new diagnoses, and chest X-ray findings.
  • 3 Dataset: The dataset represents full coded medical timelines rather than a single episode of care, using information across a health system’s departments.Records are provided as CSV files in a lightweight serialization of the OMOP-CDM format.
  • 3 Dataset: 6,739 patients contribute 41.6 million coded observations and 921,499 visits to EHRSHOT’s longitudinal benchmark.The cohort excludes patients younger than 19 or older than 88 and those with fewer than 10 clinical events.
  • 3.1 Data Source: EHRSHOT is sourced from Stanford Medicine’s STARR repository, while CLMBR-T-base uses 2.57 million source patients for foundation-model training and validation.The released benchmark preserves structured demographics, diagnoses, procedures, laboratory results, medications, and other coded observations while removing clinical notes.
  • 3.1 Data Source: FEMR supports preprocessing across EHR formats, and the authors release code for generating EHRSHOT and reproducing the benchmark results.The benchmark also visualizes its four task categories and prediction windows in Figure 2.
  • 3.2 Tasks: The benchmark includes nine binary, five 5-way multiclass, and one 14-way multilabel classification task.Some tasks assign multiple labels to individual patients, and task-specific label prevalence and cohort sizes are reported separately.

4 Baseline Models

The benchmark compares a count-based gradient boosting baseline with CLMBR-T-base, an autoregressive transformer pretrained on longitudinal structured EHRs. The models differ in representation and training, with CLMBR-T-base using causal attention and medical-code prediction.

  • 4 Baseline Models: The benchmark evaluates a count-based gradient boosting machine and CLMBR-T-base, an autoregressive language model pretrained on 2.57 million longitudinal EHRs.The authors selected these models because language modeling has performed strongly on clinical prediction while count-based features remain simple and competitive.
  • Count-based Features: Count-based features represent each patient by pre-prediction counts of medical concepts, with ontology expansion used to incorporate parent-child relationships.A LightGBM gradient boosting model is trained on the EHRSHOT training split and tuned on validation data; logistic regression and random forest are reported separately.
  • CLMBR-T-base: CLMBR-T-base predicts the next medical code from previous codes, learning representations for clinical prediction from structured EHR timelines.It uses a transformer with causally masked local attention, preserving forward-only information flow and excluding clinical text.
  • CLMBR-T-base: CLMBR-T-base has 141M trainable parameters, a hidden dimension of 768, and minute-level temporal resolution.Its transformer replaces the original CLMBR GRU formulation, and larger versions are left for future work.

5 Results

Across 15 tasks, CLMBR-T-base generally improves few-shot AUROC over the count-based GBM through 64 shots, while its advantage reverses for some diagnosis tasks at higher label counts.

  • 5 Results: The evaluation trains models with k positive and k negative examples, tunes hyperparameters on matched validation samples, and evaluates AUROC and AUPRC on the held-out test split.The benchmark uses k values from 1 through 128, except for Celiac, which is limited to k ≤64.
  • 5 Results: For CLMBR-T-base, few-shot examples fine-tune a logistic-regression head while the pretrained foundation-model weights remain frozen; GBM uses only those examples.Pretraining CLMBR-T-base took roughly four days on one Nvidia V100.
  • 5 Results: CLMBR-T-base outperforms the count-based GBM across all aggregated task categories for k ≤64.The comparison uses Macro-AUROC aggregated across subtasks within each task category.
  • 5 Results: Pretraining yields its largest advantage in data-poor settings, with gains most pronounced at intermediate k and shrinking as more labeled examples become available.At k = 1, both models struggle; the advantage generally decreases as k increases.
  • 5 Results: At k >64, the count-based GBM exceeds CLMBR-T-base on Assignment of New Diagnoses tasks, possibly because next-code pretraining is poorly matched to their one-year horizon.The authors also suggest that distinctive diagnosis signals may be easier for a tree-based model to learn directly with sufficient data.
  • 5 Results: Figure 3 reports category-level Macro-AUROC as bold lines, while blurred lines average five replicates for each subtask.The far-right All marker shows performance when trained on the full EHRSHOT training split.

6 Discussion

The authors position EHRSHOT and CLMBR-T-base as resources for reproducible, low-label healthcare ML, while acknowledging important data, model, cohort, and institutional-scope limitations.

  • 6 Discussion: Releasing pretrained weights allows researchers to replicate and build upon the work, while the results identify opportunities to improve pretrained models in few-shot settings.The authors describe this as enabling more reproducible healthcare ML research.
  • 6 Discussion: EHRSHOT enables reproducible comparison of technical approaches to clinical foundation models in few-shot settings, where labeled EHR data are costly and rare conditions may be sparsely represented.The authors acknowledge that the benchmark tasks may not themselves be the most clinically meaningful.
  • 6 Discussion: The release includes only structured data, omitting associated clinical text and images.The authors identify publication of clinical text as a continuing challenge.
  • 6 Discussion: Experiments cover only CLMBR-T-base, so the benchmark’s evidence does not compare multiple foundation-model architectures.The authors invite evaluation of additional foundation models on EHRSHOT.
  • 6 Discussion: The released cohort is less than 1% of the source EHR and was selected for the defined tasks, limiting answerable questions and representation of medical-data diversity.The full pretraining dataset was considered infeasible to release because of governance and effort constraints.
  • 6 Discussion: Because evaluation used Stanford Medicine data only, performance of the pretrained model at other institutions remains unclear.The authors anticipate some performance drop elsewhere.

7 Conclusion

The paper releases EHRSHOT, a longitudinal 15-task benchmark, together with CLMBR-T-base weights and reproducibility code to support more open clinical foundation-model research.

  • 7 Conclusion: EHRSHOT provides full longitudinal structured timelines for 6,739 patients and 15 classification tasks designed for few-shot foundation-model evaluation.The benchmark differs from prior work by covering longitudinal health data rather than a single department such as the ICU.
  • 7 Conclusion: The authors also release weights for a foundation model pretrained on more than 2.57M patient timelines and code needed to replicate the results.They frame the work as a first step toward more reproducible and open model development in healthcare ML.

Supplementary Material

The supplied supplementary-material passage contains a rights-and-licenses responsibility statement rather than substantive methodological or empirical content.

  • Supplementary Material: The authors state that they bear responsibility for violations of rights or licenses.

B.1 Dataset

EHRSHOT is distributed under research-use and access restrictions because of the dataset’s sensitive clinical content. The release includes cohort documentation, privacy protections, and CLMBR-T-base model resources.

  • Access and licensing: EHRSHOT access is gated by a research data use agreement because the dataset contains sensitive clinical information.The dataset is not uploaded to another repository because of these concerns.
  • Privacy protections: Privacy protections include deidentification, additional transformations, and exclusion of clinical notes.The release describes these measures as precautions against revealing Protected Health Information.
  • Access and licensing: EHRSHOT access is governed by Stanford’s Dataset Research Use Agreement for personal, non-commercial research.The agreement prohibits commercial use and requires individual users within an organization to register.
  • Access and licensing: The dataset is provided for non-clinical research only and must not support diagnosis or patient care.Stanford provides the dataset as-is without warranties.
  • Model release: CLMBR-T-base weights are released through the EHRSHOT website under a research data use agreement.The model release is likewise access-gated because its training data are sensitive.

C.2 Pretraining Dataset

The CLMBR-T-base pretraining dataset contains 3.67 million patient records, with 2.57 million used for model training. Demographic and timeline summaries are provided for the pretraining cohort.

  • Pretraining cohort: 2.57 million of 3.67 million pretraining patient records are used to train CLMBR-T-base.The remaining records are part of the broader pretraining dataset, whose demographics are summarized in Tables 5 and 6.

C.3 Task Definitions

EHRSHOT defines 15 prediction tasks spanning operational outcomes, new diagnoses, chest X-ray findings, and laboratory results. Each task specifies prediction timing, outcome windows, labels, and structured event or label-file representations.

  • Task coverage: EHRSHOT provides 15 tasks covering operational outcomes, new diagnoses, chest X-ray findings, and laboratory results.The task categories include binary, multiclass, and multilabel classification settings.
  • Operational Outcomes: Operational tasks predict hospital events such as long length of stay, 30-day readmission, and ICU transfer from defined admission-time points.Same-day discharges, readmissions, or transfers are excluded where specified.
  • Anticipating Lab Test Results: Laboratory tasks are multiclass predictions made immediately before results are recorded, with normal, mild, moderate, and severe outcome ranges.Baseline experiments reframe these lab tasks as binary classification, labeling abnormal results positive.
  • Assignment of New Diagnoses: New-diagnosis tasks predict first diagnoses within 365 days after inpatient discharge while excluding patients with an existing diagnosis.Examples include hypertension, hyperlipidemia, pancreatic cancer, celiac disease, lupus, and acute myocardial infarction.
  • Anticipating Chest X-ray Findings: The chest X-ray task is multilabel prediction of 14 findings at a point 24 hours before the radiology report is recorded.Labels are produced with the CheXpert NLP labeler, while the underlying unstructured reports are not released for privacy reasons.
  • Data representation: The release separates clinical events from task labels, with timestamped patient identifiers and values stored in standardized CSV schemas.Events include codes, optional values, units, and visit identifiers; labels include prediction time and task-specific values.

D Results Details

The benchmark formalizes prediction from patient timelines and compares a count-based GBM baseline with CLMBR-T-base. The provided figures examine few-shot performance and convergence toward full-data models across tasks and metrics.

  • Prediction formulation: The prediction objective uses a patient’s coded medical history through time t to predict a task-specific label at that time.Each timeline contains structured events such as diagnoses, procedures, medications, or laboratory tests.
  • Count-Based GBM: The GBM baseline represents each patient by counts of prior medical codes, producing a high-dimensional sparse matrix before ontology expansion.Ontology expansion counts each code and its OMOP parent nodes, after which LightGBM is trained for each task.
  • CLMBR-T-base: CLMBR-T-base embeds medical codes with code and rotary position embeddings, retaining the 65,536 codes contributing most to dataset entropy.Laboratory values are discretized into tokens defined by dataset-wide decile ranges.
  • CLMBR-T-base: The transformer produces cumulative event representations and scores candidate next codes through dot products with code embeddings.Training uses cross-entropy classification log-likelihood for the next event in each patient timeline.
  • Few-shot results: Few-shot figures evaluate AUROC and AUPRC across 1–128 shots plus full-data performance, with task-category macro averages and replicate averages.Figure 6 reports stronger CLMBR-T-base performance for Operational Outcomes and Anticipating Lab Test Results, while GBM is competitive for new diagnoses at higher shot counts.
  • Few-shot results: CLMBR-T-base closes the performance gap to a full-data model faster than the count-based GBM baseline in both AUROC and AUPRC.Figures 8 and 9 compare the differences between k-shot replicates and models trained on the full dataset.
Loading 2307.02028v3…