基于原型網絡的聚類式聯邦學習

張 鐵,徐林莉,周遠遠

1(中國科學技術大學 計算機科學與技術學院,合肥 230022)

2(中國科學技術大學 信息科學實驗中心,合肥 230022)

0 引 言

隨著民眾數據隱私意識的覺醒與相關法律的出臺,機器學習的數據隱私安全問題逐漸凸顯.聯邦學習(Federated Learning)[1-3]作為一個在隱私約束下的分布式機器學習訓練范式,已經被應用到機器學習與深度學習的多個領域[27,28],例如聯邦無監督學習[4],聯邦半監督學習[5],聯邦圖網絡學習[6-9],聯邦元學習[10-12]等.在聯邦學習場景中,訓練數據分布于多個用戶或移動設備處,同時存在一個用于協調各用戶分布式聯合訓練的參數服務器.由于隱私約束的限制,用戶不被允許直接發送原始數據到參數服務器或其他用戶.

傳統的聯邦學習通常假設可以訓練一個全局聯邦模型來擬合所有用戶的數據分布.然而,由于地域,文化,性別等偏向可能導致用戶之間產生數據分布異構問題,上述假設在實際分布式訓練任務中很難滿足.此外,移動設備有限的存儲與計算能力決定了用戶不能部署一個足夠大的深度神經網絡來獲得理想的模型性能.因此,近幾年一些個性化聯邦學習算法被相繼提出[13-17],旨在以聯邦學習的方式分布式地為每個用戶訓練一個個性化的模型.

作為個性化聯邦學習的一種,聚類式聯邦學習(Clustered Federated Learning)[18-22]顯式地考慮了用戶群體中由于數據分布差異而可能存在的簇狀結構.聚類式聯邦學習假設存在個不同的全局(數據)特征分布,每個用戶的局部數據采樣自其中的某一個.每個全局特征分布可以被認為是一種個性化偏向.這種潛在的數據分布差異可以被用來將用戶群體聚類為不同的簇.例如,在消費點評網站上,不同的年齡段或地域的用戶對電影、食物等商品有不同的評價習慣;在數字、字符、花草、寵物圖像等識別任務中,由于地域與身份的局限性,或采樣設備的差異,同一類別的數據可能存在不同的模式或風格,數據分布復雜多樣;此外,在一些特殊情況下,系統中也可能存在企圖通過錯誤標注樣本來影響聯邦訓練過程的惡意用戶[21].在這些場景中用戶之間數據的特征分布是異構的,因此均可以通過聚類式聯邦學習方法對用戶進行聚類,實現個性化的推薦與識別模型的訓練,以及對惡意攻擊進行防御.從樣本特征與標簽的角度來看,聚類式聯邦學習場景可以被簡單分為兩類.分別為1)多源數據:用戶局部數據分布的差異來自于地域、文化、身份背景,或采樣設備的多樣性等因素造成的不同風格.這導致同一標簽存在不同的特征.2)不相容數據:由于觀念或立場的不同,不同用戶可能對同一樣本有著不同甚至相反的看法與評價,這導致同一樣本或特征存在著不一致的標簽.

由于聯邦學習不允許用戶原始數據的共享,最近一些工作嘗試利用用戶上傳給服務器的局部模型參數對用戶群體進行聚類.例如,Ghosh等人基于模型參數之間的歐式距離提出了一個多階段的聚類式聯邦學習算法[19].然而,深度神經網絡的過參數化和排列不變性使得歐式距離并不能合理地反映模型映射關系或數據分布之間的差異程度[23],當用戶局部數據集的類別分布不一致時這個問題將愈加嚴重.此外,現有的聚類式聯邦學習方法有多階段計算,通信代價較高,難以適用于類別異構數據等問題.另一方面,由于移動設備存儲能力的有限,實際中用戶的局部數據分布通常是小樣本的,例如個人照片,消費記錄等.網絡帶寬與計算能力的限制也使得用戶很難與服務器保持長時間的持續通信以實現分布式的訓練.對于上述問題,本文從元學習方法原型網絡的角度,提出一種基于原型網絡的迭代聚類式聯邦學習框架.該框架利用原型網絡的空間映射能力為每個用戶構建用戶嵌入表示,并基于此表示,以期望最大化(EM)的算法思想交替地優化每個用戶的聚類標識與全局模型參數,進行迭代聚類式的聯邦學習,實現用戶模型的個性化訓練.

