PaperlayerView exact revision on arXiv

Source-linked AI summary

Med-BERT: pre-trained contextualized embeddings on large-scale structured electronic health records for disease prediction

Laila Rasmy, Yang Xiang, Ziqian Xie, Cui Tao, Degui Zhi

arXiv:2005.12833v1cs.CLcs.LGcs.NE

TL;DR

Deep-learning disease prediction often requires large EHR training cohorts, limiting use when labeled data are scarce. The paper introduces Med-BERT, a BERT adaptation pre-trained on 28,490,650 structured-EHR patients with domain-specific tasks. It reports improved disease-prediction performance, especially with very small fine-tuning datasets.

  • Problem

    Deep-learning EHR prediction often requires large training cohorts, while limited cases, annotation difficulty, and governance restrict access to large annotated datasets.

  • Method

    Med-BERT adapts BERT to structured diagnosis-code EHRs using contextual embeddings, code and visit representations, and code-level and patient-level pre-training tasks.

  • Results

    Med-BERT improves predictive performance across tasks, with RETAIN gains of 2.02–3.20% and performance on some small-data settings comparable to training sets almost 10 times larger.

  • Takeaways & Limitations

    Med-BERT supports disease-prediction modeling when local fine-tuning datasets are small and may reduce the training data needed to reach a given AUC.

  • Takeaways & Limitations

    The study uses only ICD diagnosis information, omits visit-interval lengths, and does not fully explore within-visit concept order.

Abstract

from arXiv · show

Deep learning (DL) based predictive models from electronic health records (EHR) deliver impressive performance in many clinical tasks. Large training cohorts, however, are often required to achieve high accuracy, hindering the adoption of DL-based models in scenarios with limited training data size. Recently, bidirectional encoder representations from transformers (BERT) and related models have achieved tremendous successes in the natural language processing domain. The pre-training of BERT on a very large training corpus generates contextualized embeddings that can boost the performance of models trained on smaller datasets. We propose Med-BERT, which adapts the BERT framework for pre-training contextualized embedding models on structured diagnosis data from 28,490,650 patients EHR dataset. Fine-tuning experiments are conducted on two disease-prediction tasks: (1) prediction of heart failure in patients with diabetes and (2) prediction of pancreatic cancer from two clinical databases. Med-BERT substantially improves prediction accuracy, boosting the area under receiver operating characteristics curve (AUC) by 2.02-7.12%. In particular, pre-trained Med-BERT substantially improves the performance of tasks with very small fine-tuning training sets (300-500 samples) boosting the AUC by more than 20% or equivalent to the AUC of 10 times larger training set. We believe that Med-BERT will benefit disease-prediction studies with small local training datasets, reduce data collection expenses, and accelerate the pace of artificial intelligence aided healthcare.

INTRODUCTION

Deep-learning disease prediction benefits from large annotated EHR datasets, but limited access and unresolved adaptations of BERT to structured EHRs constrain small-data applications. Med-BERT addresses these gaps with domain-specific structured-EHR pre-training, larger cohorts and vocabularies, and specialized clinical tasks.

  • Large, high-quality annotated datasets are typically required for deep-learning models, but rare conditions, annotation difficulty, and governance issues limit access to big EHR data.
  • Structured EHRs resemble text as sequential tokens, but clinical codes and visit ordering require domain-specific input organization.
  • Adapting BERT to structured EHRs remains non-trivial because suitable input organization and domain-specific pre-training tasks are uncertain.
  • BEHRT and G-BERT have limitations involving non-standard AUC definition, insufficient long-term context, and restricted pre-training data or inputs.
  • Med-BERT uses larger vocabulary and pre-training cohorts, longer visit sequences, code ordering, and a prolonged-length-of-stay prediction task for contextual representation learning.
  • The study evaluates Med-BERT by adding it to predictive models, comparing it with other pre-trained approaches, and testing different fine-tuning sample sizes.

METHODS Data Preparation

The study assembled one pre-training cohort and three phenotyped evaluation cohorts from Cerner and Truven databases, with cohort characteristics summarized descriptively. The pre-training cohort contained 28,490,650 unique patients after data-quality exclusions.

  • The study used Cerner HealthFacts for pre-training and evaluation, and Truven Health MarketScan only for evaluation, producing one pre-training and three evaluation cohorts.
  • Table 2 provides descriptive analysis of the cohorts used in the study.

Med-BERT

Med-BERT adapts transformer-based contextual modeling to diagnosis-code sequences in structured EHRs. Its inputs combine code identity, within-visit code order, and visit position, while omitting BERT’s sentence-separation tokens.

  • Med-BERT projects diagnosis codes, within-visit code order, and visit position into code, serialization, and visit embeddings.
  • Code embeddings represent diagnosis codes, serialization embeddings encode relative within-visit order, and visit embeddings distinguish visits in each patient sequence.
  • Unlike BERT, Med-BERT omits [CLS] and [SEP] because next-sentence prediction is not part of its tasks and visit embeddings separate visits.
  • Bidirectional transformer blocks use multihead attention to encode diagnosis-code and visit structures.

