面向數據異構的聯邦學習性能優化研究

黃聿辰,趙彥超,郝江山,陳 兵

(南京航空航天大學 計算機科學與技術學院,南京 210000)

0 引 言

聯邦學習已經成為現代大規模分布式機器學習中的一個重要方向.聯邦學習與傳統的集中式學習不同,集中式學習使用存儲在中央服務器中的大型數據集對模型進行訓練,而在聯邦學習中,訓練數據可以分布在大量的客戶端上[1],比如電話、網絡傳感器、醫院或者各類邊緣設備[2],然后通過服務器聚合成完整的全局模型,而不用傳輸用戶數據,從而保護了用戶的隱私信息.

聯邦學習本地通常采用的神經網絡,其處理各種計算機視覺任務非常有效.一般來說,在訓練過程中,神經網絡都包含所有數據樣本.然而,在聯邦學習的環境中,各個客戶端數據樣本量和樣本類別都處于變化之中,即數據異構[3].這更類似于生物系統在現實世界中的學習方式.在這種情況下,各個客戶端神經網絡只能訓練各自保有的數據集.因此中心服務器系統面臨的主要問題是當神經網絡權值適應新的任務時,可以保持對于先前學習到的數據的識別能力.

聯邦學習優化的關鍵挑戰是數據分布的異構性.由于參與聯邦學習的各個邊緣設備所處實際物理環境的不同,能為聯邦學習框架提供的數據各不相同,這種不同的數據分布會導致邊緣設備在本地訓練過程中產生模型偏移.采用這類模型進行聚合,生成的全局模型無法在全類別數據集上有優異表現.最常用的FedAvg[4]通過在與服務器通信之前對可用客戶端執行多個本地更新來解決通信瓶頸,即增加本地訓練輪數,增大本地訓練計算時間在總體任務中的時間占比,從而減少通信開銷占比.雖然它在某些應用中取得了成功,但是該方案未充分考慮聯邦學習實際應用場景,即各個客戶端數據異構,在多輪的本地訓練過程中,模型偏移隨著迭代而不斷加深,導致采用平均算法聚合的全局模型在全類別數據集上性能大幅下降,因此數據異構[5]仍是聯邦學習的一個活躍研究領域.

FedAdp[6]考慮了局部模型與全局模型的余弦相似角[7],在聚合過程中按照相似程度賦予權重,該方案一定程度上抑制了局部模型偏移,但是與FedProx相似,僅僅是直接做了全局模型趨同,未能結合各個客戶端異構性的數據分布.FEDC[8]提出對運算效率低下的客戶端進行低權重聚合,以保證全局模型的性能,但是該方案未考慮各個客戶端數據集信息量,即運算效率較低的客戶端的數據集對全局模型收斂的重要性未被考慮.Huang[9]提出一種結合本地訓練精度與訓練頻率信息進行加權的聚合方式,可以幫助服務器選擇更好的模型進行聚合,但是該權重設置屬于啟發式方案,并沒有對其收斂性進行證明.

因此,由于數據異構造成的局部模型偏移問題極大的引起了全局聚合模型識別性能的下降,針對該問題,本文提出了面向數據異構情況下的聯邦學習局部模型偏移的性能優化方案.由于神經網絡在對數據集的訓練過程中,通常僅有部分參數對數據的特征進行擬合,而另外一部分參數在訓練過程中所起到的作用較小,該現象也稱為神經網絡的過參數化.其中梯度變化量體現了該參數在模型訓練過程中起到的作用,積極擬合數據特征的相關參數會在訓練中產生更多梯度變化量.利用上述結論,本方案在聯邦學習本地訓練過程中,在其損失函數上添加待訓練模型與剩余各個客戶端模型參數之間的差值二范數,即將本地待訓練模型與其余客戶端模型參數作差并求其二范數,用剩余客戶端的模型梯度增益對應作為二范數每個參數位置的權重系數,將該整體形成懲罰項,在優化過程中,對新構建的損失函數進行梯度下降優化,從而保證了所有客戶端的模型在本地訓練過程中,依然保持著對其他客戶端數據分布的擬合能力,有效抑制局部模型偏移.所有參與聯邦學習的客戶端都能保持對其余所有客戶端數據分布的識別能力,將這些客戶端的梯度更新進行聚合,生成的全局模型性能有所保證.

