Federated Proximal Optimization for Privacy-Preserving Heart Disease Prediction: A Controlled Simulation Study on Non-IID Clinical Data

TL;DR

FedProx, a proximal-regularized extension of FedAvg, is evaluated for privacy-preserving heart disease prediction on non-IID clinical data derived from the UCI Cleveland subset. The global objective is $F(w)=\sum_{k=1}^K \frac{n_k}{n} F_k(w)$ with local loss $\ell(w;x_i^k,y_i^k)$, and FedProx adds a proximal term $\frac{\mu}{2}\|w-w_t\|^2$ to curb client drift. Across 50 runs with a grid of $\mu \in \{0.0,0.01,0.05,0.1,0.5\}$, the study shows FedProx with $\mu=0.05$ achieving 85.00% accuracy, outperforming centralized training (83.33%) and FedAvg. It also reports faster convergence (18 vs 22 rounds to 95% final accuracy) and improved cross-hospital fairness (standard deviation reduced from 3.21% to 1.42%). Limitations include simulated non-IID data from a single institution and small sample size, but the work provides a rigorous, reproducible framework for evaluating privacy-preserving federated learning in clinical settings and points toward future validation with larger, diverse datasets and differential privacy.

Abstract

Healthcare institutions have access to valuable patient data that could be of great help in the development of improved diagnostic models, but privacy regulations like HIPAA and GDPR prevent hospitals from directly sharing data with one another. Federated Learning offers a way out to this problem by facilitating collaborative model training without having the raw patient data centralized. However, clinical datasets intrinsically have non-IID (non-independent and identically distributed) features brought about by demographic disparity and diversity in disease prevalence and institutional practices. This paper presents a comprehensive simulation research of Federated Proximal Optimization (FedProx) for Heart Disease prediction based on UCI Heart Disease dataset. We generate realistic non-IID data partitions by simulating four heterogeneous hospital clients from the Cleveland Clinic dataset (303 patients), by inducing statistical heterogeneity by demographic-based stratification. Our experimental results show that FedProx with proximal parameter mu=0.05 achieves 85.00% accuracy, which is better than both centralized learning (83.33%) and isolated local models (78.45% average) without revealing patient privacy. Through generous sheer ablation studies with statistical validation on 50 independent runs we demonstrate that proximal regularization is effective in curbing client drift in heterogeneous environments. This proof-of-concept research offers algorithmic insights and practical deployment guidelines for real-world federated healthcare systems, and thus, our results are directly transferable to hospital IT-administrators, implementing privacy-preserving collaborative learning.

Figures (7)

• Figure 1: Federated Learning system architecture showing the central server coordinating with four hospital nodes. The server broadcasts the global model $w_t$ to each hospital, which performs local training and returns only gradient updates, never sharing raw patient data.
• Figure 2: End-to-end pipeline for federated heart disease prediction: from Cleveland dataset acquisition through demographic-based client partitioning, preprocessing, federated training with FedProx, and final model evaluation across all simulated hospital clients.
• Figure 3: Correlation matrix of 13 clinical features showing relationships between patient attributes. Strong correlations exist between chest pain type and disease presence, validating clinical knowledge.
• Figure 4: Distribution of key clinical features in the complete Cleveland Clinic dataset (293 patients) before partitioning. Top row: Age distribution (peak at 55-60 years), cholesterol levels (mean 246 mg/dl), and maximum heart rate achieved during exercise. Bottom row: Heart disease prevalence (45.9% positive cases), sex distribution (68% male, 32% female), and chest pain type distribution (Type 4 most common at 47% of cases). This baseline characterization informs the demographic-based stratification strategy used to create heterogeneous hospital clients.
• Figure 5: Non-IID data distribution across four simulated hospital clients created through demographic-based partitioning. Top-left: Age distributions showing Client 1 serves older patients (peak at 60 years) while Client 4 serves younger patients (peak at 50 years). Top-right: Disease prevalence ranging from 39% to 64% across clients. Bottom-left: Unbalanced sample sizes from 98 to 242 patients simulating different hospital capacities. Bottom-right: Cholesterol distributions showing client-specific patterns. These variations create substantial challenges for standard federated averaging, mimicking real-world multi-hospital heterogeneity.