PaperlayerView exact revision on arXiv

Source-linked AI summary

CEHR-BERT: Incorporating temporal information from structured EHR data to improve prediction tasks

Chao Pang, Xinzhuo Jiang, Krishna S Kalluri, Matthew Spotnitz, RuiJun Chen, Adler Perotte, Karthik Natarajan

arXiv:2111.08585v1cs.LG

TL;DR

Structured-EHR BERT adaptations do not fully incorporate temporal information across clinical domains, while common patient sequences omit intervals between visits. CEHR-BERT combines artificial time tokens, time and age embeddings, and visit-type prediction. It achieved the best AUC and PR-AUC across four prediction tasks among the evaluated models, and the authors conclude that combining its temporal components is an effective way to capture temporal information in structured EHR data.

  • Problem

    Structured-EHR BERT adaptations do not fully incorporate temporal information across clinical domains, and patient sequences can omit intervals between visits.

    Source: §1 · p. 2

  • Method

    CEHR-BERT combines artificial inter-visit time tokens and concept, age, and time embeddings with a second learning objective that predicts visit type.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    CEHR-BERT achieved the best AUC and PR-AUC across all four prediction tasks; for t2dm hf, its PR-AUC was 0.323 versus MedBert’s 0.274.

    Source: §5.2.1 · p. 5

  • Takeaways & Limitations

    The authors conclude that combining time tokens with time and age embeddings is the most effective way they tested to capture temporal information in structured EHR data.

    Source: §8 · p. 9

    The study used a context window of 300, while patient histories may exceed BERT’s 512-token sequence limit; the authors also note that VTP’s small apparent improvement may reflect random variation.

    Source: §7 · p. 9

Abstract

from arXiv · show

Embedding algorithms are increasingly used to represent clinical concepts in healthcare for improving machine learning tasks such as clinical phenotyping and disease prediction. Recent studies have adapted state-of-the-art bidirectional encoder representations from transformers (BERT) architecture to structured electronic health records (EHR) data for the generation of contextualized concept embeddings, yet do not fully incorporate temporal data across multiple clinical domains. Therefore we developed a new BERT adaptation, CEHR-BERT, to incorporate temporal information using a hybrid approach by augmenting the input to BERT using artificial time tokens, incorporating time, age, and concept embeddings, and introducing a new second learning objective for visit type. CEHR-BERT was trained on a subset of Columbia University Irving Medical Center-York Presbyterian Hospital's clinical data, which includes 2.4M patients, spanning over three decades, and tested using 4-fold cross-validation on the following prediction tasks: hospitalization, death, new heart failure (HF) diagnosis, and HF readmission. Our experiments show that CEHR-BERT outperformed existing state-of-the-art clinical BERT adaptations and baseline models across all 4 prediction tasks in both ROC-AUC and PR-AUC. CEHR-BERT also demonstrated strong transfer learning capability, as our model trained on only 5% of data outperformed comparison models trained on the entire data set. Ablation studies to better understand the contribution of each time component showed incremental gains with every element, suggesting that CEHR-BERT's incorporation of artificial time tokens, time and age embeddings with concept embeddings, and the addition of the second learning objective represents a promising approach for future BERT-based clinical embeddings.

1. Introduction

BERT adaptations for structured EHR data have shown promise, but common patient-sequence representations omit intervals between visits and leave BERT’s next-sentence objective inapplicable. CEHR-BERT addresses these gaps with temporal representations and a new visit-type learning objective.

1. Introduction

  • Few efforts have applied BERT to structured EHR data for contextualized concept embeddings, although early adaptations improved on word2vec and GloVe.Contextualized BERT representations have achieved state-of-the-art performance on NLP benchmark tasks.Source: §1 · p. 1
  • Chronologically ordered concept sequences omit time intervals between visits, limiting temporal information available to BERT for downstream prediction.Because a patient's full history is treated as one sentence, Next Sentence Prediction does not apply; replacing it with masked language modeling alone may leave other EHR information unused.Source: §1 · p. 2
  • CEHR-BERT adds artificial tokens for inter-visit intervals, combines age and time embeddings with concept embeddings, and adds Visit Type Prediction as a second learning objective.The approach adapts BERT for structured EHR data to improve disease prediction.Source: §1 · p. 2

