基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)

竇勇敢，袁曉彤

（1.南京信息工程大學(xué) 自動(dòng)化學(xué)院，江蘇 南京 210044;2.江蘇省大數(shù)據(jù)分析技術(shù)重點(diǎn)實(shí)驗(yàn)室，江蘇 南京 210044）

近些年來，隨著深度學(xué)習(xí)的興起，人們看到了人工智能的巨大潛力，同時(shí)希望人工智能技術(shù)應(yīng)用到更復(fù)雜和尖端的領(lǐng)域。而現(xiàn)實(shí)狀況是數(shù)據(jù)分散在各個(gè)用戶或行業(yè)中，用戶數(shù)據(jù)存在隱私上的敏感性和安全性。如何在保護(hù)數(shù)據(jù)隱私前提下進(jìn)行機(jī)器學(xué)習(xí)模型訓(xùn)練，讓人工智能技術(shù)發(fā)揮出更強(qiáng)大的作用成為一種挑戰(zhàn)。

為了讓這些隱私數(shù)據(jù)流動(dòng)起來，同時(shí)應(yīng)對(duì)非獨(dú)立同分布數(shù)據(jù)的影響，Google 科學(xué)家Mcmahan 等[1]提出聯(lián)邦學(xué)習(xí)(federated learning)，通過協(xié)調(diào)大量遠(yuǎn)程分布式設(shè)備在保護(hù)用戶數(shù)據(jù)隱私的前提下訓(xùn)練一個(gè)高質(zhì)量的全局模型。

目前的聯(lián)邦學(xué)習(xí)算法還存在諸多問題。首先，每個(gè)設(shè)備CPU、GPU、ISP、電池以及網(wǎng)絡(luò)連接（3G、4G、5G、WIFI）[2]等硬件差異導(dǎo)致設(shè)備間存在很大的系統(tǒng)異構(gòu)性。傳統(tǒng)的聯(lián)邦學(xué)習(xí)方法FedAvg[1]在規(guī)定時(shí)間內(nèi)將沒有訓(xùn)練結(jié)束的設(shè)備簡單丟棄，這在現(xiàn)實(shí)情況中是不可取的，浪費(fèi)了大量的計(jì)算資源。其次，每個(gè)設(shè)備的數(shù)據(jù)分布和類型存在很大的差異[3]，跨設(shè)備的數(shù)據(jù)是非獨(dú)立同分布的（non-IID），這是數(shù)據(jù)的異構(gòu)性。不同的異構(gòu)環(huán)境中模型的收斂效果差別很大，甚至無法收斂。這些系統(tǒng)級(jí)別的異構(gòu)性給聯(lián)邦學(xué)習(xí)帶來了極大的挑戰(zhàn)。

現(xiàn)有針對(duì)異構(gòu)性問題的分布式優(yōu)化算法中，大部分都是針對(duì)特定異構(gòu)環(huán)境設(shè)定的。例如:文獻(xiàn)[4-6]提出讓所有設(shè)備都參與每一輪的訓(xùn)練，雖然在異構(gòu)數(shù)據(jù)環(huán)境中的收斂性得到了保證，但是這在現(xiàn)實(shí)的聯(lián)邦環(huán)境[1]中是不可行的。這不僅增加了服務(wù)器的通信負(fù)擔(dān)，而且參與聯(lián)邦訓(xùn)練的設(shè)備也應(yīng)隨機(jī)抽取。也有方法通過共享本地?cái)?shù)據(jù)來解決數(shù)據(jù)異構(gòu)性的問題[7-8]，但這違背了聯(lián)邦學(xué)習(xí)保護(hù)用戶數(shù)據(jù)隱私的前提。在聯(lián)邦設(shè)置中，文獻(xiàn)[9]通過在服務(wù)器端設(shè)計(jì)基于動(dòng)量優(yōu)化器FEDYOGI來加快異構(gòu)數(shù)據(jù)環(huán)境中全局模型收斂速度，這雖然提高了模型的收斂速度，但卻增加了服務(wù)器的計(jì)算量，在有限的計(jì)算資源下不是好的選擇。此外，也有研究者利用二階擬牛頓法優(yōu)化模型[10]，在相同的異構(gòu)環(huán)境中，與FedAvg 相比達(dá)到相同精度下減少了通信輪數(shù)，提高了通信效率，但這潛在增加了客戶端本地的計(jì)算量。