本文的主要貢獻點總結如下:

1)首次基于元學習方法原型網絡模型對聚類式聯邦學習問題進行定義,并利用原型網絡的空間映射特點,為每個用戶構建能反映特征分布信息的用戶嵌入表示.

2)提出了一個迭代聚類式聯邦學習框架PN-ICFL.該框架基于用戶嵌入表示提出了一種新的模型距離,交替地進行用戶聚類與分布式的模型訓練.該框架可以在服務器端進行用戶聚類標識的估計,相比其他的迭代聚類式聯邦學習算法,有更低的通信代價,而且更適合于類別分布異構場景.

3)在多種異構數據場景下驗證了所提出框架的有效性.實驗結果表明,本文的PN-ICFL框架能有效地進行聚類式聯邦學習,在一些場景,特別是小樣本場景下,能取得相對于對比算法更加優異的聯邦學習模型性能.

本文的剩余部分組織如下;第1節介紹相關工作;第2節對原型網絡模型下的聚類式聯邦學習問題進行正式定義,并介紹原型網絡的預訓練;第3節介紹用戶嵌入表示,并基于期望最大化算法思想導出用戶聚類標識估計與全局模型參數聚合的近似表達式,最后提出完整的基于原型網絡的迭代聚類式聯邦學習框架PN-ICFL;第4節進行實驗驗證與結果分析;第5節為總結與展望.

1 相關工作

本文的相關工作主要包含個性化聯邦學習,聚類式聯邦學習,與原型網絡3個部分.

1.1 個性化聯邦學習

針對聯邦學習經典算法FedAvg在復雜數據場景下的局限性,一系列的個性化聯邦學習算法被陸續提出[13-17].谷歌首先通過先進行聯邦訓練得到全局模型,再在用戶端進行本地微調的方式,實現了智能手機語言模型的個性化[14].一部分研究通過為每一個用戶執行個性化的模型聚合來實現個性化的模型訓練.FedAMP算法[13]使用一個注意力誘導函數度量用戶兩兩之間的相似性,并以此作為聚合權重,進行個性化的加權模型聚合.可以有效解決跨孤島(Cross-silo)場景[2]下用戶之間類別分布不一致所導致的數據異構問題.隨后,FedFoMo算法[16]基于用戶自身個性化的目標函數,得出了一階最優的模型聚合權值,在目前大部分類別分布異構數據上達到了最好的個性化性能(SOTA).另一方面,由于深度神經網絡模型通?？梢员环纸鉃樘卣魈崛∑骱头诸惼鲀刹糠?一些工作提出聯合地訓練一個全局的特征提取器,而用戶則在本地訓練一個個性化的分類器[17].

1.2 聚類式聯邦學習

聚類式聯邦學習的目的不僅是為了基于數據分布對用戶進行聚類,而且期望能在聚類結果的基礎上實現模型的個性化訓練.Ghosh等人[19]采用多階段的訓練方式,每個用戶首先在本地訓練模型直到收斂,服務器然后根據這些模型參數之間的歐氏距離進行聚類.在得到聚類結果后,每個簇內部再執行傳統的聯邦學習.PFA算法[18]采用了一種類似的多階段思想,利用神經網絡中ReLu激活函數輸出值的稀疏性為每個用戶構造特征表示向量,并基于此度量用戶兩兩之間的相似性.考慮到不相容數據問題,CFL算法[21]基于模型梯度的相似性矩陣,實現了一種自頂向下的層級聚類.然而,以上幾種算法均需要用戶在本地訓練模型到近似收斂,計算復雜度較高,難以應用到現實大規模聯邦學習場景.此外,這些算法均未考慮類別分布不平衡,即類別分布異構問題.最近,一些工作采用周期性迭代的方式,在聯邦訓練的同時更新聚類簇.例如,IFCA算法[20]遵循FedAvg算法迭代式的整體迭代流程,用戶在每一輪與服務器進行通信時,會下載服務器端的個全局模型,并在本地計算個損失函數值,基于此決定當前的聚類劃分.相比上述其他算法,IFCA算法有實時性好,計算量較小,超參數少等優勢.然而,每一輪個全局模型的下載給算法帶來了額外的通信代價,這在模型參數較多時會造成算法效率的下降.此外,IFCA算法對全局模型的初始點的選取較為敏感,聚類過程中容易產生模型退化,算法結束時聚類簇的數量可能小于初始設定數量.