Pre-training

Med-BERT uses two pre-training tasks at different granularities: code-level masked prediction and patient-level prediction.

  • Med-BERT includes pre-training tasks at both the code level and the patient level.

(1) Masked Language Model (Masked LM)

Med-BERT uses a masked-language-model objective to learn contextualized representations from structured EHR diagnosis-code sequences. Codes are selectively masked, randomized, or left unchanged, and the model predicts the original code from its context.

  • Masked LM predicts the existence of a diagnosis code from its surrounding EHR context.
  • 80% of selected codes are replaced by [MASK], while 10% are replaced by random codes and 10% remain unchanged.
  • The masking task forms the core of Med-BERT’s contextualized embedding model.

(2) Prediction of prolonged length of stay (Prolonged LOS) in hospital

Med-BERT pre-training combines a patient-level Prolonged LOS classification task with transfer-learning evaluations on realistic disease-prediction cohorts. The study evaluates heart failure in diabetes patients and pancreatic cancer across Cerner and Truven data.

  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: Prolonged LOS is defined as any hospital stay longer than seven days, assessed from each patient’s visit history.
  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: The pre-training classification task targets patient-level Prolonged LOS to reduce pre-training complexity.
  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: The bidirectional structure is suited to Prolonged LOS because illness severity can be measured from EHR information in both temporal directions.
  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: Fine-tuning evaluates heart failure in diabetes patients and pancreatic cancer across three cohorts from two databases.
  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: The evaluation uses realistic disease phenotypes requiring diagnosis codes, prescriptions, procedures, laboratory values, and event frequencies.
  • (2) Prediction of prolonged length of stay (Prolonged LOS) in hospital: Experiments test Med-BERT with GRU, Bi-GRU, and RETAIN, compare it with t-W2V, and assess transfer learning with reduced training data.

RESULTS

Across three disease-prediction tasks, adding Med-BERT generally improves baseline-model AUC, with especially large gains in smaller-data settings. Benefits vary by base model and cohort, and uncertainty is substantial for some low-sample comparisons.

  • Adding Med-BERT improves AUCs across the evaluation tasks, including gains of over 2.5% for strong models trained on 50K samples.
  • For PaCa-Cerner, adding Med-BERT enhances base-model AUCs by 3–7%.
  • For PaCa-Truven, Med-BERT-related models produce performance gains of 2.02–4.02% despite lower average AUCs than PaCa-Cerner.
  • Med-BERT is more helpful for simpler models, following the pattern GRU > Bi-GRU > RETAIN.
  • In DHF-Cerner, Med-BERT raises GRU AUC by about 20% and reaches 0.75 with 500 training samples.
  • In PaCa-Cerner, Bi-GRU with Med-BERT reaches 0.75 using only 300 samples, while PaCa-Truven shows similar but lower overall AUC trends.

Visualization of attention patterns in Med-BERT

Med-BERT’s attention patterns vary across layers, visits, and fine-tuned tasks, revealing contextual and clinically related code connections. These visualizations provide examples of how the model represents temporal and task-specific information.

  • Med-BERT uses attention weights to visualize connections among clinical codes and inspect the model’s learned representations.The visualizations examine attention patterns from the transformer layers.
  • Across six pre-trained BERT layers, early attention is mostly syntactic, while middle layers contain medically meaningful attention patterns.Some attention heads remain within visits, while others connect the same codes across visits.
  • The same type 2 diabetes code attends differently across visits, with later attention spanning current and previous insulin-related codes.This pattern could indicate that segment embeddings capture temporal relationships between visits.
  • Fine-tuned attention patterns differ from pre-training and become task-dependent across layers in the DHF-Cerner and PaCa-Cerner cohorts.Examples show attention converging on related codes for DHF-Cerner and diagnostic-code correlations for PaCa-Cerner.
  • Attention visualizations may help researchers understand neural-network mechanisms and improve communication of health information.The paper presents this as a potential benefit of interpreting learned attention patterns.

DISCUSSION

