[Paper Review] CORE-BEHRT: A Carefully Optimized and Rigorously Evaluated BEHRT

Mikkel Odgaard, Kiril Vadimovic Klein+4 more|arXiv (Cornell University)|Apr 23, 2024

Advanced Neural Network Applications4 citations

TL;DR

CORE-BEHRT introduces a systematically optimized and rigorously evaluated BERT-based model for Electronic Health Records (EHR), improving performance through enhanced data representation (including medications and timestamps) and refined training protocols. It achieves a 0.801 AUROC average across 25 clinical tasks, with significant gains in 17 tasks and improvements in 24, demonstrating robust generalization and increased trustworthiness for clinical deployment.

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.

Motivation & Objective

• To address the lack of rigorous evaluation and optimization in BERT-based EHR models, which hinders clinical adoption.
• To isolate and quantify the impact of key design choices—data representation, architecture, and training protocol—on downstream performance.
• To improve model trustworthiness by systematically optimizing BEHRT through incremental experimentation and cross-task validation.
• To demonstrate consistent performance gains across diverse clinical prediction tasks, ensuring generalizability beyond narrow benchmarks.
• To provide a reproducible, optimized foundation for future EHR modeling research with actionable insights.

Proposed method

• Conducted incremental optimization on BEHRT by varying data representation (e.g., including medications, timestamps, ICD-10 codes) and training configurations.
• Used five-fold cross-validation with multiple pre-training runs to stabilize performance metrics and reduce convergence instability.
• Evaluated model performance on three initial tasks (death, pain treatment, infection) to guide optimization, then validated on 25 diverse clinical prediction tasks.
• Applied masking strategies and pooling methods to assess their impact on downstream AUROC, focusing on alignment with EHR task objectives.
• Employed a rigorous evaluation pipeline with statistical significance testing (p < 10⁻⁷) to validate performance improvements.
• Focused on real-world EHR data with clinical relevance, using age windows and modality constraints to simulate realistic deployment conditions.

Figure 1: This figure depicts the selection process of the best configuration. 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 predi

Experimental results

Research questions

• RQ1 How does improving data representation—particularly the inclusion of medications and timestamps—affect downstream performance in EHR-based BERT models?
• RQ2 To what extent do architectural and training protocol refinements enhance model performance beyond data representation?
• RQ3 How generalizable are the optimization gains across a broad and diverse set of clinical prediction tasks?
• RQ4 What is the impact of model size and fine-tuning data quantity on performance saturation and scalability?
• RQ5 Why do different masking ratios and pooling strategies yield inconsistent improvements across tasks, and how can this be mitigated?

Key findings

• Improving data representation increased average AUROC from 0.785 to 0.797 across three core tasks (p < 10⁻⁷), primarily due to inclusion of medication and timestamp features.
• Further optimization of architecture and training protocol raised the average AUROC to 0.801 across 25 diverse clinical tasks, with significant gains in 17 tasks and improvements in 24.
• The model achieved state-of-the-art performance in stroke prediction (AUROC 70.7%), outperforming prior work (66.9%) under more constrained data conditions.
• Performance improvements were consistent and statistically significant, indicating that the optimization process was not biased toward specific tasks.
• The inclusion of medications and timestamps had a dominant effect on performance, suggesting their critical role in capturing clinical context.
• Despite increased data size, performance saturated early for some tasks, indicating that data quantity alone is insufficient—model design and representation are more impactful.

Figure 2: Overview of the iterative optimization procedure, highlighting adapted settings in bold. 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 t