AD-BERT: Using pre-trained language model to predict the progression from mild cognitive impairment to Alzheimer's disease

Abstract

Objective: We develop a deep learning framework based on the pre-trained Bidirectional Encoder Representations from Transformers (BERT) model using unstructured clinical notes from electronic health records (EHRs) to predict the risk of disease progression from Mild Cognitive Impairment (MCI) to Alzheimer's Disease (AD).

Methods: We identified 3657 patients diagnosed with MCI together with their progress notes from Northwestern Medicine Enterprise Data Warehouse (NMEDW) between 2000 and 2020. The progress notes no later than the first MCI diagnosis were used for the prediction. We first preprocessed the notes by deidentification, cleaning and splitting into sections, and then pre-trained a BERT model for AD (named AD-BERT) based on the publicly available Bio+Clinical BERT on the preprocessed notes. All sections of a patient were embedded into a vector representation by AD-BERT and then combined by global MaxPooling and a fully connected network to compute the probability of MCI-to-AD progression. For validation, we conducted a similar set of experiments on 2563 MCI patients identified at Weill Cornell Medicine (WCM) during the same timeframe.

Results: Compared with the 7 baseline models, the AD-BERT model achieved the best performance on both datasets, with Area Under receiver operating characteristic Curve (AUC) of 0.849 and F1 score of 0.440 on NMEDW dataset, and AUC of 0.883 and F1 score of 0.680 on WCM dataset.

Conclusion: The use of EHRs for AD-related research is promising, and AD-BERT shows superior predictive performance in modeling MCI-to-AD progression prediction. Our study demonstrates the utility of pre-trained language models and clinical notes in predicting MCI-to-AD progression, which could have important implications for improving early detection and intervention for AD.

Keywords: Alzheimer's disease; Electronic health records; Mild cognitive impairment; Pre-trained language model.

Figure 1

Inclusion and exclusion criteria for the study cohorts for (a) NMEDW and (b) WCM.

Figure 2

The overview of our framework. The notes of a patient are split into sections, which are then fed to the pretrained AD-BERT model to generate a representation for each section. The patient representation is generated by global MaxPooling that aggregates all the section representations. Finally, a linear classifier combined with a sigmoid activation layer is used to predict probability of MCI-to-AD progression.

Figure 3

The illustration of samples in case and control groups. (a) For no-restrict prediction, the case and control groups are differentiated by the AD diagnosis condition after MCI diagnosis, as reflected in all diagnostic records. (b) For x-month prediction, in addition to the AD diagnosis condition within x months after MCI diagnosis, we also enforce a time constraint on the control group by requiring the last encounter to occur after x months to ensure the patients have a conversion time of at least x months.

Figure 4

Attention visualization of AD-BERT. The model pays more attention to the terms like “memory”, “MCI” and “difficulty recalling dates” than others.