FedWeight: Mitigating Covariate Shift of Federated Learning on Electronic Health Records Data through Patients Re-weighting

Federated Learning (FL) has emerged as a promising approach for research on real-world medical data distributed across different organizations, as it allows analysis of distributed data while preserving patient privacy. However, one of the prominent challenges in FL is covariate shift, where data distributions differ significantly across different clinical sites, like hospitals and outpatient clinics. These differences in demographics, clinical practices, and data collection processes may lead to significant performance degradation of the shared model when deployed for a target population. In this study, we propose a Federatively Weighted (FedWeight) framework to mitigate the effect of covariate shift on Federated Learning. Leveraging the data distribution estimated by density estimator models, we re-weight the patients from the source clinical sites, making the trained model aligned with the data distribution of the target site, thus mitigating the covariate shift between source and target sites. To make our approach also applicable to unsupervised learning, we integrate Fed-Weight into a novel federated embedded topic model (ETM), namely FedWeight-ETM. We evaluated FedWeight in cross-site FL within the eICU dataset and also cross-dataset FL between eICU and MIMIC-III data. Compared with the baseline, FedWeight-corrected FL models demonstrate superior performance for predicting patient mortality, ventilator use, sepsis diagnosis, and length of stay in the intensive care unit (ICU). Moreover, FedWeight outper-forms FedAvg in identifying important features relevant to the clinical outcomes. Leveraging Shapley Additive Explanations (SHAP), the FedWeight-corrected classifiers reveal subtle yet significant associations between drugs, lab tests, and patient outcomes. Using FedWeight-ETM, we identified known disease topics involving renal or heart failure predictive of future mortality at the ICU readmission. Together, FedWeight provides a robust FL framework to address the challenge of covariate shift from clinical silos in predicting critical patient out-comes and providing meaningful clinical features.