ICCV 2023 | 如何在资源受限下进行联邦多任务学习

摘要 · 看点

联邦学习是一种新兴的分布式机器学习方法，它允许多个边缘设备在不共享原始数据的情况下合作训练机器学习模型。然而，多个并行的联邦学习任务可能会使资源受限的设备超负荷。

在这个工作中，本文提出了第一个能够有效协调和训练多个并行任务的联邦学习系统和方法MAS（Merge and Split），以优化多个同时进行的联邦学习任务的训练性能。MAS首先通过使用多任务架构将多个任务合并成一个多合一联邦学习任务。在训练多轮之后，MAS利用在多合一训练过程中采集的任务之间的亲和力将多合一任务分成两个或多个任务，继续训练。实验表明，MAS在减少训练时间2倍和减少能耗40%的同时，取得了优于其他方法的性能。

论文名称：MAS: Towards Resource-Efficient Federated Multiple-Task Learning

 问题和挑战

联邦学习是一种新兴的分布式模型训练方法，既保护隐私，又能够协同多个边缘设备进行模型训练。它已经在计算机视觉领域得到广泛的应用，比如医学影像分析和自动驾驶。

然而，绝大多数边缘设备，算力和存储资源和云服务器比都比较有限，它们通常只能实时支持单个联邦学习任务。如果在同一个设备上同时进行多个联邦学习任务可能导致设备内存超载、算力不足和电能消耗过大。

最简单的方法是，将这些联邦学习任务视为独立的训练任务，并按顺序进行训练 (One-by-one)。这种方法可以避免设备超负荷，但完成所有任务的训练速度较慢。

另一种解决方案涉及使用多任务学习训练多个联邦学习任务，将多个联邦学习任务组合成一个多任务学习的神经网络结构 (All-in-one)。从训练时间的角度来看，这种方法更加高效，但可能会导致性能下降，因为并不是所有任务都能同等受益于一起训练。它并没有有效考虑每个任务的特性。

图1 更形象地展示了这两种方法再训练时间和性能之间的权衡关系，而我们的方法可以用较少的训练时间取得更好的性能。

图1. 现有方法在同时训练多个联邦学习任务时，往往在训练时间和性能之间存在一个权衡关系，我们的方法可以用较少的训练时间取得更好的性能

 方法

在这个工作中，我们提出了第一个能够有效协调和训练多个同时进行任务的联邦学习系统和方法MAS (Merge and Split)，它既考虑了多个任务之间的差异性，也考虑多个任务之间的相似性。图2展示了MAS的架构和流程图。

图2. MAS 的架构和工作流程

以下是MAS的工作原理：

合并多个联邦学习任务 (Merge)：MAS首先将多个联邦学习任务合并成一个多合一的联邦学习任务。该模型是一个具有多任务的模型，包括用于通用特征的共享编码器以及用于每个任务专门的解码器。因为只需要训练一个神经网络，MAS可以防止设备超负荷。合并之后，我们对这个多合一的任务进行一定轮次的训练，并在训练的过程中采集多个联邦学习任务之间的亲和度，供下一步使用。

拆分多合一任务 (Split)：在对多合一任务进行一定轮次的训练后，MAS使用在训练过程中采集的亲和度分数来评估任务之间的协同作用和差异。基于这些分数，MAS将多合一任务分成两个或更多个联邦学习任务。然后，基于多合一任务后的模型参数做为初始化，进一步训练这些任务。

图3展示了在整个训练的过程中，多个联邦学习任务的模型结构变化。

图3. MAS中任务模型结构的变化

如何拆分多合一的联邦学习任务？

1.  在客户端进行多合一联邦学习任务的训练期间，通过以下公式计算任务之间的亲和度。公式的左边为一个联邦学习任务对于另一个联邦学习任务的亲和度。

2. 对所有客户端的任务亲和度进行平均。

3. 计算单个任务的自身亲和度。

4. 选择具有最高亲和度分数的组合来拆分多合一任务。图4展示了一组多个任务之间亲和度的计算示例图，A、 B、C分别代表一个任务，任务组合 {AB,C} 取得了最高的亲和度分数，所以我们将多合一模型拆分成两个联邦学习任务，一个任务有两个解码器，分别对应任务A 和 B，另一个任务只有一个解码器 C (如图3)。

图4. 多个任务之间亲和度计算示例图，其中A、B、C分别代表一个任务

实验结果

我们使用Taskonomy数据集进行实验，该数据集是一个庞大且具有挑战性的计算机视觉数据集，包含了建筑物室内场景的图像。我们模拟了32个客户端进行实验，每个客户端包含一个建筑物的数据集，以模拟联邦学习中的数据异构性。

图5显示了这32个客户端上的数据量分布；一些客户端仅有4,000张图像，而其他客户端则拥有超过16,000张图像。

我们设计了多组多任联邦学习实验，以评估MAS在不同组合和不同数量任务下的性能，包括两组5个联邦学习任务和一组9个联邦学习任务的实验。这9个任务分别是 语义分割、深度估计、表面法线、关键点检测、边缘纹理，边缘遮挡、重塑、主曲率、自编码器。

图5. 32个客户端上的数据量分布

图6在5个联邦学习任务上，进行了性能和训练时间方面的多种方法的比较。那些能够实现更低的总测试损失和更低能耗的方法表现更佳。

一方面，多合一方法只需要最少的训练时间，但它们性能较差，测试损失最高。仅仅应用多任务学习优化 (如GradNorm[1]) 或联邦学习优化 (如FedProx[2])，几乎无法改善性能。

另一方面，一些多任务学习的方法如HOA [3]可以实现可比较的测试损失，但它需要长时间的训练和高能耗（约为我们的4-6倍）计算高阶估计的成对任务。

虽然逐个任务训练和TAG[4]在测试损失和系统指标之间取得了良好的平衡，但MAS在两个方面都表现更出色；它在测试损失上表现最佳，同时在训练时间上减少了约2倍，能耗减少了约40%。

图6. 在5个联邦学习任务上，多种方法性能和训练时间方面的比较

图7是在9个联邦学习任务上的实验结果，整体结论和图6的相似。

图7. 在9个联邦学习任务上，多种方法性能和训练时间方面的比较

结语

在本文中，我们提出了MAS，第一个能够在资源受限下有效协调和训练多个同时进行的联邦学习任务的系统。我们通过任务合并和任务拆分，考虑了多个联邦学习任务之间的协同作用和差异。大量实证验证了我们的方法可以用更低的能耗，进行更快的训练，同时取得了更好的性能。

同时训练多个联邦学习任务，有潜在的非常广泛的应用。比如在自动驾驶、机器人和智能制造等领域，需要同时执行多个资源密集型任务，MAS可以优化训练、降低能耗，提高模型性能。

相关资料

论文地址:

https://weiming.me/publication/mas/

References：

[1] Zhao Chen, Vijay Badrinarayanan, Chen-Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International Conference on Machine Learning, pages 794–803. PMLR, 2018.

[2] Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks. Proceedings of Machine Learning and Systems, 2:429–450, 2020.

[3] Trevor Standley, Amir Zamir, Dawn Chen, Leonidas Guibas, Jitendra Malik, and Silvio Savarese. Which tasks should be learned together in multi-task learning? In International Conference on Machine Learning, pages 9120–9132. PMLR, 2020.

[4] Chris Fifty, Ehsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, and Chelsea Finn. Efficiently identifying task groupings for multi-task learning. Advances in Neural Information Processing Systems, 34, 2021.

本文转自商汤学术。

加入「计算机视觉」交流群👇备注：CV