APSIPA Transactions on Signal and Information Processing > Vol 14 > Issue 1

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

Namtip Buranaburustam, University of Technology Thonburi, Thailand, Wuttipong Kumwilaisak, University of Technology Thonburi, Thailand, Chatchawarn Hansakunbuntheung, National Science and Technology Development Agency, Thailand, Nattanun Thatphithakkul, National Science and Technology Development Agency, Thailand, nattanun.tha@nstda.or.th , Kanya Kumwilaisak, Chulalongkorn University, Thailand
 
Suggested Citation
Namtip Buranaburustam, Wuttipong Kumwilaisak, Chatchawarn Hansakunbuntheung, Nattanun Thatphithakkul and Kanya Kumwilaisak (2025), "Robust ICU Mortality Prediction with Multi-Task Diffusion and Contrastive Learning Frameworks", APSIPA Transactions on Signal and Information Processing: Vol. 14: No. 1, e15. http://dx.doi.org/10.1561/116.20240085

Publication Date: 18 Jun 2025
© 2025 N. Buranaburustam, W. Kumwilaisak, C. Hansakunbuntheung, N. Thatphithakkul and K. Kumwilaisak
 
Subjects
Classification and prediction,  Applications and case studies,  Deep learning,  Information extraction
 
Keywords
Missing data imputationmortality predictiondiffusion modelmulti-task learningcontrastive learning
 

Share

Open Access

This is published under the terms of CC BY-NC.

Downloaded: 26 times

In this article:
Introduction 
Related Work 
Dataset and Preprocessing Methodology 
Imputation of Missing Data in the ICU Dataset using Diffusion Probabilistic Models 
Mortality Prediction with LSTM Networks 
The Multi-Task Diffusion Model for Mortality Prediction using Contrastive Loss 
Experimental Results 
Conclusion 
References 

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 Multi- Task 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.

DOI:10.1561/116.20240085