2. Related work

Earlier structured-EHR BERT adaptations varied in clinical scope, temporal representation, and learning objectives. CEHR-BERT takes a different approach to representing inter-visit time.

2. Related work

  • Recent structured-EHR BERT studies often focused on a single clinical domain or visit, or gave limited consideration to time.These approaches did not fully use the richness of patients’ full medical histories.Source: §2 · p. 2
  • BEHRT represented patient history with concept, visit-segment, age, and positional embeddings, but used diagnoses alone and evaluated diagnosis-code prediction.Its SEP tokens marked visit boundaries; procedures and medications were excluded.Source: §2 · p. 2
  • G-BERT used a graph neural network guided by medical ontologies, but targeted medication recommendation with single-visit input and 20K pre-training patients.Source: §2 · p. 2
  • MedBert omitted temporal information from its patient representation, used prolonged length of stay as a second objective, and evaluated two disease-prediction tasks.It was trained on 20 million patients and fine-tuned using a different data source.Source: §2 · p. 2
  • Prior temporal approaches incorporated intervals between neighboring clinical events, whereas CEHR-BERT introduced artificial time tokens into the patient sequence.Source: §2 · p. 2

3. Data and Preprocessing

The study used CUIMC-NYP’s OMOP data, drawing on three clinical domains for model training and prediction cohorts. Patient histories were represented as chronological visit sequences with artificial tokens marking inter-visit intervals.

3.1. Data

  • Training and prediction data came from CUIMC-NYP’s OMOP instance, which spans from the early 1980s to the present and includes multiple clinical domains.For pre-training and fine-tuning, the study used conditions, procedures, and medications.Source: §3.1 · p. 3

3.2. Data processing and patient representation

  • Patient histories were assembled chronologically from medical codes across three domains, with artificial time tokens marking intervals between neighboring visits.Intervals under 28 days used week tokens, intervals from 28 to 365 days used month tokens, and longer intervals used a Long Term token; visit-start and visit-end tokens marked each visit.Source: §3.2 · p. 3
  • The paper refers to this EHR representation as patient sequences: lists of visits, each containing a list of medical concepts.Source: §3.2 · p. 3

4. Methods

CEHR-BERT combines concept, age, and time representations before feeding them to BERT, while retaining masked language modeling and adding visit-type prediction. Its visit-type decoder uses contextualized concepts to predict masked visit types.

4. Methods

  • Concept, age, and time embeddings are concatenated and passed through a fully connected layer to form temporal concept embeddings for BERT.Visit-segment embeddings mark visit boundaries; time embeddings represent absolute time and age embeddings represent relative time with respect to visits.Source: §4 · p. 3, Figure 1 · p. 4
  • Visit Type Prediction adds a decoder objective that predicts masked visit types from contextualized concept embeddings, alongside masked language modeling.The visit-type sequence undergoes the same temporal transformation as the concept sequence, and multi-headed attention combines the resulting embeddings.Source: §4 · pp. 3–4

5. Experiments and Results

CEHR-BERT was evaluated on four prediction tasks and outperformed the comparison models across tasks and metrics. Few-shot results and ablations examine its performance with less training data and the contributions of its temporal components.

5.1. Experiment setup

  • 2.4M patients and 184.7M clinical data points across condition, procedure, and medication records were used for pre-training, with a 300-token context window.The context window captures more than 90% of patients’ entire medical histories.Source: §5.1 · p. 4

5.2.1. Disease prediction

  • Four prediction tasks were evaluated using a one-year observation window by default, with four-fold splits and ROC-AUC and PR-AUC reporting.Sequence models used 75:10:15 train/validation/test splits; frequency-based models used 85% for training and 15% for testing.Source: §5.2.1 · p. 5

