[论文解读] FedBone: Towards Large-Scale Federated Multi-Task Learning

Yiqiang Chen, Teng Zhang+4 more|arXiv (Cornell University)|Jun 30, 2023

Retinal Imaging and Analysis被引用 4

一句话总结

FedBone 提出了一种基于分割学习的联邦多任务学习框架，通过将服务器上的全局通用模型与客户端上的任务特定模型解耦，实现了在资源受限的边缘客户端上大规模模型的训练。它引入了 GPAggregation 方法以缓解异构任务之间的梯度冲突，在基准数据集和真实眼科数据集上实现了最先进性能，同时客户端计算量极低。

ABSTRACT

Heterogeneous federated multi-task learning (HFMTL) is a federated learning technique that combines heterogeneous tasks of different clients to achieve more accurate, comprehensive predictions. In real-world applications, visual and natural language tasks typically require large-scale models to extract high-level abstract features. However, large-scale models cannot be directly applied to existing federated multi-task learning methods. Existing HFML methods also disregard the impact of gradient conflicts on multi-task optimization during the federated aggregation process. In this work, we propose an innovative framework called FedBone, which enables the construction of large-scale models with better generalization from the perspective of server-client split learning and gradient projection. We split the entire model into two components: a large-scale general model (referred to as the general model) on the cloud server and multiple task-specific models (referred to as the client model) on edge clients, solving the problem of insufficient computing power on edge clients. The conflicting gradient projection technique is used to enhance the generalization of the large-scale general model between different tasks. The proposed framework is evaluated on two benchmark datasets and a real ophthalmic dataset. Comprehensive results demonstrate that FedBone efficiently adapts to heterogeneous local tasks of each client and outperforms existing federated learning algorithms in most dense prediction and classification tasks with off-the-shelf computational resources on the client side.

研究动机与目标

• 为解决在边缘客户端计算资源有限的联邦学习环境中训练大规模模型的挑战。
• 处理客户端之间的任务异构性，即不同客户端针对不同目标进行优化，导致模型聚合时出现梯度冲突。
• 在不依赖高客户端计算能力的前提下，提升全局模型在多样化任务上的泛化能力。
• 实现在医疗影像等隐私敏感领域（如眼科疾病诊断）中大规模模型的实际部署。

提出的方法

• 分割学习架构：全局通用模型托管在云服务器上，而客户端仅存储并计算轻量级的任务特定头。
• 客户端在本地执行图像块嵌入，将特征发送至服务器，由大规模通用模型进行特征提取。
• 服务器将提取的潜在表征返回客户端，用于任务特定的输出预测。
• GPAggregation：一种新颖的梯度聚合方法，通过历史梯度注意力对客户端梯度进行重标定，降低异构任务间的梯度冲突。
• 使用可变形卷积和自注意力机制的任务适配模块，以增强低层特征表示的交互能力。
• 模型训练通过客户端与服务器之间的交替前向和反向传播进行，最大限度减少本地参数存储与计算。

Figure 1: The overview of our proposed framework

实验结果

研究问题

• RQ1 联邦多任务学习框架是否能在保持边缘客户端低计算需求的同时，有效扩展至大规模模型？
• RQ2 在联邦聚合过程中，如何缓解由异构任务引发的梯度冲突，以提升模型泛化能力？
• RQ3 采用全局通用模型与客户端特定头的分割学习方法，在异构任务上相较于现有联邦学习基线方法，性能提升程度如何？
• RQ4 所提出的框架是否能在客户端资源极少的情况下，实现在真实世界医学视觉任务上的最先进性能？

主要发现

• 在真实世界数据集的13项眼科任务中，FedBone 在12项上优于 FedAvg、FedProx、pFedMe 和 FedEM，平均 F1 分数最高达到 95.90。
• 在 Retina-Lesions 数据集上，FedBone 在语义分割（LS 任务）中表现最佳，显示出在医学图像分割方面的强大潜力。
• 消融实验表明，GPAggregation 和任务适配模块均显著提升性能，完整 FedBone 模型在所有设置下均达到最高准确率。
• 在 CIFAR-100 多任务基准上，FedBone 达到 98.93% 的准确率，尽管使用的是现成的客户端硬件，仍优于其他联邦学习方法。
• 该框架通过将大部分模型参数和特征提取工作转移到服务器，显著降低了客户端的内存与计算需求，使低资源设备得以部署。
• GPAggregation 方法有效减少了梯度冲突，表现为在多样化任务上收敛性与泛化能力的提升。

Figure 2: The workflow of $\mathtt{FedBone}$ framework. Clients perform patch embedding locally and (1) send embeddings to the cloud server for feature extraction using the general model, and the cloud server (2) sends extracted features back to clients. Clients complete the loss computation and (3)