本文介紹的是一篇收錄于 AAAI 2023 的論文，論文由位于上海交通大學(xué)的上海市可擴(kuò)展計(jì)算與系統(tǒng)重點(diǎn)實(shí)驗(yàn)室、貝爾法斯特女王大學(xué)的華揚(yáng)老師和路易斯安那州立大學(xué)的王灝老師共同完成。

• 論文鏈接：??https://arxiv.org/abs/2212.01197?
• 代碼鏈接（包含ALA模塊的使用說(shuō)明）：??https://github.com/TsingZ0/FedALA?

該論文提出了一種用于聯(lián)邦學(xué)習(xí)的自適應(yīng)本地聚合方法，通過(guò)從全局模型中自動(dòng)捕獲客戶機(jī)所需信息的方式來(lái)應(yīng)對(duì)聯(lián)邦學(xué)習(xí)中的統(tǒng)計(jì)異質(zhì)性問(wèn)題。作者對(duì)比了 11 個(gè) SOTA 模型，并取得了超越最優(yōu)方法 3.27% 的優(yōu)異表現(xiàn)。作者將其中的自適應(yīng)本地聚合模塊應(yīng)用到其他聯(lián)邦學(xué)習(xí)方法上取得了最多 24.19% 的提升。

1 介紹

聯(lián)邦學(xué)習(xí)（FL）通過(guò)將用戶隱私數(shù)據(jù)保留在本地而不進(jìn)行傳播的方式，幫助人們?cè)诒Ｗo(hù)隱私的情況下，充分發(fā)掘用戶數(shù)據(jù)中蘊(yùn)含的價(jià)值。但由于客戶機(jī)之間數(shù)據(jù)不可見(jiàn)，數(shù)據(jù)的統(tǒng)計(jì)異質(zhì)性（數(shù)據(jù)非獨(dú)立同分布（non-IID）和數(shù)據(jù)量不平衡現(xiàn)象）便成了 FL 的巨大挑戰(zhàn)之一。數(shù)據(jù)的統(tǒng)計(jì)異質(zhì)性使得傳統(tǒng)聯(lián)邦學(xué)習(xí)方法（如 FedAvg 等）很難通過(guò) FL 過(guò)程訓(xùn)練得到適用于每個(gè)客戶機(jī)的單一全局模型。

近年來(lái)，個(gè)性化聯(lián)邦學(xué)習(xí)（pFL）方法因其應(yīng)對(duì)數(shù)據(jù)統(tǒng)計(jì)異質(zhì)性的能力受到了越來(lái)越多的關(guān)注。與尋求高質(zhì)量全局模型的傳統(tǒng) FL 不同，pFL 方法的目標(biāo)是借助聯(lián)邦學(xué)習(xí)的協(xié)同計(jì)算能力為每個(gè)客戶機(jī)訓(xùn)練適用于自身的個(gè)性化模型?，F(xiàn)有的在服務(wù)器上聚合模型的 pFL 研究可以分為以下三類：

(1) 學(xué)習(xí)單個(gè)全局模型并對(duì)其進(jìn)行微調(diào)的方法，包括 Per-FedAvg 和 FedRep；

(2) 學(xué)習(xí)額外個(gè)性化模型的方法，包括 pFedMe 和 Ditto；

(3) 通過(guò)個(gè)性化聚合（或本地聚合）學(xué)習(xí)本地模型的方法，包括 FedAMP、FedPHP、FedFomo、APPLE 和 PartialFed。

類別（1）和（2）中的 pFL 方法將全局模型中的所有信息用于本地初始化（指在每次迭代的局部訓(xùn)練之前初始化局部模型）。然而，在全局模型中，只有提高本地模型質(zhì)量的信息（符合本地訓(xùn)練目標(biāo)的客戶機(jī)所需信息）才對(duì)客戶機(jī)有益。全局模型的泛化能力較差是因?yàn)槠渲型瑫r(shí)存在對(duì)于單一客戶機(jī)來(lái)說(shuō)需要和不需要的信息。因此，研究者們提出類別（3）中的 pFL 方法，通過(guò)個(gè)性化聚合捕獲全局模型中每個(gè)客戶機(jī)所需的信息。但是，類別（3）中的 pFL 方法依舊存在（a）沒(méi)有考慮客戶機(jī)本地訓(xùn)練目標(biāo)（如 FedAMP 和 FedPHP）、（b）計(jì)算代價(jià)和通訊代價(jià)較高（如 FedFomo 和 APPLE）、（c）隱私泄露（如 FedFomo 和 APPLE）和（d）個(gè)性化聚合與本地訓(xùn)練目標(biāo)不匹配（如 PartialFed）等問(wèn)題。此外，由于這些方法對(duì) FL 過(guò)程做了大量修改，它們使用的個(gè)性化聚合方法并不能被直接用于大多數(shù)現(xiàn)有 FL 方法。