1.3 原型網絡

原型網絡(Prototypical Network)[24]是解決小樣本數據問題的一種經典的元學習方法,由Snell等人于2017年提出.原型網絡將每個類別的原始樣本映射到一個新的樣本空間中,使得在這個新空間中同類樣本距離更近,不同類別的樣本距離更遠.原型網絡使用支持集中的樣本在原型網絡空間的均值作為該類的類原型,并基于歐式距離對查詢集中的樣本進行預測.算法流程簡單,可以有效處理小樣本多分類任務.最近,Park等人[11]基于原型網絡的思想提出了一種聯邦預訓練方法,該方法旨在聯邦訓練之后,得到一個在聯邦測試階段測試用戶之間可以進行聯合訓練且能快速適應新任務的初始模型.

2 問題定義與分析

本文首先基于原型網絡模型對聚類式聯邦學習中的多分類問題進行正式定義,并提出初步解決思路.

2.1 聚類式聯邦學習定義

本文主要研究多分類任務下的聚類式聯邦學習.本文假設全局存在N個用戶和一個參數服務器;此外,存在K個全局類條件分布,稱作全局特征分布,每一個全局特征分布φk(x|y),k∈[K]代表一個個性化偏向.x與y分別代表樣本特征與對應的類標簽.用戶i∈[N]基于其局部數據集Di聯邦式地訓練一個C分類機器學習模型,C為分類任務的類別總數.數據集Di的樣本獨立采樣自聯合數據分布pi(x,y),其類別分布為pi(y),特征(類條件)分布pi(x,y)屬于K個全局特征分布中的其中一個.用戶之間除了特征分布不同之外,類別分布pi(y)也可能是不平衡且不一致的.本文將類別分布與特征分布之間存在差異分別稱為類別分布異構/類別非獨立同分布(Label Non-IID)與特征分布異構/特征非獨立同分布(Feature Non-IID).

定義原型網絡為hθ:D→M,其中D與M分別為原始樣本空間特征維度和原型網絡空間特征維度,θ為原型網絡的參數.正式地,在聚類式聯邦學習問題中,優化如下的全局目標函數:

(1)

(2)

根據上述分析,用戶聚類標識si,或聚類劃分{S1,…,SK}為全局目標函數式(1)的隱變量.因此,本文采用期望最大化(EM)算法的思想,交替地優化全局模型參數與隱變量.具體地,為了最小化式(1),本文交替地1)在聚類劃分給定的條件下,優化全局模型參數,包括全局參數原型網絡θ與個性化參數全局類原型矩陣{Ek}; 2)基于更新之后的模型參數,估計隱變量,即用戶當前聚類標識si,進行聚類劃分.

2.2 預訓練原型網絡

在實際的聯邦學習場景中,參數服務器通常為大型公司或政府部門等可信任第三方,這些機構有較高的計算與存儲能力,可以獲取到一些共有的大型數據集.基于此優勢,本文提出在聚類式聯邦學習算法執行前,服務器先利用大型公有小樣本數據集對原型網絡參數進行預訓練.本文按照傳統原型網絡的元學習訓練方法,以元學習任務為訓練單位進行預訓練,使得原型網絡成為一個在其他相關任務上具有泛化能力的空間映射函數.最終,利用預訓練得到的原型網絡θ作為下文聯邦學習框架的初始參數.類似于式(2),本文在預訓練過程中使用如下的損失函數,其中Qc為元學習任務的第c類查詢集,zc為基于相應支持集計算得到的類原型.

(3)

3 聚類式聯邦學習框架PN-ICFL

