聯邦學習：多工思想與聚類聯邦學習

1.導引

電腦科學一大定律：許多看似過時的東西可能過一段時間又會以新的形式再次迴歸。

在聯邦學習領域，許多傳統機器學習已經討論過的問題（甚至一些90年代和00年代的論文）都可以被再次被發明一次。比如我們會發現聚類聯邦學習和多工學習之間就有千絲萬縷的聯絡。

2. 多工學習回顧

我們在部落格《基於正則化的多工學習》中介紹了標準多工學習的核心是多工個性化權重+知識共享^[1]。如多工學習最開始提出的模型即為一個共享表示的神經網路：

而多工學習中有一種方法叫聚類多工學習。聚類多工學習基本思想為：將任務分為若個個簇，每個簇內部的任務在模型引數上更相似。 最早的聚類多工學習的論文^[2]為一種一次性聚類(one-shot clustering)，即將任務聚類和引數學習分為了兩個階段：第一階段，根據在各任務單獨學習得到的模型來聚類，確定不同的任務簇。第二階段，聚合同一個任務簇中的所有訓練資料，以學習這些任務的模型。這種方法把任務聚類和模型引數學習分為了兩個階段，可能得不到最優解，因此後續工作都採用了迭代式聚類(iterative clustering)，即在訓練迭代中同時學習任務聚類和模型引數的方法，

我們著重介紹Bakker等人(2003)^[3]提出了一個多工貝葉斯神經網路(multi-task Bayesian neural network)，其結構與我們前面所展現的共享表示的神經網路相同，其亮點在於基於連線隱藏層和輸出層的權重採用高斯混合模型（Gaussian mixture model）對任務進行分組。若給定資料集\(\mathcal{D} = \{D_t\}, t=1,...,T\)，設隱藏層維度為\(h\)，輸出層維度為\(T\)，\(\mathbf{W}\in \mathbb{R}^{T\times (h + 1)}\)代表隱藏層到輸出層的權重矩陣(結合了偏置)。我們假定每個任務對應的權重向量\(\mathbf{w}_t\)(\(\mathbf{W}\)的第\(t\)列)關於給定超引數獨立同分布，我們假定第\(t\)個任務先驗分佈如下：

\[\bm{w}_t \sim \mathcal{N}(\bm{w}_t | \bm{u}, \mathbf{\Sigma}) \tag{8} \]

這是一個高斯分佈，均值為\(\bm{u} \in \mathbb{R}^{h + 1}\)，協方差矩陣\(\mathbf{\Sigma} \in \mathbb{R}^{ (h+1) \times (h+1)}\)。
我們上面的定義其實假定了所有任務屬於一個簇，接下來我們假定我們有不同的簇（每個簇由相似的任務組成）。我們設有\(C\)個簇（cluster），則任務\(t\)的權重\(w_t\)為\(C\)個高斯分佈的混合分佈：

\[\bm{w}_t \sim \sum_{c=1}^C \alpha_c \mathcal{N} (\bm{w}_t | \bm{u}_c, \mathbf{\Sigma}_c) \tag{9} \]

其中，每個高斯分佈可以被認為是描述一個任務簇。式\((9)\)中的\(\alpha_c\)代表了任務\(t\)被分為簇\(c\)的先驗概率,其中task clustering(如下面左圖所示)模型中所有任務對簇\(c\)的加權\(\alpha_c\)都相同；而task-depenent模型（如下面右圖所示）中各任務對簇\(c\)的加權\(\alpha_{tc}\)不同，且依賴於各任務特定的向量\(\bm{f}_t\)。

好，大家再細細品味一下這種一個任務簇對應一個概率分佈的思想，下面我們介紹的很多聚類聯邦學習演算法都體現有這種思想的思想。

3. 潛在聚類結構的的前置假設

上面我們提到的這種來自於上個世紀90年代的多工學習思想並未過時，近年來在個性化聯邦學習中又重新煥發了生機。這篇文章我們就來歸納一下與多工學習思想關係密切的聯邦學習個性化方法之一——聚類聯邦學習。

