[Paper Review] Interpretable Survival Analysis for Heart Failure Risk Prediction

Mike Van Ness, Tomas M. Bosschieter+4 more|arXiv (Cornell University)|Oct 24, 2023

Machine Learning in Healthcare4 citations

TL;DR

This paper proposes a novel, interpretable survival analysis pipeline for predicting heart failure risk using electronic health records. By combining improved survival stacking, ControlBurn for feature selection, and Explainable Boosting Machines, the method achieves state-of-the-art performance while maintaining clinical interpretability, revealing novel insights into heart failure risk factors.

ABSTRACT

Survival analysis, or time-to-event analysis, is an important and widespread problem in healthcare research. Medical research has traditionally relied on Cox models for survival analysis, due to their simplicity and interpretability. Cox models assume a log-linear hazard function as well as proportional hazards over time, and can perform poorly when these assumptions fail. Newer survival models based on machine learning avoid these assumptions and offer improved accuracy, yet sometimes at the expense of model interpretability, which is vital for clinical use. We propose a novel survival analysis pipeline that is both interpretable and competitive with state-of-the-art survival models. Specifically, we use an improved version of survival stacking to transform a survival analysis problem to a classification problem, ControlBurn to perform feature selection, and Explainable Boosting Machines to generate interpretable predictions. To evaluate our pipeline, we predict risk of heart failure using a large-scale EHR database. Our pipeline achieves state-of-the-art performance and provides interesting and novel insights about risk factors for heart failure.

Motivation & Objective

• To address the limitations of traditional Cox models in handling non-proportional hazards and non-linear effects in survival analysis for heart failure prediction.
• To develop a machine learning-based survival pipeline that maintains high interpretability without relying on post-hoc explanation methods.
• To improve predictive accuracy over standard Cox models while preserving clinical interpretability for real-world healthcare deployment.
• To identify novel, clinically relevant risk factors for heart failure using interpretable modeling on large-scale EHR data.

Proposed method

• Uses an improved, scalable version of survival stacking to convert survival problems into binary classification tasks without estimating survival times.
• Applies ControlBurn for feature selection to reduce multicollinearity and improve model stability by selecting relevant, low-correlation features.
• Employs Explainable Boosting Machines (EBM) as the final estimator, leveraging generalized additive models for inherently interpretable predictions.
• Transforms risk sets from survival data into training samples for classification, preserving time-to-event information without pseudovalues.
• Integrates the pipeline end-to-end to handle right-censored data, optimize for discrimination, and ensure model transparency.
• Uses EHR data with temporal features (e.g., mean, max, min, last measurements) across multiple clinical variables to capture longitudinal risk patterns.

Figure 1: A summary of our interpretable survival analysis pipeline. During training (left), we use ControlBurn for feature selection (Section 3.3 ), survival stacking with subsampling to cast the survival data to classification data (Algorithm 1 ), and EBMs for generating feature importances and sh

Experimental results

Research questions

• RQ1 Can a survival analysis pipeline be designed to achieve state-of-the-art predictive performance while maintaining high interpretability for clinical use?
• RQ2 How does the proposed pipeline compare to standard Cox models and other machine learning survival models in terms of discrimination and calibration?
• RQ3 What novel clinical risk factors for heart failure emerge from the interpretable model that may be overlooked by traditional methods?
• RQ4 To what extent does feature selection via ControlBurn improve model stability and predictive performance in high-dimensional EHR data?
• RQ5 Can survival stacking without survival time estimation outperform pseudo-value-based approaches in terms of scalability and interpretability?

Key findings

• The proposed pipeline achieves state-of-the-art performance in predicting heart failure risk on a large-scale EHR database, outperforming standard Cox models and other machine learning baselines.
• The model identifies previously underappreciated risk factors such as long-term trends in body weight, creatinine, and QRS duration as strong predictors of heart failure.
• ControlBurn effectively reduces feature correlation and improves model stability, leading to more reliable and interpretable feature importance rankings.
• Explainable Boosting Machines provide high-fidelity, locally interpretable predictions that align with clinical intuition and support clinical decision-making.
• Survival stacking enables accurate modeling of time-to-event data without requiring survival time estimation, avoiding biases inherent in pseudo-value methods.
• The pipeline reveals that dynamic changes in biomarkers—such as rising creatinine and declining eGFR—carry significant predictive power beyond static baseline values.

Figure 2: Distribution of predicted survival probabilities $S(t\mid X_{i})=P(T_{i}>t)$ at $t=5$ across test patients. The left plot shows the survival probabilities using Eq. ( 2 ) as in (Craig et al., 2021 ) , while the right plot shows Eq. ( 5 ). Our survival prediction method gives a reasonable d