面向通信成本優化的聯邦學習算法

鄭 賽，李天瑞*，黃 維

（1.西南交通大學 計算機與人工智能學院，成都 611756；2.云計算與智能技術四川省高校重點實驗室（西南交通大學），成都 611756）

0 引言

在過去幾年里，機器學習在人工智能應用領域迅速發展，這些機器學習技術尤其是深度學習［1］的成功，都建立在大量數據的基礎上。然而，我們在現實世界中遇到的數據往往是小規模的、碎片化的，例如，來自移動終端設備、物聯網設備的數據和大量分布在城市中的傳感器的數據都擁有這兩個特點。出于保護用戶隱私和數據安全的要求，簡單地將這些數據聚合在一起進行模型訓練是不可行的。2018 年，歐盟開始執行《通用數據保護條例》；2021 年11 月，我國開始正式實施《中華人民共和國信息保護法》，國內數據監管法律對數據隱私保護的監管也愈發嚴格。

聯邦學習（Federated Learning，FL）是一種解決上述問題的機器學習設置。聯邦學習這一概念由McMahan 等［2］于2017 年首次提出，最近的研究者們對聯邦學習提出了一個更寬泛和準確的定義［3］：聯邦學習是一種機器學習設置，其中多個客戶端在中央服務器或服務提供商的協調下協作解決機器學習問題。每個客戶端的原始數據都存儲在本地，不進行交換或傳輸；取而代之的是，使用實時聚合的模型更新來實現學習目標。聯邦平均（Federated Averaging，FedAvg）算法是一種典型的聯邦學習算法，主要包含兩個步驟：客戶端接收服務器發送來的全局模型并進行訓練得到本地模型；服務器接收來自多個客戶端的本地模型，通過加權平均這些本地模型得到一個新的全局模型，再將其發送回客戶端。雖然FedAvg 一定程度上解決了聯邦學習中的兩個核心問題：統計異質性問題和通信成本問題，但是，作為一個樸素的聯邦學習算法，它仍然有許多地方可以改進。

本文提出一個結合生成模型和深度遷移學習的聯邦學習算法FedGT（Federated Generative Transfer）。該算法只需要客戶端和服務器進行一輪通信，能大幅降低聯邦學習的通信成本。同時，由于客戶端得到了個性化模型，所以統計異質性問題也得到了一定程度的緩解。本文工作的主要特點在于引入生成模型和模擬數據在服務器構建全局模型，并且僅需一輪通信。此前的算法都只傳輸預測模型參數并進行聚合，而本文算法則通過傳輸生成模型來生成模擬數據。具體工作如下：

1）利用生成模型在服務器生成模擬數據來建立全局預測模型，可以在一輪通信下保證聯邦學習的最終性能。

2）使用深度遷移學習中的微調來進一步適應客戶端不同分布的數據，從而緩解統計異質性的問題。

3）在不同的數據集中使用不同的模型進行實驗，結果表明本文方法具有一定的通用性。

1 相關工作

1.1 聯邦學習

作為一種能夠解決數據孤島和數據隱私安全問題的機器學習設置，相關研究者已經將聯邦學習應用到了許多領域：Hard 等［4］將FedAvg 算法用于預測智能手機鍵盤輸入法的下一個單詞；史鼎元等［5］將聯邦學習應用在信息檢索領域，提出了聯邦排序學習算法；Muhammad 等［6］將聯邦學習應用在推薦系統。然而，聯邦學習這一技術的落地仍然面臨許多挑戰。

聯邦學習環境中的客戶端通常是來自現實世界的終端設備，例如移動電話、可穿戴設備和智能設備。由于時間、地理位置和用戶習慣等因素，這些設備上的數據是非獨立同分布（non-Independent and Identically Distributed，non-IID）的，這一問題被稱為聯邦學習的統計異質性［7］問題。統計異質性導致了一些基于數據獨立同分布假設的傳統分布式機器學習算法性能低下。為了解決這一問題，聯邦學習改進算法被相繼提出，例如，FedProx［8］是一個面向統計異質性的算法，它在每個客戶端的原有優化目標上增加了一個衡量本地模型和全局模型差異的L2 正則化項，使得模型在non-IID 數據上訓練更加穩定，提高了收斂速度。