本節首先利用原型網絡為聚類式聯邦學習問題中的用戶構建特征向量表示(3.1節);隨后,根據期望最大化(EM)算法思想,交替優化隱變量(3.2節)與模型參數(3.3節),得到度量用戶特征分布差異的模型距離與全局模型參數聚合方式;最后,正式提出基于原型網絡的迭代聚類式聯邦學習框架PN-ICFL(3.4節).

3.1 基于原型網絡的用戶特征表示

聚類式聯邦學習根據用戶數據集的特征分布差異對用戶群體進行聚類.如何在真實聚類信息未知且不能直接訪問原始數據的條件下度量用戶特征分布之間的差異性是聚類式聯邦學習的首要問題.考慮到預訓練原型網絡對于未知類別的空間映射能力,本文利用用戶每類樣本在原型網絡空間中的特征表示的均值,即類原型,定義用戶的向量表示ui為:

(4)

3.2 最優聚類劃分——隱變量優化

對于全局目標函數式(1),在原型網絡參數θ與全局類原型矩陣{Ek}給定時,進行隱變量優化為用戶i判斷當前最優聚類標識si.該過程對應期望最大化(EM)算法中的E步.

為了最小化式(1),本文基于式(2)計算每個用戶在給定個全局類原型矩陣條件下的用戶局部損失函數.具體地,傳統原型網絡在訓練中使用負對數似然作為損失函數,樣本距離度量d(,)為歐式距離的平方[24],每個樣本x的后驗概率定義為:

(5)

其中zc為第c類的類原型.將全局類原型矩陣Ek中相應的全局類原型代入上式,則關于用戶的局部數據集Di的負對數似然損失函數可以寫為:

(6)

考慮到歐式距離的非負性,上述損失函數有如下上界:

(7)

(8)

上式中右端的求和項可以進一步按類別寫為C部分子項之和:

(9)

(10)

圖1 基于原型網絡的模型距離Fig.1 Prototypical network-based model distance

3.3 全局模型參數聚合——模型參數優化

本文假設在隱私約束下不允許用戶原始數據的共享,但用戶可以將自己的經驗類別分布發送給可信任的服務器,同時假設每類原始樣本在經過原型網絡映射且均值化后的向量表示,即類原型,不會造成隱私的泄露.基于以上假設,在計算得到用戶嵌入ui表示后,用戶i將其與類別分布pi(y)發送給服務器,服務器基于式(9)和式(10)為該用戶估計聚類表示si.隨后,本文在當前用戶的聚類標識給定的條件下,分別更新全局原型網絡參數θ與全局類原型矩陣{Ek}.該過程對應期望最大化算法的M步.

3.3.1 更新全局類原型矩陣

在得到當前輪聚類劃分之后,對于一個聚類簇Sj,本文認為其中的用戶應擁有相同的特征分布.基于式(2)定義,本文采用交替優化策略,首先固定全局原型網絡參數θ.假設每輪所有用戶參與聯邦訓練,則該簇的最優全局類原型矩陣Ek應最小化如下的目標函數:

(11)

由式(7)中的不等式可知,以上的簇內損失函數有上界:

(12)

(13)

(14)

式(14)即為服務器端按不同聚類簇進行的個性化的全局類原型矩陣的模型聚合.

3.3.2 更新全局原型網絡參數

(15)

3.4 基于原型網絡的迭代聚類式聯邦學習框架PN-ICFL

基于上述分析,本文正式提出基于原型網絡的迭代聚類式聯邦學習框架PN-ICFL(Prototypical Network-based Iterative Clustered Federated Learning),PN-ICFL框架遵循聯邦學習經典FedAvg算法的“周期性通信-局部更新”的基本流程,具體細節見算法1.

算法1.PN-ICFL

輸入：聚類簇個數K,用戶個數N,通信輪數T,局部更新輪數E,局部學習率μ,用戶元學習任務數量B,公有數據集G

1.服務器用公有數據集G與式(3)預訓練原型網絡,得到θ(0)

3.fort=0toT-1do

5.for用戶i∈t并行do

6. 基于局部數據集Di計算當前用戶嵌入表示ui;

8. end for

9. for 用戶i∈tdo

