Learning Graphical Structure of Electronic Health Records with Transformer for Predictive Healthcare

Effective modeling of electronic health records (EHR) is rapidly becoming an important topic in both academia and industry. A recent study showed that utilizing the graphical structure underlying EHR data improves the performance of prediction tasks such as heart failure diagnosis prediction. However, EHR data do not always contain the complete structural information. Moreover, when it comes to claims data, they do not have any structural information to begin with. Under such circumstances, can we still do better than just treating EHR data as a flat-structured bag-of-features? In this paper, we study the possibility of utilizing the implicit structure of EHR by using Transformer for prediction tasks on public electronic health records. Specifically, we make a connection between graph networks and Transformer, then use a variant of Transformer on encounter-based prediction tasks such as medication prediction and masked node prediction. Our model empirically demonstrates superior prediction performance to previous approaches on two publicly available EHR datasets, indicating that it can serve as an effective general-purpose representation learning algorithm for EHR data.