由于分布在客戶端的數據是non-IID 的，這使訓練單個全局模型難以適用于所有客戶端，所以為每個客戶端構建個性化模型十分重要。個性化聯邦學習［9］算法通常與遷移學習［10］、知識蒸餾［11］、元學習［12］、多任務學習［13］等其他機器學習技術相結合。遷移學習使深度學習模型能夠利用在解決一個問題時獲得的知識來解決另一個相關問題，Wang 等［14］利用客戶端的本地數據對全局模型參數進行再次更新，從而得到客戶端的個性化模型。知識蒸餾通過讓學生網絡模仿教師網絡，將大型教師網絡中的知識提取到小型學生網絡中，減少了網絡的參數量。過擬合是聯邦學習個性化模型的一個重要挑戰，Yu 等［15］提出將全局模型視為教師、將個性化模型視為學生，通過知識蒸餾來減輕個性化過程中過擬合的影響。Li 等［16］提出了FedMD 算法，這是一個基于知識蒸餾和遷移學習的聯邦學習框架，它允許客戶端使用本地私有數據集和全球公共數據集單獨訓練個性化模型。在多任務學習中，多個相關任務被同時解決，模型通過利用各個任務的共性和差異性來達到更好的訓練效果。Smith 等［13］表明多任務學習是構建個性化聯邦模型的一種合理方式，提出了一個聯邦多任務學習框架來應對聯邦學習中與通信、掉隊和容錯相關的挑戰。

聯邦學習中的另一個核心挑戰是高昂的通信成本。現實世界中終端設備數量龐大，通信環境復雜等因素導致了聯邦學習的通信壓力非常大。為了減少在這種復雜環境下的通信量，面向通信成本優化的聯邦學習算法主要從這兩方面進行研究：減少客戶端和服務器的通信次數；減小每次通信中傳輸的數據規模。例如，Yao 等［17］在客戶端的原有優化目標上增加了基于最大均值差異（Maximum Mean Discrepancy，MMD）距離的加權差異項，通過MMD 距離來衡量全局模型和本地模型的差異，從而加速全局模型的收斂，減少訓練過程中的通信次數。而Caldas 等［18］受到常被用于防止模型過擬合的隨機失活算法［19］的啟發，提出了聯邦隨機失活算法。在每個全連接層上，該算法丟棄固定數量的全連接層參數，但保證相鄰兩層失活后的輸出矩陣的維度仍然能夠進行矩陣運算；而在每個卷積層上，該算法通過丟棄固定數量的卷積核來減少參數。在傳統的隨機失活算法中，失活后的模型仍然具有和失活前模型一樣的大小，而在聯邦隨機失活算法中，因為只傳輸激活的參數，所以能顯著減少每輪的通信量。除了減少通信次數和減小通信數據規模這兩個方面，異步通信［20-21］和通信拓撲優化［22］也是兩個重要的面向通信成本優化的研究方向，但是由于實現技術難度大等原因，目前這兩個方向的相關研究還較少。

1.2 生成模型

自動編碼器（AutoEncoder，AE）和生成對抗網絡（Generative Adversarial Net，GAN）［23］是兩種著名的生成模型。AE 和GAN 都有許多變體，如變分自動編碼器（Variational AutoEncoder，VAE）［24］、WGAN（Wasserstein Generative Adversarial Network）［25］、條件生成對抗網絡（Conditional Generative Adversarial Net，CGAN）［26］等。VAE由編碼器和解碼器組成，編碼器將數據樣本x編碼為隱層表示z，解碼器將隱層表示z解碼回數據空間，這兩個過程可以分別表示為：

VAE 的訓練目標是使重建誤差盡可能小，即使x和盡可能接近。VAE 損失函數如下所示：

其中DKL指的是KL 散度。

GAN 同樣也包含編碼器和解碼器，通常被稱為生成器網絡G(z)的解碼器將隱層表示z映射到數據空間，同時通常被稱為判別器網絡D(x)的編碼器將訓練一個代表數據真實性的概率y=D(x) ∈[0，1]，其中：y越接近1，代表x是真實數據的概率越大；y越接近0，則代表x來自生成器網絡G(z)的概率越大。