12. end for

15.end for

算法2.局部更新(Local update)

輸入：原型網絡參數θ,學習率μ,迭代輪數E,用戶的元學習任務數量B;

輸出：更新后的原型網絡參數θ,類別分布向量pi(y);

1.{Tl}←將數據集D劃分成B個小樣本元學習任務;

2.fott=0 to E-1 do

3. forl=0 to B-1 do

7. 更新原型網絡θ=θ-μL(θ);

8. end for

9.end for

在聯邦訓練的第t輪中,服務器首先進行用戶采樣,得到當前輪參與訓練的用戶集合t.t中的每一個用戶在本地基于局部數據集Di計算用戶嵌入表示ui,并以原型網絡的訓練方式進行局部更新(第5～8行).用戶隨后將更新之后的原型網絡θ,用戶嵌入表示ui與類別分布向量pi(y)發送給服務器.在收到用戶的上述數據后,服務器基于式(10)計算模型距離{dik},并為每個用戶估計當前聚類標識si(第9～12行).最后,服務器為每個簇執行簇內的全局類原型矩陣的聚合(第13行),并為原型網絡執行全局模型參數聚合(第14行),得到下一輪的全局類原型矩陣和原型網絡參數θ(t+1).值得注意的是,一方面,模型距離{dik}的計算與用戶聚類標識si的估計可以在服務器端直接進行,因此用戶不再需要如IFCA算法中每輪下載所有的全局模型,避免了K倍的模型下載通信量.另一方面,用戶在本地計算得到用戶嵌入表示ui之后,即可立即將其與類別分布向量pi(y)發送給服務器,因此,服務器在計算模型距離,并對全局類原型矩陣以最新聚類劃分進行參數聚合的同時,各用戶可以在本地對全局原型網絡進行局部更新.這一特點可以提高算法的并行性,在一定程度上節省時間開銷.

此外,意識到全局類原型矩陣Ek可以視為當前聚類簇Sk內所有用戶的全局聚合信息,因此,在用戶的局部更新的過程中,全局類原型矩陣可以作為全局監督信息,指導原型網絡的本地更新.具體地,對于用戶i,本文利用全局類原型矩陣Esi替換式(3)中基于支持集計算得到的當前的類原型zc,得到如下的損失函數:

(16)

L總(θ)=L(θ)+λ·Lg(θ)

(17)

其中,λ為控制全局類原型信息對局部更新影響的超參數.全局類原型矩陣Esi可以在服務器端得到最新聚類劃分之后,由用戶下載獲得.用戶也可以直接利用上一輪模型測試時下載的全局類原型矩陣作為當前Esi的估計.本文將采用全局信息增強后的損失函數L總(θ)進行原型網絡局部更新的框架稱為PN-ICFL+.

在模型測試時,用戶i可以從服務器端下載第si個簇的全局類原型矩陣Esi和全局原型網絡參數θ.測試樣本x經過原型網絡映射得到特征表示hθ(x),然后基于Esi和最短距離原則進行類別預測.對于未參與聯邦訓練的新用戶,則可以首先使用局部數據集計算用戶嵌入表示,再由服務器基于最短距離原則為其估計聚類標識,并發送回相應的全局類原型矩陣.該新用戶隨后基于此為測試樣本預測類別.可以注意到,由于從服務器端下載的全局類原型矩陣包含所有類別的類原型,因此,在測試階段用戶也可以對局部數據集中缺失類別的樣本進行預測.

4 實驗設計與分析

本文的實驗部分將在小樣本,類別分布異構數據等場景下對前文提出的聚類式聯邦學習框架PN-ICFL進行驗證.

4.1 數據集構建

為了模擬聯邦學習中可能出現的類別分布異構問題,本文基于Lin等人提出的用戶數據集構建方法[25],在聚類式聯邦學習場景為用戶構建局部數據集.具體地,本文通過一個狄利克雷分布(Dirichlet Distribution)Dir(α)為每個用戶生成類別分布向量pi(y),其中α是控制類別非獨立同分布,即類別不平衡程度的超參數,α越大時,類別分布向量pi(y)越接近均勻分布.本文中選取α=100與α=3分別代表類別獨立同分布與類別非獨立同分布場景.在實驗所采用的C=10分類任務中,α=100時每個用戶包含所有類別的樣本,類別分布接近均勻分布,而α=3時每個用戶僅擁有3～5個類別的樣本.

