Robust ICU Mortality Prediction with Multi-Task Diffusion and Contrastive Learning Frameworks

Abstract

Predicting death in the intensive care unit (ICU) plays an important role in clinical decision-making and patient care to increase hospital performance and help to communicate with patients and families about treatment decisions on time. Machine learning and deep learning have been used. Widely used in ICU patient data to predict mortality. The data are usually time series data, which have common data problems such as missing values and imbalance of classification. This paper presents a Multi-Task Diffusion Model (MTDM) designed to address the dual challenges of missing data and mortality prediction in ICU settings. The Multi-Task Diffusion Model (MTDM) introduces an innovative approach by integrating diffusion models for high-fidelity imputation of incomplete clinical time-series data and an LSTM network for mortality prediction, capturing temporal dependencies. By unifying imputation and prediction tasks, the MTDM ensures seamless optimization, addressing challenges such as noisy and missing data. Furthermore, the Siamese network with contrastive loss enhances feature representation by distinguishing between patient profiles with similar and dissimilar outcomes, enabling nuanced clinical insights. A feedback mechanism between the imputation and prediction models ensures joint optimization, improving overall performance even in the presence of noisy or incomplete data. The proposed MultiTask Diffusion Model (MTDM) demonstrated superior imputation accuracy across varying missing data rates and achieved state-of-the-art performance in mortality prediction when evaluated on the Medical Information Mart for Intensive Care III (MIMIC-III) dataset, Medical Information Mart for Intensive Care IV (MIMIC-IV), and eICU Collaborative Research Database, underlining its robustness and efficacy for critical care applications. The experimental results confirm that integrating diffusion-based imputation with predictive modeling enhances the robustness and reliability of outcomes. The MTDM framework offers a comprehensive solution for ICU mortality prediction, addressing both data quality issues and predictive accuracy to support critical care decision-making. © 2024 N. Buranaburustam, W. Kumwilaisak, C. Hansakunbuntheung, N. Thatphithakkul and K. Kumwilaisak.