基于梯度統計的無線聯邦學習的功率控制優化

薛曉慧， 郭志華， 芮光輝， 厲娜， 馬曉琴

(國網青海省電力公司, 西寧 810008)

聯邦學習(federated learning, FL)是一種新的邊緣學習框架，使許多邊緣設備能夠在無線網絡中協同訓練機器學習模型[1]。與傳統的集中式數據中心學習相比，FL具有許多獨特的優點，如保護隱私、減少網絡擁塞、分布式計算能力等[2]。FL的主要瓶頸是通信成本，因為在每一次的模型訓練中，大量的參與邊緣設備都會向邊緣服務器發送更新模型，因此如何實現可靠通信是FL應用的關鍵[3]。

通過利用無線波形的疊加特性，空中計算(AirComp)可以加速模型聚合，從而提高通信效率[4]。考慮到信道的衰落，設備的選擇和功率控制是實現可靠通信的關鍵。文獻[5]通過聯合優化邊緣設備的發射功率和邊緣服務器的接收比例因子，從而使聚合信號的均方誤差(mean square error, MSE)最小化。結果表明，靜態信道中的最優發射功率呈現基于閾值的切換結構。即，如果每個設備的質量指標超過最優閾值，則應用信道反轉功率控制，否則應用全功率傳輸。然而，這項工作純粹是針對基于空中計算的信號聚合，忽略了上下文學習。文獻[6]引入了一種基于截斷的方法，用于排除具有深度衰落信道的邊緣設備，從而在學習性能和聚集誤差之間取得良好的平衡。文獻[7]一種聯合設備選擇和接收機波束形成的設計方法，以找到滿足均方誤差要求的最大選擇設備數，從而提高學習性能。文獻[8]通過利用梯度向量中的稀疏模式，將每個設備中的梯度估計投影到一個低維向量中，并且僅傳輸重要的梯度條目，同時累積先前迭代的誤差。另外，文獻[9]設計了類似于文獻[10]中的信道反轉的功率控制方案，使得從所選設備發送的梯度向量在邊緣服務器處對齊。然而，上述方法的空中飛行功率控制研究都忽略了梯度的統計特性：①梯度在訓練迭代過程中的分布不一定相同，即使在同一次迭代中，梯度向量的每個入口的分布也可能是不相同的。一個普遍情況實梯度分布隨著迭代而變化，并且在每個特征維度上是不同的。②如果每個設備的梯度分布未知，將梯度歸一化為零均值和單位方差的分布是不可行的。因此，由于忽略了梯度分布的上述特性，現有的空中飛行層功率控制方法在實際應用中可能效果較差。

基于上述問題，在考慮梯度統計量的基礎上，研究衰落信道中無線信道的最優功率控制問題。目標是在給定梯度的一階和二階統計量的情況下，通過聯合優化每個設備的發射功率和邊緣服務器的去噪因子，在每次迭代時最小化聚合模型的均方誤差，從而提高FL的精度。

1 系統模型

1.1 無線網絡的聯邦學習

考慮如圖1所示的無線FL框架，其中通過邊緣服務器的協調跨K個邊緣設備協作地訓練共享人工智能模型。設K={1,2,…,K}表示邊緣設備的集合，每個設備k∈K通過與自己的用戶交互收集一部分標記的訓練數據，構成一個本地數據集，表示為Sk，邊緣服務器是未知的。w∈RD表示要學習的D維模型參數。測量模型誤差的損失函數定義為

(1)

(2)

gk(t)、pk(t)、w(t)分別表示本地梯度、控制量、模型向量圖1 無線聯邦學習框架Fig.1 Wireless federated learning framework

1.2 梯度聚集的空中計算

梯度元素{gk,d(t)}(?k∈K)在設備k上獨立且服從同一分布，在聯邦學習中，本地數據集由于隱私問題，邊緣服務器通常不知道它們的實際分布。因此，在邊緣服務器中，從這些本地數據集訓練的本地梯度{gk,d(t)}的分布在設備之間權重一致。

梯度元素{gk,d(t)}(?t∈N)在迭代t上是不同分布的，另外梯度分布隨時間是非平穩的，非平穩分布是有效的，因為梯度值通常在開始時迅速變化，然后隨著訓練的進行逐漸接近零。