隨后,為了決定用戶局部數據集的類條件分布,即特征分布,本文首先對每個用戶隨機分配聚類標識si.當聚類式聯邦學習問題中的K個全局特征分布φk(x|y)給定時,用戶的特征分布則為φsi(x|y).下文的實驗基于常用數據集,通過各種方法構造不同的全局特征分布.用戶的局部數據集則通過從其全局特征分布φsi(x|y)對應的數據集中采樣一定數量的樣本得到,具體細節見4.4節.

本文實驗中,服務器使用小樣本數據集omniglot作為公有數據集對原型網絡進行預訓練.

4.2 基準算法

本文實驗中,除了提出的PN-ICFL框架之外,選擇4個相關算法作為對比算法.它們分別是:

FedAvg:McMahan[3]等人2017年提出的聯邦學習基準算法,為所有用戶學習一個全局唯一的聯邦模型.

Local model:局部模型,每個用戶僅在本地基于局部數據集訓練一個個性化的模型.

FedFoMo:個性化聯邦學習算法,由Zhang等人提出[16],可以在大部分異構數據上達到SOTA的性能.

IFCA:迭代聯邦聚類算法[20],基于損失函數數值迭代地執行用戶聚類,實現個性化的訓練.

為了遵循FedAvg算法的迭代式更新過程,考慮實驗對比的公平性,本文未對層級聚類式聯邦學習方法如CFL[21]進行驗證.此外,每個用戶隨機地選擇20%的局部數據作為IFCA算法和FedFoMo算法中的驗證集,用于估計損失函數值.由于聚類初始點所造成的聚類式聯邦學習算法的隨機性,對于IFCA算法和本文的PN-ICFL框架,本文利用5個不同的隨機種子所決定的聚類初始中心執行5次實驗,取5次算法的平均指標作為最終評價指標.具體地,PN-ICFL框架從第一輪參與用戶的用戶嵌入表示集合中隨機選取個元素并轉化為相應的矩陣形式,作為初始全局類原型矩陣.

4.3 評價指標

本文采用所有用戶在測試集上分類模型的平均測試準確率作為算法性能的評價指標.此外,對于聚類式聯邦學習算法,本文基于用戶真實聚類標識,計算其聚類結果的調整蘭德系數(Adjusted Rank Index,ARI),評價其聚類性能.調整蘭德系數的值域為[-1,1],值越大時表明聚類算法取得了更優更接近真實聚類標簽的聚類結果.

4.4 小樣本數據場景

在跨設備(Cross-device)[2]聯邦學習場景中,由于移動設備存儲能力、職業身份等因素的限制,用戶的本地訓練樣本數量有限,且用戶不能與服務器保持長期的通信.例如,消費點評軟件僅擁有單個用戶對極少部分商品的評價記錄;個人手機中花草或寵物等其他類型照片的儲備量通常不會達到幾百甚至更大規模.考慮到以上問題,本文首先在小樣本數據場景中設置類別獨立同分布與類別非獨立同分布兩種環境對PN-ICFL框架的有效性進行驗證.

4.4.1 類別獨立同分布(Label IID)

為了模擬K個全局特征分布φk(x|y),本文首先構建一個多源數字數據集,稱作Multi-source Digits數據集,其部分樣本如圖2所示.Multi-source Digits數據集由4種常用數字相關的數據集組成,包括MNIST,USPS,SVHN與SIGN[26]數據集.其中MNIST與USPS是兩種不同風格的手寫體數字數據集,SVHN是采樣自現實街景門牌號碼的數字數據集,SIGN是一個手勢數字數據集,每個數據集代表分別一種全局特征分布.類別總數C=10.設置狄利克雷分布分布超參數α=100代表類別獨立同分布,并設置用戶局部數據集的大小|Di|=100模擬小樣本數據場景,此時每個用戶對于每類只有大約10個樣本.由于 PN-ICFL 框架采用原型網絡與類原型作為模型結構,為了模型的一致性,對于其他算法本文在 PN-ICFL 原型網絡的基礎上加入一個全連接層作為分類器,得到端到端的多分類模型.