本文主要貢獻有:

1)通過增加本地待訓練模型與其余客戶端模型參數差值的二范數,并在該二范數上添加模型梯度權重,將整體作為本地訓練懲罰項,抑制本地模型對各自數據集的模型偏移,來達到提升全局模型識別精度的效果.

2)在不改變聚合方式的情況下,有效解決了數據異構的問題,使得模型框架更好的適配實際應用環境,實現了優化框架的即插即用.

3)對于提出的優化方案進行了系統的證明,從理論上保證了收斂性和收斂速度.

1 相關工作

大規模機器學習對數據中心服務器的性能需求非常巨大,因此眾多分布式學習方案應運而生.隨著手機、傳感器等邊緣計算設備算力的增長和普及,以及將數據移動到數據中心的通信開銷較為龐大的問題,在分布式設備組成的網絡中進行各自的本地學習來進行統計模型生成的方案變得越來越有吸引力.這個方案被稱為聯邦學習,它需要解決隱私、數據設備異構以及大規模分布式網絡通信開銷等新問題.

最近提出的一些優化方法,對上述問題提出了不少有創意的解決方案.如Boyd[10]所提方案以及小批量優化方法[11].在大型網絡中,允許不精確的本地更新來平衡通信與計算,和在一個小部分集合的設備組之間進行即時通信.Fedpaq[12]提出了一種周期平均和量化功能的高效通信聯邦學習方法,它通過多任務學習框架為每個設備學習獨立但是相關的模型并進行周期性更新和發送.盡管該方法在理論上有保證,但這種方法并不能推廣到所有問題上.在非凸設置下,一種基于平均原始局部隨機梯度下降(SGD)更新的方法[13],實驗證明性能良好.SkewScout[14]提出了通過本地訓練輪數的限制來控制局部模型偏移的程度,從而提升全局模型的識別精度,但是該方案僅僅考慮了模型性能,并未考慮在聯邦學習實際應用場景中,通信開銷也是一個重要的問題,壓縮本地訓練輪數會導致通信開銷占比顯著增大,從而導致計算與數據傳輸比例大幅減小.IFCA[15]通過按照模型參數相似程度進行簇劃分,在各個簇內進行模型共享,從而使得簇內模型更新方向的一致性,抑制簇內模型偏移.該方案提升了簇間的全局聚合模型性能,但是該方案在實驗過程中發現,需要每輪訓練過程中保持較高的客戶端參與率,否則各個簇內客戶端數量不夠,無法保持該方案的優化效果,但是聯邦學習實際應用場景中,邊緣設備的不可用性也是需要考慮的一部分因素.

盡管,最近的一些工作[16,17]探索了數據異構環境下的收斂性保證問題,但他們要求所有設備參與每一輪通信,這在現實聯邦學習網絡中往往是不可行的.還有一些啟發式方法旨在通過共享本地設備數據或服務器端代理數據來解決數據異構[18,19].如Liam[20]該作者通過各個客戶端將自己本地數據進行多維度映射,并在該過程中保存數據特征,然后將映射結果傳輸給服務器,由服務器對所有客戶端上報數據進行整合,再下發給各個客戶端,客戶端進行本地訓練時,先在服務器下發的多個客戶端映射數據集上進行訓練,然后在本地數據集上進行訓練,從而減少數據異構帶來的全局性能損失.該方案的全局模型性能較高,但是即使對客戶端數據進行的映射操作,依然存在暴露本地數據的行為.

因此,這些方法是不現實的.除了增加網絡帶寬負擔之外,根據模型權重推測出數據分布[21]的方式違反了聯邦學習的隱私保護要求.將各個客戶端的數據進行周期性局部分享的方式[22]則需要詳細地生成或收集此類局部輔助數據,計算開銷較大,局部輔助數據生成難度較大.

除了數據上的異構性之外,系統上的異構性也是聯邦學習中的一個關鍵問題.由于硬件(CPU、內存)、網絡連接(3G、4G、5G、WIFI)和功率(電池水平)的變化,系統屬性也有所不同.這些系統級的異構性極大地引發了諸如掉隊和容錯等問題.在異構網絡中,聯邦學習優化的一種策略是忽略不受約束的設備,即在訓練框架中,放棄聚合無法完成一定訓練輪數的設備所提供的梯度更新信息[23].然而,這會對收斂產生負面影響,因為它減少了有助于訓練的有效設備的數量,如果被丟棄的設備具有特定的數據特征,則會在全局模型收斂的過程中產生偏差.

