[论文解读] 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 AIMedicine被引用 3

一句话总结

本文提出 FedFBN，一种联邦学习框架，通过冻结预训练模型中的批量归一化层，缓解数据异质性和部分标签导致的收敛问题，适用于医学图像分类。通过结合迁移学习与冻结的 BN 层，FedFBN 在多样化的非独立同分布（non-iid）和部分标注数据集上，性能优于 FedAvg 和 FedBN，展现出在真实世界医学影像场景中更强的泛化能力和鲁棒性。

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.

研究动机与目标

• 解决在数据分布、非独立同分布且存在部分标签标注的情况下，训练准确全局模型的挑战。
• 克服 FedAvg 和 FedBN 在处理医学影像中机构间领域偏移和标签重叠方面的局限性。
• 通过使用预训练模型并冻结批量归一化层，利用迁移学习稳定训练过程并提升收敛性能。
• 在标签完整度和数据异质性各异的合成与真实世界数据集上，评估 FedFBN 的性能表现。
• 证明在联邦学习设置中，冻结 BN 层可实现比自适应或聚合 BN 统计量更优的泛化性能。

提出的方法

• FedFBN 通过在训练过程中全程冻结批量归一化层统计量，对 FedBN 进行改进，防止在本地和全局聚合过程中更新。
• 该框架采用预训练的深度神经网络作为模型主干，利用迁移学习初始化权重，以提升收敛速度。
• 批量归一化层在本地训练和全局聚合过程中均不更新，避免因领域偏移导致的性能下降。
• 全局模型通过 FedAvg 方式聚合非 BN 层权重，而 BN 统计量在所有客户端间保持固定。
• 该方法在基于 NIH 胸部 X 光 14 数据集生成的合成 iid 和非 iid 数据集上进行评估，同时使用外部 MIMIC-CXR 测试集。
• 性能通过多种标签配置下的 AUROC 进行衡量：共享标签、全部标签以及部分标签重叠。

(a) 4 Shared Disease Labels

实验结果

研究问题

• RQ1 在联邦学习设置中冻结批量归一化层，是否能提升在非独立同分布、部分标注的医学影像数据集上的模型收敛性和性能？
• RQ2 在标签完整度和数据异质性变化的情况下，FedFBN 与 FedAvg 和 FedBN 相比，在泛化能力和鲁棒性方面表现如何？
• RQ3 在联邦医学图像分类中，使用预训练模型并冻结 BN 层是否优于更新或聚合 BN 统计量的策略？
• RQ4 FedFBN 是否能在保护数据隐私的分布式设置下，实现与集中训练模型相当的性能？
• RQ5 与基线方法和现有联邦学习方法相比，FedFBN 在存在领域偏移的外部测试集（如 MIMIC-CXR）上的表现如何？

主要发现

• 在合成 iid 数据集的 4 个共享疾病标签上，FedFBN 的平均 AUROC 达到 0.89，显著优于基线模型（p<0.001）。
• 在 MIMIC 外部测试集上，FedFBN 的平均 AUROC 为 0.75，优于所有对比模型（p<0.001）。
• 在非 iid 数据集且标签完整的情况下，FedFBN 的性能与基线 NIH 模型相当（平均 AUROC 0.84，p=0.1），并优于 FedAvg 和本地 FedBN（p<0.001）。
• 在非 iid 数据集且存在部分标签的情况下，FedFBN 在 7 个共享标签上的平均 AUROC 为 0.84（p=0.03），在全部标签上的平均 AUROC 为 0.81（p=0.06），优于 FedAvg 和本地 FedBN（p<0.001）。
• 在所有测试集（包括 MIMIC-CXR）上，FedFBN 的泛化能力均优于 FedAvg 和 FedBN，平均 AUROC 达到 0.75（p<0.001）。
• FedFBN 在避免数据共享的前提下，性能与集中训练的 NIH 基线模型相当（平均 AUROC 0.81，p=0.06），凸显其临床实用性。

(b) All Disease Labels