[论文解读] Auto-FedAvg: Learnable Federated Averaging for Multi-Institutional Medical Image Segmentation

Yingda Xia, Dong Yang+16 more|arXiv (Cornell University)|Apr 20, 2021

Privacy-Preserving Technologies in Data参考文献 50被引用 36

一句话总结

Auto-FedAvg 学习数据驱动的聚合权重用于联邦学习，使用Dirichlet分布、可微分的方法来组合客户端模型，在CIFAR-10和多机构医学分割任务上提升FedAvg。

ABSTRACT

Federated learning (FL) enables collaborative model training while preserving each participant's privacy, which is particularly beneficial to the medical field. FedAvg is a standard algorithm that uses fixed weights, often originating from the dataset sizes at each client, to aggregate the distributed learned models on a server during the FL process. However, non-identical data distribution across clients, known as the non-i.i.d problem in FL, could make this assumption for setting fixed aggregation weights sub-optimal. In this work, we design a new data-driven approach, namely Auto-FedAvg, where aggregation weights are dynamically adjusted, depending on data distributions across data silos and the current training progress of the models. We disentangle the parameter set into two parts, local model parameters and global aggregation parameters, and update them iteratively with a communication-efficient algorithm. We first show the validity of our approach by outperforming state-of-the-art FL methods for image recognition on a heterogeneous data split of CIFAR-10. Furthermore, we demonstrate our algorithm's effectiveness on two multi-institutional medical image analysis tasks, i.e., COVID-19 lesion segmentation in chest CT and pancreas segmentation in abdominal CT.

研究动机与目标

• 解决跨机构医疗影像中联邦客户端之间的非独立同分布数据问题。
• 开发可微分的聚合权重学习机制，而不是使用固定、基于数据规模的权重。
• 确保隐私保护的联邦学习，尽量减少额外通信。
• 在医疗影像任务上展示相对于FedAvg和最先进的FL方法的改进。
• 探索网络级和层级聚合变体，以及Dirichlet与Softmax参数化。

提出的方法

• 将本地模型参数与全局聚合参数解耦以进行迭代更新。
• 将聚合权重表示为可学习参数（softmax或Dirichlet），以强制单纯形约束。
• 使用Dirichlet分布来建模聚合权重并实现基于梯度的优化。
• 引入一种通信高效的LearnAggWeight步骤，在不共享原始数据的情况下更新聚合权重。
• 提出网络级和层级聚合方案，以按层或整个网络灵活加权客户端贡献。
• 分析聚合权重学习区间t0以在额外通信与性能之间取得平衡。

实验结果

研究问题

• RQ1 可学习的、数据驱动的聚合权重是否能在跨机构FL的非i.i.d.数据下改进联邦平均？
• RQ2 Dirichlet分布的聚合权重是否比基于Softmax的权重具有更好的性能和泛化能力？
• RQ3 网络级与层级聚合对性能与通信成本的影响是什么？
• RQ4 聚合权重学习区间t0在跨机构场景下如何影响性能与效率？

主要发现

• 采用网络级Dirichlet权重的Auto-FedAvg（Auto-FedAvg-N-Dirichlet）在CIFAR-10的变体中取得最佳结果，最终准确度为88.98%，比FedMA（87.53%）高1.45%。
• 在多国新冠病灶分割任务中，Auto-FedAvg-N-Dirichlet变体在全局模型性能、平均本地模型性能和本地泛化方面优于FedAvg，例如在报告的实验中全局测试平均水平较FedAvg提高约2个百分点。
• Auto-FedAvg在CIFAR-10及两个医学影像任务（COVID-19病灶分割与胰腺分割）上持续优于FedAvg、FedProx与FedMA基线。
• 在多数任务中，Dirichlet聚合权重通常比Softmax更有效，且网络级聚合在大多数情形下优于层级聚合。
• 在训练早期对聚合权重的积极学习可加速收敛，后期稳定权重则可进一步提升泛化。
• 在胰腺分割实验中，Auto-FedAvg-N-Dirichlet* 在全局和局部性能上均达到最佳，明显优于FedAvg。