[论文解读] Application of Federated Learning in Building a Robust COVID-19 Chest X-ray Classification Model

Amartya Bhattacharya, Manish Gawali+2 more|arXiv (Cornell University)|Apr 22, 2022

COVID-19 diagnosis using AI被引用 4

一句话总结

本文提出了一种联邦学习（FL）框架，用于在不集中患者数据的情况下，训练一个用于从胸部X光片中分类COVID-19的鲁棒深度学习模型。通过在三个非独立同分布（non-IID）数据源上进行训练，并使用FedAvg聚合本地模型更新，联邦学习模型的性能与在全部数据上训练的集中式模型相当，同时保护了数据隐私，并提升了在多样化数据分布上的泛化能力。

ABSTRACT

While developing artificial intelligence (AI)-based algorithms to solve problems, the amount of data plays a pivotal role - large amount of data helps the researchers and engineers to develop robust AI algorithms. In the case of building AI-based models for problems related to medical imaging, these data need to be transferred from the medical institutions where they were acquired to the organizations developing the algorithms. This movement of data involves time-consuming formalities like complying with HIPAA, GDPR, etc.There is also a risk of patients' private data getting leaked, compromising their confidentiality. One solution to these problems is using the Federated Learning framework. Federated Learning (FL) helps AI models to generalize better and create a robust AI model by using data from different sources having different distributions and data characteristics without moving all the data to a central server. In our paper, we apply the FL framework for training a deep learning model to solve a binary classification problem of predicting the presence or absence of COVID-19. We took three different sources of data and trained individual models on each source. Then we trained an FL model on the complete data and compared all the model performances. We demonstrated that the FL model performs better than the individual models. Moreover, the FL model performed at par with the model trained on all the data combined at a central server. Thus Federated Learning leads to generalized AI models without the cost of data transfer and regulatory overhead.

研究动机与目标

• 通过避免敏感患者医学影像数据的集中化，解决医疗人工智能中的数据隐私和合规性挑战。
• 减少因在单一来源、分布多样性有限的数据上训练而导致的模型偏差。
• 通过利用多个非独立同分布（non-IID）数据源，提升模型在多样化临床数据分布上的泛化能力。
• 证明联邦学习可在不传输原始数据的情况下，实现与集中式训练相当的性能。
• 通过比较单个、联邦和联合训练设置下的模型性能，验证联邦学习在医学影像中的有效性。

提出的方法

• 在来自不同来源的三个独立、非独立同分布（non-IID）的胸部X光片数据集上，分别训练了三个基于卷积神经网络（CNN）的个体深度学习模型。
• 将所有数据合并为一个集中式“联合”模型作为基准比较，该模型在单台服务器上、拥有完整数据访问权限的情况下进行训练。
• 应用联邦平均（FedAvg）算法，在三个客户端之间训练一个全局联邦学习（FL）模型，每个客户端本地训练并仅共享模型更新。
• 重复进行5轮联邦训练，共20个总训练周期，每轮后进行全局模型聚合。
• 使用与阈值无关的评估指标——曲线下面积（AUROC）和精确率-召回率曲线下面积（AUPRC），在每个客户端的测试集上评估模型，以确保公平比较。
• 通过从不将原始图像传输至中心服务器，确保了数据隐私；所有训练和聚合操作均在本地安全完成。

实验结果

研究问题

• RQ1 联邦学习能否提升在多样化、非独立同分布（non-IID）数据分布上进行COVID-19胸部X光片分类模型的泛化能力？
• RQ2 联邦模型的性能与在单个数据源上训练的模型以及整合全部数据的集中式模型相比如何？
• RQ3 在无需传输敏感医学影像数据的情况下，联邦学习在多大程度上能保持模型的鲁棒性？
• RQ4 与在单一数据源上训练的模型相比，FL框架是否能减少数据偏差？
• RQ5 FL能否在不损害患者数据隐私的前提下，实现与集中式训练相当的性能？

主要发现

• 联邦学习（FL）模型在不同来源的测试集上均优于所有单个客户端模型，表现出更优的泛化能力。
• 在Client1的测试集上，FL模型的AUROC达到99.61%，优于单独的Client1模型（99.96%），并显著超过Client2（37.85%）和Client3（43.85%）的模型。
• 在Client2的测试集上，FL模型的AUROC达到92.31%，优于单独的Client2模型（93.94%），并大幅优于Client1（37.85%）和Client3（43.85%）的模型。
• 在Client3的测试集上，FL模型的AUROC达到88.31%，优于单独的Client3模型（97.73%），并显著优于Client1（43.85%）和Client2（37.85%）的模型。
• FL模型在三个客户端测试集上的AUROC得分分别为99.61%、92.31%和88.31%，表明其在多样化数据分布上均保持了稳定且优异的性能。
• FL模型的性能几乎与集中式“联合”模型相当，后者在相同测试集上的AUROC得分分别为99.86%、99.92%和93.51%，证实了FL可在不传输数据的情况下实现与集中式训练相当的性能。