除了數(shù)據(jù)異構(gòu)性，每個(gè)參與聯(lián)邦訓(xùn)練的客戶端的硬件存在差異，這導(dǎo)致設(shè)備間存在很大的系統(tǒng)異構(gòu)性[11]。例如：在文獻(xiàn)[12-15]中，介紹了在異構(gòu)環(huán)境中目前最新的聯(lián)邦學(xué)習(xí)研究進(jìn)展，在全局模型聚合階段的更新方式同F(xiàn)edAvg[1]一樣，在指定的時(shí)間窗口內(nèi)，服務(wù)器將未完成訓(xùn)練的設(shè)備直接丟棄，不允許上傳本輪訓(xùn)練的模型參數(shù)。各參與訓(xùn)練的設(shè)備不能根據(jù)自己硬件性能在本地執(zhí)行可變數(shù)量的本地工作，缺乏自主調(diào)節(jié)能力。

在解決聯(lián)邦學(xué)習(xí)異構(gòu)性的問題上，近鄰優(yōu)化的更新方式廣泛地用于研究，包括高效通信分布式機(jī)器學(xué)習(xí)[16]、聯(lián)邦學(xué)習(xí)中公平性和魯棒性的權(quán)衡[17]。近鄰優(yōu)化在原理上與有偏正則化相同，其中文獻(xiàn)[18]中考慮有偏正則化的方法對(duì)FedAvg進(jìn)行重新參數(shù)化，提出FedProx，通過有偏正則化約束每個(gè)設(shè)備學(xué)習(xí)的本地模型更加接近于全局模型，并允許各參與訓(xùn)練的設(shè)備在本地執(zhí)行可變數(shù)量的工作，在異構(gòu)環(huán)境中提供了收斂的保證。由于FedProx 在優(yōu)化全局模型參數(shù)w時(shí)和FedAvg 方式相同，通過簡單平均本地上傳的模型參數(shù)來更新全局模型參數(shù)，導(dǎo)致全局模型收斂速度慢，缺乏直接對(duì)全局模型參數(shù)的優(yōu)化。

受小批量近似更新的元學(xué)習(xí)機(jī)制[19]的啟發(fā)，本文提出了基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法，在本地模型更新階段通過近鄰優(yōu)化約束本地模型更新更加接近于全局模型，在全局模型聚合階段通過求解近似全局梯度，利用梯度下降來更新全局模型參數(shù)。最終實(shí)現(xiàn)全局模型能夠在較少的通信輪數(shù)下達(dá)到更快更穩(wěn)定的收斂結(jié)果。

本文的貢獻(xiàn)主要體現(xiàn)在以下3 個(gè)方面：

1）區(qū)別于已有的方法，不在對(duì)全局模型參數(shù)進(jìn)行簡單平均。在全局模型聚合階段，通過利用本地上傳的模型參數(shù)近似求出平均全局梯度，同時(shí)也避免求解一階導(dǎo)數(shù)。

2）針對(duì)異構(gòu)性導(dǎo)致的全局模型收斂慢甚至無法收斂的問題，區(qū)別于現(xiàn)有的聯(lián)邦學(xué)習(xí)算法，本文提出基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法，通過隱式隨機(jī)梯度下降來更新全局模型參數(shù)，能夠使全局模型參數(shù)實(shí)現(xiàn)更加高效的更新，從而可以在有限的計(jì)算資源下加快模型的收斂速度。

3）和現(xiàn)有的工作相比，本文的算法在高度異構(gòu)的合成數(shù)據(jù)集上，30 輪左右就可以達(dá)到FedAvg 的收斂效果，40 輪左右可以達(dá)到FedProx 的收斂效果。在相同收斂效果的前提下，本文的算法比FedProx 減少了近50%的通信輪數(shù)。