聚類聯邦學習基本上都基於這樣一個假設：雖然聯邦學習中各節點的資料是Non-IID的，但是我們可以假定某些節點的資料可以歸為一個聚類簇（簇內節點的分佈近似IID）。實際上，這個思想讓我們聯想到高斯混合分佈。高斯混合分佈就假設每個節點的資料取樣自高斯混合分佈中的一個成分(對應一個簇)，而經典的高斯混合聚類就是要確定每個節點和簇的對應關係。如果採用這種假設，那麼在做實驗時就可以考慮基於混合分佈的劃分方法，參見我的部落格《基於混合分佈的多工學習》。

注意，雖然這裡假定每個聚類簇對應一個成分分佈，但現實情況的資料分佈是模糊的，不同簇之間的資料仍然可能會有關聯。而這就為我們進行簇間的多工知識共享提供了動機。

按照其聚類的時間，我們按照上文聚類多工學習的分發，可將聚類聯邦學習的方法也分為一次性聚類和迭代式聚類。

一次性聚類即在模型訓練開始之前就先對client資料進行聚類。因為client資料不出庫，常常將\(T\)個client初步訓練後的引數\(\{\widetilde{\theta}_t\}^T_{t=1}\)發往server，然後由server將\(T\)個client聚為\(S_1,...,S_k\)。然後對每一個類\(k(k\in [K])\)，執行傳統聯邦學習方法（如FedAvg^[5]或二階方法^[6][7]）求解出最終的\(\hat{\theta}_k\)。注意, 此種方式預設每個聚類簇得到一個引數\(\hat{\theta}_k\)。同樣，類似我們上面對一次性聚類多工學習的分析，如果一旦開局的聚類演算法產生了錯誤的估計，那麼接下來的時間裡演算法將無法對其進行更正。此外，這種方法常常使用類似於K-means的距離聚類方法，而從理論上講對於神經網路兩個模型擁有相近的引數距離，但其輸出可能大不相同（由於模型對於隱藏層單元的置換不變性）。

動態式的聚類為則為在訓練的過程中一邊訓練，一邊根據模型引數\(\widetilde{\theta}_t\)的情況來動態調整聚類結果。這種方式既可以簇內任務直接廣播引數\(\theta_k\)，也可以僅僅共享引數的變化量\(\Delta \theta_k\)。

按照其任務劃分粒度的方式。我們可將方法分為節點間多工和簇間多工。節點間多工類似聚類多工學習，它為將每個節點視為一個任務(訓練一個個性化模型)，然後將節點劃分為多個簇，簇內和簇間都會進行知識共享。簇間多工為將節點劃分為多個簇，假定簇內資料IID。然後簇內節點直接廣播引數，此時將每個簇視為一個任務，然後再在簇間進行知識共享。

4. 聚類聯邦學習經典論文閱讀

4.1 TNNLS 2020:《Clustered Federated Learning: Model-Agnostic Distributed Multitask Optimization Under Privacy Constraints》^[8]

本文采用的是迭代式聚類和節點間多工，它在每輪迭代時，都會根據節點的引數相似度進行一次任務簇劃分，同一個任務簇共享引數的變化量 \(g\)，以此既能達到每個節點訓練一個個性化模型，又能完成知識共享和相似的節點相互促進的目的。本論文的亮點是聚類簇的個數可以隨著迭代變化，無需實現指定任務簇個數\(K\)做為先驗。

本篇論文演算法的每輪通訊描述如下：

(1) 第\(t\)個client節點執行

• 從server接收對應的簇引數\(g_k\)。
• 令\(\theta_{old}=\theta_t=\theta_t + g_k\)。
• 執行\(E\)個區域性epoch的SGD：
  \[\theta_t = \theta_t - \eta \nabla \mathcal{l}(\theta_t; b) \]
  (此處將區域性資料\(D_t\)劃分為多個\(b\))