梯度元素{gk,d(t)}(?d∈{1,2,…,D})是獨立的，但在梯度向量維數D上是不同分布的，只要數據樣本中的特征是獨立的但分布不完全相同。

每次迭代t時邊緣服務器處的興趣梯度為

(3)

為了獲得式(3)，所有設備以模擬方式同時傳輸其梯度向量gk(t)，如圖1所示的空中計算原理。本文中考慮塊衰落信道，其中信道系數在FL中的每個迭代的持續時間內保持不變，但是可以從一個迭代到另一個迭代獨立地改變。本文中將一次迭代的持續時間定義為一個時間塊，以t∈N為索引，每個邊緣設備和邊緣服務器都配備有單個天線。設hk(t)表示在第t時間塊從設備k到邊緣服務器的復數信道系數，假設它是由平穩遍歷過程生成的。每個傳輸塊采用D個時隙的持續時間，一個時隙表示D維梯度向量中的一個條目，將每個梯度向量gk(t)乘以表示為bk(t)的預處理因子。邊緣服務器處的接收信號向量表達式為

(4)

(5)

將每個設備k∈K作為峰值功率預算Pk，即

pk(t)≤Pk, ?k∈K；?t∈N

(6)

在接收到y(t)時，邊緣服務器的去噪因子表示為η(t)，從而恢復梯度為

(7)

1.3 梯度統計

(8)

通常，由于每個設備上的峰值功率預算，梯度聚集式(8)的單個失調誤差和以MSE為單位的復合失調誤差不能同時為0。直接使用兩個誤差的權重，即梯度方差和梯度均值，很難在這兩個誤差之間找到平衡點。為了解決這個問題，引入了兩個可選的梯度統計參數。

(9)

令β(t)表示g(t)的平方多元變異系數(square multivariate coefficient of variation, SMCV)，表達式為

(10)

通過式(9)和式(10)，可以將式(8)中的MSE改寫為

(11)

由式(11)可知，當梯度MSNα(t)在單個和復合不對稱誤差的權重中，梯度SMCVβ(t)在MSE表達式中的作用更顯著。尤其當模型訓練剛開始時，即β(t)→0，如圖2所示，可以忽略單個信號未對準誤差。

ⅡD、non-ⅡD表示2維區域和非2維區域圖2 交替梯度統計的實驗結果Fig.2 Experimental results of alternating gradient statistics

2 已知梯度統計量的最優功率控制

當梯度統計量α(t)以及β(t)已知時，就解決了最小均方誤差的最優功率控制問題。為了方便，在本節中省略了迭代參數t。對于每個設備k∈K，用能量p和去噪因子η定義了聚集水平，即

(12)

C1≤…≤Ck≤…≤CK

(13)

2.1 一般情況下的功率控制問題

本節討論考慮一般情況下最小均方誤差的最優功率控制問題，將這個問題表述為

(14)

基于引理2，求解問題P1可以等價于在全局冪次區域的K個專屬次區域中最小化目標函數，表示為{Ml}l∈K，并比較其相應的最優解，得到全局最優解為

Ml={[p1,p2,…,pK]∈RK|pk=Pk,

?k∈{1,2，…,l};

0≤pk

(15)

?k∈{1,2,…,l};

pk≥0,?k∈{l+1,l+2,…,K}}

(16)

k∈{l+1,l+2,…,K}

(17)

(18)

將式(17)代入式(14)①并使其有關η的導數為0，可以得到最優的去噪因子η在第一松弛子區域中定義的問題P1的封閉形式為

(19)

(20)

定理1：解決問題P1的每個設備的最優發射功率的表達式為

?k∈{l*+1,l*+2,…,K}}

(21)

邊緣服務器的最優去噪因子的表達式為

(22)

(23)

(24)

將式(24)替換式(23)，可以得到

(25)

定理1：表明這些裝置k∈{1，2,…,l*}聚合能力并不比設備l*的高，應以全功率傳輸其梯度，即pk=Pk，而設備k∈{l*+1,l*+2,…,K}聚合能力高于設備l*，使用部分功率進行傳輸，以便它們具有相同的聚合級別，表達式為

(26)

2.2 最優發射功率函數

對于l∈K，可以找到每個域Xl的邊界和相應的最優發射功率p*。為此，需要以下關于最優發射功率值上下界的引理。

最后，可以得出最優發射功率函數在梯度統計量和噪聲方差方面的性質。

