Compound Density Networks for Risk Prediction using Electronic Health Records

TL;DR

This paper introduces CDNet, an integrated model combining imputation and prediction for EHR data, which improves mortality prediction accuracy and uncertainty estimation by jointly tuning imputation and prediction components.

Contribution

The paper presents a novel end-to-end framework, CDNet, that jointly optimizes imputation and prediction in EHR data using a combination of GRU, MDN, and RAN, outperforming existing methods.

Findings

Outperforms state-of-the-art models on MIMIC-III mortality prediction.

Regularizing imputed values enhances prediction accuracy.

Captures both aleatoric and epistemic uncertainties.

Abstract

Electronic Health Records (EHRs) exhibit a high amount of missing data due to variations of patient conditions and treatment needs. Imputation of missing values has been considered an effective approach to deal with this challenge. Existing work separates imputation method and prediction model as two independent parts of an EHR-based machine learning system. We propose an integrated end-to-end approach by utilizing a Compound Density Network (CDNet) that allows the imputation method and prediction model to be tuned together within a single framework. CDNet consists of a Gated recurrent unit (GRU), a Mixture Density Network (MDN), and a Regularized Attention Network (RAN). The GRU is used as a latent variable model to model EHR data. The MDN is designed to sample latent variables generated by GRU. The RAN serves as a regularizer for less reliable imputed values. The architecture of CDNet enables GRU and MDN to iteratively leverage the output of each other to impute missing values, leading to a more accurate and robust prediction. We validate CDNet on the mortality prediction task on the MIMIC-III dataset. Our model outperforms state-of-the-art models by significant margins. We also empirically show that regularizing imputed values is a key factor for superior prediction performance. Analysis of prediction uncertainty shows that our model can capture both aleatoric and epistemic uncertainties, which offers model users a better understanding of the model results.