Personalized Federated Learning for Heart Failure Mortality Prediction under Non-IID Clinical Data

Abstract

This research explores how federated learning (FL) can be applied to predict heart failure mortality, emphasizing the protection of patient data privacy while maintaining model accuracy in realistic clinical environments. Using the Heart Failure Clinical Records dataset, I simulate 30 medical institutions and compare three FL algorithms—FedAvg, FedProx, and pFedMe—under conditions where data is not independently and identically distributed (non-IID). FedAvg serves as a baseline for centralized aggregation, FedProx introduces a proximal term to address client drift, and pFedMe employs Moreau envelopes for personalized model updates. The results demonstrate that both FedProx and pFedMe outperform FedAvg, achieving a final test accuracy of 70% versus 66.67%, with pFedMe further exhibiting the most stable training dynamics and minimal loss drift. These findings underscore the critical role of regularization and personalization in federated healthcare systems, particularly in heterogeneous environments where data distributions vary across clients. The work provides practical guidance for FL deployment in real-world medical settings, highlighting a trade-off between accuracy, model stability, and patient-specific adaptation—all while maintaining strict data privacy compliance. This contributes to the broader effort of enabling secure, collaborative AI in healthcare.