An explainable knowledge distillation method with XGBoost for ICU mortality prediction

Abstract

Background and objective: Mortality prediction is an important task in intensive care unit (ICU) for quantifying the severity of patients' physiological condition. Currently, scoring systems are widely applied for mortality prediction, while the performance is unsatisfactory in many clinical conditions due to the non-specificity and linearity characteristics of the used model. As the availability of the large volume of data recorded in electronic health records (EHRs), deep learning models have achieved state-of-art predictive performance. However, deep learning models are hard to meet the requirement of explainability in clinical conditions. Hence, an explainable Knowledge Distillation method with XGBoost (XGB-KD) is proposed to improve the predictive performance of XGBoost while supporting better explainability.

Methods: In this method, we first use outperformed deep learning teacher models to learn the complex patterns hidden in high-dimensional multivariate time series data. Then, we distill knowledge from soft labels generated by the ensemble of teacher models to guide the training of XGBoost student model, whose inputs are meaningful features obtained from feature engineering. Finally, we conduct model calibration to obtain predicted probabilities reflecting the true posterior probabilities and use SHapley Additive exPlanations (SHAP) to obtain insights about the trained model.

Results: We conduct comprehensive experiments on MIMIC-III dataset to evaluate our method. The results demonstrate that our method achieves better predictive performance than vanilla XGBoost, deep learning models and several state-of-art baselines from related works. Our method can also provide intuitive explanations.

Conclusions: Our method is useful for improving the predictive performance of XGBoost by distilling knowledge from deep learning models and can provide meaningful explanations for predictions.

Keywords: Explainable machine learning; Intensive care units; Knowledge distillation; Mortality prediction.