SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the count-based GBM, these few-shot examples are the only training examples seen by the model. For the pretrained CLMBR-T-base model, we use these few-shot examples to fine-tune a logistic regression head appended to the top of the model, while keeping the weights of the pretrained CLMBR-T-base model frozen. Pretraining the CLMBR-T-base model took roughly 4 days on a single Nvidia V100 hosted in an on-premise compute cluster.
The AUROC of each model across all 4 task categories is presented in Figure 3. In the Appendix, we show this grouping for AUPRC in Figure 5. We also break down each individual task's AUROC in Figure 6 and AUPRC in Figure 7 of the Appendix. We also include results for additional baselines in the Appendix in Figures 10 and 11. The bolded lines are the Macro-AUC for each model within a task category, averaged across all subtasks at each  $k$ . We include the performance of each model trained on the entire EHRSHOT training split on the far right of every plot as All.
As shown in Figure 3, the pretrained foundation model CLMBR-T-base (blue) outperforms the count-based GBM (red) across all aggregated task categories for  $k \leq 64$ . This demonstrates the benefits of pretraining in few-shot settings, as the model can leverage patterns learned across millions of patients to derive more accurate representations out-of-the-box than a model trained from scratch. CLMBR-T-base outperforms the count-based GBM across all  $k$  on the Operational Outcomes and the majority of Anticipating Lab Test Results and Anticipating Chest X-ray Findings tasks. For these three task groups, the advantage of CLMBR-T-base seems most pronounced at intermediate levels of  $k$  between 8 and 128. At extremely low  $k$  (i.e.  $k = 1$ ), both models struggle to learn anything, while as  $k$  increases the advantage of the pretrained model tends to shrink, a trend noted elsewhere [21]. This is most visible in the far-right of the plot at the All marker, which represents the performance of each model when trained on the full EHRSHOT training dataset.
In fact, the count-based GBM exceeds the performance of CLMBR-T-base on the Assignment of New Diagnoses tasks at  $k > 64$ . This suggests that the advantage of pretraining comes primarily from improved initialization of patient representations, and that the largest gains are achieved in the most data poor regimes.
There are several possible reasons for CLMBR-T-base's underperformance at higher values of  $k$  for the Assignment of New Diagnoses tasks. First, the CLMBR-T-base model's training objective is next code prediction, which makes it ill-suited for predictive tasks with long time horizons (which for these tasks is 1 year). Second, if a simple tree-based model exists for a task (i.e. a few medical concepts tightly correlate with a diagnosis), then it may be more difficult for a pretrained model to coerce patient representations learned over millions of patients to that specific task than training a model from scratch with enough data to learn those distinctive signals. We believe that this reversal in model rankings demonstrates a key strength of EHRSHOT – namely, the diversity of its predictive tasks can help identify opportunities for improving pretraining and few-shot strategies.
We release all of our model weights, evaluation tasks, and data processing code to fully reproduce our results. To the best of our knowledge, the release of our pretrained CLMBR-T-base model is one of the first examples of such a clinical FM having its pretrained weights made publicly available [48].
We believe that EHRSHOT represents a useful contribution to the ML community by enabling more reproducible healthcare ML research. The release of our pretrained CLMBR-T-base model's weights will allow the community to replicate and build upon our work. Our results identify opportunities for improving pretrained models in few-shot settings.
Acquiring labeled EHR data is expensive and time-consuming. Additionally, certain rare conditions may only be present in a small cohort of patients out of millions within a health system [28]. Thus, model performance in low-label settings is of paramount importance in healthcare. As our results in Section 5 demonstrate, pretrained FMs can yield large performance gains in few-shot settings. While we acknowledge that the tasks themselves may not be the most clinically meaningful, we believe that EHRSHOT offers a valuable contribution by providing a reproducible and rigorous point of comparison for different technical approaches to developing clinical FMs.
Limitations. There are several limitations to this work. First, we only release structured data – i.e., we do not publish any of the clinical text or images associated with our patients. While many datasets for medical images exist [3], publishing clinical text remains a challenge [40]. Second, we only
Figure 3: Aggregated AUROC across all subtasks within each of the 4 task categories for  $k \in \{1,2,4,8,12,16,24,32,48,64,128\}$  shots. We show performance on the full training set as All. The bolded lines are the Macro-AUROC for each model, averaged across all subtasks within a task category for each value of  $k$ . The blurred lines are the average AUROC across 5 replicates for each subtask within a task category. CLMBR-T-base (blue) consistently outperforms the count-based GBM (red) at  $k \leq 64$ , but lags in higher label settings for the Assignment of New Diagnoses tasks.
consider one type of foundation model (CLMBR-T-base) for our experiments [41]. We look forward to seeing the additional foundation models that the community applies to our benchmark. Third, we release a very small cohort of patients  $(< 1\%)$  from our source EHR database, and specifically select these patients for the tasks that we define. Releasing our full pretraining dataset would be infeasible from a governance and effort perspective. Thus, while necessary in order to publish our EHR dataset and still broader than existing ICU-specific datasets, our cohort selection process limits the types of questions we can answer and does not reflect the full diversity of medical data. Fourth, as we only were able to evaluate our pretrained FM on Stanford Medicine data, it is unclear how well our pretrained model will perform at other institutions. We anticipate there will be some drop in performance, but the extent is unclear. Fifth, several of our tasks are "low label" in the most extreme sense - for example, the Celiac task only has 13 positive patients in its test set. This makes obtaining low variance estimates of model performance difficult. We aim to mitigate this by adding additional patients to our benchmark in future releases.
Societal Implications. We believe that the release of this dataset can help spur positive innovations for improving clinical care with ML. However, we recognize that there are patient privacy concerns anytime EHR data is released. We believe we sufficiently mitigate this risk through the rigorous deidentification process on which our data is subjected [5]. Additionally, we gate access to the dataset through a research data use agreement. Another concern is that models trained on biased data will reflect those biases [7]. Thus, the pretrained FM that we release may propagate biases in care delivery or outcomes present in our source EHR database [7]. However, we hope that by encouraging the full release of models, we can help the community better identify and mitigate these issues [25].
We publish EHRSHOT, a benchmark containing the structured data of 6,739 patients' full longitudinal medical timelines specifically geared towards few-shot evaluation of foundation models for clinical data. Unlike most prior work, EHRSHOT contains longitudinal health data rather than a single department (e.g. ICU). We define a set of 15 tasks ranging from well-studied outcomes like 30-day readmission to lesser explored settings such as anticipating abnormal lab values. Finally, we release the weights of a foundation model pretrained on over 2.57M patient timelines and publish the code needed to replicate our results. We hope that this work represents a first step towards moving the field of ML for healthcare towards more reproducible and open model development.
We thank the Stanford AIMI Center and Stanford Medicine Research IT for their assistance in publishing this dataset. This work was supported in part by the Mark and Debra Leslie Endowment for AI in Healthcare, the Clinical Excellence Research Center at Stanford Medicine, and Technology and Digital Solutions at Stanford Healthcare. MW is supported by an NSF Graduate Research Fellowship. JF was supported in part by a Stanford AIMI-HAI Partnership Grant.