Heart Failure Mortality Prediction
An empirical machine learning investigation on the UCI Heart Failure Clinical Records cohort using transfer learning strategies and SHAP explainability to evaluate mortality risk discrimination. Published and presented in IEEE ICIMCIS 2025.
Problem Statement & Clinical Motivation
Cardiovascular diseases, particularly heart failure, remain a leading cause of mortality globally. Early and accurate mortality risk prediction can help clinicians allocate intensive monitoring and interventions.
Clinical tabular datasets frequently suffer from small sample sizes, severe class imbalance (substantially fewer mortality events compared to survivors), and non-linear physiological interactions.
Standard deep neural networks often underperform gradient-boosted decision trees on tabular domains, and black-box models lack the transparency demanded by medical professionals.
Investigation Goals
- 01Investigate whether TabNet, an attention-based neural architecture designed specifically for tabular data, can achieve competitive discrimination when combined with transfer learning.
- 02Evaluate four distinct experimental transfer learning scenarios to determine the most effective knowledge transfer mechanism for clinical mortality prediction.
- 03Apply SHAP (SHapley Additive exPlanations) to interpret model predictions and validate learned risk factors against recognized cardiology evidence.
Research Approach & Core Decisions
Employed TabNet with sequential attention mechanisms and sparse feature selection masks. Preprocessed clinical data with standardized scaling and addressed severe class imbalance using SMOTE strictly inside training folds. Formulated four progressive experimental scenarios to evaluate transfer learning effects.
TabNet as Core Architecture
TabNet uses sequential attention steps to select sparse subsets of features at each decision step, combining neural representation learning with tree-like decision paths.
Fold-Isolated SMOTE Resampling
Applied Synthetic Minority Over-sampling Technique (SMOTE) strictly within training folds to avoid data leakage into validation and testing partitions.
Clinical Metric Priority (Recall & PR-AUC)
Prioritized Recall and Precision-Recall AUC over standard Accuracy, because missing a high-risk mortality patient (false negative) carries critical clinical consequences.
Global and Local SHAP Attribution
Used SHAP values to explain both overall biomarker importance rankings and patient-level risk trajectories for clinical transparency.
Multi-Scenario Empirical Evaluation
Four TabNet-based experimental scenarios were evaluated across accuracy, precision, recall, F1 score, ROC AUC, and PR AUC. Each scenario was independently executed 20 times using identical data splits and hyperparameters, with 95% confidence intervals estimated on an 80/20 train-test partition.
Metric Comparison Across Scenarios (S1 β S4)
Mean Performance Matrix (20 Independent Runs)
Source: Table III (ICIMCIS 2025)| Scenario | Architecture Pipeline | Accuracy | Precision | Recall | F1 Score | ROC AUC | PR AUC |
|---|---|---|---|---|---|---|---|
| S1 β Baseline TabNet | UCI β TabNet | 0.704 | 0.538 * | 0.658 | 0.585 | 0.763 | 0.629 |
| S2 β In-Domain Pretraining | UCI β Reconstruction Pretraining β TabNet | 0.596 | 0.443 | 0.813 | 0.563 | 0.737 | 0.604 |
| S3 β Cross-Domain Transfer | CVD Dataset β Feature Alignment β Pretraining β UCI Fine-tuning β TabNet | 0.683 | 0.500 | 0.761 | 0.597 | 0.757 | 0.628 |
| S4 β Feature-Based Transfer | UCI β LR / RF / XGB Models β Probabilistic Meta-Features β TabNet | 0.662 | 0.487 | 0.937 * | 0.639 * | 0.869 * | 0.745 * |
In mortality risk screening, Recall (Sensitivity) is prioritized over raw precision because missing an at-risk mortality event carries severe clinical consequences. Scenario 4 (Feature-Based Transfer) achieves the highest Recall (0.937), ROC AUC (0.869), PR AUC (0.745), and F1 (0.639), successfully identifying the vast majority of positive mortality cases.
Conversely, Scenario 1 (Baseline TabNet) maintains the highest Precision (0.538) with a more conservative decision threshold, resulting in fewer false alarms at the expense of lower sensitivity (Recall 0.658).
Transfer Learning Pipeline & Scenarios
Four distinct transfer learning configurations were designed and benchmarked to determine the optimal mechanism for transferring tabular representation knowledge to a small-sample clinical mortality dataset.
Baseline TabNet
Standard TabNet architecture trained from scratch directly on the target clinical cohort with fold-isolated SMOTE resampling.
In-Domain Pretraining
Self-supervised masked tabular reconstruction on the same clinical cohort to pretrain encoder representations before supervised fine-tuning.
Cross-Domain Pretraining
Pretraining on the external 70,000-record Cardiovascular Disease dataset across aligned common attributes, then transferring encoder weights to UCI fine-tuning.
Auxiliary Teacher Meta-Features + TabNet
Out-of-fold probabilistic predictions from Logistic Regression, Random Forest, and XGBoost models are concatenated with raw clinical features as inputs into TabNet.
p_lr, p_rf, and p_xgb are model-derived probabilistic meta-features generated via out-of-fold estimation, not primary clinical measurements.Dual-Domain Datasets & Feature Alignment
To evaluate cross-domain pretraining, the research combined a target clinical heart failure cohort with a large-scale cardiovascular disease source dataset across strictly aligned common attributes.
UCI Heart Failure Clinical Records
Cohort of heart failure patients monitored during follow-up period, recording key biomarkers, physiological indicators, and mortality event outcomes.
Cardiovascular Disease (CVD) Dataset
Large-scale cardiovascular cohort from Kaggle used for self-supervised cross-domain pretraining to learn generalized tabular representations.
Cardiovascular Cross-Domain Feature Alignment Schema
Only overlapping features were mapped into the pretraining schema. Non-shared CVD variables were excluded.
Data Preprocessing & Fold-Isolated SMOTE
A three-stage preprocessing pipeline ensured numerical stability, outlier resistance, and balanced class distributions without introducing data leakage into evaluation folds.
Median Imputation
- β’Missing numerical observations imputed using training fold medians
- β’Robust to outliers in clinical biomarker distributions
StandardScaler
- β’Zero-mean and unit-variance normalization: z = (x - u) / s
- β’Prevents large-scale attributes from dominating gradient updates
Fold-Isolated SMOTE
- β’Sampling strategy: 1.0 (balanced 1:1 training ratio)
- β’k-Neighbours: 5, Random state: 42
- β’Executed strictly within training folds to prevent validation leakage
TabNet Architecture & Hyperparameters
TabNet uses sequential attention to select sparse feature subsets at each decision step, enabling neural networks to mimic decision trees while preserving gradient descent optimization.
SHAP Attribution & Feature Importance Rankings
SHAP (SHapley Additive exPlanations) was applied to compute feature attribution across experimental configurations, verifying that model decision pathways aligned with established clinical biomarkers while revealing how meta-features dominate Scenario 4.
Primary Clinical Predictors
In baseline and cross-domain models, clinical variables directly drive predictions. Follow-up duration (time), ejection_fraction, and serum_creatinine emerge as top mortality predictors.
Meta-Features + Clinical Biomarkers
In Scenario 4, the model-derived probabilistic meta-features (p_rf, p_xgb, p_lr) become top-ranked dominant predictors, while primary clinical variables serve as secondary guidance.
Research Insights & Scientific Conclusions
Key conclusions drawn from the empirical comparison of TabNet architectures and transfer learning strategies for clinical mortality risk prediction.
Feature-based transfer produced the highest recall, ROC AUC, PR AUC, and F1 among the four scenarios.
By augmenting TabNet input with out-of-fold probabilistic predictions from diverse teacher models (Logistic Regression, Random Forest, XGBoost), Scenario 4 achieved a recall of 0.937 and ROC AUC of 0.869.
Baseline TabNet achieved the highest precision, reflecting a more conservative decision boundary.
Scenario 1 yielded a precision of 0.538. While it generated fewer false alarms, its lower recall (0.658) means a larger fraction of high-risk patients would be missed compared to the transfer configurations.
Cross-domain transfer increased recall but was affected by distribution differences between the CVD and UCI populations.
Pretraining on 70,000+ general cardiovascular records (Scenario 3) improved sensitivity (recall 0.761 vs 0.658 in baseline), though differing feature distributions between broad CVD and specific heart failure cohorts limited global precision.
SHAP analysis consistently identified time, ejection fraction, and serum creatinine as important clinical predictors, while Scenario 4 introduced task-aligned probabilistic meta-features.
Attribution patterns confirmed established cardiology literature: lower ejection fraction and elevated serum creatinine correlate with higher mortality risk, while Scenario 4 prioritized probabilistic meta-features as dominant composite predictors.
Experimental research using public clinical datasets. Results should not be interpreted as a clinically validated diagnostic or decision-support system.