1 客戶端-服務(wù)器的聯(lián)邦學(xué)習(xí)更新架構(gòu)

聯(lián)邦學(xué)習(xí)更新架構(gòu)主要有客戶端-服務(wù)器和去中心化對(duì)等計(jì)算架構(gòu)。其中最常用的是客戶端-服務(wù)器的聯(lián)邦學(xué)習(xí)更新架構(gòu)。訓(xùn)練過程主要分為兩個(gè)階段：本地模型更新階段和全局模型聚合階段。具體更新過程如圖1 所示。

圖1 客戶端-服務(wù)器聯(lián)邦學(xué)習(xí)架構(gòu)Fig.1 Federated learning architecture of client and server

1）本地模型更新

在本地模型更新階段，服務(wù)器首先隨機(jī)選取K個(gè)客戶端，然后服務(wù)器發(fā)送全局模型參數(shù)[[wt]]給被選客戶端，客戶端利用本地?cái)?shù)據(jù)并行執(zhí)行E個(gè)epoch 的隨機(jī)梯度下降，然后將更新后的模型參數(shù)經(jīng)過同態(tài)加密算法[20]加密，之后再上傳至服務(wù)器。

2）全局模型聚合

2 隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法設(shè)計(jì)

在本節(jié)中，主要介紹聯(lián)邦近鄰優(yōu)化算法和隱式隨機(jī)梯度下降優(yōu)化算法的關(guān)鍵要素。由于聯(lián)邦學(xué)習(xí)是通過大量設(shè)備與中央服務(wù)器協(xié)同學(xué)習(xí)一個(gè)最優(yōu)的全局模型，因此我們的最終目標(biāo)是最小化：

式中：wk是設(shè)備k在本地迭代過程中所得的近似最優(yōu)解；w是需要求解全局模型的最優(yōu)解；Fk(wk):=,每個(gè)設(shè)備本地?cái)?shù)據(jù)xk服從不同的分布 Dk,損失函數(shù)是預(yù)測值與真實(shí)值之間的差。式(1)包含兩方面的優(yōu)化過程：1）在本地模型訓(xùn)練階段，每個(gè)設(shè)備通過全局模型參數(shù)w學(xué)習(xí)一個(gè)本地近似最優(yōu)wk；2）在全局模型聚合階段，服務(wù)器通過各設(shè)備上傳的wk利用隱式隨機(jī)梯度下降來調(diào)整全局模型參數(shù)w，使w與所有wk的平均距離較小。具體的算法流程為：

在算法1 中，步驟4)～6)為本地模型訓(xùn)練階段，7)～9)為Server 全局模型更新階段，然后將更新后的模型參數(shù)發(fā)送給下一輪參與訓(xùn)練的設(shè)備。不斷重復(fù)以上過程，直至模型損失收斂。

2.1 聯(lián)邦近鄰優(yōu)化

在本地模型訓(xùn)練階段，主要在本地模型更新時(shí)引入帶參數(shù)的近鄰算子約束本地模型更新更加接近于全局模型，這種本地優(yōu)化算法被稱為Fed-Prox 算法[18]，每個(gè)設(shè)備k的本地目標(biāo)函數(shù)被重新定義為

式中：λ是一個(gè)約束本地模型和全局模型差異的超參數(shù)；wt表示在第t輪服務(wù)器聚合更新之后的全局模型參數(shù)。

2.2 基于隱式隨機(jī)梯度下降的全局模型更新優(yōu)化

由鏈?zhǔn)椒▌t可以得到：

所以?Gk(wt)=，式（4）展現(xiàn)了全局模型的梯度估計(jì)可以通過求解當(dāng)前任務(wù)的近似更新來計(jì)算。在第t輪，所選設(shè)備在本地?cái)?shù)據(jù)集上利用隨機(jī)梯度下降更新E輪后，求出近似最優(yōu)解。服務(wù)器通過式(4)可以計(jì)算出平均的全局梯度:

