Federated learning for predicting clinical outcomes in patients with COVID-19

Abstract

Federated learning (FL) is a method used for training artificial intelligence models with data from multiple sources while maintaining data anonymity, thus removing many barriers to data sharing. Here we used data from 20 institutes across the globe to train a FL model, called EXAM (electronic medical record (EMR) chest X-ray AI model), that predicts the future oxygen requirements of symptomatic patients with COVID-19 using inputs of vital signs, laboratory data and chest X-rays. EXAM achieved an average area under the curve (AUC) >0.92 for predicting outcomes at 24 and 72 h from the time of initial presentation to the emergency room, and it provided 16% improvement in average AUC measured across all participating sites and an average increase in generalizability of 38% when compared with models trained at a single site using that site's data. For prediction of mechanical ventilation treatment or death at 24 h at the largest independent test site, EXAM achieved a sensitivity of 0.950 and specificity of 0.882. In this study, FL facilitated rapid data science collaboration without data exchange and generated a model that generalized across heterogeneous, unharmonized datasets for prediction of clinical outcomes in patients with COVID-19, setting the stage for the broader use of FL in healthcare.

Extended Data Fig. 1. Test performance of models predicting 72h oxygen treatment trained on local data only versus the performance of the best global model available on the server.

Test performance of models predicting 72h oxygen treatment trained on local data only (Local) versus the performance of the best global model available on the server (FL (gl. best)). b, Generalizability (average performance on other sites’ test data) as a function of a site’s dataset size (# cases). The average performance improved by 18% (from 0.760 to 0.899 or 13.9 percentage points) compared to locally trained models alone, while average generalizability of the global model improved by 34% (from 0.669 to 0.899 or 23.0 percentage points).

Extended Data Fig. 2. Confusion Matrices at a site with unbalanced data and mostly mild cases.

a, Confusion matrices on the test data at site 16 predicting oxygen treatment at 72h using the locally trained model. b, Confusion matrices on the test data at site 16 predicting oxygen treatment at 72h using the best Federated Learning global model. We show the ROCs for two different cut-off values t of the EXAM risk score.

Extended Data Fig. 3. Effect of data set size on model performance.

ROCs of the best global model in comparison to the mean ROCs of models trained on local datasets to predict 24/72-h oxygen treatment devices for COVID positive/negative patients respectively, using the test data of 5 large datasets from sites in the Boston area. The Mean ROC is calculated based on 5 locally trained models, with the gray-area showing the standard deviation of the ROCs. We show the ROCs for three different cut-off values t of the EXAM risk score.

Extended Data Fig. 4. Failures cases at an independent test site.

CXRs from two failure cases at CDH. The above is noisy data where each available value has been anonymized by adding a zero-mean Gaussian noise with the standard deviation of 1/5 of the standard deviation of the cohort distribution.

Extended Data Fig. 5. Safety enhancing features used in EXAM.

Additional data-safety-enhancing features were assessed by only sharing a certain percentage of weight updates with the largest magnitudes before sending them to the server after each round of learning. We show that by using partial weight updates during FL, models can be trained that reach a performance comparable to training while sharing the full information. This differential privacy technique decreases the risk for model inversion or reconstruction of the training image data through gradient interception.

Extended Data Fig. 6. Characteristics of EMR data used in EXAM.

Min. and max. values (asterisks) and mean and standard deviation (length of bars) for each EMR feature used as an input to the model. n specifies the number of sites that had this particular feature available. Missing values were imputed using a MissedForest algorithm. .

Extended Data Fig. 7. Distribution of oxygen treatments between EXAM sites.

The boxplots show the quartiles of the minimum, the maximum, the sample median, and the first and third quartiles (excluding outliers) of the oxygen treatments applied at different sites at time of Emergency Department admission and after 24 and 72- hour periods. The types of oxygen treatments administered are ‘room air’, ‘low-flow oxygen’, ‘high-flow oxygen (non-invasive)’, and ‘ventilator’.

Extended Data Fig. 8. Site variations in oxygen usage.

Normalized distributions of oxygen devices at different time points, comparing the site with largest dataset size (site 1) and a site with unbalanced data, including mostly mild cases (site #16).

Extended Data Fig. 9. Description of the EXAM Federated Learning study.

a, Previously developed model, CDS, to predict a risk score that corresponds to respiratory outcomes in patients with SARS-COV-2. b, Histogram of CORISK results at MGB, with an illustration of how the score can be used for patient triage, in which ‘A’ is an example threshold for safe discharge that has 99.5% negative predictive value, and ‘B’ is an example threshold for Intensive Care Unit (ICU) admission that has 50.3% positive predictive value. For the purpose of the NPV calculation (threshold A), we defined the Model Inference to be Positive if it predicted oxygen need as LFO or above (COVID risk score ³0.25) and Negative if it predicted oxygen need as RA (<0.25). We defined the Disease to be Negative if the patient was discharged and not readmitted, and Positive if the patient was readmitted for treatment. For the purpose of PPV calculation (threshold B), we defined the Model Inference to be Positive if it predicted oxygen need as MV or above (^³0.75) and Negative if it predicted oxygen need as HFO or less (<0.75). We defined the disease to be Positive if the patient required MV or if they died, and we defined the disease as Negative if the patient survived and did not require MV. The EXAM score can be used in the same way. c, Federated Learning using a client-server setup.

Extended Data Fig. 10. Calibration Plots for the MGB data and the new independent dataset, CDH, used for model validation.

Fig. 1 |

Data used in the EXAM FL study. a, World map indicating the 20 different client sites contributing to EXAM study. b, Number of cases contributed by each institution or site contributed (client #1 represents the site contributing the largest number of cases). c, Chest X-ray intensity distribution at each client site. d, Age of patients at each client-site showing the minimum and maximum ages (asterisks), mean age (triangle) and standard deviation (horizontal bar). The number of samples of each client site is shown in Supplemental Table 1.

Fig. 2 |

Performance of federated learning versus local models. a, Performance on each client’s test set for predicting 24h oxygen treatment for models trained on local data only (Local) versus the performance of the best global model available on the server (FL (gl. best)). “Avg.” stands for the average test performance across all sites. b, Generalizability (average performance on other sites’ test data, as represented by average AUC) as a function of a client’s dataset size (# cases). The green horizontal line shows the generalizability performance of the best global model. The performance for 18 of 20 clients is shown, as client 12 had outcomes only for 72 hours (see Extended Data Fig. 1) and client 14 had cases only with room air treatment, such that the evaluation metric (avg. AUC) was not applicable in either of these cases (see Methods). Data for client 14 was also excluded from the computation of the average generalizability of the local models.

Fig. 3 |

Comparison of federated learning-trained to locally-trained models. a, ROC at a site with unbalanced data and mostly mild cases (client-site #16). b, ROC of the local model at client-site #12 (a small dataset), mean ROC of models trained on larger datasets corresponding to the 5 client-sites in the Boston area (#1, #4, #5, #6, #8), and ROC of the best global model to predict oxygen treatment at 72 h for different thresholds of the EXAM score (left, middle, right). The mean ROC is calculated based on 5 locally-trained models, and the gray-area shows the standard deviation of the ROCs. ROCs for three different cut-off values t of the EXAM risk score are shown. “Pos” and “Neg” stand for the number of positive and negative cases defined by this range of the EXAM score, respectively.

Fig. 4 |

Performance of the best global model on the largest independent data set. a, Performance (ROC) (top) and confusion matrices (bottom) of the EXAM FL model on the CDH dataset for predicting oxygen treatment at 24 h. b, Performance (ROC) (top) and confusion matrices (bottom) of the EXAM FL model on the CDH dataset for predicting oxygen treatment at 72 h. ROCs for three different cut-off values t of the EXAM risk score are shown. “Pos” and “Neg” stand for the number of positive and negative cases defined by this range of the EXAM score, respectively.