TOO-BERT: A Trajectory Order Objective BERT for self-supervised representation learning of temporal healthcare data

Abstract

Healthcare data accumulation over time, particularly in Electronic Health Records (EHRs), plays a pivotal role by offering a vast repository of patient data with the potential to enhance patient care and predict health outcomes. While Bert-inspired models have shown promises in modeling EHR trajectories, the challenge lies in capturing intricate disease-intervention relationships over time. This study introduces TOO-BERT, enhancing MLM representations by explicitly leveraging sequential patient trajectory information at code and visit levels. TOO-BERT excels in learning frequent sequential patterns by refining the TOO self-supervised objective through two proposed methods, Conditional Code Swapping (CCS) and Conditional Visit Swapping (CVS) weighting functions. Evaluation on MIMIC-IV and Malmö Diet cohort datasets demonstrates TOO-BERT's performance in predicting Heart Failure (HF), Alzheimer's disease (AD), and Prolonged Length of Stay (PLS). Notably, TOO-BERT outperforms Bert in HF prediction, even with limited fine-tuning data. Our findings illustrate the potency of integrating TOO objectives in MLM-based models, enabling intricate EHR data relationships to be captured. Attention analysis highlights the model's capability to learn complex structural patterns.