[论文解读] FedMix: Mixed Supervised Federated Learning for Medical Image Segmentation

Jeffry Wicaksana, Zengqiang Yan+5 more|arXiv (Cornell University)|May 4, 2022

Privacy-Preserving Technologies in Data被引用 10

一句话总结

FedMix 提出了一种用于医学图像分割的新型联邦学习框架，将像素级、边界框和图像级等多种标签类型统一到单一训练过程中，通过伪标签优化和动态样本选择提升模型质量。该方法引入自适应客户端加权机制，显著提升全局模型收敛性，在乳腺肿瘤和皮肤病变分割任务上大幅超越当前最先进方法。

ABSTRACT

The purpose of federated learning is to enable multiple clients to jointly train a machine learning model without sharing data. However, the existing methods for training an image segmentation model have been based on an unrealistic assumption that the training set for each local client is annotated in a similar fashion and thus follows the same image supervision level. To relax this assumption, in this work, we propose a label-agnostic unified federated learning framework, named FedMix, for medical image segmentation based on mixed image labels. In FedMix, each client updates the federated model by integrating and effectively making use of all available labeled data ranging from strong pixel-level labels, weak bounding box labels, to weakest image-level class labels. Based on these local models, we further propose an adaptive weight assignment procedure across local clients, where each client learns an aggregation weight during the global model update. Compared to the existing methods, FedMix not only breaks through the constraint of a single level of image supervision, but also can dynamically adjust the aggregation weight of each local client, achieving rich yet discriminative feature representations. To evaluate its effectiveness, experiments have been carried out on two challenging medical image segmentation tasks, i.e., breast tumor segmentation and skin lesion segmentation. The results validate that our proposed FedMix outperforms the state-of-the-art method by a large margin.

研究动机与目标

• 解决现有联邦学习方法假设客户端标签监督水平一致的局限性。
• 实现在医学图像分割中有效利用异构标签类型——强监督（像素级）、弱监督（边界框）和最弱监督（图像级）的能力。
• 通过利用易于获取的弱标签，减少医学专家的标注负担，同时不共享原始数据。
• 通过在训练过程中动态选择高一致性伪标签样本，提升模型鲁棒性与收敛性。
• 设计一种自适应聚合机制，根据本地数据质量和数量对客户端进行加权，从而提升全局模型性能。

提出的方法

• 提出统一的 FedMix 框架，支持在单一联邦学习流程中集成像素级、边界框和图像级标签的混合监督。
• 应用针对不同监督类型的优化策略，从弱标签数据中生成高质量伪标签，提升标签准确性。
• 通过两台独立训练模型之间的预测一致性，实现动态样本选择，以过滤低置信度的伪标签样本。
• 使用一致性阈值 ε，基于模型间预测的一致性选择评估样本，确保仅使用可靠的伪标签进行训练。
• 设计自适应聚合函数，根据本地数据质量与数量计算客户端特定权重，提升全局模型更新的稳定性。
• 使用标注数据与伪标签数据的组合训练本地模型，并通过客户端模型的加权聚合更新全局模型。

Figure 1 : Examples of different levels of medical image labels, where the image-level class labels in (b) contain only the lesion category. The bounding box labels in (c) contain not only the lesion category, but also a coarse location. The pixel-level labels in (d) contain both the lesion category

实验结果

研究问题

• RQ1 联邦学习框架能否在不共享数据的前提下，有效利用多种图像监督级别（像素级、边界框、图像级）进行医学图像分割？
• RQ2 基于两模型预测一致性的动态样本选择机制，如何影响伪标签质量与模型收敛性？
• RQ3 与均匀聚合相比，自适应客户端加权在异构联邦设置下能在多大程度上提升全局模型性能？
• RQ4 在完全监督与半监督联邦学习场景下，FedMix 相较于当前最先进方法（如 FedAvg 和 FedST）表现如何？
• RQ5 所提出的框架能否在保持或提升分割精度的同时，显著降低标注负担，适用于多样化的医学影像客户端？

主要发现

• 在 BUS 数据集（C1）上，FedMix 的 Dice 系数（DC）达到 66.92%；在 BUSIS（C2）上为 88.49%；在 UDIAT（C3）上为 86.95%，显著优于 FedAvg 和 FedST。
• 在皮肤病变分割任务中，FedMix 在 C1 上取得 81.30% 的 DC，C2 上为 78.10%，C3 上为 94.43%，C4 上为 91.11%，平均 DC 达 86.24%，展现出优异的跨客户端泛化能力。
• 动态样本选择机制显著提升性能：当 ε = 0.9 时，BUS（C1）上 DC 达 86.07%，BUSIS（C2）上为 89.1%，UDIAT（C3）上为 79.6%，验证了其在过滤不准确伪标签中的关键作用。
• 自适应聚合使乳腺肿瘤分割任务中平均 DC 相较 FedAvg 提升 1.00%，在皮肤病变分割中整体提升 0.69%，在低质量数据客户端上提升更高（达 1.44%）。
• 消融实验表明，若移除动态样本选择，性能显著下降（如平均 DC 从 80.78% 降至 37.67%），证明其对稳定训练的必要性。
• FedMix 在半监督设置下优于 FedST，在完全监督场景中也持续优于 FedAvg，证实其鲁棒性与可扩展性。

Figure 2 : An illustration of our proposed Mixed Supervised Federated Learning (FedMix) framework. The local client update utilizes every available supervision for training. Based on which, an adaptive weight aggregation procedure is used for the global federated model update. Compared to the existi