[论文解读] Robust Split Federated Learning for U-shaped Medical Image Networks

Ziyuan Yang, Yingyu Chen+5 more|arXiv (Cornell University)|Dec 13, 2022

Privacy-Preserving Technologies in Data被引用 9

一句话总结

本文提出鲁棒分割联邦学习（RoS-FL），一种结合分割学习与联邦学习的混合框架，用于实现U型医学图像网络的隐私保护与计算高效训练。通过将模型拆分为三部分——头部（客户端）、主体（计算服务器）和尾部（客户端），并分布部署于客户端、计算服务器与聚合服务器之间，RoS-FL在保护输入、标签、输出及模型隐私的同时，支持并行训练。其核心贡献是一种动态权重校正策略（DWCS），可在数据异构环境下稳定训练，减少模型漂移，提升客户端间性能的一致性。

ABSTRACT

U-shaped networks are widely used in various medical image tasks, such as segmentation, restoration and reconstruction, but most of them usually rely on centralized learning and thus ignore privacy issues. To address the privacy concerns, federated learning (FL) and split learning (SL) have attracted increasing attention. However, it is hard for both FL and SL to balance the local computational cost, model privacy and parallel training simultaneously. To achieve this goal, in this paper, we propose Robust Split Federated Learning (RoS-FL) for U-shaped medical image networks, which is a novel hybrid learning paradigm of FL and SL. Previous works cannot preserve the data privacy, including the input, model parameters, label and output simultaneously. To effectively deal with all of them, we design a novel splitting method for U-shaped medical image networks, which splits the network into three parts hosted by different parties. Besides, the distributed learning methods usually suffer from a drift between local and global models caused by data heterogeneity. Based on this consideration, we propose a dynamic weight correction strategy ( extbf{DWCS}) to stabilize the training process and avoid model drift. Specifically, a weight correction loss is designed to quantify the drift between the models from two adjacent communication rounds. By minimizing this loss, a correction model is obtained. Then we treat the weighted sum of correction model and final round models as the result. The effectiveness of the proposed RoS-FL is supported by extensive experimental results on different tasks. Related codes will be released at https://github.com/Zi-YuanYang/RoS-FL.

研究动机与目标

• 解决U型医学图像网络中的隐私问题，其中原始数据、标签和输出均属敏感信息。
• 通过将重型模型组件卸载至服务器，减轻联邦学习中的客户端计算负载。
• 缓解分布式训练环境中因数据异构性导致的模型漂移问题。
• 在保护客户端与服务器间数据与模型隐私的前提下，实现并行训练。
• 通过新颖的动态权重校正机制，提升在非独立同分布（non-IID）数据下客户端间的训练稳定性和性能一致性。

提出的方法

• 将U型网络拆分为三部分：头部（客户端侧）、主体（服务器侧）和尾部（客户端侧），以实现计算分担并增强隐私保护。
• 采用联邦学习与分割学习相结合的混合架构，支持并行训练并减轻客户端计算负担。
• 仅在客户端与服务器之间传输中间特征图，避免传输原始输入、标签或最终输出。
• 实施一种动态权重校正策略（DWCS），通过计算当前与前一全局模型的加权和来校正模型漂移。
• 定义权重校正损失以量化连续通信轮次间模型漂移，并优化校正模型以最小化该损失。
• 在中心化聚合服务器上聚合客户端梯度与服务器端模型更新，以更新全局模型。

实验结果

研究问题

• RQ1 混合分割联邦学习框架能否有效保护U型医学图像网络中的输入、标签、输出及模型参数？
• RQ2 在医学图像模型的分布式训练中，如何缓解由数据异构性引起的模型漂移？
• RQ3 所提出的动态权重校正策略（DWCS）在多大程度上提升了非独立同分布（non-IID）数据下客户端间的训练稳定性和性能一致性？
• RQ4 与标准联邦学习相比，所提出的RoS-FL框架是否能在降低客户端计算成本的同时实现具有竞争力的性能？
• RQ5 RoS-FL的性能对超参数选择（如校正权重μ）的敏感度如何？

主要发现

• RoS-FL在客户端间实现一致的性能表现，消除了基线方法中因模型漂移导致的显著性能差距。
• 在恢复任务中，最佳PSNR达到42.81，仅需1次通信轮次与500轮本地训练，优于其他配置。
• 在分割任务中，最高Dice分数达88.21，对应2次通信轮次与150轮本地训练，展现出对数据异构性的强鲁棒性。
• 该方法对超参数μ不甚敏感，在宽范围（1×10⁻⁶至100）内表现稳定，最优性能出现在μ = 1×10⁻⁴。
• 结果表明，当本地训练轮次超过某一阈值后，性能在图像任务中开始下降，归因于模型漂移加剧，凸显DWCS在稳定训练中的关键作用。
• 视觉对比显示，与其它方法相比，RoS-FL生成的输出边缘更清晰，结构更完整，伪影更少。