Optimizing Federated Learning for Medical Image Classification on Distributed Non-iid Datasets with Partial Labels
This addresses a specific bottleneck in federated learning for medical image classification, offering an incremental improvement for scenarios with non-iid data and partial labels.
The paper tackled the problem of federated learning convergence on distributed non-iid medical image datasets with partial labels by proposing FedFBN, a framework that freezes batch normalization layers, and demonstrated it outperforms current aggregation strategies.
Numerous large-scale chest x-ray datasets have spearheaded expert-level detection of abnormalities using deep learning. However, these datasets focus on detecting a subset of disease labels that could be present, thus making them distributed and non-iid with partial labels. Recent literature has indicated the impact of batch normalization layers on the convergence of federated learning due to domain shift associated with non-iid data with partial labels. To that end, we propose FedFBN, a federated learning framework that draws inspiration from transfer learning by using pretrained networks as the model backend and freezing the batch normalization layers throughout the training process. We evaluate FedFBN with current FL strategies using synthetic iid toy datasets and large-scale non-iid datasets across scenarios with partial and complete labels. Our results demonstrate that FedFBN outperforms current aggregation strategies for training global models using distributed and non-iid data with partial labels.