• 將\(\widetilde{g}_t = \theta_t-\theta_{old}\)發往server。

(2) server節點執行

• 從\(T\)個client接收\(\widetilde{g}_1、\widetilde{g}_2,...\widetilde{g}_T\)。

• 對每一個簇\(k\)，計算簇內平均引數變化：

\[\hat{g}_{k} = \frac{1}{|S_k|}\sum_{t\in S_k}\widetilde{g}_t \]

• 將每個簇的簇內變化\(\hat{g}_k\)發往對應的client。

• 根據不同節點引數變化量的餘弦距離\(\alpha_{i,j}=\frac{\langle \widetilde{g}_i, \widetilde{g}_j\rangle}{||\widetilde{g}_i||||\widetilde{g}_j||}\)重新劃分聚類簇。

4.2 NIPS 2020:《An Efficient Framework for Clustered Federated Learning》^[8]

本文采用的是迭代式聚類和簇間多工。本論文需要事先指定聚類簇的個數\(K\)。它在每輪迭代時，都會根據節點的引數相似度重新進行一次任務簇劃分，第\(k\)個任務簇對其所有節點廣播引數\(\theta_k\)，也就是說每個聚類簇構成一個學習任務(負責學習屬於該任務簇的個性化模型)。

本篇論文演算法的每輪通訊描述如下：

(1) 第\(t\)個client節點執行

• 從server接收簇引數\(\theta_1,\theta_2,...\theta_{K}\)

• 估計其所屬的簇：\(\hat{k}=\underset{k\in[K]}{\text{argmin }}\mathcal{l}(\theta_k;D_t)\)

• 對簇引數執行\(E\)個區域性epoch的SGD：

  \[\theta_k = \theta_k - \eta \nabla \mathcal{l}(\theta_k; b) \]

  (此處將區域性資料\(D_t\)劃分為多個\(b\))

• 將最終得到的簇引數做為\(\widetilde{\theta}_t\)，並和該client對應的簇劃分一起發往server。

(2) server節點執行

• 從\(T\)個client接收\(\widetilde{\theta}_1、\widetilde{\theta}_2,..., \widetilde{\theta}_T\)和各client的簇劃分情況

• 根據節點引數的平均值來更新簇引數：

\[\hat{\theta}_k = \frac{1}{|S_k|} \sum_{t \in S_k}\widetilde{\theta}_t \]

由於現實的聚類關係很模糊，該演算法在具體實現時效仿多工學習中的權值共享^[1]機制，允許不同簇(任務)之間共享部分引數。具體地，即在訓練神經網路模型時，先使用所有client的訓練資料學習一個共享表示，然後再執行聚類演算法為每個簇學習神經網路的最後一層（即多工層）。

4.3 WWW 2021:《PFA: Privacy-preserving Federated Adaptation for Effective Model Personalization》^[9]

本文采用的是一次性聚類。本論文需要事先指定聚類簇的個數\(K\)，然後在訓練迭代開始前，每個節點先訓練一個稀疏表徵模型模型，然後將稀疏表徵模型得到的稀疏向量傳到server，server再根據各client稀疏向量之間的相似度進行任務簇劃分。之後，各任務簇分別執行傳統FedAvg演算法學習每個簇對應的個性化模型。值得一提的是，雖然該演算法屬於聚類聯邦學習，但我個人認為嚴格來說它與標準的多工學習還有一定距離，因為它的每個簇之間缺少知識共享。 不過它確實體現了“多工”的思想，而且與我們前面講的聚類演算法關聯十分密切，最終我還是將其加入本專題。

4.4 INFOCOM 2021:《Resource-Efficient Federated Learning with Hierarchical Aggregation in Edge Computing》

這篇文章嚴格來說不屬於上面所說的聚類聯邦學習，也不是多工學習/個性化聯邦學習的範疇。其文中提到的clustered應該翻譯成"（區域網構成的）叢集的"。但我覺得這篇文章算是從另一個工程的角度來應用了分組聚合的思想，我在這裡還是覺得將其介紹一下。