Med-BERT improves disease-prediction performance across training regimes and is especially effective when fine-tuning with only hundreds of samples. The authors also identify instability and input limitations that constrain interpretation and scope.

  • Med-BERT improves prediction across multiple configurations and is particularly effective for extreme transfer learning with only several hundred fine-tuning samples.The discussion attributes this to pre-trained representations capturing complex input semantics.
  • Masked LM predicts masked codes bidirectionally, while Prolonged LOS predicts whether a patient has a prolonged hospital stay.Both tasks use labels generated without human annotations to reinforce contextual and sequential representation learning.
  • 2.02–3.20%: adding Med-BERT further improves RETAIN across all three tasks, while gains are larger for the simpler GRU and Bi-GRU models.The simpler models reach performance comparable to RETAIN, and Med-BERT_only also achieves good performance.
  • t-W2V helps RETAIN more than Bi-GRU or GRU and can hurt GRU performance, possibly because its shallow structure limits sequential modeling.The explanation is presented as probable rather than established.
  • Under 500 samples, RETAIN curves can heavily overlap, while large standard deviations indicate uneven distributions and insufficiently stable results.These overlaps make it difficult to determine which model is absolutely better in those settings.
  • 500 samples: GRU+Med-BERT reaches AUC 0.75 in PaCa-Cerner, compared with over 5,000 samples for GRU only.The authors describe this as performance on par with a training set almost 10 times larger, potentially reducing annotation costs.
  • Med-BERT’s ICD-9 and ICD-10 vocabulary has broader coverage than BEHRT and G-BERT and is evaluated across Cerner and Truven data.The authors connect this vocabulary choice to potential transferability.
  • The current model uses only ICD diagnosis information, omits visit-time intervals, and does not fully explore within-visit concept ordering.The authors identify adding time, medications, procedures, and laboratory tests as future work.

CONCLUSION

The paper proposes Med-BERT as a contextualized embedding model for structured EHRs and evaluates it on disease-prediction tasks. Experiments and visualizations show promising predictive and interpretive performance, including with few training samples.

  • Med-BERT is a contextualized embedding model pre-trained on large-scale structured EHRs and evaluated on disease-prediction tasks.
  • Domain-specific input formats and pre-training tasks were designed for modeling structured EHR information.
  • Experiments show that Med-BERT can improve baseline deep-learning models across training-set sizes and obtain promising results with very few samples.
  • The visualization module exposed underlying EHR semantics and model mechanisms, with meaningful examples verified by clinical experts.
  • The authors believe Med-BERT may also benefit other clinical problems.

Appendix A: Data Extraction and Preparation

The appendix describes how structured diagnosis data were extracted, ordered, cleaned, and divided into cohorts for pre-training and evaluation. The preparation produced a 28,490,650-patient pre-training dataset and defined disease-specific cohorts using diagnosis and clinical criteria.

  • Cerner Health Facts contains longitudinal patient-level EHR data from over 600 U.S. hospitals and clinics and more than 68 million unique patients.Available data include diagnoses, procedures, laboratory results, medications, encounters, and other clinical observations.
  • Visits were ordered temporally, and diagnosis codes within visits were ranked by admission status, capture phase, and database-provided priority.
  • Diagnosis codes were represented using ICD-9-CM and ICD-10-CM, with within-visit order determined by the ranking criteria.Adding code order during evaluation produced only limited gains compared with randomly scattering codes.
  • 28,490,650 patients remained after removing records with fewer than three diagnosis codes or incorrect time information, then splitting data 7:1:2 into training, validation, and testing sets.
  • The diabetes-heart-failure cohort excluded several diabetes subtypes and required repeated diabetes encounters for type II diabetes-focused analysis.
  • Type II diabetes eligibility was additionally verified using A1C readings above 6.5 or use of specified antidiabetic agents.
  • Heart-failure cases were identified using ICD-9/10 codes and corroborating evidence such as diuretics, high BNP, or relevant procedures.
  • The final DHF cohort contained 39,727 cases and 632,920 controls after exclusions for incorrect data, intervening death records, and age below 18 at diabetes diagnosis.

A.3. Pancreatic cancer cohort (PaCa)

The pancreatic cancer cohort was defined from ICD-9 and ICD-10 codes in Cerner HealthFacts, with further details provided in a supplemental flowchart. A separate evaluation cohort was extracted from Truven MarketScan, while supplemental attention examples illustrate Med-BERT’s learned patterns.

  • Cohort definition: 11,486 pancreatic cancer cases aged 45 years or older, without a prior cancer diagnosis, were eligible from the Cerner HealthFacts cohort.The cohort was initially identified using ICD-9 codes beginning with 157 and ICD-10 codes beginning with C25.
  • Cohort definition: Supplemental Figure 2 presents the flowchart for defining the pancreatic cancer cohort.
  • Evaluation cohort: The Truven Health MarketScan Research Databases were used to extract a separate cohort for evaluation.The 2015 databases integrate de-identified patient-level medical, drug, dental, productivity, laboratory, risk-assessment, hospital-discharge, and electronic-medical-record data.
  • Attention patterns: Attention patterns progress from short-range within-visit links in layer one to token correspondences in layer two and more interpretable relevance patterns in layer three.Supplemental Figure 3 illustrates attention connections from the first three transformer layers for a sample patient sequence.
Loading 2005.12833v1…