圖3 最優發射功率與與梯度的關系Fig.3 Relationship between optimal transmission power and gradient

2.3 特殊情況下的功率控制問題

在這一節中，證明了基于閾值的功率控制和全功率傳輸是最優功率控制策略的兩種特殊情況，其中梯度SMCVβ→∞、β→0。

(1)β→∞: 具體來說，在定理1中令β→∞，有如下推論。

推論1：最優發射功率β→∞具有基于閾值的結構，即

(27)

其中，最優去噪因子表達式為

(28)

l*的表達式見式(20)。

推論2：最優發射功率β→0表示全功率傳輸，即

(29)

最優去噪因子的表達式為

(30)

當β→0時，基于上述證明有l*=K。從邊緣服務器處的每個設備接收的梯度向量的方向獨立于發送設備的功率，因此，當復合信號失調誤差固定時，增加所有器件的功率可以減小噪聲引起的誤差。

3 未知梯度統計量的自適應功率控制

在本節中，討論梯度統計α(t)以及β(t)是未知的實際情況。估算每一個時間段的α(t)以及β(t)，通過估計P1問題的α(t)以及β(t)的最優解，設計了一種基于最優解的自適應功率控制方案。

3.1 參數估計

在本節中，本文提出了一種估算α(t)以及β(t)在每個時刻，時間塊t分別基于式(9)和式(10)中的定義。

(1)估計α(t)：設Bk(t)=‖gk(t)‖表示迭代t時設備k的梯度范數。在每一輪局部訓練結束時，讓每個設備在發送梯度向量gk(t)之前將其梯度范數Bk(t)傳輸到邊緣服務器。與發送梯度gk(t)的通信成本相比，發送Bk(t)的通信成本可以忽略不計，因為梯度范數Bk(t)是一個標量，而梯度gk(t)是一個維數D非常大的向量。根據定義式(9)，可以將梯度MSN估計為所有參與設備的梯度范數平方的平均值，即

(31)

(2)估計β(t)：根據式(10)中的定義，梯度SMCVβ(t)取決于md(t)和σd(t)。在每個設備在時間塊t發送其梯度之前，無法估計β(t)。然而，從圖2中真實數據集的實驗結果可以看出β(t)相鄰迭代之間高度相關。因此，本文中β(t)使用時間塊t處的聚集梯度t-1為

(32)

3.2 具有自適應功率控制的FL

本節提出了算法1中的具有自適應功率控制的FL過程。首先，每個設備在本地使用其本地數據集(算法1的第5步)在當前模型上執行一步SGD。之后，每個設備計算其局部梯度的范數，并通過常規數字傳輸(算法1的第6～7步)將其上傳到邊緣服務器。進一步邊緣服務器估計基于在時間塊t處接收到的梯度范數和歷史聚集梯度參數α(t)以及β(t)(算法1的第9步和第16步)，然后分別基于式(21)和式(22)得到最優發射功率和去噪因子(算法1的第10步)。最后邊緣服務器向每個設備通知最優發射功率，并且每個設備使用空中計算的模擬方式同時向邊緣服務器發射具有所分配功率的本地梯度(算法1的第12～第14步)。

在算法1中尋找最優功率控制的計算復雜度主要包括按聚合能力排序設備的時間復雜度，即O(KlgK)，并通過引理4尋找l*的時間復雜度，即O(K)。

4 實驗結果

4.1 實驗配置

(1)基線方法：將本文自適應功率控制方案與以下基線方法進行比較。

①無錯誤傳輸：聚合梯度完全更新，沒有任何傳輸錯誤，這相當于集中式SGD算法[15]。

②具有已知統計信息的功率控制：本文假設梯度統計信息在訓練開始時是已知的，然后應用所本文的功率控制。在本實驗中，實際的梯度統計是從1 000個梯度樣本中獲得的，沒有傳輸誤差[16]。

③在文獻[11]中基于閾值的功率控制：在文獻[11]中給出的功率控制方案中，它假設信號是標準化的。

④全功率傳輸：所有設備以全功率Pk傳輸，邊緣服務器應用式(22)中的最優去噪因子，其中l*=K。

