Explainable Transformer Models for Clinical Prediction Tasks on Structured Electronic Health Records

Jun Ni Du; Lukas Adamek; Maxim Kryukov; Flavio Dormont; Ziv Bar-Joseph; Sven Jager; Brandon Rufino

cs.LGAug 20, 2026

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

Abstract

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 $\geq 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 $\geq 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
Model	Name	Values	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.

Explore similar work

DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning
Mar 10, 2026Deyi Li, Zijun Yao, Qi Xu +4Longitudinal EHR Modeling Predictive Representation Learning
University of Florida, Gainesville, Florida, USA · University of Kansas, Lawrence, Kansas, USA · MD Anderson Cancer Center, Houston, Texas, USA +1
ChatHealthAI: Aligning Electronic Health Record Representations with Large Language Models for Grounded Clinical Reasoning
Jun 1, 2026Bo-Hong Wang, Baicheng Peng, Ruilin Wang +3LLM Grounding LLM Alignment
School of Computer Science, McGill University, Montreal, QC, Canada · Mila - Quebec AI Institute, Montreal, QC, Canada
INTERVenE: Temporal-Abstraction-Interval Based Transformers for Short-Horizon Medical Event Prediction
Aug 30, 2026Shahar Oded, Yuval ShaharEvent Time Prediction Interpretable ML
The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Beer-Sheva, Israel

Mar 10, 2026J/K move · Enter open · S save

DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

Deyi Li, Zijun Yao, Qi Xu, Muxuan Liang+3

University of Florida, Gainesville, Florida, USA · University of Kansas, Lawrence, Kansas, USA · MD Anderson Cancer Center, Houston, Texas, USA+1

Longitudinal EHR Modeling Predictive Representation Learning Graph Transformer Language Model Pretraining

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

DT-BEHRT: Disease Trajectory-aware Transformer for Interpretable Patient Representation Learning

ChatHealthAI: Aligning Electronic Health Record Representations with Large Language Models for Grounded Clinical Reasoning

INTERVenE: Temporal-Abstraction-Interval Based Transformers for Short-Horizon Medical Event Prediction