[论文解读] SelfFed: Self-Supervised Federated Learning for Data Heterogeneity and Label Scarcity in Medical Images

Sunder Ali Khowaja, Kapal Dev+2 more|arXiv (Cornell University)|Jul 4, 2023

Privacy-Preserving Technologies in Data被引用 5

一句话总结

本文提出 SelfFed，一种用于医疗影像的自监督联邦学习框架，旨在解决物联网医疗（IoMT）中数据异质性和标签稀缺性的问题。该框架采用两阶段方法：首先利用 Swin Transformer 通过掩码自编码（MAE）进行预训练，以应对数据异质性；随后通过一种新颖的聚合策略进行对比微调，以缓解标签稀缺问题，在非独立同分布（non-IID）数据上实现最高 8.8% 的准确率提升，并且在仅使用 10% 标注样本的情况下仍优于基线模型。

ABSTRACT

Self-supervised learning in the federated learning paradigm has been gaining a lot of interest both in industry and research due to the collaborative learning capability on unlabeled yet isolated data. However, self-supervised based federated learning strategies suffer from performance degradation due to label scarcity and diverse data distributions, i.e., data heterogeneity. In this paper, we propose the SelfFed framework for medical images to overcome data heterogeneity and label scarcity issues. The first phase of the SelfFed framework helps to overcome the data heterogeneity issue by leveraging the pre-training paradigm that performs augmentative modeling using Swin Transformer-based encoder in a decentralized manner. The label scarcity issue is addressed by fine-tuning paradigm that introduces a contrastive network and a novel aggregation strategy. We perform our experimental analysis on publicly available medical imaging datasets to show that SelfFed performs better when compared to existing baselines and works. Our method achieves a maximum improvement of 8.8% and 4.1% on Retina and COVID-FL datasets on non-IID datasets. Further, our proposed method outperforms existing baselines even when trained on a few (10%) labeled instances.

研究动机与目标

• 解决在物联网医疗（IoMT）背景下，联邦学习用于医疗影像时面临的数据异质性与标签稀缺性双重挑战。
• 克服传统联邦学习方法在数据分布偏斜导致的非独立同分布（non-IID）数据设置下性能下降的问题。
• 通过在去中心化环境中利用有限标注数据实现有效学习，减少对大规模标注数据集的依赖。
• 开发一种鲁棒且保护隐私的框架，确保在多样化、去中心化的医疗影像客户端中保持模型性能。
• 通过整合自监督预训练与对比微调及自适应聚合策略，提升联邦医疗人工智能的泛化能力与收敛性。

提出的方法

• 提出一种两阶段 SelfFed 框架：(1) 利用掩码自编码（MAE）与 Swin Transformer 编码器进行预训练，以去中心化方式学习鲁棒表征。
• 在预训练过程中对 60% 的 16x16 图像块应用增强型掩码，以促进鲁棒特征学习。
• 在微调阶段引入对比网络，从同一输入的两个视图中学习，从而在标签稀缺条件下提升表征质量。
• 设计一种新颖的客户端聚合策略，对在微调阶段多次参与的客户端进行降权处理，以减少过拟合客户端带来的偏差。
• 在预训练与微调阶段均使用 AdamW 优化器配合余弦衰减与学习率预热策略，确保联邦训练与集中式训练之间保持一致的超参数设置。
• 采用全局模型聚合协议，客户端本地训练并将其更新后的权重发送至服务器进行聚合，从而保护数据隐私。

实验结果

研究问题

• RQ1 使用掩码自编码进行自监督预训练是否能缓解联邦医疗影像学习中因数据异质性导致的性能下降？
• RQ2 所提出的对比微调结合新颖聚合策略，在非独立同分布（non-IID）设置下，如何提升标签稀缺条件下的性能表现？
• RQ3 当在联邦设置中对有限标注数据进行微调时，SelfFed 在多大程度上能超越 ImageNet 预训练模型（如 ViT、BEiT、MAE）？
• RQ4 当仅使用 10% 的标注数据时，所提出框架在非独立同分布（non-IID）与独立同分布（IID）数据分布下是否仍能保持鲁棒性与一致性？
• RQ5 所提出的新颖聚合策略是否能有效降低频繁参与微调的客户端的影响，从而提升模型公平性与收敛性？

主要发现

• 在非独立同分布（non-IID）数据条件下，SelfFed 相较于 ImageNet 预训练基线模型，在 Retina 数据集上实现最高 8.8% 的准确率提升，在 COVID-FL 数据集上实现 4.1% 的准确率提升。
• 即使仅使用 10% 的标注样本，SelfFed 在 IID 与非独立同分布（non-IID）设置下仍优于在 100% 标注数据上训练的基线模型。
• 所提出的 SelfFed-MAE 预训练阶段显著降低了由数据异质性引起的表现下降，优于从零开始训练的 Swin Transformer 及 ImageNet 预训练模型。
• 结合新颖聚合策略的对比微调方法有效缓解了频繁参与客户端的影响，提升了模型稳定性与公平性。
• 在 CXR14 数据集上，SelfFed 的表现优于预训练的 ViT、BEiT 与 MAE 模型，证实其在多样化医疗影像基准上的鲁棒性。
• 该方法在多个数据集与不同设置下均表现出一致的高性能，证明其在真实 IoMT 场景中面对有限标签与异质数据时的有效性。