分布式多任务学习及联邦学习个性化

1 导引

现在多任务学习根据数据的收集方式可以粗略地被分为两种，一个是集中化的计算方法，即假定数据被事先收集到一个中心节点上然后再运行模型， 大多数基于神经网络的多任务学习应用，比如CV和NLP，主要都用的这种方法^[1][2][3][4]。

另外还有一种是分布式的计算方法，这种方法假定异构的（heterogeneous）数据分别由各个任务分别以分布式的方式收集。这种方法常常被建模为在多智能体(multi-agent)系统（比如移动设备、无人驾驶汽车、智慧城市等）中的分布式学习。这种方式通常有两个假定，一是各任务节点和中心节点之间的网络通信代价很高，二是数据只能在任务节点存放，不能拷贝到中心节点（由于隐私性和通信代价问题）。近年来由于联邦学习的火热，该方法得到了很多的重视。

2 分布式多任务学习

在分布式多任务学习中，传统的处理方式^[5][6][7]仍然是多个任务节点分摊任务，然后将信息交给主节点汇总（比如在分布式近端映射算法中，任务节点进行梯度计算，主节点负责近端映射）。

近年来，随着去中心化优化算法的研究发展，越来越朝着去中心化的路线发展^[8][9][10]，也就是尽量满足使任务节点直接相互通信，而减少任务节点与主节点的通信。同时，随着联邦学习的发展，也越来越注重联邦学习中的经典问题，比如拜占庭容错等。

3 多任务学习和联邦学习之恋

我们前面提到，分布式多任务学习朝着联邦学习的路线发展。但其实联邦学习和多任务学习原本是很不一样的。在标准的联邦学习中，每个节点任务不共享数据，但是可以共享参数，以此联合训练出各一个全局的模型（可能是主从client-server结构，也可能是去中心化结构）。也就是说，联邦学习下每个节点的任务是一样的。

而多任务学习是要针对不同的任务协同训练出多个不同的模型。

但是，为什么分布式多任务学习会走向联邦学习呢？其实，不是分布式多任务选择了联邦学习，而是联邦学习选择了多任务学习。 原来，联邦学习由于数据不独立同分布（Non-IID） ，每个模型训练出的局部模型差异会很大，就会使得构建一个全局的、通用的模型难度很大。形式化地，传统联邦学习的优化目标函数可以写为^[13]：

\[\begin{aligned} f(w) &= \sum_{k=1}^K \frac{n_k}{n} F_k(w) \\ F_k(w) &= \frac{1}{n_k}\sum_{i = 1}^{n_k}\mathcal{l}(h(x_i; w), y_i) \end{aligned} \]

其中\(K\)为总节点个数，\(n_k\)为低\(k\)个节点的样本个数。联邦学习的训练过程中，会现将数据按照Non-IID划分到各client节点，然后再各client节点的数据划分train/test/val。而对于传统联邦学习而言，每个client都会使用全局模型\(w\)进行测试。

我们知道，在IID条件下，在分布式优化中我们常常假定\(f(w)=\mathbb{E}_{D_k}[F_k(w)]\)，其中\(D_k\)为第\(k\)个节点的数据集。然而，在数据Non-IID条件下，\(F_k\)就不是一个对\(f\)的良好近似。所以这意味着我们想训练一个全局的模型\(w\)满足所有节点的要求难度很大。

而个性化联邦学习不求构建一个全局的通用模型\(w\)，而是为每个节点分别构建一个个性化的模型\(w_k\)。这样，同样一个下一个单词预测的任务，同样给定"I love eating,"，但对于下一个单词每个client会给出不同的答案。联邦学习个性化的常见手段有元学习、多任务学习、迁移学习等。

具体到如何为各任务节点构建各不相同的模型。 有论文^[11][12]提出使每个节点采用知识共享的方式直接训练各不相同的模型这样一种训练方式，这就被冠名为联邦多任务学习了。多任务学习的目标函数常常会写为正则项的形式：

\[\begin{aligned} f(w_1,w_2, ...,w_K) &= \sum_{k=1}^K F_k(w_k) + \lambda g(w_1,w_2,...,w_k) \end{aligned} \]

不过，基于正则项的多任务学习正则项捕获个性化模型间的复杂关系，但由于正则项的复杂性只能优化简单的模型（线性模型或线性模型的组合）^[14][15]。而有些联邦多任务学习方法^[16][17][18]虽然牺牲了正则项的复杂性以训练更复杂的模型，但是又丧失了捕获任务间复杂关系的能力。

