Federated Learning Performance on Early ICU Mortality Prediction with Extreme Data Distributions

Abstract

Federated Learning (FL) is a novel machine learning technique that allows multiple parties to collaboratively train a global model without sharing their local data, thus addressing data privacy and security concerns. These issues are of paramount importance to the healthcare domain, due to strict regulations regarding patient data. By utilizing data from multiple medical institutions, FL could lead to generalizable models for one of the most prominent medical tasks, the early prediction of mortality risk in the ICU setting, where patients in critical condition are treated. This paper evaluates the performance of various FL algorithms in a realistic FL scenario and, also, in the presence of FL clients with ‘extreme’ data distributions, using real world data from a collaborative research database. Overall, the FL models perform, in general, substantially better than the local models of the participating hospitals and slightly worse than the ‘ideal’ model, which is trained on the centralized data. FedProx, a client-side optimization FL algorithm, regulates more effectively the contribution of a large FL client, with ‘extreme’ bias against the underrepresented class, while the server-side optimization FL algorithms incorporate the beneficial information of a smaller FL client into the global model more efficiently.