本文介紹的是一篇收錄於 AAAI 2023 的論文，論文由位於上海交通大學的上海市可擴充套件計算與系統重點實驗室、貝爾法斯特女王大學的華揚老師和路易斯安那州立大學的王灝老師共同完成。
論文連結：https://arxiv.org/abs/2212.01197
程式碼連結（包含ALA模組的使用說明）：https://github.com/TsingZ0/FedALA
該論文提出了一種用於聯邦學習的自適應本地聚合方法，透過從全域性模型中自動捕獲客戶機所需資訊的方式來應對聯邦學習中的統計異質性問題。作者對比了 11 個 SOTA 模型，並取得了超越最優方法 3.27% 的優異表現。作者將其中的自適應本地聚合模組應用到其他聯邦學習方法上取得了最多 24.19% 的提升。
1 介紹
聯邦學習（FL）透過將使用者隱私資料保留在本地而不進行傳播的方式，幫助人們在保護隱私的情況下，充分發掘使用者資料中蘊含的價值。但由於客戶機之間資料不可見，資料的統計異質性（資料非獨立同分布（non-IID）和資料量不平衡現象）便成了 FL 的巨大挑戰之一。資料的統計異質性使得傳統聯邦學習方法（如 FedAvg 等）很難透過 FL 過程訓練得到適用於每個客戶機的單一全域性模型。
近年來，個性化聯邦學習（pFL）方法因其應對資料統計異質性的能力受到了越來越多的關注。與尋求高質量全域性模型的傳統 FL 不同，pFL 方法的目標是藉助聯邦學習的協同計算能力為每個客戶機訓練適用於自身的個性化模型。現有的在伺服器上聚合模型的 pFL 研究可以分為以下三類：
(1) 學習單個全域性模型並對其進行微調的方法，包括 Per-FedAvg 和 FedRep；
(2) 學習額外個性化模型的方法，包括 pFedMe 和 Ditto；
(3) 透過個性化聚合（或本地聚合）學習本地模型的方法，包括 FedAMP、FedPHP、FedFomo、APPLE 和 PartialFed。
類別（1）和（2）中的 pFL 方法將全域性模型中的所有資訊用於本地初始化（指在每次迭代的區域性訓練之前初始化區域性模型）。然而，在全域性模型中，只有提高本地模型質量的資訊（符合本地訓練目標的客戶機所需資訊）才對客戶機有益。全域性模型的泛化能力較差是因為其中同時存在對於單一客戶機來說需要和不需要的資訊。因此，研究者們提出類別（3）中的 pFL 方法，透過個性化聚合捕獲全域性模型中每個客戶機所需的資訊。但是，類別（3）中的 pFL 方法依舊存在（a）沒有考慮客戶機本地訓練目標（如 FedAMP 和 FedPHP）、（b）計算代價和通訊代價較高（如 FedFomo 和 APPLE）、（c）隱私洩露（如 FedFomo 和 APPLE）和（d）個性化聚合與本地訓練目標不匹配（如 PartialFed）等問題。此外，由於這些方法對 FL 過程做了大量修改，它們使用的個性化聚合方法並不能被直接用於大多數現有 FL 方法。
為了從全域性模型中精確地捕獲客戶機所需資訊，且相比於 FedAvg 不增加每一輪迭代中的通訊代價，作者提出了一種用於聯邦學習的自適應本地聚合方法（FedALA）。如圖 1 所示，FedALA 在每次本地訓練之前，透過自適應本地聚合（ALA）模組將全域性模型與本地模型進行聚合的方式，捕獲全域性模型中的所需資訊。由於 FedALA 相比於 FedAvg 僅使用 ALA 修改了每一輪迭代中的本地模型初始化過程，而沒有改動其他 FL 過程，因此 ALA 可被直接應用於大多數現有的其他 FL 方法，以提升它們的個性化表現。
圖 1：在第次迭代中客戶機上的本地學習過程
2 方法
2.1 自適應本地聚合（ALA）
圖 2：自適應本地聚合（ALA）過程
自適應本地聚合（ALA）過程如圖 2 所示。相比於傳統聯邦學習中直接將下載的全域性模型
覆蓋本地模型
得到本地初始化模型
的方式（即
），FedALA 透過為每個引數學習本地聚合權重，進行自適應本地聚合。
其中，作者把
這一項稱之為“更新”。此外，作者透過逐元素權重剪枝方法
實現正則化並將
中的值限制在 [0,1] 中。
因為深度神經網路（DNN）的較低層網路相比於較高層傾向於學習相對更通用的資訊，而通用資訊是各個本地模型所需資訊，所以全域性模型中較低層網路中的大部分資訊與本地模型中較低層網路所需資訊一致。為了降低學習本地聚合權重所需的計算代價，作者引入一個超引數 p 來控制 ALA 的作用範圍，使得全域性模型中較低層網路引數直接覆蓋本地模型中的較低層網路，而只在較高層啟用 ALA。
其中，
表示
中的神經網路層數（或神經網路塊數），
與
中的低層網路形狀一致，
則與中剩下的 p 層高層網路形狀一致。
作者將
中的值全部初始化為 1，且在每一輪本地初始化過程中基於舊的更新。為了進一步降低計算代價，作者採用隨機取樣 s
其中，
是更新的學習率。作者在學習的過程中，將除之外的其他可訓練引數凍結。
圖 3：在 MNIST 和 Cifar10 資料集上 8 號客戶機的學習曲線
透過選擇較小的 p 值，在幾乎不影響 FedALA 表現的情況下，大幅度地降低 ALA 中訓練所需的引數。此外，如圖 3，作者觀察到：一旦在第一次訓練將其訓練到收斂，即使在後續迭代中訓練，其對本地模型質量也沒有很大影響。也就是說，每個客戶機可以複用舊的實現對其所需資訊的捕獲。作者採取在後續迭代中微調的方式，降低計算代價。
2.2 ALA 分析
在不影響分析的情況下，為了簡便起見，作者忽略
並設
。根據上述公式可以得到
，其中
代表
。作者可以把在 ALA 中更新
看成更新
。
梯度項
在每一輪中都做了逐元素地縮放。不同於本地模型訓練（或 fine-tuning）方法，上述對
的更新過程能夠感知到全域性模型中的通用資訊。在不同的迭代輪數之間，動態變化的
為 ALA 模組引入了動態資訊，使其 FedALA 容易適應複雜的環境。
3 實驗
作者在實際（practical）資料異質環境下的 Tiny-ImageNet 資料集上用 ResNet-18 進行了對超引數 s 和 p 的對 FedALA 影響的研究，如表 1 所示。對於 s 來說，採用越多的隨機取樣的本地訓練資料用於 ALA 模組學習可以使個性化模型表現更好，但也同時增加了計算代價。在使用 ALA 的過程中，可以根據每個客戶機的計算能力調整 s 的大小。從表中可以得知，即使使用極小的 s（如 s=5），FedALA 依舊具有傑出的表現。對於 p 來說，不同的 p 值對個性化模型的表現幾乎沒有影響，在計算代價方面卻有著巨大的差別。這一現象也從一個側面展示了 FedRep 等方法，將模型分割後保留靠近輸出的神經網路層在客戶機不上傳的做法的有效性。使用 ALA 時，我們可以採用較小且合適的 p 值，在保證個性化模型表現能力的情況下，進一步降低計算代價。
表 1：關於超引數和對 FedALA 影響的研究
作者在病態（pathological）資料異質環境和實際（practical）資料異質環境下，將 FedALA 與 11 個 SOTA 方法進行了對比和詳細分析。如表 2 所示，資料顯示 FedALA 在這些情況下的表現都超越了這 11 個 SOTA 方法，其中 “TINY” 表示在 Tiny-ImageNet 上使用 4-layer CNN。例如，FedALA 在 TINY 情況下比最優基線（baseline）高了 3.27%。
表 2：病態和實際資料異質環境下的實驗結果
此外，作者也在不同異質性環境和客戶機總量情況下評估了 FedALA 的表現。如表 3 所示，FedALA 在這些情況下依舊保持著優異的表現。
表 3：其他實驗結果
根據表 3 的實驗結果，將 ALA 模組應用於其他方法可以獲得最多 24.19% 的提升。
最後，作者還在 MNIST 上視覺化了 ALA 模組的加入對原本 FL 過程中模型訓練的影響，如圖 4 所示。不啟用 ALA 時，模型訓練軌跡與使用 FedAvg 一致。一旦 ALA 被啟用，模型便可以透過全域性模型中捕獲的其訓練所需資訊徑直朝著最優目標最佳化。
圖 4：4 號客戶機上模型訓練軌跡的視覺化圖