[论文解读] FedIIC: Towards Robust Federated Learning for Class-Imbalanced Medical Image Classification

Nannan Wu, Li Yu+3 more|arXiv (Cornell University)|Jun 28, 2022

Imbalanced Data Classification Techniques被引用 5

一句话总结

FedIIC 提出了一种隐私保护的联邦学习框架，通过联合优化特征学习与分类器校准，以解决医学图像分类中的类别不平衡问题。该框架采用两级对比学习（客户端内与客户端间）以提升特定类别的表征能力，并引入基于难度感知的logit调整机制，动态平衡决策边界，在真实世界与模拟的医学数据集上实现了最先进的性能，平衡准确率显著提升。

ABSTRACT

Federated learning (FL), training deep models from decentralized data without privacy leakage, has shown great potential in medical image computing recently. However, considering the ubiquitous class imbalance in medical data, FL can exhibit performance degradation, especially for minority classes (e.g. rare diseases). Existing methods towards this problem mainly focus on training a balanced classifier to eliminate class prior bias among classes, but neglect to explore better representation to facilitate classification performance. In this paper, we present a privacy-preserving FL method named FedIIC to combat class imbalance from two perspectives: feature learning and classifier learning. In feature learning, two levels of contrastive learning are designed to extract better class-specific features with imbalanced data in FL. In classifier learning, per-class margins are dynamically set according to real-time difficulty and class priors, which helps the model learn classes equally. Experimental results on publicly-available datasets demonstrate the superior performance of FedIIC in dealing with both real-world and simulated multi-source medical imaging data under class imbalance. Code is available at https://github.com/wnn2000/FedIIC.

研究动机与目标

• 为解决联邦学习在医学图像分类中面临的类别不平衡关键挑战，即少数疾病类别代表性不足且识别效果差。
• 开发一种隐私保护的联邦学习框架，无需在机构间共享原始数据。
• 在全局模型训练中同时缓解因背景属性不平衡导致的属性偏差，以及因先验概率差异导致的类别偏差。
• 在去中心化、多源医学影像环境下，提升少数类别性能，同时不牺牲多数类别的准确率。

提出的方法

• 引入两级对比学习：客户端内对比学习增强每个客户端本地数据中的特征判别能力，而客户端间对比学习则通过跨客户端对齐特征，提升泛化性能。
• 使用温度调度的对比损失对特征提取器进行校准，使其适应类别先验，降低类别间相似性，提升表征质量。
• 提出一种难度感知的logit调整（DALA）机制，根据样本实时难度和类别先验概率动态调整分类器logits，以平衡决策边界。
• 将对比特征学习与动态分类器校准两个模块整合进统一的联邦学习流程中，联合减少属性偏差与类别偏差。
• 采用全客户端参与的全局模型聚合策略，确保在异构、不平衡数据源下实现稳健收敛。
• 该方法使用标准深度学习骨干网络实现，兼容现有联邦学习训练协议，确保在临床环境中的实际部署可行性。

实验结果

研究问题

• RQ1 在医学影像中全局类别不平衡的联邦学习场景下，联合优化特征学习与分类器校准是否能提升性能？
• RQ2 在全局数据不平衡条件下，两级对比学习（客户端内与客户端间）在学习鲁棒、类别特定表征方面是否有效？
• RQ3 基于样本难度与类别先验的概率的动态logit调整是否能带来更优的决策边界并提升少数类别识别能力？
• RQ4 FedIIC在处理真实世界与模拟的多源、类别不平衡医学图像数据集时，与当前最先进方法相比表现如何？

主要发现

• FedIIC在所有评估数据集上均取得最高的平衡准确率（BACC），在Real数据集上领先第二名方法（CReFF）3.99%，在ISIC上领先7.32%，在ICH上领先2.01%。
• 消融实验表明，结合客户端内与客户端间对比学习及DALA机制可获得最佳性能，ISIC数据集上BACC达78.84%，ICH数据集上达84.22%。
• 将温度参数设置为t=0.5（基于类别先验）相比t=0显著提升性能，验证了动态温度调度的重要性。
• 原型向量正交化可降低类别间相似性并提升性能，应用后BACC从72.64%提升至78.84%。
• 引入难度感知logit调整（d=0.25）可提升BACC，但略微降低F1，证实了平衡准确率与精确率/召回率调和平均值之间的权衡。
• 完整FedIIC框架在ISIC上取得最高F1得分78.05，在ICH上达84.73，表明在所有指标上均实现一致提升。