A Three-Stage Prediction Model Based on Transformer for Chronic Kidney Disease

Chronic kidney disease (CKD) is a serious global health threat. At the terminal stage, kidney function is nearly completely lost. Therefore, predicting the development of CKD based on a patient's visits can enable doctors to intervene early and delay the disease's progression. In this paper, we propose a three-stage prediction model named Imputation-Capture-Prediction (ICP) and based on the Transformer architecture, for chronic kidney disease (CKD) using electronic health records (EHRs). The first stage is to address the missing data problem in EHR, and ICP employs a two-stage imputation method, using the deep learning method SAITS module after recent padding. The second stage is designed to better capture this temporal dependency and the relationships between features, where ICP incorporates a two-branch architecture and introduces two modules: Time-Aware Convolution (TC) and Dynamic-Static-Medical Graph Attention Network (DSMGAT), to extract diverse feature information. The TC module is designed to capture the relationships within visit records, accounting for the unequal lengths of visit intervals while emphasizing the importance of recent records. The DSMGAT module, on the other hand, considers various categories of record features, using a Graph Attention Network (GAT) with learnable weights to model the relationships among them. Then we use a Feed-Forward Network to predict the estimated glomerular filtration rate (eGFR). To evaluate the effectiveness of our method, we compared it with several advanced approaches using a real EHR dataset, TFHCKD. The Mean Absolute Error (MAE) and Mean Squared Error (MSE) were 0.0344 and 0.0028, respectively, demonstrating a significant improvement over existing methods.