生成器網絡和判別器網絡被同時訓練：更新G的網絡參數來最小化ln( 1-D(G(z)))，更新D的網絡參數來最小化ln(D(x))，二者進行著一種兩方最大最小博弈（Two-player min-max game），其值函數為：

1.3 深度遷移學習模型

遷移學習可以解決機器學習中訓練數據不足的問題，它試圖通過放寬訓練數據和測試數據必須是獨立同分布的假設，將知識從源域轉移到目標域。遷移學習的基本方法可以分為四類：基于樣本的遷移、基于模型的遷移、基于特征的遷移和基于關系的遷移。基于特征和基于模型的遷移通常表現得更好，這也是目前大多數遷移學習工作的研究熱點。由于深度學習在許多研究領域取得了主導地位，研究如何通過深度神經網絡有效地轉移知識也變得至關重要，這類方法被稱為深度遷移學習。深度遷移學習可以分為四類：基于實例的深度遷移學習、基于映射的深度遷移學習、基于網絡的深度遷移學習和基于對抗的深度遷移學習。基于網絡的深度遷移學習主要通過重復使用在源域中預訓練的部分網絡來遷移知識，微調［27］就是一種基于網絡的深度遷移學習方法，它的主要思想如下：

深度神經網絡的淺層通常學習數據的一般性特征，但隨著網絡的深入，深層網絡更注重學習特定性特征。因此，當有一個完成訓練的模型時，可以通過凍結淺層網絡的參數并更新深層網絡的參數，將該模型快速應用于新的數據集或者訓練任務。

類似地，在聯邦學習中，全局模型學習一般性特征，而局部模型學習特定性特征。可以使用全局模型學習一般性特征，通過微調在每個客戶端上學習特定性特征，從而快速訓練本地模型。

2 單輪通信聯邦學習算法設計

在介紹本文提出的單輪通信聯邦學習算法FedGT 之前，先通過介紹FedAvg 算法來了解聯邦學習算法的基本流程。

FedAvg 算法的參與者有1 個中心服務器和N個客戶端，首先中心服務器要初始化模型參數w0，然后進行T輪迭代：服務器將模型參數發送給隨機選出的K個客戶端，客戶端接收到模型參數對其更新后發回服務器，最后服務器聚合各個客戶端的模型參數得到新的模型參數。

算法1 FedAvg 算法。

不同于FedAvg 算法的多輪通信，FedGT 算法僅在客戶端和服務器之間進行一輪通信。FedGT 算法主要包括三個步驟：各個客戶端利用本地數據訓練一個用于生成數據樣本的生成模型和一個用于推斷標簽的局部預測模型，然后將這兩個模型的參數發送給服務器；服務器利用各個客戶端的生成模型生成數據樣本，然后再用客戶端的預測模型給這些樣本打標簽，從而得到一個模擬數據集，服務器再利用該模擬數據集訓練一個全局的預測模型并發送給客戶端；各個客戶端收到全局預測模型后再次利用全局預測模型和本地真實數據訓練出個性化本地預測模型。FedGT算法流程如圖1所示。

圖1 FedGT算法流程Fig.1 Flowchart of FedGT algorithm

2.1 客戶端生成模型和預測模型

FedGT 算法會在服務器生成一個模擬數據集，并通過該模擬數據集訓練出能代表數據一般性特征的全局預測模型。據我們所知，生成模擬數據的方式通常分為兩類：一類為不考慮數據各個維度相關性的分布擬合法；一類為考慮數據各個維度相關性的神經網絡生成模型。

分布擬合法可以分為兩個步驟：選擇不同的數據分布，使用統計方法估計這些分布的參數值；確定哪個分布更加符合數據樣本，或者說，選出p值最大的分布作為最終分布。具體地，分布一般由四個參數定義：位置、規模、形狀和閾值。這些參數定義了不同的分布：位置參數規定了分布在X軸上的位置；規模參數決定了分布中的擴散程度；形狀參數使分布具有不同的形狀；閾值參數則定義了分布在X軸上的最小值。分布的參數可以用各種統計方法來估計。例如最大似然估計法通過最小化負對數似然函數值來求得對分布參數的估計值。然后可以采用例如Kolmogorov-Smirnov 檢驗的方式計算出各個分布在該數據上的p值，選擇出p值最大的分布即為最終確定的該數據的分布。