圖2 多源數字數據集示例Fig.2 Examples of multi-source digits dataset

設置聯邦訓練通信輪數T=20,用戶數量N=48,用戶局部更新輪數E=1.對于PN-ICFL框架,采用原型網絡原始文獻[24]中的網絡結構.設置學習率μ=0.01,小樣本元學習任務數量B=2.對于其他算法,設置學習率μ=0.1,批量樣本大小B=50.本文分別選取用戶參與率R=1.0與R=0.5模擬全部用戶參與和部分用戶參與的場景,并分別設置λ=1.0與λ=2.0作為PN-ICFL+框架中損失函數Lg項的系數.實驗測試結果見表1(括號內數值為調整蘭德系數).

表1 多源數字數據集測試結果(Label IID)Table 1 Test results on multi-source digit (Label IID)

由表1可知在存在多種全局特征分布的情況下,傳統聯邦學習算法FedAvg很難有效地為所有用戶訓練出一個全局模型,因此此時需要進行用戶個性化的訓練.同時,當用戶本地數據量較少時,個性化聯邦學習算法FedFoMo與迭代聯邦聚類算法IFCA對于測試準確率雖有提升,卻不能有效地發現并利用用戶的聚類結構.例如,表1中IFCA的調整蘭德系數ARI僅為0.37(R=1.0)與0.00(R=1.0),這是因為IFCA在小樣本數據場景下不能對損失函數進行準確的估計,而且對初始聚類中心,即初始全局類原型矩陣的選取更敏感.本文的PN-ICFL框架實現了更優的聚類結果,其ARI值均大于等于0.7,同時取得了更高的平均測試準確率;此外,PN-ICFL+框架加入全局信息增強項,進一步提高了測試準確率.

另一方面,為了進一步展示PN-ICFL框架在此場景下的優勢,本文在圖3中展示出了在通信輪數分別為5、10、15、20時各種對比算法的平均測試準確率.由圖3可知,本文的PN-ICFL框架在聯邦訓練早期(T=5時)即可得到較好的模型性能,而其他對比算法則需要較長期的訓練,這是因為本文的框架利用了預訓練的原型網絡,可以快速地根據類原型信息為用戶群體進行聚類.這一特點使得PN-ICFL框架能更好的適用于通信能力有限的移動設備場景.

圖3 平均測試準確率隨通信輪數變化曲線Fig.3 Curve of averaged test accuracy with the number of communication rounds

4.4.2 類別非獨立同分布(Label Non-IID)

在實際的聯邦學習場景中,用戶可能只擁有部分類別的樣本,而且不同用戶之間的類別分布是不一致的.例如在多源數字數據集場景中,某一用戶僅擁有類別標簽“1”-“6”的樣本,而另一用戶僅擁有類別標簽“4”-“9”的樣本.為了模擬這種類別分布異構的場景,本文設置狄利克雷分布分布超參數α=3;此外,采用IFCA算法文獻[20]中的翻轉手寫體數據集Rotated MNIST.該數據集基于MNIST手寫體數據集構造,具體地,對于MNIST數據集中的每一個圖像樣本,分別將其翻轉0,90,280,270度,每個角度被視為一個全局特征分布.因此,聚類簇的個數K=4.

設置聯邦訓練通信輪數T=20,用戶數量N=48,用戶局部更新輪數E=2.對于PN-ICFL框架,本部分實驗采用3層卷積層和2層全連接層作為原型網絡的網絡結構.設置學習率μ=0.01,小樣本元學習任務數量B=2.對于其他算法,設置學習率μ=0.1,批量樣本大小B=50.本文分別選取用戶參與率R=1.0與R=0.5模擬全部用戶參與和部分用戶參與的場景,并分別設置λ=0.2與λ=0.1作為PN-ICFL+框架中損失函數Lg項的系數.實驗測試結果見表2(括號內數值為調整蘭德系數).

