[论文解读] FedDropoutAvg: Generalizable federated learning for histopathology image classification

Gozde N. Gunesli, Mohsin Bilal+2 more|arXiv (Cornell University)|Nov 25, 2021

AI in cancer detection被引用 12

一句话总结

该论文提出FedDropoutAvg，一种用于组织病理学图像分类的新型联邦学习框架，通过在联邦平均过程中引入随机客户端选择和随机参数丢弃，提升模型泛化能力。该方法在包含120万张图像的多中心数据集上实现了与集中式训练相当的性能，优于标准FedAvg和FedProx在未见中心上的泛化表现。

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.

研究动机与目标

• 为解决多中心组织病理学图像分类中数据去中心化和异质性带来的模型泛化挑战。
• 克服标准联邦平均（FedAvg）和FedProx的局限性，后者依赖数据规模加权，可能因数据不平衡和质量差异而出现欠拟合。
• 开发一种无需事先了解各中心数据质量或性能信息即可提升鲁棒性和泛化能力的方法。
• 在训练中心的保留测试集和未参与联邦训练的独立中心测试集上评估所提方法，以评估其在真实场景中的可部署性。

提出的方法

• FedDropoutAvg在每轮联邦训练开始时引入随机客户端选择，降低通信和计算负载。
• 在聚合阶段，中央服务器在平均前对选定客户端的模型参数应用随机丢弃，引入随机性以提升泛化能力。
• 该方法在各客户端本地进行完整模型训练，保留模型容量，仅在全局聚合阶段应用丢弃。
• 聚合过程为丢弃后剩余参数的加权平均，权重基于本地数据规模，与FedAvg类似但增加了随机性。
• 该方法设计为与现有联邦学习优化技术兼容，例如梯度平均中的动量。
• 未被选中的客户端在该轮不参与通信，提升大规模部署中的通信效率。

实验结果

研究问题

• RQ1 在多中心组织病理学图像分类中，联邦平均过程中的随机客户端选择和参数丢弃是否能提升模型泛化能力？
• RQ2 FedDropoutAvg在训练中心的保留测试集和未见中心上的性能与FedAvg和FedProx相比如何？
• RQ3 在客户端参与和参数聚合中引入随机性，是否能提升对数据异质性和质量差异的鲁棒性？
• RQ4 FedDropoutAvg能否在不共享原始数据的前提下实现接近集中式训练的性能？
• RQ5 该方法在真实联邦设置中的通信效率和计算负载方面有何影响？

主要发现

• 在来自21个中心的120万张图像切片数据集上，FedDropoutAvg的分类性能比FedAvg和FedProx更接近集中式训练。
• 使用FedDropoutAvg训练的模型在未参与联邦训练的中心的独立测试集上泛化能力显著更强，表现出更优的分布外性能。
• 该方法在训练中心的保留测试集和未见中心上的表现均优于FedAvg和FedProx，表明其对数据异质性具有更强的鲁棒性。
• 随机客户端选择降低了通信和计算负载，因每轮仅部分客户端参与，提升了可扩展性。
• 该方法是首个在医学图像分析联邦学习中应用随机客户端和参数选择的方法，为泛化能力提供了新范式。
• 该方法通过在本地训练完整模型并在聚合时仅应用丢弃，保持了高模型容量，与以往采用子模型训练的方法不同。