分布擬合法在采用不同分布擬合數據樣本時，是對數據的每一個維度進行單獨擬合，所以它并沒有考慮數據各維度的相關性；而現實世界的數據樣本各個維度存在極大的相關性，例如圖像數據的鄰近像素點，文本數據的上下文都說明數據各維度存在極大相關性。基于深度學習的一些生成模型考慮了數據各維度的相關性，例如典型的全連接層輸出特征各個維度在計算時都將輸入特征的各個維度乘以權重參數。

這些生成模型通常由編碼器和解碼器兩個部分組成，模型訓練的目標為最小化解碼器的重建誤差。在數據生成階段，只需將噪聲輸入到解碼器，解碼器即可輸出生成的模擬數據。VAE 和GAN 主要在編碼器和解碼器的設計上具有明顯差別：VAE 模型參數少，模型結構簡單，易于調試，但是在復雜數據上的表現不佳；而GAN 通常可以用于更加復雜的數據的生成，但其模型訓練過程參數的調試比較具有挑戰性。所以FedGT 算法使用VAE 在簡單的數據集上進行實驗，而在較復雜的數據集上則采用GAN 進行實驗。

建立生成模型可以得到模擬數據樣本X，然而僅通過沒有標簽的數據樣本無法建立預測模型，所以客戶端需要建立預測模型發送到服務器，從而得到服務器生成的模擬數據的標簽Y。本文選取了兩個圖像分類任務的數據集分別進行實驗，客戶端的預測模型分別為簡單卷積網絡（Simple Convolutional Neural Network，Simple-CNN）和修改后的ResNet-18［28］。

2.2 服務器全局預測模型

對于一個客戶端，服務器會接收到的內容包括：生成模型的解碼器Deci(z)，本地預測模型，數據數量numi。生成模型解碼器Deci(z)輸入噪聲后可以得到模擬數據樣本，生成的數量為numi，這些模擬數據會被輸入客戶端的本地預測模型得到對應的預測值。于是服務器在接受到來自一個客戶端的內容后，可以產生一個樣本數量為numi的子數據集。當全部N個客戶端都發送內容到服務器后，服務器得到N個子數據集后，將這些數據集合并為，使用該數據集訓練出全局預測模型Pglobal并將其發送給各個客戶端。這些步驟在算法2 中進一步介紹。

2.3 客戶端個性化預測模型

FedAvg 和FedProx 等算法在得到全局預測模型后算法就結束了，所以每個客戶端最終得到的模型是相同的全局模型；但是由于每個客戶端上的數據通常是non-IID 的，同樣的模型在某些客戶端上表現良好，但在另一些客戶端上會表現得很糟糕。對于這個問題，一個更好的解決方案是對于每個客戶端單獨訓練個性化模型。在FedGT 算法中，客戶端在接收到全局預測模型后，會利用本地的數據通過微調來得到個性化模型。具體地，客戶端會凍結網絡模型的淺層網絡參數，利用本地數據對深層網絡參數進行調整。

算法2 FedGT 算法。

3 實驗與結果分析

本文將FedGT 與FedAvg、FedProx 和集中式學習進行了比較。集中式學習是指所有客戶端將它們的數據傳輸到一個中央服務器，然后服務器使用這些數據建立一個全局模型，并將全局模型參數發回給客戶端，這種方法的一個重要問題是傳輸客戶端真實數據的同時伴隨著數據泄露的風險。

為了說明算法的通用性，分別在CIFAR-10 和MNIST 數據集上采用了不同的生成模型和預測模型進行實驗。MNIST 數據集包含了7 × 104張手寫數字的灰度圖像，所有圖像被分為10 類，分別為手寫數字0～9，并且每張圖像已經被標準化為28 × 28，訓練集包含6 × 104張圖像，測試集包含1 × 104張圖像。CIFAR-10 數據集由6 × 104張32 × 32 的彩色圖像組成，所有圖像被分為10 類，分別為飛機、汽車、鳥、貓、鹿、狗、青蛙、馬、船和卡車，每類有6 × 103張圖像，包括5 × 103張訓練圖像和1 × 103張測試圖像，數據集總共有5 × 104張訓練圖像和1 × 104張測試圖像。

3.1 在MNIST數據集上的實驗