5.2.2. Performance of few-shot learning

  • 5% of training data gave CEHR-BERT an AUC of nearly 0.78 and PR-AUC of almost 0.29 for t2dm hf, while comparison models using up to 80% achieved AUCs of 0.60–0.76 and PR-AUCs of 0.12–0.26.The same overall trend held across other prediction tasks, with hospitalization as an exception in the relative performance of BERT models.Source: §5.2.2 · p. 6

5.3. Ablation studies

  • Ablations removed one component at a time; Table 9 also reports network size and pre-training time to inform model choice by use case.The ablation variants were evaluated using the disease-prediction procedure.Source: §5.3 · pp. 6–7

5.2.1. Disease prediction

  • CEHR-BERT achieved the best AUC and PR-AUC across all four prediction tasks, while sequence models outperformed frequency-based models overall.MedBert and BEHRT were the next top-performing models; hospitalization was an exception to their usual ordering, with Bi-LSTM second after CEHR-BERT.Source: §5.2.1 · pp. 5–6, Table 3 · p. 7

5.3.1. Assessment of time tokens

  • Artificial time tokens improved performance over MedBERT- and BEHRT-style representations in all tasks except discharge home death PR-AUC, where M-BERT was slightly better.Removing visit-start and visit-end tokens reduced performance to a level comparable to the SEP-token variant, suggesting that visit boundaries matter for the time tokens’ role.Source: §5.3.1 · p. 7

5.3.2. Assessment of visit type prediction

  • Visit-type prediction improved performance across tasks except hospitalization, where NS-BERT slightly outperformed CEHR-BERT and both remained largely similar.The objective assumes that concepts associated with different visit types follow different distributions; the paper reports distinct frequent-concept distributions by visit type.Source: §5.3.2 · p. 7

5.3.3. Assessment of time/age embeddings

  • 32.3±1.0% to 28.2±0.2% PR-AUC in t2dm hf for CEHR-BERT versus NT-BERT; CEHR-BERT also outperformed the summation-based ALT-BERT.NT-BERT was slightly better for HF readmission, while CEHR-BERT outperformed it on the other tasks.Source: §5.3.3 · p. 8

6. Discussion

The discussion examines what ATT embeddings represent and whether sequence inflation explains CEHR-BERT’s performance gain.

6.1. Understanding artificial time tokens

  • PCA placed ATT tokens in increasing interval order from right to left, with week tokens clustered and month tokens roughly linear.VS and VE were far from time tokens, consistent with encoding visit boundaries rather than time intervals.Source: §6.1 · p. 8
  • The analysis suggests BERT learns meaningful ATT relationships, including associations between different time tokens and groups of concepts.Source: §6.1 · p. 9

6.2. Inflated patient sequence

  • In HF readmission, shortening other models to a 180-day window lowered performance versus the original evaluation, suggesting CEHR-BERT’s gain was attributable to ATT tokens.ATTs lengthened patient sequences and caused some records to be cut off.Source: §6.2 · p. 9

7. Future work

The authors identify context length and future model extensions as areas for further work.

  • The experiments used a 300-token context window, while some patient histories may exceed BERT’s 512-token sequence limit.The authors propose sliding windows to create regional contextualized representations, which could be combined using a 1D convolution or LSTM.Source: §7 · p. 9
  • Future work includes investigating VTP’s small, possibly random-variation gain, embedding continuous lab data with a different strategy, and testing whether time methods generalize to other models.Source: §7 · p. 9

8. Conclusion

The conclusion presents CEHR-BERT as a way to incorporate temporal information into structured-EHR BERT models, combining artificial time tokens with time and age embeddings. The authors report that it outperforms existing state-of-the-art BERT-based approaches across several prediction tasks.

  • CEHR-BERT outperforms existing state-of-the-art BERT-based approaches across several prediction tasks, and the authors identify both time tokens and combined time/age embeddings as the most effective approach.Source: §8 · p. 9

