Source-linked AI summary
Hi-BEHRT: Hierarchical Transformer-based model for accurate prediction of clinical events using multimodal longitudinal electronic health records
Yikuan Li, Mohammad Mamouei, Gholamreza Salimi-Khorshidi, Shishir Rao, Abdelaali Hassaine, Dexter Canoy, Thomas Lukasiewicz, Kazem Rahimi
TL;DR
Longitudinal EHR risk prediction is constrained by models’ difficulty processing very long, multimodal histories. The paper introduces Hi-BEHRT, a hierarchical Transformer with adapted contrastive pretraining, and reports improved prediction across four disease-risk tasks, including for patients with longer histories.
Problem
Existing deep-learning models struggle to process long, multimodal EHR sequences, risking omission of historical information needed for prediction.
Method
Hi-BEHRT uses hierarchical local feature extraction and global aggregation for long multimodal EHR sequences, with BYOL-based self-supervised pretraining adapted to EHR.
Results
Hi-BEHRT outperformed BEHRT by 1%-5% AUROC and 3%-8% AUPRC across four risk-prediction tasks, while pretraining matched a 5% training-data model using 1%.
Takeaways & Limitations
Hi-BEHRT supports risk prediction from long, comprehensive multimodal EHR histories and can transfer effectively when training data are relatively limited.
Takeaways & Limitations
Performance nearly saturates when 50% of the training dataset is used, indicating a remaining accuracy ceiling in the tested setting.
Abstract
from arXiv · showhide
Electronic health records represent a holistic overview of patients' trajectories. Their increasing availability has fueled new hopes to leverage them and develop accurate risk prediction models for a wide range of diseases. Given the complex interrelationships of medical records and patient outcomes, deep learning models have shown clear merits in achieving this goal. However, a key limitation of these models remains their capacity in processing long sequences. Capturing the whole history of medical encounters is expected to lead to more accurate predictions, but the inclusion of records collected for decades and from multiple resources can inevitably exceed the receptive field of the existing deep learning architectures. This can result in missing crucial, long-term dependencies. To address this gap, we present Hi-BEHRT, a hierarchical Transformer-based model that can significantly expand the receptive field of Transformers and extract associations from much longer sequences. Using a multimodal large-scale linked longitudinal electronic health records, the Hi-BEHRT exceeds the state-of-the-art BEHRT 1% to 5% for area under the receiver operating characteristic (AUROC) curve and 3% to 6% for area under the precision recall (AUPRC) curve on average, and 3% to 6% (AUROC) and 3% to 11% (AUPRC) for patients with long medical history for 5-year heart failure, diabetes, chronic kidney disease, and stroke risk prediction. Additionally, because pretraining for hierarchical Transformer is not well-established, we provide an effective end-to-end contrastive pre-training strategy for Hi-BEHRT using EHR, improving its transferability on predicting clinical events with relatively small training dataset.
I. INTRODUCTION
Linked longitudinal EHRs offer comprehensive patient histories for risk prediction, but their complexity and length can exceed existing models’ capacity. The study therefore targets models that can use longer, multimodal records without losing historical information.
- Linked longitudinal EHRs provide comprehensive patient trajectories for developing risk prediction models.
- Existing EHR studies often use only disease and medication records, typically limiting patients to a few hundred records.
- Longer and more comprehensive records can create sequences with thousands of events, making missing important historical information a critical bottleneck.
- Hierarchical Transformer architectures are considered a solution for extracting local temporal features before higher-level sequence processing.
- The study compares four disease-risk tasks using diagnoses, medications, procedures, tests, blood pressure, drinking, smoking, and BMI information.
- Pretraining is also challenging because hierarchical feature extraction makes masked-record labels less clearly defined.
B. Case identification
The study defines four 5-year incident-risk outcomes and represents multimodal primary-care records in a hierarchical sequence model. Hi-BEHRT segments medical histories with overlapping sliding windows to capture local temporal interactions before global aggregation.
- Case identification: The study estimates 5-year incident risk for heart failure, diabetes, chronic kidney disease, and stroke.
- Case identification: Input records include diagnoses, medications, hospital procedures, GP tests, blood pressure, drinking status, smoking status, and BMI.
- Case identification: Diagnostic records are mapped from primary-care Read codes to ICD-10 level 4 using NHS Digital and SNOMED-CT vocabularies.
- Representations: BEHRT represents records with token, age, segmentation, and position embeddings across available clinical codes and categorical variables.
- Representations: Hi-BEHRT uses sliding windows and a local Transformer feature extractor to learn temporal interactions within smaller segments.
- Representations: The segment representations appear non-overlapping but are built from overlapping tokens in the sliding window.
D. Training and validation of the models
Models are evaluated on four risk-prediction tasks using AUROC and AUPRC on held-out validation data. The self-supervised strategy adapts BYOL to Hi-BEHRT through EHR and segment augmentations.
- Training and validation: AUROC and AUPRC are reported on each task’s validation set after training on the combined training and tuning data.
- Model training: Hi-BEHRT globally aggregates local segment features and uses a pooling layer to make risk predictions.
- Pretraining: BYOL pretraining is adapted to Hi-BEHRT by combining its representation-learning framework with masked-language modeling.
- Pretraining: BYOL uses online and target networks, with the online predictor matched to the target projector through a mean squared loss.
- Pretraining: The target-network weights follow 𝜁←𝜏𝜁+ (1 −𝜏)𝜃, where 𝜏 is a decay factor between 0 and 1.
- Pretraining: EHR augmentation uses random cropping and token masking, while segment augmentation modifies selected latent segment representations.
G. Implementation details
The implementation uses large multimodal EHR cohorts and compares BEHRT with Hi-BEHRT under shared training settings. Sequence length, augmentation probabilities, and cohort characteristics define the experimental setup.
- Pretraining: Pretraining uses 50% random EHR cropping, 20% token masking, and 50% segment augmentation probability.
- Pretraining: Segment augmentation masks representations to zero with 85% probability or adds Gaussian noise with 15% probability.
- Optimization: BEHRT is limited to a sequence length of 256, whereas Hi-BEHRT’s local-extractor complexity grows linearly with the number of segments.
- Cohorts: The pretraining cohort contains 2,844,733 patients, with 2,438,352 training, 406,381 tuning, and 1,219,078 validation patients.
- Cohorts: The model includes 1,995 diagnosis, 378 medication, 275 test, 960 procedure, 24 systolic-BP, 17 diastolic-BP, and 34 BMI categories.
- Cohorts: The selected cohort excludes patients with less than 3 years of learning period, producing task-specific patient-count variability.
B. Model performance evaluation
Hi-BEHRT was evaluated against BEHRT across four clinical risk-prediction tasks, including settings with limited training data. The evaluation also examined transfer from self-supervised pretraining.
- Hi-BEHRT was evaluated on four risk-prediction tasks and compared with BEHRT and a pretrained Hi-BEHRT variant.The evaluation used validation-set performance and also examined subgroups by EHR length.
- 1% to 5% and 3% to 8% absolute improvement were reported for AUROC and AUPRC, respectively, over BEHRT across all risk-prediction tasks.The comparison used a smaller model size and less model complexity for Hi-BEHRT.
REPRESENTATIONS TO RISK PREDICTION TASK
The study examined whether pretrained Hi-BEHRT representations transfer to risk prediction with limited data and whether hierarchical processing benefits patients with longer EHR histories. Pretraining improved small-data performance, while Hi-BEHRT retained advantages over BEHRT across EHR-length subgroups.
- Transfer from pretrained representations: Pretraining enabled Hi-BEHRT fine-tuned on 1% of the training data to achieve performance similar to training on 5% without pretraining.The comparison was made between fine-tuned pretrained representations and Hi-BEHRT trained from scratch.
- Performance by EHR length: Hi-BEHRT outperformed BEHRT across all subgroups and all four risk-prediction tasks.The subgroups were defined by learning-period EHR length of 0–256 records versus more than 256 records.
- Performance by EHR length: 3%–6% AUROC and 3%–10% AUPRC improvements were reported for patients whose EHR length exceeded BEHRT’s 256-record capacity.Comparisons also examined groups with similar positive-case proportions to assess the effect of including more records.
- Performance by EHR length: Including almost the entire EHR sequence and using a sliding window helped Hi-BEHRT focus on local temporal features across long and short sequences.The passage links this design to more comprehensive risk estimation under both circumstances.
C. Ablation analysis
The ablation analysis evaluates Hi-BEHRT with different training-data sizes and modality combinations. Adding modalities generally improves performance, with gains varying by recording frequency and importance.
- Training size and performance: Hi-BEHRT was evaluated with 1%, 5%, 10%, 20%, 50%, and 100% of the training dataset for heart-failure risk prediction.Models were averaged over three random seeds.
- Ablation of modality: Adding procedure, test, BP, BMI, drinking, and smoking information alongside diagnosis and medication assessed how EHR richness affected performance.Each additional modality was included one at a time.
- Ablation of modality: Performance in terms of AUPRC and AUROC generally improves as more modalities are included.The cited ablation passage reports a trend rather than a specific numerical gain.
- Ablation of modality: The modality ablation evaluated diagnosis, medication, procedure, test, BP measurement, BMI, smoking status, and drinking status.The figure denotes these modalities as D, M, P, T, B, I, S, and A.
IV. DISCUSSION
Hi-BEHRT extends BEHRT to process long, multimodal EHR histories and improves risk prediction across four diseases. Contrastive pretraining improves performance with limited data, while modality scarcity and training-data saturation remain important boundaries.
- Model architecture and performance: Hi-BEHRT incorporates long EHR sequences from multiple modalities and outperformed BEHRT across heart failure, diabetes, CKD, and stroke risk prediction.The reported improvements were 1%-5% for AUROC and 3%-8% for AUPRC.
- Multimodal EHR: Including medication, procedure, test, BP, BMI, drinking, and smoking information extends BEHRT beyond diagnosis-only inputs.These modalities provide additional information for modelling.
- Multimodal EHR: Modality contribution depends on recording frequency and importance, with scarce modalities contributing less and potentially being overpowered by frequent modalities.The authors identify low-frequency modality handling as a future-work direction.
- Pretraining: With BYOL and MLM pretraining, Hi-BEHRT achieved similar performance using 1% of training data as a model without pretraining using 5%.The authors report that pretraining can expand the power-law region and reach it with less data.
- Limitations and future work: Model performance almost saturated when using 50% of the training dataset.The authors characterize this as reaching the irreducible error region and call for more robust architectures.
- Implications: The framework targets risk prediction with long, comprehensive EHR and provides self-supervised pretraining for settings with limited training data.The stated motivation is to support risk prediction using the full medical record.
Supplementary
The supplementary section introduces procedure and test data as additional EHR modalities used in the modelling.
- Procedure and test: Test data are recorded in the CPRD test table using Read codes for examinations, diagnostic procedures, and laboratory procedures.The experiment used Read-code-level information and excluded more detailed quantitative values.
A. Procedure
The supplementary material describes procedure and test data sources, outcome-code tables, and age-stratified model evaluation. Hierarchical BEHRT performs better across age groups, especially for younger patients on heart-failure and diabetes tasks.
- A. Procedure: Procedure records come from CPRD-linked HES admitted-care data and use OPCS 4.6 coding with start, end, and event dates.They are recorded at NHS healthcare-provider admissions or attendances.
- A. Procedure: Test records use CPRD Read codes for history, examinations, diagnostic procedures, and laboratory procedures, without detailed quantitative information.Only Read-code-level information was used in the experiment.
- Supplementary outcome definitions: Supplementary ICD-10 tables identify patients with heart failure, diabetes, CKD, and stroke in hospital or general-practice records.The supplied tables list disease-specific code resources rather than performance results.
- Age-stratified evaluation: Hierarchical BEHRT performs better than BEHRT across baseline-age subgroups and substantially outperforms it for heart failure and diabetes, especially among younger patients.The evaluated groups were ages 35-50, 50-70, and 70-90 years.
IX. SIZE AND OVERLAP OF SLIDING WINDOW
Sliding-window configuration affects both model complexity and predictive performance. Larger windows are not uniformly better, and overlap becomes important when windows are large.
- Window size and overlap: AUPRC decreases 4% with window size 100 and no overlap versus stride size 50.The analysis indicates that some overlap is necessary for larger windows.
- Window size and overlap: AUPRC for window size 100 without overlap decreases 2% compared with window size 50 without overlap.Thus, increasing window size alone does not guarantee better performance.
- Window size and overlap: Larger stride sizes can reduce the number of segments and model complexity when the window size is 50.For small windows, stride size has little effect on predictive performance.
- Hyper-parameter tuning: Hidden size 150 and intermediate size 108 achieve similar performance to larger configurations in 5-year heart failure risk prediction.These settings were selected through hierarchical BEHRT hyper-parameter tuning.