由表2可知,用戶類別分布的不均衡與不一致,即類別分布異構現象給聚類式聯邦學習問題增加了新的難度.此時,FedFoMo算法不能根據局部的損失函數值與模型參數信息實現有效的個性化聯邦訓練.而IFCA算法則幾乎不能發現有效的聚類結構.然而,PN-ICFL框架仍然可以取得有效的用戶聚類結果,有著最佳的平均測試準確率.這是由于本文中提出的模型距離,即式(10),可以看作不同的類別模型距離的加權和,考慮到了用戶數據的類別分布pi(y),因此也更適合于類別分布異構的場景.

4.5 一般數據場景

上一節重點驗證了PN-ICFL框架在小樣本數據場景中的優異性.在另一些聯邦學習場景中,用戶端可能為一些機構,如醫院或公司,此時,每個用戶擁有較多的本地數據.另外,由于具有高效且穩定的通信能力,用戶與服務器之間可以保持較長期的持續的聯合訓練.本小節在4.4.2節中的Rotated MNIST數據集的基礎上,增大用戶局部數據集大小與聯邦訓練通信輪數.具體地,設置局部數據集大小|Di|=1000,通信輪數T=40.為了驗證算法對不同類別異構場景的有效性,本文分別在α=3與α=100的條件下進行測試.設置用戶參與率R=1.0,局部更新輪數E=1.其他超參數與4.4.2節相同.實驗結果見表3.

表3 翻轉手寫體數據集測試結果(R=1.0)Table 3 Test results on multi-source digit (R=1.0)

由表3實驗結果可知,在上述的一般數據場景下,由于用戶擁有較多的訓練數據,而且可以與服務器保持多輪的通信,基準算法FedAvg可以取得較高的測試準確率.此外,在α=3時,FedFoMo算法受到類別分布異構的影響,不能實現個性化的模型訓練.由于用戶數據量的提升,IFCA算法在此場景中可以發現有效的聚類結構,對于α=3與α=100兩種類別分布場景均有最優的測試準確率.同時可以看出,本文的PN-ICFL框架可以取得與IFCA算法相近的結果,同時在用戶聚類結果上表現更好,其調整蘭德系數ARI大于ICFA算法,例如,當α=100時,本文的方法可以實現真實的用戶聚類結構,即ARI=1.0.這一現象主要是由于IFCA算法中采用了K個端到端的神經網絡,其學習能力大于PN-ICFL框架中采用的基于全局唯一的原型網絡的分類模型.

因此,除小樣本數據場景之外,本文的PN-ICFL框架在一般數據場景下也可以實現較優異的個性化聯邦模型訓練.

5 總結與展望

本文從元學習方法原型網絡的角度,對于聚類式聯邦學習問題提出了一種基于原型網絡的迭代聚類式聯邦學習框架PN-ICFL.與現有的聚類式聯邦學習算法如IFCA相比,PN-ICFL框架遵循FedAvg算法迭代式的訓練流程,通信代價較低,更適用于小樣本與類別分布異構數據場景等.具體地,本文基于原型網絡模型對聚類式聯邦學習問題進行正式定義,并從期望最大化算法的思想推導出了一種迭代式的“服務器端聚類-用戶端訓練”的框架PN-ICFL.服務器使用由用戶局部類原型組成的用戶嵌入表示計算用戶與聚類簇之間的模型距離.此模型距離是PN-ICFL框架的核心,由基于類別分布的類別模型距離加權和組成,從形式上考慮了類別分布異構的影響.預訓練原型網絡的使用使得PN-ICFL在小樣本數據場景中表現優異,可以較快速地發現有效的聚類結構,實現個性化的訓練.相關的實驗結果證明了PN-ICFL框架的有效性.

本文假設弱隱私約束條件成立,服務器不能直接訪問用戶原始數據,但允許訪問用戶的類別分布信息,并且類原型向量不會泄露出用戶的個人隱私信息.這個假設在一些嚴格的隱私約束條件下將不再成立.因此,未來的工作應以隱私保護為研究重點,嘗試從多方安全計算等隱私保護技術[27,28]的角度對本文的PN-ICFL框架提出改進措施,使其適用于更多的聯邦學習場景.