Appendix A. Prediction Tasks

Prediction tasks are framed by target and outcome cohorts, with an index event and prediction window determining who is labeled positive. The paper describes observation, hold-off, and prediction windows and validates on four downstream prediction problems.

Appendix A. Prediction Tasks

  • A prediction is positive when the outcome index date falls between the target cohort’s index date and the end of its prediction window.Target cohorts define the initial group; outcome cohorts are subsets meeting the event criteria.Source: Appendix A · p. 12
  • Features come from the observation window, followed by a hold-off window, and the prediction window follows the index event.Patient histories are chronological sequences of medical engagements; the paper focuses on condition, medication, and procedure events.Source: Appendix A · p. 12
  • The model is validated on four downstream prediction problems, with cohort-definition parameters listed in Table 5.Figure 2 visualizes the cohort definition against events in patient medical histories.Source: Appendix A · p. 12, Figure 2 · p. 12

A.1. Feature Engineering

Baseline frequency features are constructed by rolling up clinical codes where possible and counting their occurrences within the observation window.

A.1. Feature Engineering

  • Diagnosis and procedure ICD-9 or ICD-10 codes are rolled up to two- or three-digit codes, while CPT procedure codes are rolled up to the second level from the root.Medication codes are rolled up to the ingredient level.Source: §A.1 · p. 12
  • Codes that cannot be rolled up retain their original form.This applies after the specified roll-up rules for diagnosis, procedure, and medication codes.Source: §A.1 · p. 12
  • The baseline features count rolled-up codes within the observation window.The passage specifies frequency-based features for logistic regression and XG-boost.Source: §A.1 · p. 12

A.2. Type 2 Diabetes Mellitus patients who developed Heart Failure

The appendix defines cohorts for heart failure prediction, readmission, death after discharge, and hospitalization, with additional figures and tables documenting related analyses.

A.2. Type 2 Diabetes Mellitus patients who developed Heart Failure

  • The T2DM target cohort includes patients with T2DM history, while the outcome cohort meets heart-failure diagnosis, high-BNP lab, or specified treatment criteria during the prediction window.The cohort definition also excludes several other diabetes types; related OMOP concept sets are listed in Tables 6 and 7.Source: §A.2 · p. 12

A.3. Heart Failure patients who were readmitted within 30 days

  • For heart-failure readmission, the target cohort has an inpatient visit with a heart-failure diagnosis, and readmission within 30 days defines the outcome.The prediction window is 30 days.Source: §A.3 · p. 13

A.4. Patients who were discharged and died within one year

  • For death after discharge, the target cohort includes inpatient visits discharged home, and death within 360 days is the outcome.The index visit uses visit concept id 9201 or 262, with discharge to home or other nursing facilities.Source: §A.4 · p. 13

A.5. Hospitalization

  • Hospitalization uses a distinct cohort structure: the first hospital visit is the index event, observation follows it, and an inpatient visit is the outcome.Patients must have 2–30 visit occurrences in the observation window.Source: §A.5 · p. 13, Figure 3 · p. 13

Appendix B. Additional Figures and Analyses

  • Appendix B includes tables of T2DM and heart-failure cohort concepts, readmission AUC and PR-AUC, and parameter counts across BERT representations.The cited table titles identify these topics but provide no cell values here.Source: Table 6 · p. 14, Table 7 · p. 15, Table 8 · p. 15
  • Appendix B also includes an EHR representation figure and a PCA visualization of artificial time-token embeddings.The PCA figure describes a two-dimensional visualization of the tokens’ base embeddings.Source: Figure 4 · p. 16, Figure 5 · p. 17
  • Appendix B plots AUC and PR-AUC across training percentages for few-shot tasks, including HF readmission and hospitalization.The listed figures also cover discharge-home death; separate tables report visit-type concepts and patient counts or percentages.Source: Figure 6 · p. 18, Figure 7 · p. 18, Figure 8 · p. 19
Loading 2111.08585v1…