式中：St為K個(gè)設(shè)備的子集；t為當(dāng)前訓(xùn)練輪數(shù)；為按固定輪數(shù)衰減的學(xué)習(xí)率；ηgi為初始化學(xué)習(xí)率，在訓(xùn)練模型初期用較大的學(xué)習(xí)率對(duì)全局模型進(jìn)行優(yōu)化，隨著通信輪數(shù)的不斷增加學(xué)習(xí)率逐步減小，有效保證了全局模型在訓(xùn)練過程中能以較快的速度逐步趨于穩(wěn)定。更新后的wt+1作為下一輪訓(xùn)練的全局模型參數(shù)。

從式(3)～(6)推導(dǎo)過程很容易看出，本文提出基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法是直接對(duì)全局模型參數(shù)進(jìn)行優(yōu)化，而不是簡單平均所有設(shè)備上傳的本地模型參數(shù)作為更新后的全局模型參數(shù)。因?yàn)??Gk(wt)=，所以在服務(wù)器端只需通過就可以得到平均全局模型梯度，因此避免了求解一階導(dǎo)數(shù)，然后利用隨機(jī)梯度下降對(duì)全局模型參數(shù)進(jìn)行更新。相比于FedProx，本算法在信息比較冗余的情況下能更高效地利用有效信息。其次，在迭代的過程中也會(huì)很快收斂到最小值附近，加快模型的收斂速度。

3 實(shí)驗(yàn)與結(jié)果

為了驗(yàn)證本文提出的隱式隨機(jī)梯度下降優(yōu)化算法的有效性，本文在3 個(gè)真實(shí)數(shù)據(jù)集和3 個(gè)合成數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)，在分類和回歸任務(wù)上進(jìn)行評(píng)估，并與當(dāng)前具有代表性的解決異構(gòu)性問題的方法FedProx[18]以及經(jīng)典的FedAvg[1]算法進(jìn)行比較。

3.1 實(shí)驗(yàn)設(shè)置

在Linux 系統(tǒng)下，包括2 塊GeForce GTX 1 080 Ti 和1 塊GeForce GTX TITAN X 的服務(wù)器上進(jìn)行仿真實(shí)驗(yàn)，代碼使用Tensorflow 框架實(shí)現(xiàn)，基于Python3 來實(shí)現(xiàn)基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法。其中，訓(xùn)練輪數(shù)、每輪迭代次數(shù)、選擇設(shè)備數(shù)量、學(xué)習(xí)率等超參數(shù)設(shè)置如表1 所示。

表1 超參數(shù)設(shè)置Table 1 Setting of Hyperparameters

為了保證評(píng)估方法與結(jié)果的公平性，本文提出的方法與FedProx、FedAvg 使用了相同的本地求解器，在模擬系統(tǒng)異構(gòu)設(shè)置時(shí)，掉隊(duì)的設(shè)備數(shù)量分別設(shè)置為0%、50%、90%。生成合成數(shù)據(jù)集本文使用了和FedProx 類似的方法，通過式（7）生成本地?cái)?shù)據(jù)：

式中：W∈10×60；x∈60；b∈10。通過式（7）生成30 個(gè)設(shè)備的數(shù)據(jù)集，同樣每輪隨機(jī)抽取10 個(gè)參與訓(xùn)練。

3.2 3 個(gè)真實(shí)數(shù)據(jù)集和模型

Sent140[21]是一個(gè)Twitter 帶有表情的文本信息情感分類數(shù)據(jù)集，該任務(wù)使用的是一個(gè)兩層LSTM，包含256 個(gè)隱藏層單元，每個(gè)Twitter 帳戶對(duì)應(yīng)一個(gè)設(shè)備。該模型以25 個(gè)字符序列作為輸入，通過兩個(gè)LSTM 層和一個(gè)全連接層，每個(gè)訓(xùn)練樣本輸出一個(gè)字符。

