Explainable Transformer Models for Clinical Prediction Tasks on Structured Electronic Health Records
Authors: Jun Ni Du, Lukas Adamek, Maxim Kryukov, Flavio Dormont, Ziv Bar-Joseph, Sven Jager, Brandon Rufino
Organizations: Data & Computational Science, R&D, Sanofi Toronto, Canada · Data & Computational Science, R&D, Sanofi Barcelona, Spain · Clinical Real World Evidence, R&D, Sanofi Cambridge, USA · School of Computer Science, Carnegie Mellon University Pittsburgh, USA · Data & Computational Science, R&D, Sanofi Frankfurt, Germany
Predictive models over structured electronic health records (EHRs) remain central to machine learning for healthcare, but few have jointly emphasized quantitative laboratory information and interpretability with respect to input medical events. We present BERT-LER, a BERT-style model for coded EHR timelines pretrained and fine-tuned from a de-identified EHR dataset of 75 million patients, that encodes laboratory test results as discrete tokens while retaining graded information through percentile-based binning, paired with Integrated Gradients for token-level attributions grounded in the input EHR sequence. We benchmark our approach on the public EHRShot benchmark suite and on an asthma severity progression study based on real-world data. This addresses a methodological gap in EHR foundation-style modeling by unifying laboratory value representation and explainability in a single framework, while assessing whether both predictive performance and explanations generalize beyond standard clinical prediction tasks. Across EHRShot and asthma tasks, BERT-LER achieves predictive performance that is competitive with, and on laboratory-related tasks often exceeds, publicly available benchmark models, and provides attributions that align with clinically known risk factors. Our architecture and explainability approach can be applied to many therapeutic areas and prediction tasks using language models trained on structured EHRs.
Figures & tables
Ontology
Purpose
BERT-LER
MED-BERT
ICD10CM
Dx
65231
64497
NDC
Rx
43230
0
RXNORM
Rx
22866
0
ICD9CM
Dx
14136
13861
LOINC
Lx
4680
0
CPT
Pr
1790
0
Table 1: Number of unique medical codes per ontology for the diagnosis-only BERT (MED-BERT) and BERT-LER models. All coding systems are from UMLS ( Bodenreider, 2004 ) , except demographics, which include age, race, ethnicity, and gender as defined in the TriNetX dataset.
Prediction Task
Split
Patients (class counts)
Loss-of-control
Train
19575 (11612/7963)
Validation
4143 (2422/1721)
Test
4108 (2472/1636)
Multiclass
Train
19575 (11612/1168/6795)
Validation
4143 (2422/270/1451)
Test
4108 (2472/239/1397)
Table 2: Number of patients in each split for each prediction task. Parentheses show class counts: negative/positive for binary tasks and controlled/exacerbation/symptomatic for multiclass.
Figure 1: Distribution of attribution scores for a hypothetical patient in the loss-of-control model. The model predicted a 60% probability of loss of control. Input tokens are shown on the right, with darker green highlights indicating stronger positive contributions to the prediction.
BERT-LER
BERT-LER (no lab value emb.)
CLMBR
Med-BERT
RETAIN
Prediction Task
0.80 (0.85, 0.75)
0.76 (0.81, 0.70)
0.84 (0.88, 0.80)
0.69 (0.75, 0.63)
0.65 (0.72, 0.58)
ICU Admission
0.80 (0.83, 0.77)
0.79 (0.82, 0.76)
0.79 (0.82, 0.76)
0.72 (0.75, 0.69)
0.70 (0.74, 0.66)
Readmission to ICU
0.78 (0.80, 0.76)
0.77 (0.79, 0.75)
0.80 (0.82, 0.78)
0.69 (0.71, 0.67)
0.71 (0.74, 0.69)
Long Stay in ICU
0.910 (0.912, 0.907)
0.787 (0.790, 0.783)
0.891 (0.893, 0.889)
0.688 (0.692, 0.684)
0.800 (0.804, 0.796)
Anemia Lab
0.79 (0.80, 0.78)
0.70 (0.71, 0.69)
0.77 (0.78, 0.76)
0.68 (0.69, 0.66)
0.63 (0.65, 0.61)
Hyperkalemia Lab
0.77 (0.78, 0.76)
0.65 (0.67, 0.63)
0.75 (0.76, 0.74)
0.63 (0.64, 0.61)
0.58 (0.60, 0.57)
Hypoglycemia Lab
Table 3: Average ROC-AUC values of various models on the EHRShot prediction tasks. Bold values indicate the best-performing model for each task and metric. Parentheses show the 95% confidence interval bounds derived from bootstrapping, reported as (upper, lower). We additionally report a BERT-LER lab-ablation variant with lab tokens retained but percentile-value embeddings removed.
Figure 2: Attribution score distribution for normal lab-anemia test results on the EHRShot dataset. Each row corresponds to a unique token. Black bars show the 68% interval, the white line indicates the median, and the numbers on the right (N) show the token counts in the test set.
Prediction Task
PR-AUC
ROC-AUC
Loss-of-control
BERT-LER
0.599 (0.572, 0.624)
0.682 (0.665, 0.698)
BERT-LER (no lab value emb.)
0.586 (0.558, 0.611)
0.676 (0.660, 0.692)
Med-BERT
0.509 (0.484, 0.533)
0.622 (0.604, 0.640)
RETAIN
0.476 (0.452, 0.501)
0.593 (0.577, 0.612)
XGBoost
0.554 (0.528, 0.581)
0.647 (0.630, 0.665)
Table 4: Performance of the models across three prediction tasks for asthma. Bold values indicate the best-performing model for each task and metric. Values in parentheses represent 95% confidence intervals derived from bootstrapping.
Figure 3: Attribution score distribution for predicting asthma loss of control. Top 30 tokens with minimum frequency ≥30 are shown, ordered by adjusted scores.
Figure 4: Top 30 laboratory test tokens with the highest regression coefficients for asthma loss of control. Each dot represents a lab test from a patient, with values indicated by color in decile groups. Negative regression coefficients are labeled in red to distinguish them from positive regression coefficients. Tokens have minimum frequency ≥40 .
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Figure A.1 : The study index date of January 1, 2018 was selected to ensure that all patients had one year of baseline data and one year of follow-up data prior to the COVID-19 pandemic, which impacted asthma care and management ( Philip et al., 2022 ; Lin and Cerrone, 2022 ) . The figure illustrates the index date along with the baseline and follow-up observation windows used for cohort construction.
Name
Values
Learning Rate
5e-5 , 1e-5, 5e-6
Effective Batch Size
16, 32
Appendix
Table A.1 : Pretraining hyperparameter sweep, with the selected parameters highlighted. The same parameters were used for both Med-BERT and BERT-LER. Pretraining was performed on a single A10 GPU with 24 GB of VRAM. The best-performing model saw 20 million unique patient records during training, slightly more than one epoch of data.
Name
Values
number of layers
12
number of attention heads
12
attention dropout
0.15
positional encoding
relative key query
embedding dimension in embedding layer
768
embedding dimension in hidden layers
768
Appendix
Table A.2 : Other training parameters for the model. Early stopping monitored the validation loss, stopping training after no improvement over 3 validation steps and at least one completed epoch. During hyperparameter tuning, models stopped between 1 and 2 epochs, having seen 20 to 40 million unique patient medical histories.
Source ontology
Purpose
# codes
LOINC
Lx
37,590
SNOMED
Dx / clinical
18,174
RxNorm
Rx
4,678
CPT4
Pr
3,730
CARE_SITE
Other
396
RxNorm Extension
Rx
255
Appendix
Table A.3 : Composition of the released CLMBR-T-base vocabulary by OMOP source ontology, reproduced from the model card ( Wornow et al., 2023 ) . Purpose labels are approximate and parallel Table 1 ; SNOMED in OMOP can span conditions, procedures, and observations.
Model
Name
Values
Task & Best Value
Loss-of-control
Multiclass
Exacerbation
BERT-LER
Learning Rate
5e-5, 1e-5, 5e-6
1e-5
5e-6
5e-6
Gradient Accumulation Steps
1, 2
1
1
1
Attention Dropout Probability
0.0, 0.2, 0.4
0.4
0.2
0.0
Med-BERT
Learning Rate
5e-5, 1e-5, 5e-6
5e-6
1e-5
5e-6
Gradient Accumulation Steps
1, 2
1
1
1
Appendix
Table A.4 : Hyperparameters used for the models trained on the asthma prediction tasks.
Figure A.2 : Top contributing factors from the RETAIN model for lab anemia prediction.
Figure A.3 : Top contributing factors identified by the RETAIN model for predicting asthma loss of control.
Figure A.4 : Top contributing factors identified by the XGBoost model for predicting asthma loss of control.
The growing adoption of electronic health record (EHR) systems has provided unprecedented opportunities for predictive modeling to guide clinical decision making. Structured EHRs contain longitudinal observations of patients across hospital visits, where each visit is represented by a set of medical codes. While sequence-based, graph-based, and graph-enhanced sequence approaches have been developed to capture rich code interactions over time or within the same visits, they often overlook the inherent heterogeneous roles of medical codes arising from distinct clinical characteristics and contexts. To this end, in this study we propose the Disease Trajectory-aware Transformer for EHR (DT-BEHRT), a graph-enhanced sequential architecture that disentangles disease trajectories by explicitly modeling diagnosis-centric interactions within organ systems and capturing asynchronous progression patterns. To further enhance the representation robustness, we design a tailored pre-training methodology that combines trajectory-level code masking with ontology-informed ancestor prediction, promoting semantic alignment across multiple modeling modules. Extensive experiments on multiple benchmark datasets demonstrate that DT-BEHRT achieves strong predictive performance and provides interpretable patient representations that align with clinicians' disease-centered reasoning. The source code is publicly accessible at https://github.com/GatorAIM/DT-BEHRT.git.
Deyi Li, Zijun Yao, Qi Xu +4
University of Florida, Gainesville, Florida, USA · University of Kansas, Lawrence, Kansas, USA · MD Anderson Cancer Center, Houston, Texas, USA +1
Large language models (LLMs) exhibit strong natural-language reasoning abilities for clinical decision support, but struggle to effectively model structured longitudinal electronic health records (EHRs). In contrast, EHR foundation models can learn predictive patient representations, yet lack interpretable language-based reasoning. To bridge this gap, we propose ChatHealthAI, a multimodal reasoning framework that aligns structured EHR representations from a pretrained EHR foundation model with the semantic space of a frozen LLM through a task-aware resampler. By integrating longitudinal patient representations with refined clinical event descriptions, ChatHealthAI enables clinically grounded natural-language reasoning while maintaining accurate patient prediction. We evaluated ChatHealthAI on three clinical predictive tasks from the EHRSHOT benchmark. Results show that ChatHealthAI improves reasoning quality and interpretability while preserving competitive predictive performance. These findings highlight the potential of integrating EHR foundation models with pretrained LLMs for interpretable clinical prediction.
Bo-Hong Wang, Baicheng Peng, Ruilin Wang +3
School of Computer Science, McGill University, Montreal, QC, Canada · Mila - Quebec AI Institute, Montreal, QC, Canada
Electronic Health Record (EHR) prediction models in the intensive care unit must learn from sparse and irregular measurements while preserving the clinical meaning of time and supporting transparent decision-making. We present INTERVenE, a family of Transformer architectures whose input is an interval-based, knowledge-based temporal abstraction (KBTA), a token stream of named clinical concepts (states, trends, events, contexts) drawn from a curated medical ontology, rather than an unnamed bin index or a raw measurement triplet. This naming layer is what we ask KBTA to do: it makes the model's per-token attributions resolve to clinical concepts by construction. INTERVenE offers two complementary variants: an auto-regressive decoder that generates future abstraction trajectories with a per-step risk readout (localizing when and after which events risk rises), and a bidirectional encoder for single-pass joint risk and time-to-event prediction. Evaluated on 57,078 MIMIC-IV admissions against GRU-D, STraTS, and KarmaLego, INTERVenE-Enc reaches a support-weighted AUPRCw of 0.672, improving by 0.041 over the strongest neural baseline with non-overlapping 95% bootstrap CIs, while also taking the best AUROCw (0.901) and length-of-stay MAE (44.4,h). INTERVenE-Ar (AUROCw0.854, AUPRCw0.587 under the same evaluation contract - a strictly harder generative readout) provides a complementary token-level risk trajectory. An input-representation ablation confirms the lift transfers across structured discretizations, positioning KBTA-based intervals as the interpretable substrate that makes per-token attributions resolve to meaningful clinical concepts within the deployed model.
Shahar Oded, Yuval Shahar
The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Beer-Sheva, Israel