[论文解读] Closing the Generalization Gap of Cross-silo Federated Medical Image Segmentation

An Xu, Wenqi Li+8 more|arXiv (Cornell University)|Mar 18, 2022

Privacy-Preserving Technologies in Data被引用 7

一句话总结

该论文提出FedSM，一种新型联邦学习框架，通过引入个性化模型和模型选择器，弥合了跨硅联邦学习与集中式训练在医学图像分割任务中的泛化差距。该框架采用一种新型优化方法SoftPull，平衡本地与全局知识，在避免客户端漂移的同时，实现与集中式训练相当的性能。

ABSTRACT

Cross-silo federated learning (FL) has attracted much attention in medical imaging analysis with deep learning in recent years as it can resolve the critical issues of insufficient data, data privacy, and training efficiency. However, there can be a generalization gap between the model trained from FL and the one from centralized training. This important issue comes from the non-iid data distribution of the local data in the participating clients and is well-known as client drift. In this work, we propose a novel training framework FedSM to avoid the client drift issue and successfully close the generalization gap compared with the centralized training for medical image segmentation tasks for the first time. We also propose a novel personalized FL objective formulation and a new method SoftPull to solve it in our proposed framework FedSM. We conduct rigorous theoretical analysis to guarantee its convergence for optimizing the non-convex smooth objective function. Real-world medical image segmentation experiments using deep FL validate the motivations and effectiveness of our proposed method.

研究动机与目标

• 为解决由于非独立同分布数据和客户端漂移导致的跨硅联邦学习在医学图像分割任务中产生的泛化差距。
• 通过为每个客户端学习个性化模型，消除对必须适配所有客户端的单一全局模型的需求。
• 设计一种模型选择器，基于数据分布相似性，动态选择最适合给定测试输入的个性化模型。
• 提出一种新的个性化联邦学习优化公式，并开发SoftPull方法，以平衡本地与全局模型更新，提升泛化能力。
• 对收敛性进行理论分析，并通过实证验证，证明FedSM的性能可与集中式训练相媲美。

提出的方法

• 提出FedSM框架，为每个客户端训练一组个性化模型，并共享一个模型选择器，用于将测试输入路由到最合适的个性化模型。
• 引入一种新颖的个性化联邦学习目标：基于可学习系数λ，对本地与全局模型更新进行加权插值。
• 提出SoftPull方法，通过施加软约束优化个性化目标，以促进与本地数据和全局模型的一致性。
• 基于分割头最后一层的特征嵌入构建模型选择器，实现对测试输入的最佳个性化模型的动态选择。
• 在模型选择器中引入阈值γ，当无个性化模型足够置信时，可回退至全局模型，提升对未见客户端的鲁棒性。
• 在标准联邦学习假设下，对SoftPull所优化的非凸光滑目标函数提供严谨的收敛性理论分析。

Figure 1 : The proposed FedSM framework with “super model”.

实验结果

研究问题

• RQ1 联邦学习框架能否在跨硅医学图像分割任务中弥合联邦学习与集中式训练之间的泛化差距？
• RQ2 通过个性化模型避免客户端漂移，是否能带来优于传统联邦学习方法（如FedAvg）的泛化性能？
• RQ3 模型选择器能否基于数据分布相似性，有效识别最适合给定测试输入的个性化模型？
• RQ4 SoftPull优化方法是否能实现优于现有个性化联邦学习基线方法的泛化性能？
• RQ5 个性化目标中可调的插值系数λ是否能通过防止在小规模本地数据集上过拟合，从而提升泛化能力？

主要发现

• FedSM在真实世界医学图像分割任务中实现了与集中式训练相当的测试Dice分数，有效弥合了泛化差距。
• 在视网膜分割任务中，FedSM优于FedAvg和本地训练，尤其在小样本客户端上表现更优，当使用最优阈值γ时，性能提升最高达3%。
• SoftPull优于其他个性化方法（如APFL、Per-FedAvg和FT），其全局Dice系数比次优基线高出0.62%。
• 在低数据相似性任务（如视网膜）中，最优λ值为0.7；在高相似性任务（如前列腺）中，最优λ值为0.3，验证了SoftPull公式的自适应能力。
• 损失曲面分析表明，SoftPull的损失景观比FedAvg更平坦，表明其泛化能力更优，尽管存在轻微收敛误差。
• 模型选择器能成功识别未见客户端中最相似的数据分布，其选择频率与特征可视化（t-SNE）及实证性能一致。

Figure 2 : Training curves comparison. The curves are non-decreasing because we record the best result during training.