[Paper Review] Optimizing Federated Learning for Medical Image Classification on Distributed Non-iid Datasets with Partial Labels

Pranav Kulkarni, Adway Kanhere+2 more|arXiv (Cornell University)|Mar 10, 2023

COVID-19 diagnosis using AIMedicine3 citations

TL;DR

This paper proposes FedFBN, a federated learning framework that freezes batch normalization layers in pretrained models to mitigate convergence issues caused by data heterogeneity and partial labels in medical image classification. By combining transfer learning with frozen BN layers, FedFBN achieves superior performance over FedAvg and FedBN across diverse non-iid and partially labeled datasets, demonstrating improved generalization and robustness in real-world medical imaging scenarios.

ABSTRACT

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.

Motivation & Objective

• To address the challenge of training accurate global models in federated learning when data is distributed, non-iid, and annotated with partial labels.
• To overcome the limitations of FedAvg and FedBN in handling domain shift and label overlap across institutions in medical imaging.
• To leverage transfer learning by using pretrained models with frozen batch normalization layers to stabilize training and improve convergence.
• To evaluate FedFBN’s performance across synthetic and real-world datasets with varying label completeness and data heterogeneity.
• To demonstrate that frozen BN layers in a federated setting can yield better generalization than adaptive or aggregated BN statistics.

Proposed method

• FedFBN modifies FedBN by freezing batch normalization layer statistics throughout the training process, preventing updates during local and global aggregation.
• The framework uses a pretrained deep neural network as the model backbone, leveraging transfer learning to initialize weights and improve convergence.
• Batch normalization layers are not updated during local training or global aggregation, avoiding domain shift-induced performance degradation.
• The global model is updated via FedAvg-style aggregation of non-BN layer weights, while BN statistics remain fixed across all clients.
• The method is evaluated on synthetic iid and non-iid datasets derived from the NIH Chest X-Ray 14 dataset, as well as external MIMIC-CXR test sets.
• Performance is measured using AUROC across multiple label configurations: shared labels, all labels, and partial label overlaps.

(a) 4 Shared Disease Labels

Experimental results

Research questions

• RQ1 Can freezing batch normalization layers in a federated learning setup improve model convergence and performance on non-iid, partially labeled medical imaging datasets?
• RQ2 How does FedFBN compare to FedAvg and FedBN in terms of generalization and robustness across varying label completeness and data heterogeneity?
• RQ3 Does the use of pretrained models with frozen BN layers outperform strategies that update or aggregate BN statistics in federated medical image classification?
• RQ4 Can FedFBN achieve performance comparable to centrally trained models while preserving data privacy in a distributed setting?
• RQ5 How does FedFBN perform on external test sets with domain shift, such as MIMIC-CXR, compared to baseline and existing FL methods?

Key findings

• FedFBN achieved a mean AUROC of 0.89 on 4 shared disease labels in synthetic iid datasets, significantly outperforming the baseline model (p<0.001).
• On the MIMIC external test set, FedFBN achieved a mean AUROC of 0.75, outperforming all compared models (p<0.001).
• In non-iid datasets with complete labels, FedFBN matched the performance of the baseline NIH model (mean AUROC 0.84, p=0.1) and outperformed FedAvg and local FedBN (p<0.001).
• On non-iid datasets with partial labels, FedFBN achieved a mean AUROC of 0.84 on 7 shared labels (p=0.03) and 0.81 on all labels (p=0.06), outperforming FedAvg and local FedBN (p<0.001).
• FedFBN generalized better than FedAvg and FedBN across all test sets, including MIMIC-CXR, with a mean AUROC of 0.75 (p<0.001).
• FedFBN demonstrated comparable performance to the centrally trained NIH baseline model (mean AUROC 0.81, p=0.06) while avoiding data sharing, highlighting its clinical utility.

(b) All Disease Labels