随着深度學習大獲成功,保護用戶數據隐私變得越來越重要。
聯邦學習(Federated Learning)應運而生,這是一種基于隐私保護的分布式機器學習框架。
它可以讓原始數據保留在本地,讓多方聯合共享模型訓練。
但它有一個問題——數據的異質化(data heterogeneity),即不同的參與方的本地數據來自不同的分布,這将嚴重影響全局模型的最終性能,背後原因也十分複雜。
字節跳動、新加坡國立大學及中科院自動化所的學者們首次發現了關鍵影響因素。
即:數據異質化導緻了表征的維度坍縮(dimensional collapse),由此大大限制了模型的表達能力,影響了最終全局模型的性能。
爲了緩解這一問題,研究人員提出了一個新聯邦學習正則項:FedDecorr。
結果表明,使用該方法後,數據異質化帶來的維度坍縮問題被有效緩解,顯著提升模型在該場景下的性能。
同時這一方法實現簡單,幾乎不會帶來額外計算負擔,可以很容易地加入到多種聯邦學習算法上。
如何影響?觀察一:更嚴重的數據異質化會爲全局模型(global model)帶來更嚴重的維度坍縮
首先,爲了更好地理解數據異質化是如何影響全局模型輸出表征的,研究人員探索了随着數據異質化越來越嚴重,全局模型輸出表征是如何而變化的。
基于模型輸出的表征,估計其表征分布的協方差矩陣(covariance matrix),并且按照從大到小的順序可視化了該協方差矩陣的特征值。結果如下圖所示。α 越小,異質化程度越高,α 爲正無窮時爲同質化場景。k 爲特征值的 index。
對于該曲線,如果特征值大部分相對較大,即意味着表征能夠更加均勻地分布在不同的特征方向上。而如果該曲線隻有前面少數特征值較大,而後面大部分特征值都很小,就意味着表征分布被壓縮在少數特征方向上,即維度坍縮現象。
因此,從圖中可以看到,随着數據異質化程度越來越高(α 越來越小),維度坍縮的現象就越來越嚴重。
觀察二:全局模型的維度坍縮來自聯邦參與各方的局部模型的維度坍縮
由于全局模型是聯邦參與各方的局部模型融合的結果,因此作者推斷:全局模型的維度坍縮來源于聯邦參與各方的局部模型的維度坍縮。
爲了進一步驗證該推斷,作者使用與觀察 1 類似的方法,針對不同程度數據異質化場景下得到的局部模型進行了可視化。結果如下圖所示。
從圖中可以看到,對于局部模型,随着數據異質化程度的提升,維度坍縮的現象也越來越嚴重。因此得出結論,全局模型的維度坍縮來源于聯邦參與各方的局部模型的維度坍縮。
怎麽解決?
受到以上兩個觀察的啓發,由于全局模型的維度坍縮來源于本地局部模型的維度坍縮,研究人員提出在本地訓練階段來解決聯邦學習中的表征維度坍縮問題。
首先,一個最直觀的可用的正則項爲以下形式:
其中爲第個特征值。該正則項将約束特征值之間的方差變小,從而使得較小的特征值不會偏向于 0,由此緩解維度坍縮。
然而,直接計算特征值往往會帶來數值不穩定,計算時間較長等問題。因此借助以下 proposition 來改進方法。
爲了方便處理,需要對表征向量做 z-score 歸一化。這将使得協方差矩陣變成相關系數矩陣(對角線元素都是 1)。
基于這個背景,可以得到以下 proposition:
這一 proposition 意味着,原本較爲複雜的基于特征值的正則化項,可以被轉化爲以下易于實現且計算方便的目标:
該正則項即是簡單的約束表征的相關系數矩陣的 Frobenius norm 更小。研究人員将該方法命名爲FedDecorr。
因此,對于每個聯邦學習參與方,本地的優化目标爲:
其中爲分類的交叉熵損失函數,β 爲一個超參數,即 FedDecorr 正則項的系數。
實驗結果
首先,驗證使用 FedDecorr 是否可以有效緩解維度坍縮。
在 α =0.01/0.05 這兩個強數據異質化的場景下,觀察使用 FedDecorr 對模型輸出表征的影響。
結果如下圖所示。
可以看到,使用 FedDecorr 可以有效地緩解本地局部模型的維度坍縮,從而進一步緩解全局模型的維度坍縮。
在 CIFAR10/100 兩個數據集上驗證方法。研究團隊發現 FedDecorr 可以很方便的加入到之前提出的多個聯邦學習方法,并且帶來顯著提升:
同時,爲了展示方法的可擴展性,作者在較大規模數據集(TinyImageNet)上進行了實驗,并且也觀察到了顯著提升:
此外還基于 TinyImageNet,驗證了 FedDecorr 在更大規模聯邦參與方的場景下的有效性。
結果如下表。通過實驗結果展示了 FedDecorr 可以被用于較大規模聯邦參與方的場景。
FedDecorr 對正則項系數(超參數 β)的魯棒性結果如下圖所示。
通過實驗,發現 FedDecorr 對于其超參數 β 有較強的魯棒性。
同時發現将 β 設爲 0.1 是一個不錯的默認值。
最後,研究人員驗證了在聯邦學習時,使用不同的 local epoch 下 FedDecorr 也可以帶來普遍的提升:
論文地址:
https://arxiv.org/abs/2210.00226
代碼鏈接:
https://github.com/bytedance/FedDecorr