MNIST[22]是一個(gè)0～9 手寫體數(shù)字識(shí)別數(shù)據(jù)集，在這個(gè)任務(wù)上利用邏輯回歸的方法研究手寫數(shù)字圖像分類問題。為了生成非獨(dú)立同分布數(shù)據(jù)，本文將數(shù)據(jù)隨機(jī)分布在1 000 個(gè)設(shè)備中，每個(gè)設(shè)備只有2 種數(shù)字。模型的輸入是28×28 維的圖像，輸出是0～9 這10 個(gè)數(shù)字的標(biāo)簽。

EMNIST[23]是MNIST 數(shù)據(jù)集的擴(kuò)展，包含0～9 數(shù)字和26 個(gè)英文字母的大小寫，構(gòu)成了更大難度的62 類手寫字符圖像分類任務(wù)，但在實(shí)驗(yàn)中只隨機(jī)抽取10 個(gè)小寫字母，每個(gè)設(shè)備分配5 個(gè)類，在這個(gè)任務(wù)上利用邏輯回歸的方法研究圖像分類問題。模型的輸入是28×28 維的圖像，輸出是a～j 這10 個(gè)類的標(biāo)簽。

對(duì)于以上所有數(shù)據(jù)集，客戶端的本地?cái)?shù)據(jù)分配遵循冪律分布[24]。本文在本地分配80%為訓(xùn)練集，20%為測試集。各設(shè)備數(shù)據(jù)集組成如表2 所示。

表2 設(shè)備數(shù)據(jù)集分布Table 2 Datasets distribution on devices

3.3 合成數(shù)據(jù)集實(shí)驗(yàn)結(jié)果分析

首先在第1 個(gè)實(shí)驗(yàn)中，為了驗(yàn)證本文的算法在異構(gòu)數(shù)據(jù)集上有更快的收斂速度，本文在3 組合成數(shù)據(jù)集上進(jìn)行實(shí)驗(yàn)，分別是Synthetic_0_0、Synthetic_0.5_0.5、Synthetic_1_1，從左到右數(shù)據(jù)異構(gòu)性逐漸增強(qiáng)，異構(gòu)性越強(qiáng)，對(duì)模型收斂影響越大。本文通過損失的減小速度和梯度方差[25]的變化來衡量模型的收斂速度，結(jié)果如圖2 所示。為了證明本文方法的公平性和有效性，約束項(xiàng)λ統(tǒng)一設(shè)置成相同的值。由圖2 訓(xùn)練損失和梯度方差可以看出，本文的方法在第30 輪左右達(dá)到了FedAvg 的收斂效果，在第40 輪左右達(dá)到了FedProx 的收斂效果，并且40 輪以后還在繼續(xù)收斂。梯度方差（variance of local gradient,VLG）越小表示越穩(wěn)定，收斂性越好。VLG 可表示為

圖2 合成數(shù)據(jù)集實(shí)驗(yàn)結(jié)果分析Fig.2 Analysis of experimental results of synthetic datasets

實(shí)驗(yàn)中，通過使所有設(shè)備執(zhí)行相同的工作量來模擬不存在系統(tǒng)異構(gòu)性的情況，隨著數(shù)據(jù)異構(gòu)性增強(qiáng)，全局模型收斂結(jié)果最終會(huì)趨于某個(gè)區(qū)間，因此本文取最后一半通信輪數(shù)的平均測試精度作為模型好壞的評(píng)判標(biāo)準(zhǔn)，在合成數(shù)據(jù)集上平均測試精度如表3 所示，可以看出本文提出的算法平均測試精度普遍高于FedProx 和FedAvg。

表3 合成數(shù)據(jù)集上平均測試精度Table 3 Average test accuracy on synthetic datasets %

3.4 真實(shí)數(shù)據(jù)集實(shí)驗(yàn)結(jié)果分析

