本文提出了聯邦匹配平均(FedMA)算法。FedMA通過對提取到的具有相似特征的隱元素(即卷積層的通道,LSTM的隱狀態,全連接層的神經元)進行匹配和平均,按層構建共享全局模型。FedMA訓練的CNN和LSTM模型在數據集MNIST,CIFAR-10,Shakespeare上優於最新聯邦學習算法FedProx,同時提高了通信效率。
按我的理解,由於模型可辨識性(model identifiability)問題,神經網絡和任意具有多個等
效參數化潛變量的模型都會具有多個局部極小值。如果一個足夠大的訓練集可以唯一確定一組模型參數,那么該模型被稱為可辨認的。帶有潛變量的模型通常是不可辨認的,因為通過相互交換潛變量我們能得到等價的模型。例如,考慮神經網絡的第一層,我們可以交換單元 i 和單元 j 的傳入權重向量、傳出權重向量而得到等價的模型。如果神經網絡有 m 層,每層有 n 個單元,那么會有種排列隱藏單元的方式。這種不可辨認性被稱為權重空間對稱性(weight space symmetry)。根據這種對稱性任何給定的神經網絡,它的許多變化只在參數的順序上不同,構成了實際等價的局部最優。概率聯邦神經匹配(PFNM)通過在對神經網絡的參數求平均值之前找到其參數的排列來解決這個問題。本文在PFNM的基礎上提出了FedMA算法。
本文討論了為什么FedAvg中對所有客戶端的模型參數直接平均是不合理的,例如:在最簡單的單隱藏層全連接神經網絡中,FC神經網絡可以表示為:
根據對稱性可以寫為:
L是隱含層的節點數目,置換矩陣∏將W1的L列與W2的L行進行置換。
假設是最佳權重,在同質數據集
上訓練的權重為
找不到任何置換矩陣使得
成立
所以不能在一開始就做
(即兩個模型參數直接平均就可以得到最優模型)
聯邦匹配平均優化公式:
表示從數據集j中學習得到的第l個神經元
表示全局模型中的第i個神經元,
是這一對神經元的相似函數
給定J個客戶端提供的權重計算得到聯邦神經網絡權重:
基於優化公式與最大二分匹配問題的關系,所以本文將此方法稱為匹配平均。
在每次迭代中,先根據給定的權值矩陣估計找到對應的全局模型,然后根據匈牙利算法將全局模型和數據集j'上的局部神經元進行匹配,得到新的拓展全局模型,由於數據異質性,局部模型j'可能存在部分神經元,它們不存在於由其他局部模型構成的全局模型中。因此,我們希望避免“差”匹配,即如果最優匹配的代價大於某個閾值,從相應的局部神經元創建一個新的全局神經元。還需要一個中等大小的全局模型,因此用一些遞增函數f(L')來懲罰它的大小。其中,全局模型大小記為L。
FedMA算法:
首先,中央服務器只從客戶端收集第一層的權重,並執行前面描述的單層匹配以獲得聯邦模型的第一層權重。然后中央服務器將這些權重廣播給客戶端,客戶端繼續訓練其數據集上的所有連續層,同時保持已經匹配的聯邦層凍結。然后,將此過程重復到最后一層,根據每個客戶端數據的類比例對其進行加權平均。FedMA 方法要求通信輪數等於網絡中的層數。
實驗分析
數據集:MNIST,CIFAR-10,Shakespeare
模型:VGG-9,LSTM
下圖展示了層匹配 FedMA 在更深的 VGG-9和 LSTM 上的性能。在異構環境中,FedMA 優於 FedAvg、FedProx和其他基線模型(即客戶端 CNN模型及其集成)訓練得到的 FedProx。
FedMA 的優點之一是它比 FedAvg 更有效地利用了通信輪次,即 FedMA 不是直接按元素平均權重,而是識別匹配的卷積濾波器組,然后將它們平均到全局卷積濾波器中。
最后,作者研究了 FedMA 的通信性能。通過將 FedMA 與 FedAvg、FedProx 進行比較,FedMA 在所有情況下都優於 FedAvg 和 FedProx。
總結:
本文提出了一種利用概率匹配和模型大小自適應的分層聯邦學習算法FedMA,通過實驗驗證了FedMA的收斂速度和通信效率。本文證明了 FedMA 可以有效地利用訓練后的局部模型。在后續工作中,作者考慮利用近似二次分配解(Approximate Quadratic Assignment Solutions)的方法引入其他的神經網絡層,例如殘差連接和批量歸一化層,從而進一步改進LSTM的聯邦學習效果。此外,探索 FedMA 的容錯性並研究其在更大數據集上的性能非常重要,特別是針對那些即使在數據可以聚合的情況下也無法進行有效訓練的數據集。