最接近本文的工作的是Prox[24],該文章的作者提出的FedProx算法,就像本文的算法一樣,也使用參數差值的二范數作為懲罰項.然而,與本文的算法不同的是,在FedProx中,懲罰項是各向同性的,其直接將本地模型與全局模型更新方向進行同步,并未考慮全局模型本身已經是各個局部模型在本地異構數據上的模型偏移之后的聚合結果,所以此時的全局模型更新方向并不能代表最優解方向.Ohad[25]提高了處理二次目標的收斂速度,證明了其可以通過提升數據容量大小,從而達到加速收斂.最近的一項工作[26]證明了FedAvg變體的對數據異構的收斂速度,它還為實踐中已知的一種現象提供了一個理論解釋,即當局部迭代次數過高時全局性能會下降,因為在數據異構情況下,過多的本地迭代會導致局部模型偏移,所以之后的全局聚合模型無法在全類別數據集上有優異性能.

2 系統模型

本節將詳細介紹聯邦學習中存在的數據異構問題帶來的影響,分為兩個小節對系統建模以及問題定義進行介紹.

2.1 系統建模

在概率論與統計學中,獨立同分布(Independent and identically distributed,縮寫為IID)是指一組隨機變量中每個變量的概率分布都相同,且這些隨機變量互相獨立.

在機器學習任務中,訓練數據集的獨立同分布是常規假設,常用的標準數據集如MNIST,共包含6萬張圖片,分為10類,每一類6千張圖片,即均勻劃分,在集中式訓練過程中,每一個mini-batch都是均分采樣,符合獨立同分布的假設.

但是在聯邦學習實際應用場景中,由于邊緣設備本地所包含數據集各不相同,各客戶端數據集Dk,其中k為客戶端編號.Di∪Dj=0或者Di∩Dj→0,即客戶端數據的樣本交集為0或者交集很少.在這種數據分布下,本地客戶端的訓練會產生模型更新的偏移,如圖1所示.

圖1 模型更新偏移Fig.1 Model update offset

因此,參與訓練的各個客戶端本地模型更新的聚合方向與全局理論最優解方向不一致,就會導致最終生成全局模型識別精度的下降.

2.2 問題定義

為了修正本地模型對各自數據集在訓練過程中產生的模型偏移,即保證模型在本地訓練過程中,能減少模型在其余客戶端模型關鍵參數位置的變化.該問題可以表述為:

θi為客戶端i的本地待訓練模型.θk為其余客戶端模型,其中k為除i以外的客戶端.Gk為其余客戶端本地更新梯度增益,用來標識其余客戶端模型的關鍵參數,使得關鍵參數在客戶端i的本地優化過程中得到更高的權重,從而保留該參數的更新方向.

數學表述為:

argminθiL(θi)+λ∑k∈S(θi-θk)Gk(θi-θk)

(1)

將該參數作為懲罰項加入到本地訓練過程中的損失函數中進行梯度下降優化,在優化過程中就可以使得每一個客戶端在本地進行模型的訓練與更新的同時,保持其余客戶端的模型更新方向,使得每一個客戶端都不會產生對自己本地數據分布的嚴重模型偏移,從而在后續的全局聚合中,能夠有效提高生成的全局模型對于全類數據集的識別精度.

3 局部模型偏移抑制

該章節介紹了優化方案在本地的實現過程,算法的整體流程,以及對優化方案的理論證明.

3.1 局部優化

神經網絡對于數據樣本的訓練普遍存在過參數化現象,即對于某些特定數據分布的樣本集合,神經網絡中的神經元參數只有部分對于該數據分布的擬合起到作用,可將其稱為關鍵參數.因此在各個客戶端本地訓練過程中,應當減少對于其余客戶端關鍵參數的更新,從而保證各個客戶端本地訓練完畢之后對于其余客戶端數據依然可以有效識別.

本文通過限制局部更新使其接近其余客戶端模型,在局部子問題中添加一個約束項.客戶端不僅僅最小化常規局部損失函數Ls,而是使用以下新構建損失函數Lt, s進行約束更新:

Lt, s(θ)=Ls(θ)+λ∑j∈S(θ-θt-1,j)Gt-1,j(θ-θt-1,j)

(2)

其中,Ls(θ)為客戶端集合,S為客戶端集合,t為全局訓練輪數,j為其余客戶端標識,θt-1,j為上一輪全局訓練的各個客戶端模型,其中j∈S,Gt-1,j為上一輪全局訓練的各個客戶端梯度,其中j∈S.用本地模型與、其余客戶端模型參數作差并取其二范數,在該值各個參數位置用模型梯度作為權重,將其加入損失函數進行SGD,可以減少訓練過程中其余客戶端關鍵參數位置的變化幅度,從而保留其余客戶端信息量.

本文將該優化方案簡稱為FedOC.維護FedOC所需的歷史數據的存儲和傳輸,存在數據量大以及數據傳輸成本高的問題.然而,通過對損失函數的計算整合,可以避免這些潛在的問題.損失函數Lt,s可以被重新整理為:

(3)

即應用多項式乘法規則,將兩個多項式逐個相乘,其中常數C為θt-1,j和Gt-1,j兩個項的相乘結果,因為這兩個項為上一輪運算結果,因此其結果為常數,故用常數C表示.

服務器只需要維護并向邊緣設備傳輸除θ之外的2個與θ維度相同的元素:

ut-1=∑j∈S(Gt-1,j),vt-1=∑j∈S(Gt-1,j)θt-1,j

(4)

按照公式(3),每個客戶端本地訓練過程,僅需要公式(4)的兩個模型系數即可.

需要注意的是,該方案只需要從邊緣設備發送與本地梯度相關的聚合信息到服務器.在隱私性方面,它與經典的FedAvg算法沒有區別.圖2給出了框架結構圖.

圖2 局部模型偏移抑制框架Fig.2 Local model offset suppression framework

3.2 算法流程與描述

3.2.1 算法流程

算法.局部模型偏移抑制(FedOC)

輸入：K為客戶端總數;T為全局通信輪數;θ為訓練模型;N為所有客戶端訓練樣本總數;nk為客戶端k的樣本數量;

1.ForkinKdo:

2. 服務器發送訓練模型θ給客戶端k

4.服務器側運算參數u0,v0

5.End for;

6.Fort=0…T-1 do:

7. 服務器隨機選擇的K子集kt

8.服務器向kt發送全局模型θt,參數ut,vt

9.本地客戶端優化新構建損失函數

10.θt+1,j=argminθFk(θ)

11.=L(θ)+λθTutθ-2λθTvt+C

12.客戶端上傳本地梯度Gt,j

13.服務器接收各個客戶端上傳梯度,生成本地模型

14.運算新一輪的ut+1,vt+1

16.End for;

3.2.2 算法描述

正式訓練之前,進行預訓練的數據收集,由參與聯邦學習訓練的客戶端首先完成1輪初步訓練,并上傳各自訓練梯度,由服務器進行對應模型更新.在正式的聯邦學習訓練過程中,服務器選擇客戶端集合的任意子集參與本輪聯邦學習訓練,并向子集中的客戶端傳遞預訓練過程中所有客戶端上報的參數θ0,u0,v0,各個參與本輪訓練的客戶端用指定參數構建損失函數,完成SGD優化,上傳本輪更新梯度,由服務器更新對應本地模型和全局模型,并運算出下一輪需要提供給客戶端進行損失函數構建的參數進行下發,不斷迭代,直到訓練達到全局輪數為止.

最后,由于FedOC只對FedAvg進行輕量級修改,僅僅在本地訓練過程中采用了新構建的損失函數進行優化,對于聚合方面并未有所改動.因此可以使用在FedAvg成立的推理過程,并且使FedOC輕松地集成到現有的系統中.特別是,本文提出,FedAvg是FedOC的一個特殊情況,即λ=0.

3.3 收斂性及收斂速度證明

3.3.1 定義,假設以及定理介紹

接下來對框架的收斂性以及收斂速度進行證明,首先,將介紹一些定義,假設和定理.

定義1.(光滑性)函數f具有Lipschitz連續梯度,其中常數L>0(即f是L-平滑的):

(5)

定義2.(強凸性)函數f是μ強凸的,其中μ>0

(6)