具体在数据分布方面，论文^[11][12]都保持了经典多任务学习的假设（参见我的博客《多任务学习中的数据分布》），不过有些许区别。论文^[11]中每个任务的训练数据分布和损失函数都不同。但是论文^[12]中假定每个任务不同之处只有训练数据的分布。

4 分布式多任务学习和联邦多任务学习的区别

此二者非常相似，但是联邦多任务学习可以看做是分布式多任务学习在特殊条件下的限制版，即联邦多任务学习中可能更关注节点的容错性，以及节点数据集隐私（节点之间的数据不能共享），单纯的分布式多任务学习一般没这几个需求。此外还有一点就是，按照最初的传统联邦多任务学习一般是有中心节点的(如论文^[11]中所说)，而分布式多任务学习是可以去中心化的（如论文^[10]中所说）。但是也有论文把联邦多任务学习也去中心化了(^[12])，所以这个应该算不上主要依据。

5 我的研究

我的研究现在关注的是分布式/联邦的多任务学习方法。而分布式的多任务学习方法其思想常常来源于基于正则化的多任务学习，这是一种非神经网络的多任务学习方法，已经得到了充分的研究，大家可以参见我的博客《基于正则化的多任务学习》回顾一下这种方法。

参考

• [1] Long M, Cao Z, Wang J, et al. Learning multiple tasks with multilinear relationship networks[J]. arXiv preprint arXiv:1506.02117, 2015.

• [2] Misra I, Shrivastava A, Gupta A, et al. Cross-stitch networks for multi-task learning[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2016: 3994-4003.

• [3] Hashimoto K, Xiong C, Tsuruoka Y, et al. A joint many-task model: Growing a neural network for multiple nlp tasks[J]. arXiv preprint arXiv:1611.01587, 2016.

• [4] Kendall A, Gal Y, Cipolla R. Multi-task learning using uncertainty to weigh losses for scene geometry and semantics[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2018: 7482-7491.

• [5] Baytas I M, Yan M, Jain A K, et al. Asynchronous multi-task learning[C]//2016 IEEE 16th International Conference on Data Mining (ICDM). IEEE, 2016: 11-20.

• [6] Liu S, Pan S J, Ho Q. Distributed multi-task relationship learning[C]//Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. 2017: 937-946.

• [7] Dinuzzo F, Pillonetto G, De Nicolao G. Client–server multitask learning from distributed datasets[J]. IEEE Transactions on Neural Networks, 2010, 22(2): 290-303.

• [8] Zhang C, Zhao P, Hao S, et al. Distributed multi-task classification: A decentralized online learning approach[J]. Machine Learning, 2018, 107(4): 727-747.

• [9] Yang P, Li P. Distributed primal-dual optimization for online multi-task learning[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2020, 34(04): 6631-6638.

• [10] Li J, Abbas W, Koutsoukos X. Byzantine Resilient Distributed Multi-Task Learning[J]. arXiv preprint arXiv:2010.13032, 2020.

• [11] Smith V, Chiang C K, Sanjabi M, et al. Federated multi-task learning[J]. Advances in Neural Information Processing Systems, 2017.

• [12] Marfoq O, Neglia G, Bellet A, et al. Federated multi-task learning under a mixture of distributions[J]. Advances in Neural Information Processing Systems, 2021, 34.

• [13] McMahan B, Moore E, Ramage D, et al. Communication-efficient learning of deep networks from decentralized data[C]//Artificial intelligence and statistics. PMLR, 2017: 1273-1282.

• [14] Paul Vanhaesebrouck, Aurélien Bellet, and Marc Tommasi. “Decentralized Collaborative Learning of Personalized Models over Networks”. In: AISTATS. 2017.

• [15] Valentina Zantedeschi, Aurélien Bellet, and Marc Tommasi. “Fully Decentralized Joint Learning of Personalized Models and Collaboration Graphs”. In: ed. by Silvia Chiappa and Roberto Calandra. Vol. 108. Proceedings of Machine Learning Research. Online: PMLR, Aug. 2020, pp. 864–874.

• [16] Filip Hanzely, Slavomıér Hanzely, Samuel Horváth, and Peter Richtárik. “Lower bounds and optimal algorithms for personalized federated learning”. In: 34th Conference on Neural Information Processing Systems (NeurIPS 2020). 2020.

• [17] Yutao Huang et al. “Personalized cross-silo federated learning on non-iid data”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 35. 9. 2021, pp. 7865–7873.

• [18] Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. “Ditto: Fair and robust federated learning through personalization”. In: International Conference on Machine Learning. PMLR. 2021, pp. 6357–6368.