為了從全局模型中精確地捕獲客戶機(jī)所需信息，且相比于 FedAvg 不增加每一輪迭代中的通訊代價(jià)，作者提出了一種用于聯(lián)邦學(xué)習(xí)的自適應(yīng)本地聚合方法（FedALA）。如圖 1 所示，F(xiàn)edALA 在每次本地訓(xùn)練之前，通過(guò)自適應(yīng)本地聚合（ALA）模塊將全局模型與本地模型進(jìn)行聚合的方式，捕獲全局模型中的所需信息。由于 FedALA 相比于 FedAvg 僅使用 ALA 修改了每一輪迭代中的本地模型初始化過(guò)程，而沒(méi)有改動(dòng)其他 FL 過(guò)程，因此 ALA 可被直接應(yīng)用于大多數(shù)現(xiàn)有的其他 FL 方法，以提升它們的個(gè)性化表現(xiàn)。

圖 1：在第次迭代中客戶機(jī)上的本地學(xué)習(xí)過(guò)程

2 方法

2.1 自適應(yīng)本地聚合（ALA）?

圖 2：自適應(yīng)本地聚合（ALA）過(guò)程

自適應(yīng)本地聚合（ALA）過(guò)程如圖 2 所示。相比于傳統(tǒng)聯(lián)邦學(xué)習(xí)中直接將下載的全局模型覆蓋本地模型得到本地初始化模型的方式（即），F(xiàn)edALA 通過(guò)為每個(gè)參數(shù)學(xué)習(xí)本地聚合權(quán)重，進(jìn)行自適應(yīng)本地聚合。

其中，作者把這一項(xiàng)稱之為“更新”。此外，作者通過(guò)逐元素權(quán)重剪枝方法實(shí)現(xiàn)正則化并將中的值限制在 [0,1] 中。

因?yàn)樯疃壬窠?jīng)網(wǎng)絡(luò)（DNN）的較低層網(wǎng)絡(luò)相比于較高層傾向于學(xué)習(xí)相對(duì)更通用的信息，而通用信息是各個(gè)本地模型所需信息，所以全局模型中較低層網(wǎng)絡(luò)中的大部分信息與本地模型中較低層網(wǎng)絡(luò)所需信息一致。為了降低學(xué)習(xí)本地聚合權(quán)重所需的計(jì)算代價(jià)，作者引入一個(gè)超參數(shù) p 來(lái)控制 ALA 的作用范圍，使得全局模型中較低層網(wǎng)絡(luò)參數(shù)直接覆蓋本地模型中的較低層網(wǎng)絡(luò)，而只在較高層啟用 ALA。

其中，表示中的神經(jīng)網(wǎng)絡(luò)層數(shù)（或神經(jīng)網(wǎng)絡(luò)塊數(shù)），與中的低層網(wǎng)絡(luò)形狀一致，則與中剩下的 p 層高層網(wǎng)絡(luò)形狀一致。

作者將中的值全部初始化為 1，且在每一輪本地初始化過(guò)程中基于舊的更新。為了進(jìn)一步降低計(jì)算代價(jià)，作者采用隨機(jī)采樣 s

其中，是更新的學(xué)習(xí)率。作者在學(xué)習(xí)的過(guò)程中，將除之外的其他可訓(xùn)練參數(shù)凍結(jié)。

圖 3：在 MNIST 和 Cifar10 數(shù)據(jù)集上 8 號(hào)客戶機(jī)的學(xué)習(xí)曲線

通過(guò)選擇較小的 p 值，在幾乎不影響 FedALA 表現(xiàn)的情況下，大幅度地降低 ALA 中訓(xùn)練所需的參數(shù)。此外，如圖 3，作者觀察到：一旦在第一次訓(xùn)練將其訓(xùn)練到收斂，即使在后續(xù)迭代中訓(xùn)練，其對(duì)本地模型質(zhì)量也沒(méi)有很大影響。也就是說(shuō)，每個(gè)客戶機(jī)可以復(fù)用舊的實(shí)現(xiàn)對(duì)其所需信息的捕獲。作者采取在后續(xù)迭代中微調(diào)的方式，降低計(jì)算代價(jià)。