在現實工業環境下，聯邦學習常常是基於引數伺服器(parameter server)的，但引數伺服器位於遠端雲平臺上，邊緣節點與它之間的通訊經常不可用，而邊緣節點的數量龐大，這就導致通訊代價可能會很高。本文的基本思想即是減少邊緣節點和PS之間的通訊，加強邊緣節點之間的協作。論文將不同的client劃分為了多個區域網/叢集，每個叢集都有一個leader node（LN）做為叢集頭。每個叢集先分別執行FedAvg演算法將引數聚合到LN(叢集內是同步的)，然後再由PS非同步地蒐集各LN的引數並進行聚合。最後再將新的引數廣播到各個邊遠client。

本篇論文還考慮了每個client的資源消耗等複雜資訊，此處為了簡單起見，我們將簡化後演算法的每輪通訊描述如下：

(1) 第\(t\)個client節點執行

• 從\(LN_k\)接收引數\(\theta\)做為本地的\(\theta_t\)
• 執行\(E\)個區域性epoch的SGD：
  \[\theta_t = \theta_t - \eta \nabla \mathcal{l}(\theta_t; b) \]
  (此處將區域性資料\(D_t\)劃分為多個\(b\))
• 將新的引數\(\widetilde{\theta}_t\)發往所在簇的\(LN_k\)。

(2) \(LN_k\)節點執行

• 從叢集內的\(|S_k|\)個client接收引數\(\{\widetilde{\theta}_j\}(j\in S_k)\)

• 根據簇內節點引數的加權平均值來更新簇引數：

\[\theta_k = \frac{\sum_{j \in S_k}{ |D_j|}\widetilde{\theta}_j}{\sum_{j\in S_k} |D_j|} \]

• 將簇引數\(\theta_k\)發往PS。
• 從PS接收引數\(\theta\)並將其發往client節點。

(2) PS節點執行

• 從\(k\)個\(LN_k\)節點接收\({\theta}_k\)（非同步地）。

• 根據一階指數平滑來更新引數：

\[\hat{\theta} = (1-\alpha)\theta + \alpha \theta_k \]

\(\quad\quad\)(在實際論文中\(\alpha\)是一個在迭代中變化的量，此處為了簡化省略)

• 將新的引數\(\hat{\theta}\)發往\(LN_k\)

引用

• [1] Caruana R. Multitask learning[J]. Machine learning, 1997, 28(1): 41-75

• [2] Thrun S, O'Sullivan J. Discovering structure in multiple learning tasks: The TC algorithm[C]//ICML. 1996, 96: 489-497.

• [3] Bakker B J, Heskes T M. Task clustering and gating for bayesian multitask learning[J]. 2003.

• [4] Sattler F, Müller K R, Samek W. Clustered federated learning: Model-agnostic distributed multitask optimization under privacy constraints[J]. IEEE transactions on neural networks and learning systems, 2020, 32(8): 3710-3722.

• [5] 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.

• [6] Wang S, Roosta F, Xu P, et al. Giant: Globally improved approximate newton method for distributed optimization[J]. Advances in Neural Information Processing Systems, 2018, 31.

• [7] Ghosh A, Maity R K, Mazumdar A. Distributed newton can communicate less and resist byzantine workers[J]. Advances in Neural Information Processing Systems, 2020, 33: 18028-18038.

• [8] Ghosh A, Chung J, Yin D, et al. An efficient framework for clustered federated learning[J]. Advances in Neural Information Processing Systems, 2020, 33: 19586-19597.

• [9] Liu B, Guo Y, Chen X. PFA: Privacy-preserving Federated Adaptation for Effective Model Personalization[C]//Proceedings of the Web Conference 2021. 2021: 923-934.

• [10] Wang Z, Xu H, Liu J, et al. Resource-Efficient Federated Learning with Hierarchical Aggregation in Edge Computing[C]//IEEE INFOCOM 2021-IEEE Conference on Computer Communications. IEEE, 2021: 1-10.