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

This research presents the Multi-Task Diffusion Model (MTDM), which addresses issues of missing data and mortality prediction by utilizing diffusion models for data imputation and Long Short-Term Memory (LSTM) for outcome prediction. To enhance feature extraction, the framework employs a Siamese network with contrastive loss, distinguishing between patient profiles with similar and dissimilar outcomes. Additionally, 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 model is evaluated on the Medical Information Mart for Intensive Care III (MIMIC-III) dataset, achieving imputation accuracy across multiple missing data rates and performance in mortality prediction, with a ROC-AUC score of 0.92. 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.