2.2 ALA 分析?

在不影響分析的情況下，為了簡(jiǎn)便起見(jiàn)，作者忽略并設(shè)。根據(jù)上述公式可以得到，其中代表。作者可以把在 ALA 中更新看成更新。

梯度項(xiàng)在每一輪中都做了逐元素地縮放。不同于本地模型訓(xùn)練（或 fine-tuning）方法，上述對(duì)的更新過(guò)程能夠感知到全局模型中的通用信息。在不同的迭代輪數(shù)之間，動(dòng)態(tài)變化的為 ALA 模塊引入了動(dòng)態(tài)信息，使其 FedALA 容易適應(yīng)復(fù)雜的環(huán)境。

3 實(shí)驗(yàn)

作者在實(shí)際（practical）數(shù)據(jù)異質(zhì)環(huán)境下的 Tiny-ImageNet 數(shù)據(jù)集上用 ResNet-18 進(jìn)行了對(duì)超參數(shù) s 和 p 的對(duì) FedALA 影響的研究，如表 1 所示。對(duì)于 s 來(lái)說(shuō)，采用越多的隨機(jī)采樣的本地訓(xùn)練數(shù)據(jù)用于 ALA 模塊學(xué)習(xí)可以使個(gè)性化模型表現(xiàn)更好，但也同時(shí)增加了計(jì)算代價(jià)。在使用 ALA 的過(guò)程中，可以根據(jù)每個(gè)客戶機(jī)的計(jì)算能力調(diào)整 s 的大小。從表中可以得知，即使使用極小的 s（如 s=5），F(xiàn)edALA 依舊具有杰出的表現(xiàn)。對(duì)于 p 來(lái)說(shuō)，不同的 p 值對(duì)個(gè)性化模型的表現(xiàn)幾乎沒(méi)有影響，在計(jì)算代價(jià)方面卻有著巨大的差別。這一現(xiàn)象也從一個(gè)側(cè)面展示了 FedRep 等方法，將模型分割后保留靠近輸出的神經(jīng)網(wǎng)絡(luò)層在客戶機(jī)不上傳的做法的有效性。使用 ALA 時(shí)，我們可以采用較小且合適的 p 值，在保證個(gè)性化模型表現(xiàn)能力的情況下，進(jìn)一步降低計(jì)算代價(jià)。

表 1：關(guān)于超參數(shù)和對(duì) FedALA 影響的研究

作者在病態(tài)（pathological）數(shù)據(jù)異質(zhì)環(huán)境和實(shí)際（practical）數(shù)據(jù)異質(zhì)環(huán)境下，將 FedALA 與 11 個(gè) SOTA 方法進(jìn)行了對(duì)比和詳細(xì)分析。如表 2 所示，數(shù)據(jù)顯示 FedALA 在這些情況下的表現(xiàn)都超越了這 11 個(gè) SOTA 方法，其中 “TINY” 表示在 Tiny-ImageNet 上使用 4-layer CNN。例如，F(xiàn)edALA 在 TINY 情況下比最優(yōu)基線（baseline）高了 3.27%。

表 2：病態(tài)和實(shí)際數(shù)據(jù)異質(zhì)環(huán)境下的實(shí)驗(yàn)結(jié)果

此外，作者也在不同異質(zhì)性環(huán)境和客戶機(jī)總量情況下評(píng)估了 FedALA 的表現(xiàn)。如表 3 所示，F(xiàn)edALA 在這些情況下依舊保持著優(yōu)異的表現(xiàn)。

表 3：其他實(shí)驗(yàn)結(jié)果

根據(jù)表 3 的實(shí)驗(yàn)結(jié)果，將 ALA 模塊應(yīng)用于其他方法可以獲得最多 24.19% 的提升。

最后，作者還在 MNIST 上可視化了 ALA 模塊的加入對(duì)原本 FL 過(guò)程中模型訓(xùn)練的影響，如圖 4 所示。不激活 ALA 時(shí)，模型訓(xùn)練軌跡與使用 FedAvg 一致。一旦 ALA 被激活，模型便可以通過(guò)全局模型中捕獲的其訓(xùn)練所需信息徑直朝著最優(yōu)目標(biāo)優(yōu)化。

圖 4：4 號(hào)客戶機(jī)上模型訓(xùn)練軌跡的可視化圖