Uncertainty-Aware Pre-Trained Foundation Models for Patient Risk Prediction via Gaussian Process

Jiaying Lu, Shifan Zhao, Wenjing Ma, Hui Shao, Xiao Hu, Yuanzhe Xi, Carl Yang

Abstract: Patient risk prediction models are crucial as they enable healthcare providers to proactively identify and address potential health risks. Large pre-trained foundation models offer remarkable performance in risk prediction tasks by analyzing multimodal patient data. However, a notable limitation of pre-trained foundation models lies in their deterministic predictions (i.e., lacking the ability to acknowledge uncertainty). We propose Gaussian Process-based foundation models to enable the generation of accurate predictions with instance-level uncertainty quantification, thus allowing healthcare professionals to make more informed and cautious decisions. Our proposed approach is principled and architecture-agnostic. Experimental results show that our proposed approach achieves competitive performance on classical classification metrics. Moreover, we observe that the accuracy of certain predictions is much higher than that of the uncertain ones, which validates the uncertainty awareness of our proposed method. Therefore, healthcare providers can trust low-uncertainty predictions and conduct more comprehensive investigations on high-uncertainty predictions, ultimately enhancing patient outcomes with less expert intervention.