[论文解读] Federated Contrastive Learning for Dermatological Disease Diagnosis via On-device Learning

Yawen Wu, Dewen Zeng+6 more|arXiv (Cornell University)|Feb 14, 2022

Cutaneous Melanoma Detection and Management被引用 7

一句话总结

本文提出了一种用于皮肤科疾病诊断的设备端联邦对比学习框架，通过有限的标注数据提升模型性能。通过在自监督预训练期间实现设备间特征共享并移除本地负样本，该方法提升了数据多样性与表征质量，在皮肤科数据集上的召回率和精确率均优于现有的联邦学习与对比学习基线方法，达到最先进水平。

ABSTRACT

Deep learning models have been deployed in an increasing number of edge and mobile devices to provide healthcare. These models rely on training with a tremendous amount of labeled data to achieve high accuracy. However, for medical applications such as dermatological disease diagnosis, the private data collected by mobile dermatology assistants exist on distributed mobile devices of patients, and each device only has a limited amount of data. Directly learning from limited data greatly deteriorates the performance of learned models. Federated learning (FL) can train models by using data distributed on devices while keeping the data local for privacy. Existing works on FL assume all the data have ground-truth labels. However, medical data often comes without any accompanying labels since labeling requires expertise and results in prohibitively high labor costs. The recently developed self-supervised learning approach, contrastive learning (CL), can leverage the unlabeled data to pre-train a model, after which the model is fine-tuned on limited labeled data for dermatological disease diagnosis. However, simply combining CL with FL as federated contrastive learning (FCL) will result in ineffective learning since CL requires diverse data for learning but each device only has limited data. In this work, we propose an on-device FCL framework for dermatological disease diagnosis with limited labels. Features are shared in the FCL pre-training process to provide diverse and accurate contrastive information. After that, the pre-trained model is fine-tuned with local labeled data independently on each device or collaboratively with supervised federated learning on all devices. Experiments on dermatological disease datasets show that the proposed framework effectively improves the recall and precision of dermatological disease diagnosis compared with state-of-the-art methods.

研究动机与目标

• 解决在标注数据稀缺且分布于移动设备上的情况下，训练高精度皮肤科疾病诊断模型的挑战。
• 克服对比学习在联邦设置中的局限性，即每个设备的数据多样性较低会降低性能。
• 通过避免原始数据共享，保护患者隐私，同时通过设备间特征交换增强表征学习。
• 通过使用预训练模型进行有效微调，提升标注效率，仅需极少的标注数据。
• 开发一种两阶段框架，结合联邦自监督预训练与协作或本地微调，用于医学图像诊断。

提出的方法

• 提出两阶段框架：首先在无标注数据上使用对比学习进行联邦自监督预训练，随后使用有限标注数据进行微调。
• 在预训练阶段提出设备间特征交换（FE），在不共享原始图像的前提下提升数据多样性与对比学习质量。
• 在对比学习过程中移除本地负样本来增强学习表征的聚类效果并减少噪声。
• 在预训练与微调两个阶段均使用 FedAvg 作为聚合策略，以保持设备间模型的一致性。
• 采用孪生网络架构，并结合数据增强技术（如随机裁剪、颜色抖动）生成正负样本对以计算对比损失。
• 支持本地与联邦微调，根据数据可用性与隐私约束灵活部署。

实验结果

研究问题

• RQ1 当每个设备仅有少量且无标注数据时，联邦对比学习能否有效学习到用于皮肤科疾病诊断的鲁棒表征？
• RQ2 在客户端数据多样性较低的联邦设置中，设备间特征共享如何提升对比学习的性能？
• RQ3 移除本地负样本对学习表征质量与下游诊断准确率有何影响？
• RQ4 所提框架在保持或提升诊断性能的前提下，能在多大程度上减少标注需求？
• RQ5 在皮肤科数据集上，所提方法与最先进联邦学习与对比学习基线相比，在召回率与精确率方面表现如何？

主要发现

• 在 10% 标注数据下，所提方法在联邦微调中达到 48.03% 的召回率与 42.87% 的精确率，优于最佳基线方法 FedSimCLR，召回率提升 2.14 个百分点，精确率提升 1.43 个百分点。
• 在 40% 标注数据比例下，所提方法达到 55.13% 的召回率与 48.73% 的精确率，性能与使用两倍标注数据训练的基线方法相当或更优。
• 消融实验表明，特征交换使召回率从 44.72% 提升至 47.45%，进一步移除本地负样例后，召回率提升至 48.03%（10% 标注数据）。
• 在所有标注比例（10%、20%、40%、80%）下，所提方法在本地与联邦微调设置中均持续优于所有基线方法，显著提升召回率与精确率。
• 在包含不同肤色与疾病表现的皮肤科数据集上评估显示，该框架对多种皮肤类型具有鲁棒性。
• 结果表明，特征共享在低数据、隐私受限环境中显著增强表征学习，使其适用于真实世界中的移动皮肤科应用。