分布式多任務學習及聯邦學習個性化

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.