在MNIST 數據集上使用VAE 作為生成模型，使用Simple-CNN 作為預測模型。Simple-CNN 的結構如表1 所示，除了表1 中提到的卷積層和池化層，使用ReLU 函數作為網絡的激活函數，并增加了隨機失活層防止過擬合。

表1 Simple-CNN的結構Tab.1 Structure of Simple-CNN

對于數據集在客戶端上的劃分，采用了IID 和non-IID 兩種方式。首先將訓練集和測試集合并得到7 × 104張圖像，然后設置了20 個客戶端，IID 的數據劃分方式為：每個客戶端互不重復地隨機選取3.5 × 103張圖像作為本地數據集，其中6/7 被作為本地訓練集，1/7 被作為本地測試集。而non-IID 的劃分方式為：將所有數據切分成40 份，其中每份數據只包含10 個類別中的1 個類別，每個客戶端互不重復地隨機選取2 份數據作為本地數據集，這保證了每個客戶端上至多有2 種類別的數據，同樣地，non-IID 劃分的訓練集-測試集比例為6∶1。

在FedGT 算法的步驟3 中，客戶端接收到全局預測模型進行微調時，凍結Conv1、Conv2 和FC1 這三層的參數，只更新FC2 這一全連接層的參數。使用Adam 作為優化器，使用交叉熵作為損失函數，每批次數據量為64，訓練輪數為500，學習率為5 × 10-4。

實驗選取了集中式學習，FedAvg、FedProx、FedMD 和FedDyn［29］作為基準算法與FedGT 算法進行比較。其中FedAvg 和FedProx 作為聯邦學習中較早發表的經典算法，常被用于各類聯邦學習算法實驗的基準算法。FedDyn 是目前最新的聯邦學習算法，它在客戶端本地模型更新的損失函數中加入了一個動態的正則器，使客戶端每輪的損失函數動態更新，同時使全局經驗損失和局部經驗損失的最小值保持漸進一致，FedDyn 算法的通信效率和準確率均優于FedAvg 和FedPox。FedMD 是個性化聯邦學習中的一個重要算法，它將知識蒸餾和遷移學習同時運用在聯邦學習訓練過程中，允許客戶端使用本地私有數據集和全球公共數據集單獨訓練個性化模型。

基準算法在訓練過程中使用隨機梯度下降（Stochastic Gradient Descent，SGD）作為優化器，使用交叉熵作為損失函數，每批次數據量為64，訓練輪次為200，初始學習率為0.1，每50 輪衰減80%。

表2 列出了所選基準算法和FedGT 算法在MNIST 數據集上的準確率，由于集中式學習是將數據集中到服務器上進行訓練，所以它在non-IID 劃分和IID劃分下的準確率是相同的。FedGT 算法在IID 和non-IID 數據劃分下的準確率超過了實驗對比的所有基準算法，然而不同于基準算法，FedGT 算法在non-IID數據劃分下的準確率比IID數據劃分的準確率高。

表2 不同算法在MNIST數據集上的準確率 單位：%Tab.2 Accuracies of different algorithms on MNIST dataset unit：%

3.2 在CIFAR-10數據集上的實驗

在CIFAR-10 數據集上使用GAN 作為生成模型，采用修改后的ResNet-18 作為預測模型。修改后的ResNet-18 殘差單元的結構如圖2 所示，圖中Conv1、Conv2 和Conv3 即為表3中的Conv1、Conv2 和Conv3。

表3 修改后的ResNet-18網絡結構Tab.3 Modified ResNet-18 network structure

圖2 修改后的ResNet-18殘差單元結構Fig.2 Structure of residual unit of modified ResNet-18

由于ResNet-18 網絡輸入圖像的尺寸為3 × 224 × 224，而CIFAR10 的圖像尺寸為3 × 32 × 32，所以對ResNet-18 進行了一定調整，主要改動為：將第一個卷積核為7 × 7 的卷積層以及一個最大池化層替換為一個卷積核為7 × 7 的卷積層，以此來適應CIFAR-10 的圖像尺寸，具體網絡結構參數如表3 所示。其中：Conv2d（3，1）代表卷積核為3 × 3、步長為1的二維卷積層；ResUnit(n，m，k1，k2，k3)代表一個如圖2 所示的殘差單元，Conv1、Conv2 和Conv3 的卷積核都為n×n，輸出通道都為m，步長分別為k1、k2 和k3；AvgPool（4，4）代表核為4 × 4、步長為4 的平均池化層；Linear（512，10）代表輸入為512 維向量、輸出為10 維向量的全連接層。