在本實(shí)驗(yàn)中，為了驗(yàn)證本文提出的算法在高度系統(tǒng)異構(gòu)性和數(shù)據(jù)異構(gòu)性環(huán)境下的整體效果，本節(jié)在3 個(gè)聯(lián)邦學(xué)習(xí)常用真實(shí)數(shù)據(jù)集和一個(gè)合成數(shù)據(jù)集上比較不同算法的穩(wěn)定性和收斂效果，其中Synthetic_1_1 客戶端本地類別設(shè)置為5，實(shí)現(xiàn)在數(shù)據(jù)異構(gòu)性基礎(chǔ)上模擬不同系統(tǒng)異構(gòu)性的聯(lián)邦設(shè)置。

本文通過約束設(shè)備的本地工作量，使每個(gè)設(shè)備訓(xùn)練指定的E來模擬系統(tǒng)的異構(gòu)性，對(duì)于不同的異構(gòu)設(shè)置，隨機(jī)選擇不同的E(E<20)分配給0%、50%和90%當(dāng)前參與訓(xùn)練的設(shè)備。當(dāng)?shù)絷?duì)者為0% 時(shí)，代表所有設(shè)備執(zhí)行相同的工作量(E=20)。在指定的全局時(shí)間周期內(nèi)，當(dāng)E<20 時(shí)，F(xiàn)edAvg 會(huì)丟掉這些掉隊(duì)者，本文的算法和Fed-Prox 會(huì)合并這些掉隊(duì)者，不同的是本文在全局模型聚合階段會(huì)有效地使用合并掉隊(duì)者的模型參數(shù)，利用隱式隨機(jī)梯度下降對(duì)全局模型進(jìn)一步優(yōu)化。真實(shí)數(shù)據(jù)集上的訓(xùn)練損失如圖3 所示，從上到下3 行圖片分別代表0%、50%和90%的掉隊(duì)者。隨著迭代輪數(shù)的不斷增加，平均損失逐漸趨于穩(wěn)定，從圖3 中可以看出本文提出的算法的收斂速度明顯優(yōu)于Fedavg 和FedProx。

圖3 真實(shí)聯(lián)邦數(shù)據(jù)集實(shí)驗(yàn)結(jié)果分析Fig.3 Analysis of experimental results of realistic federated datasets

表4 給出了在高度異構(gòu)環(huán)境下模型的平均測試精度，從表中可以看出掉隊(duì)者為90%時(shí)，本文提出的算法的平均測試精度最高，其次是Fed-Prox。本文算法在MNIST 數(shù)據(jù)集上比FedProx高5%。實(shí)驗(yàn)中，在Sent140 數(shù)據(jù)集上通過設(shè)置相同超參數(shù)進(jìn)行比較不同算法運(yùn)行時(shí)間，在通信輪數(shù)為200 的情況下，F(xiàn)edAvg、FedProx 和本文所提算法運(yùn)行時(shí)間分別為67 min、108 min、108 min。

表4 高度異構(gòu)環(huán)境各算法平均測試精度Table 4 Average test accuracy of each algorithm in highlyheterogeneous environment %

4 結(jié)束語

本文提出了一種基于隱式隨機(jī)梯度下降優(yōu)化的聯(lián)邦學(xué)習(xí)算法。全局模型聚合階段不再是簡單的平均各設(shè)備上傳的模型參數(shù)，而是利用本地上傳的模型參數(shù)近似求出全局梯度，同時(shí)避免求解一階導(dǎo)數(shù)。利用隨機(jī)梯度下降對(duì)全局模型參數(shù)進(jìn)行更新，在信息冗余的情況下能更準(zhǔn)確地利用有效信息，隨著通信輪數(shù)不斷增加，全局模型會(huì)很快收斂到最小值附近。在3 個(gè)合成數(shù)據(jù)集和3 個(gè)真實(shí)數(shù)據(jù)集上的實(shí)驗(yàn)結(jié)果充分表明：該算法能夠在不同異構(gòu)環(huán)境中均表現(xiàn)出更快更穩(wěn)健的收斂結(jié)果，顯著提高了聯(lián)邦學(xué)習(xí)在實(shí)際應(yīng)用系統(tǒng)中的穩(wěn)定性和魯棒性。