FedDropoutAvg: Generalizable federated learning for histopathology image classification

TL;DR

This study proposes FedDropoutAvg, a new federated learning approach for training a generalizable model that takes advantage of randomness, both in client selection and also in federated averaging process, and is the first study to use a randomized client and local model parameter selection procedure in a federated setting for a medical image analysis task.

Abstract

Federated learning (FL) enables collaborative learning of a deep learning model without sharing the data of participating sites. FL in medical image analysis tasks is relatively new and open for enhancements. In this study, we propose FedDropoutAvg, a new federated learning approach for training a generalizable model. The proposed method takes advantage of randomness, both in client selection and also in federated averaging process. We compare FedDropoutAvg to several algorithms in an FL scenario for real-world multi-site histopathology image classification task. We show that with FedDropoutAvg, the final model can achieve performance better than other FL approaches and closer to a classical deep learning model that requires all data to be shared for centralized training. We test the trained models on a large dataset consisting of 1.2 million image tiles from 21 different centers. To evaluate the generalization ability of the proposed approach, we use held-out test sets from centers whose data was used in the FL and for unseen data from other independent centers whose data was not used in the federated training. We show that the proposed approach is more generalizable than other state-of-the-art federated training approaches. To the best of our knowledge, ours is the first study to use a randomized client and local model parameter selection procedure in a federated setting for a medical image analysis task.

Figures (8)

• Figure 1: Workflow diagram of the proposed FedDropoutAvg approach. Training phase: At the beginning of each training round, Central Server sends global model weights to some clients which are randomly selected from all the clients participating in the training. Using the received model, local training takes place at each client on their local datasets. At the end of local training epochs, each Client i sends the parameters of their locally trained model to Central Server. Then, Central Server randomly drops out some of the parameters of the received models and aggregates the models into a global model by averaging. This training process continues for some number of rounds. Testing phase: After the federated training is over, Central Server sends the final model to independent clients for use on their own test sets, simulating a real life scenario.
• Figure 2: Sample tumor (first three rows) and non-tumor (last two rows) image patches from different data centers. It can be observed that there is large intra-class variation between the two classes. Some image patches contain artifacts; stain color variation can also be seen in images from different centers.
• Figure 3: Heatmap of F1 scores of the locally trained models (not federated) on the held-out test sets. Mean F1 of each locally trained model is given on the rightmost column.
• Figure 4: Heatmap of F1 scores of the locally trained models (not federated) on the independent test centers. Mean F1 of each locally trained model is given on the rightmost column.
• Figure 5: F1 scores of the federated approaches on (a) the held-out test sets of the centers participated in training, (b) the datasets of the independent test centers (Note that these centers did not participate in the training of the models).