Contrastive learning with time-aware transformer on EHR data

Abstract

Patient representation learning has been a hot topic in the medical fields in recent years due to the fact that personal indexes of each person can be different from the population levels. Therefore, researchers have been trying contrastive learning using EHR data to achieve a personalized representation of each patient, resulting in more individual and specific recommendations.

Earlier contrastive learning work has made use of scan images or clinical notes during hospitalization to analyze the survival status or predict the upcoming events of patients. However, the impact of both categorical and numerical data in a multimodal setting has not yet been focused on, even though hospitals and clinical sites have a considerable amount of such information collected over years. Therefore, in this paper, we concentrate on contrastive learning using multimodal data, also including special types with hierarchical structure such as ICD-10 and medication codes.

We conduct exhaustive experiment with different types of data to see how each of them can be augmented for a contrastive learning framework. We also come up with a novel pretrained model schema for understanding the hierarchy of ICD-10 and medication codes, which is integrated into a unified multimodal model. For learning the temporal and semantic information of patients event sequences in a unified manner, we follow SimCLR framework with TAAT being the main encoder module and InfoNCE loss acting as the objective function.

The obtained results show that our pretrained models can perfectly understand the inherent hierarchy within ICD-10 and medication codes, and our augmentation methods work very well for both normal data types and special data types with hierarchical structure. The preliminary findings from the multimodal model suggest that it has learned meaningful latent representations for patients and can form clusters of the ones with similar features, and these representations can be utilized for downstream tasks such as prediction and classification. We believe this would contribute a further step toward the goal of personalized healthcare with artificial intelligence.