(2)數據集：本文在MNIST、CIFAR-10和SVHN 3個數據集上評估了性能。MNIST數據集由數字0到9的10個類別組成，共有70 000個標記數據樣本(60 000個樣本用于訓練，10 000個樣本用于測試)。CIFAR-10數據集包括10種不同類型物體的60 000幅彩色圖像(50 000幅用于訓練，10 000幅用于測試)。SVHN是一個用于開發機器學習和目標識別算法的真實圖像數據集，對數據預處理和格式化的要求最低，包括99 289個標記數據樣本(73 257個用于訓練，26 032個用于測試)。

(3)數據分布：研究梯度對SMCVβ的影響為了獲得最優的發射功率，本實驗在移動設備之間模擬了兩種類型的數據集分區，即IID設置和非IID設置。對于前者，將訓練樣本隨機分成100個相等的碎片，每個碎片分配給一個特定的設備。對于后者，首先對數據進行數字標簽排序，將數據分成200個相等的分片，然后隨機分配2個分片給每個設備。

(4)神經網絡與超參數：在所有的實驗中，采用了卷積神經網絡(層數為11層)。超參數設置如下：動量優化參數為0.5，兩個全局聚合之間的局部更新數為1，將局部批大小為10，學習率為γ=0.01。

4.2 結果分析

圖4比較了3個數據集的測試精度，分別是IID數據集分區和非IID數據集分區，其中所有設備的平均接收SNR設置為10 dB。結果表明，與已知梯度統計量的方案相比，性能差距很小，說明本文的梯度統計量估計方法是有效的。實驗還發現，該功率控制方法的模型精度優于基于閾值的功率控制和全功率傳輸。從圖2中，可以發現平均梯度SMCVβ(t)在IID數據集分區中，它小于非IID數據集分區中的分區，并且隨著迭代次數的增加而增加。基于閾值的功率控制在IID分區或訓練開始時會發生嚴重的精度下降，這是因為在這種情況下，梯度SMCV很小，因此MSE由復合失調誤差控制。因此，僅考慮單個失調誤差的基于閾值的功率控制性能要差得多。此外，在非IID分區或訓練結束時，全功率傳輸的精度下降也很大，這是因為梯度SMCV很大，因此全功率傳輸方案無法最小化在這種情況下主導MSE的個別失調誤差。

圖4 3個數據集的測試精度Fig.4 Test accuracy of three data sets

圖5說明了在平均接收SNR=5 dB時具有非IID數據分區的MNIST的測試精度，結果表明，在低信噪比下，該方案的總體性能仍優于兩種基線方法。具體而言，全功率傳輸比基于閾值的功率控制方案具有更好的性能。這主要是因為當噪聲方差較大時，全功率傳輸能夠強烈抑制控制均方誤差的噪聲誤差。

圖5 非IID數據分區的MNIST的測試精度Fig.5 Test accuracy of MNIST for non IID data partition

最后，圖6比較了不同功率控制方案在不同設備數K下的測試精度。這里使用非IID分區的MNIST數據集，將所有設備的平均接收信噪比設置為SNRk=10 dB，并對50個模型訓練的結果進行平均。首先，可以觀察到，當參與裝置的數目K增加時，所有四種方案所達到的測試精度都增加，但當K足夠大時，不能得到進一步改善。特別是當K≥30時，所有方案的測試精度保持不變，因為邊緣服務器可以聚合足夠的數據進行平均。其次，該方案在整個K域內的性能優于基于門限的功率控制和全功率傳輸。當K較小時，即圖6中的K=4，全功率傳輸接近基于閾值的功率控制，但是由于缺乏用于減少失調誤差的功率自適應，性能隨著K的增加而降低。

圖6 設備數量的性能比較Fig.6 Performance comparison of equipment quantity

5 結論

綜合考慮梯度統計量的影響，本文提出了一種衰落信道條件下基于無線聯邦學習的最優功率控制方法。通過3個數據集上的實驗結果可以得出如下結論。

(1)本文控制方法能夠峰值功率限制下，使空中計算的聚合誤差最小化，具有較好的控制性能。

(2)因為當噪聲方差較大時，全功率傳輸能夠強烈抑制控制均方誤差的噪聲誤差，因此全功率傳輸比基于閾值的功率控制方案具有更好的性能。全功率傳輸接近基于閾值的功率控制，但是由于缺乏用于減少失調誤差的功率自適應，性能隨著的K增加而降低。

(3)各器件的最優發射功率隨梯度SMCV的增大而減小，隨噪聲方差的增大而增大。