Heart Failure Prediction & Risk Stratification using Machine Learning

This study demonstrates that a carefully calibrated stacked ensemble machine learning model, trained on readily available electronic medical record data from the All of Us Research Program, can effectively predict heart failure with high accuracy (ROC-AUC 0.927) and provide clinically actionable risk stratification for early diagnosis and proactive care.

The Gist

Imagine your heart is the engine of a car. Heart Failure (HF) is like that engine starting to sputter, lose power, or overheat. It's a very common problem that affects millions of people, but often, the warning signs are subtle. People might think they are just getting older, stressed, or a little out of shape, so they ignore the symptoms until it's too late.

This paper is about building a "Digital Mechanic"—a smart computer program—that can look at a patient's routine medical records and say, "Hey, this engine is showing early signs of trouble," long before a major breakdown happens.

Here is how the researchers built this Digital Mechanic, explained in simple terms:

1. The Ingredients: What Did They Feed the Computer?

Usually, to diagnose heart failure, you need expensive, high-tech tools like MRI machines or specialized heart scans. But the researchers wanted to build a tool that works in any doctor's office, even a small one without fancy equipment.

So, they fed their computer only the "grocery list" of data that is already available in almost every patient's file:

• Basic Stats: Age, gender, and where they live (specifically, how wealthy or poor their neighborhood is).
• Vital Signs: Blood pressure and body weight (BMI).
• Lab Results: Simple blood tests for things like sugar, salt, kidney function, and blood cell counts.
• History: Do they have high blood pressure? Do they smoke? Do they have an irregular heartbeat (Atrial Fibrillation)?

The Analogy: Think of it like a mechanic who doesn't need to take the car apart to know it's sick. They just listen to the engine noise, check the oil color, and look at the mileage. If the oil is dark and the mileage is high, they know there's a problem, even without a full teardown.

2. The Training: Teaching the Computer to Spot Patterns

The researchers used data from the "All of Us" program, a massive database of medical records from over 37,000 people.

• The Class: They had about 13,500 people who did have heart failure and 23,500 who didn't.
• The Challenge: The computer is like a student. If you show it 100 pictures of cats and only 1 picture of a dog, it will just guess "cat" every time to get a high score. The researchers had to teach the computer to pay special attention to the "dog" (the heart failure patients) so it wouldn't miss them.

They tried many different types of computer brains (algorithms), from simple math formulas to complex "neural networks" that mimic the human brain.

• The Winner: They built a "Stacked Ensemble." Imagine a team of experts. One is a specialist in blood pressure, another in blood sugar, another in age trends. Instead of asking just one expert, they asked all of them, and then a "Team Captain" (a final algorithm) listened to all their opinions to make the final decision. This team approach worked the best.

3. The Results: How Good is the Mechanic?

The computer became incredibly good at its job:

• Accuracy: It correctly identified heart failure about 85-90% of the time.
• The "Why": The researchers used a tool called SHAP (which is like a magnifying glass) to see why the computer made its decisions. It turned out the computer was looking at the right things:
  • Atrial Fibrillation (an irregular heartbeat) was the biggest red flag.
  • Age and High Blood Pressure were close seconds.
  • Sodium levels and Poverty Index (socioeconomic status) were also major clues.
  • Interestingly, having low levels of certain good things (like healthy blood cells or albumin) was a warning sign.

4. The Real-World Test: Adjusting for the "Average" Person

Here is the tricky part. The group of people the computer learned from had a very high rate of heart failure (about 36%). But in the real world, only about 2.5% of people have it.

• The Problem: If you use the computer on the general public without fixing this, it will be too paranoid. It will scream "Heart Failure!" at almost everyone because it's used to seeing it so often in its training data.
• The Fix: The researchers "recalibrated" the computer. They told it, "Remember, in the real world, this is rare. Don't panic unless you are really sure."
• The Result: After this adjustment, the computer's predictions became realistic. If you screened 1,000 random people, it would predict about 25 cases, which matches reality perfectly.

5. The Superpower: Sorting the Risk