該定義旨在允許本地客戶在靈活地執行訓練,以便每個本地目標都可以收斂.局部計算與通信的數量可以通過調整局部迭代的數量來調整,即更多的局部迭代表明更精確的局部解.

此外,本文介紹以下假設:

假設1.目標函數f(x)是有界的,即最優解的有界性minf(θ)=f(θ*)>-∞.

假設1顯然滿足,因為目標函數f(x)存在一個最小值.

假設2中關于隨機梯度平方范數的條件是常規的.與隨機梯度的期望的假設相比,這是一個更加合理的假設.

(7)

假設3確保了局部梯度的可估計性.

3.3.2 定理以及其證明

結合上述定義和假設,本文將繼續介紹并在后續證明了一些有用的定理.

定理1.根據定義3,新構建本地損失函數具有γ-不確定解,因此可得:

(8)

其中c為客戶端數量.

證明：

對于定理1,即全局梯度的有界性.

根據定義3可得:

全局梯度可表示為:

將級數展開可得:

根據假設2,得證:

證明完畢.

定理2.根據定義2,如果f(x)是μ-強凸,則:

(9)

證明：

對于定理2,即本地損失相比于最優解的有界性.

由于根據強凸函數線性結合的性質可得,新構建損失函數同樣為μ-強凸.

定義：

令Γ′(θ′)=0

可得:

因此可得:

帶入公式可得:

整理結果,得證:

證明完畢.

定理3.根據定義1,定義2,定義3,假設2,以及假設3,可得,在T輪全局迭代之后,FedOC收斂于全局最優解模型θ*.

(10)

證明：

對于定理3,即FedOC收斂于全局最優解模型θ*,全局損失的穩定下降.

根據定義1,可得:

根據假設3,可得:

不等式兩邊同時減去:

得證:

證明完畢.

定理4.根據定理3結果可得,全局輪數可由相關常數提供有界性保證,即:

(11)

其中ξ為初始損失與最優損失差值的線性相關值.

證明：

根據定理3的結果,即:

令t+1=T,1-2μησ=q

存在常數ξ,使得:

不等號兩邊取對數,并令:

可得:

等式兩邊同時除以logq,

得證:

證明完畢.

4 實驗與分析

本節介紹了優化方案所采用的實驗設置,各項對比實驗、消融實驗和實驗結果的分析.

4.1 實驗基礎設置

本文的算法框架基于開源pytorch框架,采用RTX2080作為模型運算單元,并且實現了FedProx,sharedata兩種框架進行對比實驗.在聯邦學習實際應用場景中,各個邊緣設備算力參差不齊,因此選擇對算力要求相對較低的全連接網絡作為訓練模型.

采用MNIST,FMNIST,CIFAR10這3種常用數據集進行實驗.訓練樣本為50000個,測試樣本為10000個.實際訓練過程中由于MNIST和FMNIST的樣本維度為(28,28,1)而CIFAR10為(32,32,3),因此為了方便模型進行數據輸入,將CIFAR10所有樣本數據進行灰度化.這樣僅需在模型輸入階段將輸入維度由784修改為1024即可.

首先將客戶端設置為100,在MNIST,FMNIST,CIFAR10這3種數據集上進行與FedProx,sharedata的對比實驗.各個客戶端在訓練開始之前,隨機分配總類別的其中3種,每一種類別隨機選取3000個樣本,形成總計9000個,3大類的數據異構分布場景.

實驗過程中,對數據分布不斷做隨機劃分,進行多次實驗,選取平均值作為實驗結果.

4.2 實驗結果分析

客戶端數量和數據集的不同都會對實驗結果產生影響,因此設計了相關對比實驗.消融實驗為未采用優化方案的常規FedAvg和采用了優化方案的FedOC進行對比.

4.2.1 客戶端數量,數據集對比實驗

客戶端數量為100的3種數據集下的各個優化方案對比實驗結果見圖3.可以發現,各個優化算法在CIFAR10數據集上的表現相對MNIST,FMNIST要略差一些,這是由于CIFAR10的數據樣本相對復雜,簡單的全連接網絡對于其特征的擬合沒有對于MNIST以及FMNIST那么完善.

圖3 客戶端數量為100且在3種數據集下對比實驗Fig.3 Comparative experiments on three datasets with 100 clients

