Federated Learning of Electronic Health Records Improves Mortality Prediction in Patients Hospitalized with COVID-19

Abstract

Machine learning (ML) models require large datasets which may be siloed across different healthcare institutions. Using federated learning, a ML technique that avoids locally aggregating raw clinical data across multiple institutions, we predict mortality within seven days in hospitalized COVID-19 patients. Patient data was collected from Electronic Health Records (EHRs) from five hospitals within the Mount Sinai Health System (MSHS). Logistic Regression with L1 regularization (LASSO) and Multilayer Perceptron (MLP) models were trained using local data at each site, a pooled model with combined data from all five sites, and a federated model that only shared parameters with a central aggregator. Both the federated LASSO and federated MLP models performed better than their local model counterparts at four hospitals. The federated MLP model also outperformed the federated LASSO model at all hospitals. Federated learning shows promise in COVID-19 EHR data to develop robust predictive models without compromising patient privacy.

Figure 1.. Study Design and Model Workflow.

(A) Criteria for patient inclusion in study. (B) Overview of local and pooled models. Local models only utilize data from the site itself while pooled models incorporate data from all sites. Both local and pooled MLP and LASSO models were utilized. (C) Overview of federated model. Parameters from a central aggregator are shared with each site, and sites do not have direct access to clinical data from others. After models are trained locally at a site, parameters with and without added noise are sent back to the central aggregator to update federated model parameters. A federated LASSO and federated MLP model were utilized.

Figure 2.. Federated Model Training.

Performance of (A) federated MLP and (B) federated LASSO models as measured by area under the receiver-operating characteristic (AUC-ROC) versus the number of training epochs. Binary Cross-Entropy Loss of (C) Federated MLP and (D) Federated LASSO versus the number of training epochs.

Figure 3.. Model Performance by Site.

Performance of all models (LASSO local, LASSO pooled, LASSO federated, MLP local, MLP pooled, MLP federated (no noise) by area under the receiver-operating characteristic (AUC-ROC) at (A) Mount Sinai Brooklyn (MSB) (n=611) (B) Mount Sinai West (MSW) (n=485), (C) Mount Sinai Morningside (MSM) (n=749), (D) Mount Sinai Hospital (MSH) (n=1644), and (E) Mount Sinai Queens (MSQ) (n=540). Averages of receiver-operating characteristic after 10-fold cross validation are shown. Average performance of each model across all five sites is presented in (F).