CORE-BEHRT: A Carefully Optimized and Rigorously Evaluated BEHRT

TL;DR

CORE-BEHRT systematically optimizes BEHRT for EHR data by dissecting the impact of data representation, architectural components, and training procedures. The study shows that enriching inputs with medications, full-depth codes, background context, and precise timestamps, combined with Time2Vec embeddings, RoPE/SwiGLU refinements, and diversified pooling, yields robust AUROC gains (e.g., from 0.785 to 0.801, $p<10^{-7}$) that generalize across 25 clinical tasks and withstand temporal shifts in out-of-time evaluations. Across optimization rounds, 17 of 25 tasks showed significant improvements, with 24 tasks benefitting overall, underscoring the approach’s generalizability and potential for broader adoption in clinical decision support. Limitations include dataset size, reliance on medication-coded data, and challenges of replicating private EHR setups, but the work provides a concrete, evidence-backed foundation for trustworthy, scalable BERT-based EHR models.

Abstract

The widespread adoption of Electronic Health Records (EHR) has significantly increased the amount of available healthcare data. This has allowed models inspired by Natural Language Processing (NLP) and Computer Vision, which scale exceptionally well, to be used in EHR research. Particularly, BERT-based models have surged in popularity following the release of BEHRT and Med-BERT. Subsequent models have largely built on these foundations despite the fundamental design choices of these pioneering models remaining underexplored. Through incremental optimization, we study BERT-based EHR modeling and isolate the sources of improvement for key design choices, giving us insights into the effect of data representation, individual technical components, and training procedure. Evaluating this across a set of generic tasks (death, pain treatment, and general infection), we showed that improving data representation can increase the average downstream performance from 0.785 to 0.797 AUROC ($p<10^{-7}$), primarily when including medication and timestamps. Improving the architecture and training protocol on top of this increased average downstream performance to 0.801 AUROC ($p<10^{-7}$). We then demonstrated the consistency of our optimization through a rigorous evaluation across 25 diverse clinical prediction tasks. We observed significant performance increases in 17 out of 25 tasks and improvements in 24 tasks, highlighting the generalizability of our results. Our findings provide a strong foundation for future work and aim to increase the trustworthiness of BERT-based EHR models.

Figures (6)

• Figure 1: This figure depicts the incremental optimization of the configurations. Non-selected configurations are depicted in a lighter shade, indicating they were considered but not chosen based on their impact on the average Area Under the Receiver Operating Characteristic (AUROC), across three prediction tasks, whose prediction windows are denoted in parentheses. Models were trained in five-fold cross-validation and evaluated on one test set per task. An extension of these experiments for masking ratios and pooling strategies is demonstrated in \ref{['appendix:mr_and_heads_results']}. * Model performance metrics are averaged not only across cross-validation folds but also over multiple pre-training runs to mitigate instabilities in convergence.
• Figure 2: Overview of the iterative optimization procedure, highlighting adapted settings in bold. Symbols denote either addition (+) or removal (-) of a setting with no symbol indicating a direct change. We augmented the input data with medication codes, patient sex, event timestamps, full International Classification of Diseases (ICD-10), and Anatomical Therapeutic Chemical (ATC) codes. Improvements to the model include integrating time2vec embeddings, Rotary Position Embeddings (RoPE), and the SwiGLU activation function. For pre-training (PT), we investigated various masking ratios and Med-BERT's secondary PT task: predicting prolonged length of stay in the hospital. We investigated various pooling strategies for fine-tuning (FT), finalizing on Bidirectional Gated Recurrent Units (BiGRU) based on its superior results.
• Figure 3: This figure illustrates the out-of-time fine-tuning and evaluation process. Training used data only up to 01/2020. The remaining three years were split for validation and testing: validation labels covered outcomes from 03/2020 to 06/2021, and testing labels covered outcomes from 09/2021 onward. Patients with outcomes before these dates in their respective sets were excluded.
• Figure 4: Selection pipeline for model selection and final evaluation
• Figure 5: Age distribution for positive patients in the cross-validation set, before age group selection, for each of the 25 downstream tasks.