分析整體實驗結果得出結論,共享數據的優化方案在各個數據集上的表現,都要優于其他兩個優化方案.這是因為共享數據方案使得聯邦學習系統擁有趨于獨立同分布的數據集分布,因此模型可以在訓練中充分擬合數據特征而不會產生局部模型偏移,更加接近集中式學習方案,從而提升識別精度.通常,數據共享型優化方案可以作為其他優化方案的目標值.但是由于聯邦學習對于隱私保護的需求,這類對客戶端樣本信息進行公開或者部分公開的做法并不嚴謹.

再將客戶端數量調整為200進行實驗,觀察各優化算法在客戶端數量變化之后的性能,實驗結果如圖4所示.

圖4 客戶端數量為200且在3種數據集下對比實驗Fig.4 Comparative experiments on three datasets with 200 clients

根據實驗結果得出結論,客戶端數量增加之后,各個優化手段在3種數據集上的識別精度都有所提升.在實驗室環境中,公共數據集樣本種類有限,隨著客戶端數量增加,隨機分發數據時,擁有相同種類甚至同一樣本的客戶端數量會成比例上升,各個客戶端的數據異構性產生弱化,因此隨著客戶端數量上升,全局模型擬合能力會隨之提升.

將本方案與現有方案FedProx進行比較,可以看出,在識別精度方面,要優于FedProx 5個百分點.FedProx 在本地損失函數上僅添加了本地待訓練模型參數與當前全局輪數下發的全局模型的差值二范數,此時的全局模型是由上一輪的各個本地客戶端模型按照平均算法聚合生成的,然而各個本地模型因為其數據異構性質,會產生較大的局部模型偏移,采用該本地模型進行聚合生成的全局模型和理論全局模型最優解存在差異,籠統地將本地待優化模型參數與全局模型參數做同步的方案并不精確,只是在各向同性上抑制了本地客戶端的偏移,而FedOC充分考慮各個客戶端在本地訓練過程中其余客戶端的更新方向,因此在最終的識別精度上FedOC也要更高.

4.2.2 消融實驗

該小節實驗中加入了未采用任何針對非獨立同分布數據集進行優化的常規FedAvg訓練框架作為消融實驗進行對比,用戶數為100,3種數據集對比實驗結果如圖5所示.

圖5 客戶端數量為100且在3種數據集下消融實驗Fig.5 Ablation experiments on three datasets with 100 clients

用戶數為200,3種數據集對比實驗結果如圖6所示.

圖6 客戶端數量為200且在3種數據集下消融實驗Fig.6 Ablation experiments on three dataset with 200 clients

可以看出,未采用優化手段的FedAvg識別精度上升緩慢,這是由于局部模型偏移引起的全局模型更新方向偏移,最終遠離理論全局最優解方向,因此識別性能提升緩慢.

4.2.3 結果分析

1)本方案不涉及客戶端數據的共享,在隱私保護方面與常規FedAvg相同,可采用所有適用于FedAvg的隱私保護方案.

2)本方案僅僅對FedAvg進行輕量化改造,并未對聚合過程進行變化,僅在本地訓練過程中對損失函數進行優化,因此適用于FedAvg以及其變體場景,實現即插即用.

3)本方案充分考慮系統中所有客戶端的數據分布,并在訓練中保持其更新方向,因此在全局模型性能上優于如FedProx等現有算法.

4)常規FedAvg算法在非獨立同分布數據集上性能較差,模型偏移導致其收斂緩慢.實驗中其余優化算法均可在規定全局訓練輪數中收斂.

5)在訓練總樣本類別固定的情況下,增加系統內客戶端數量,可以提升數據分布相似性濃度,從而提高全局模型性能.

表1為實驗對比數據.

表1 實驗對比數據Table 1 Experimental comparision data

5 結 語

本文主要提升了聯邦學習在數據異構條件下的全局模型性能.基于局部梯度信息量來保存各個客戶端上對擬合起重要作用的關鍵模型參數,從而抑制本地客戶端模型訓練時的模型偏移.雖然本地訓練過程中加入了額外的懲罰項,并且服務器端需要進行梯度的矩陣運算,一定程度上延長了訓練時間,但是全局模型性能的提升是顯著的.下一步可以考慮在矩陣運算部分進行輕量化,以提升訓練效率.