對于數據集在客戶端上的劃分，采用和在MNIST 數據集上相同的IID 劃分和non-IID 劃分策略，首先合并訓練集和測試集，客戶端數量為20，如果為IID 劃分則隨機選取，如果為non-IID 劃分則使每個客戶端至多有兩個類別的樣本，客戶端上訓練集和測試集的比例保持和原數據集一致，都為5∶1。

客戶端接收到全局預測模型進行微調時，凍結ConvIn、Layer1、Layer2 和Layer3 層的參數，只更新Layer4 和Linear。在FedGT 和基準算法上使用的實驗參數與在MNIST 上的實驗參數一致，表4 列出了實驗結果。基準算法在CIFAR-10上的準確率整體上低于MNIST 上的準確率，FedGT 算法的準確率在IID 和non-IID 數據劃分上都優于基準算法。

表4 算法在CIFAR-10數據集上的準確率 單位：%Tab.4 Accuracies of algorithms on CIFAR-10 dataset unit：%

3.3 通信效率和計算開銷

FedGT 算法的主要目的是減少聯邦學習中的通信輪數和通信數據量。通常一個聯邦學習模型需要進行E輪通信，每輪通信均傳輸預測模型參數；而FedGT 算法只需要一輪通信，客戶端發送給服務器生成模型解碼器和本地預測模型的參數，服務器發送給客戶端全局預測模型參數。表5 中對比了FedGT 和基準算法在兩個數據集上的通信量。由表5 可以看出：在MNIST 數據集上，FedGT 算法的通信量約為FedAvg、FedProx、FedDyn 算法通信量的1/10，約為FedMD 算法的1/100；在CIFAR-10 數據集上，FedGT 算法的通信量約為FedAvg、FedProx、FedDyn 算法的通信量的1/100，約為FedMD算法的1/10。

表5 MNIST和CIFAR-10數據集上的通信效率對比Tab.5 Comparison of communication efficiency on MNIST and CIFAR-10 datasets

聯邦學習通信成本的減少會帶來計算量的增加，本文以每秒浮點運算次數（Floating-Point Operations Per Second，FLOPS）為單位計算了模型在前向和反向傳播時的計算量，并考慮了損失函數不同對反向傳播過程中計算量的影響和服務器端聚合模型的計算開銷，最終得到的計算量如表6 所示，其中客戶端這一欄為單個客戶端的計算開銷。計算結果顯示，FedGT 算法增加了聯邦學習在服務器上的計算開銷，但減少了客戶端上的計算開銷。聯邦學習的客戶端通常為移動智能設備，只有當設備空閑時才能進行模型訓練，而這些設備計算能力遠不如計算能力強大的服務器，所以FedGT算法減少客戶端計算開銷、增加服務器計算開銷有利于聯邦學習的落地。

表6 MNIST和CIFAR-10數據集上的計算開銷對比 單位：FLOPsTab.6 Comparison of computing overhead on MNIST and CIFAR-10 datasets unit unit：FLOPs

4 結語

本文提出了一種新的聯邦學習算法——FedGT 算法。該算法采用服務器生成模擬數據來訓練全局模型，這一做法能夠將通信的輪數減少至一輪，并且還在客戶端通過模型的微調實現模型個性化來解決客戶端異質性問題。FedGT 算法在不同的數據集上使用不同的架構進行了實驗，結果表明它在數據以IID 和non-IID 方式分布時均優于FedAvg、FedProx、FedDyn 和FedMD 算法。

FedGT 算法為聯邦學習的研究提出了一個新的方向：除了傳輸目標訓練模型的參數，還可以傳輸其他信息（例如生成模型）來加快模型訓練，減少通信輪數，這是一個新穎且具有挑戰的研究方向。我們的下一步研究工作是在更廣泛的數據集上檢驗FedGT 算法，這樣才能真正捕捉到聯邦學習真實場景下的大規模分布的復雜性。