The most exciting part is how this helps doctors. Instead of treating everyone the same, the computer sorts patients into 10 risk groups (like a ladder from 1 to 10).

• The Magic Stat: If a clinic only has time to check the top 10% of people on this ladder (the highest risk), they will catch 75% of all the actual heart failure cases in that group.
• The Analogy: Imagine a lighthouse. You can't shine the light on the whole ocean. But if you aim it at the specific area where the ships are most likely to be, you save the most boats. This tool tells doctors exactly where to shine the light.

The Bottom Line

This paper shows that we don't always need expensive MRI machines or genetic tests to find heart failure early. By using a smart computer program trained on the simple, everyday data already sitting in our medical files, we can:

1. Spot heart failure earlier.
2. Understand why a patient is at risk (it's not a "black box").
3. Prioritize care for the people who need it most, saving money and lives.

It's a step toward a future where your doctor can check your "engine health" with a quick glance at your routine blood work, catching problems before they become emergencies.

Technical Summary

1. Problem Statement

Heart Failure (HF) is a leading cause of morbidity, mortality, and healthcare costs, affecting approximately 6.7 million adults in the U.S. Early diagnosis is challenging because symptoms are often non-specific or misinterpreted as normal aging. While Machine Learning (ML) models exist for HF prediction, many suffer from:

• Dependency on expensive modalities: Reliance on imaging (MRI, CT, Echocardiography) or invasive monitoring limits scalability in primary care or resource-constrained settings.
• Lack of generalizability: Many models are trained on single-center datasets with specific demographic biases.
• Poor calibration: Models often optimize for discrimination (ROC-AUC) without ensuring predicted probabilities reflect real-world absolute risk, which is critical for clinical decision support.
• Class imbalance: Medical datasets often have skewed distributions (fewer HF cases than non-HF), leading to models that perform well on accuracy but poorly on detecting the minority class (HF).

Objective: To develop a scalable, interpretable, and well-calibrated ML framework for HF prediction and risk stratification using only routinely available Electronic Medical Record (EMR) data (demographics, vitals, labs, comorbidities) without specialized imaging or genomic data.

2. Methodology

Dataset

• Source: NIH "All of Us" Research Program (de-identified EMR data).
• Cohort: 37,070 adults (13,577 HF cases, 23,493 non-HF controls).
• Labels: Derived from ICD-9/ICD-10 and SNOMED codes.
• Subtypes: A secondary experiment categorized HF into Heart Failure with Reduced Ejection Fraction (HFrEF) and Preserved Ejection Fraction (HFpEF), resulting in a 3-class problem (No HF, HFrEF, HFpEF).
• Exclusions: Patients with cancer, terminal illness, or deceased status were excluded.

Feature Engineering

The study utilized 18 routinely available features across four domains:

1. Measurements: Albumin, BMI, Creatinine, Glucose, Hemoglobin, Potassium, Sodium, HDL, Systolic/Diastolic BP.
2. Conditions: Anemia, Atrial Fibrillation, Hypertensive disorder, Obesity.
3. Lifestyle: Smoking history.
4. Demographics: Age, Sex at birth, Deprivation Index (socioeconomic status).

Data Preprocessing

• Missingness: Median imputation.
• Outliers: IQR-based winsorization (capping at 1.5x IQR).
• Normalization: Quantile Transformer to map features to a standard normal distribution ($N(0,1)$), addressing non-Gaussian EMR data.
• Encoding: One-hot encoding for categorical variables.
• Class Imbalance: Addressed via three strategies: Random Undersampling, SMOTE (Oversampling), and Class Reweighting on the raw unbalanced data.

Model Architecture

The study benchmarked several algorithms and proposed a Custom Stacked Ensemble:

• Base Learners: XGBoost, LightGBM, CatBoost, and Multi-Layer Perceptron (MLP).
• Meta-Learner: Logistic Regression (LR).
• Training Strategy: Stratified 5-fold Cross-Validation and a 70/30 hold-out test split.
• Hyperparameters: Optimized via grid search (e.g., XGBoost depth=8, learning rate=0.02).

Calibration and Risk Stratification

• Calibration: Applied Platt Scaling to correct probability outputs.
• Prevalence Adjustment: Since the study cohort had a 36.6% HF prevalence (vs. ~2.5% in the real world), the authors applied prior probability shift in log-odds space to adjust predictions to a realistic 2.5% base rate.
• Stratification: Patients were grouped into 10 risk deciles to identify high-risk populations for targeted screening.

3. Key Results

Binary Classification (HF vs. Non-HF)

The Stacked Ensemble trained on the unbalanced (raw) data with class weights achieved the best performance:

• ROC-AUC: 0.927
• PR-AUC: 0.895 (Critical for imbalanced datasets)
• Accuracy: 0.856
• Comparison: Outperformed individual models like Logistic Regression (0.893 AUC) and SVM (0.910 AUC). XGBoost alone was close (0.929 AUC) but the ensemble provided better stability and calibration.

Multiclass Classification (No HF, HFrEF, HFpEF)

• Performance: Lower discrimination (Macro-AUC ~0.87) and significantly lower recall for subtypes (HFrEF recall ~48%, HFpEF recall ~15%).
• Analysis: The difficulty is attributed to overlapping phenotypes in EMR data and the lack of echocardiographic data (LVEF) required to definitively distinguish subtypes.

Explainability (SHAP Analysis)

SHAP analysis on the XGBoost component identified the top 5 influential features:

1. Atrial Fibrillation
2. Age
3. Hypertensive Disorder
4. Sodium
5. Deprivation Index

• Interpretation: Higher values of Atrial Fibrillation, Age, Hypertension, and Deprivation Index increased HF risk, while higher Albumin, Potassium, and HDL were protective. These findings align with established clinical pathophysiology.

Risk Stratification & Calibration

• Calibration: After adjusting for the 2.5% real-world prevalence, the model showed excellent calibration (Weighted Brier Score: 0.0178; Calibration Slope: ~1.0).
• Screening Efficiency:
  • Screening the top 10% highest-risk patients captured 74.7% of all HF cases.
  • Screening the top 20% captured 86.8% of cases.
  • This demonstrates high efficiency for resource-constrained clinical settings.

4. Key Contributions

1. Scalable AI-CDSS Framework: Demonstrated that high-performance HF prediction is possible using only low-cost, routinely collected EMR variables, removing the barrier of expensive imaging.
2. Rigorous Calibration: Moved beyond simple discrimination metrics to provide clinically interpretable absolute risk probabilities adjusted for real-world prevalence, a critical step often missing in clinical ML literature.
3. Handling Imbalance: Systematically compared undersampling, oversampling, and reweighting, finding that reweighting on raw data yielded the best balance of precision and recall for this specific dataset.
4. Explainability: Integrated SHAP analysis to ensure model decisions align with medical knowledge, fostering clinician trust.
5. Risk Deciles: Proposed a practical risk stratification tool that allows clinicians to prioritize a small subset of the population for further diagnostic testing while capturing the vast majority of true cases.

5. Significance and Limitations

• Significance: The study provides a blueprint for deploying AI in population health screening. By proving that a calibrated ensemble model can identify ~75% of HF cases by screening only the top 10% of the population, it offers a pathway for proactive care and early intervention in primary settings.
• Limitations:
  • Lack of External Validation: The model was validated only on internal splits of the "All of Us" dataset; external validation on independent health systems is needed to confirm generalizability.
  • Subtype Discrimination: The model struggled to distinguish HFrEF from HFpEF without imaging data, highlighting the need for future integration of echocardiographic features.
  • Retrospective Nature: The study relies on historical EMR data; prospective validation in a real-time clinical workflow is required.

Conclusion: The authors successfully developed a robust, interpretable, and calibrated machine learning system for heart failure prediction using standard EMR data. The model's ability to accurately stratify risk and provide actionable probability estimates makes it a strong candidate for integration into Clinical Decision Support Systems (AI-CDSS) for early HF detection.