基于脈沖神經網絡的無線空中聯邦學習

楊瀚哲，游家偉，文鼎柱，石遠明

（上海科技大學,上海 201210）

0 引言

隨著移動通信、物聯網技術的發展，無線網絡中的移動設備數量呈現指數級的增長，這些設備所擁有的海量計算和數據資源極大地推動了邊緣人工智能在網絡邊緣的部署，從而為移動用戶提供具有高可靠性、低延遲的智能服務，由此形成了一個新的研究領域，稱為邊緣機器學習[1-2]。

在眾多的邊緣機器學習框架中，應用最為廣泛的是分布式的聯邦邊緣學習（FEL,Federated Edge Learning）[3-7]，它不需要用戶上傳自己的數據進行集中式訓練，而是在每一輪的訓練中，由一個中心服務器將全局機器學習模型廣播給所有的設備，設備在本地訓練該模型，得到局部梯度或新的局部模型并將該局部梯度/ 模型傳回給中心服務器更新全局模型。這很大程度上保護了用戶的隱私安全。但是，當前的聯邦學習訓練的往往是傳統的深度神經網絡模型，比如卷積神經網絡。這些模型的訓練需要在移動設備上消耗大量的能量、占據相當的計算資源。這大大限制了聯邦邊緣學習在邊緣網絡中的應用，尤其是當邊緣設備為計算和能量資源受限的物聯網設備時。除此之外，由于機器學習模型往往具有大量的參數需要傳輸，聯邦學習需要消耗龐大的通信資源。

一方面，為了解決邊緣設備能量和計算資源受限的問題，本文考慮在聯邦學習系統中訓練一種新的網絡模型——脈沖神經網絡（SNN,Spiking Neural Network）。脈沖神經網絡是一種更加仿生的網絡，它比起傳統神經網絡更加的節能[8-9]，并且有專屬的硬件設備為其訓練加速[10]。文獻[11]的作者研究了基于脈沖神經網絡的聯邦學習，該文獻指出，在聯邦學習框架下，脈沖神經網絡是可行的，而且比傳統神經網絡表現更加良好。另一方面，為了降低聯邦學習訓練的通信開銷，本文將利用空中計算技術進行模型參數聚合。

為了提升聯邦學習的性能，本文致力于最小化由空中計算所造成的模型參數失真，該問題為一個非凸的二次約束二次規劃問題[12]。通過使用分支定界算法，本文提出了一種最優解決方案。最后，在基于CIFAR10 數據集的實驗表明，所提出算法的性能優于其他已有算法，例如SDR[13]。

1 模型介紹

1.1 脈沖神經網絡

考慮一個有J個隱藏層的脈沖神經網絡，用Lj表示該網絡的第j個隱藏層，其中的第i個神經元表示為ni,j。脈沖神經網絡中每個神經元的輸入與輸出均為脈沖信號，如圖1 所示。在這個網絡中，每個輸入數據被編碼成長度為S的脈沖信號，并且逐個地輸入進脈沖神經網絡中。當2 ≤s≤S時，每個神經元的第s個時刻的脈沖信號由該神經元前一層的第s-1 個時刻的脈沖信號求得。把第s個脈沖信號所處時刻稱為時刻s。在本文中，使用帶泄漏整合發放模型（LIF,Leaky-Integrate-and-Fire）節點當作脈沖神經網絡的神經元。

圖1 脈沖神經元示意圖

在一個訓練輪次中，脈沖神經網絡使用S個脈沖信號進行訓練。在前向傳播中的任意時刻，比如時刻s，每個神經元聚合所有的上一層的神經元在時刻s-1 的輸出脈沖信號，并與該神經元本身在時刻s-1 的神經元隱狀態（通常稱作膜電位，用ui,j(s)表示神經元ni,j在時刻s的膜電位）聚合，得到該神經元在時刻s的脈沖輸出（用oi,j(s) 表示神經元ni,j在時刻s的脈沖輸出）。從而，神經元ni,j在時刻s的膜電位可以表示為：

其中，λ是衰減因子，i′代表Lj-1層的神經元的索引，wi′,i是對應關聯的權重參數。對應的脈沖輸出可以表示為：

如果在s時刻，該神經元的輸出脈沖信號為1，則該時刻該神經元的膜電位會被重置為某一設定好的定值ur。

脈沖神經網絡的反向傳播與RNN 相似。考慮任意一個權重wi′,i，它連接了神經元ni′,j-1與ni,j。每個訓練輪次中，它的梯度可以通過將所有時刻的梯度累加起來得到，即可以表示為：

其中L是損失函數。式(3) 中的主要的困難在于oi,j(s) 是不可微分的函數。為了解決這個問題，使用了替代梯度法[11]，即用閾值函數的逼近函數的梯度替代該梯度：

其中ξ是衰減因子，取決于脈沖信號的總長度S。對于很大的S，ξ應該是足夠小的，以防止梯度爆炸問題，反之亦然。

1.2 聯邦學習模型

考慮一個擁有多天線的邊緣服務器與k個單天線設備的聯邦學習系統。該系統的目的是通過共同訓練一個脈沖神經網絡來優化全局損失函數L，用公式表示為：

其中，w是脈沖神經網絡的參數向量，L(w) 是用于分類任務的交叉熵損失函數，Dk與Lk(w)分別為設備k的本地數據集與局部模型的交叉熵損失函數，表示設備k的本地數據集的數量。具體而言，考慮任一訓練輪次t，邊緣服務器首先廣播全局模型的參數wt-1給每個邊緣設備。然后第k個設備基于收到的模型參數wt-1，用本地數據集Dk根據式(3) 與(4) 更新本地模型，并利用式(6) 計算本地梯度更新。

其中，gk(w) 為t輪次的設備k的本地梯度更新，wk,t為t訓練輪次設備k完成本地訓練后的本地模型參數，wt-1為t訓練輪次開始時邊緣服務器廣播給各個邊緣設備的全局模型的參數。最終，邊緣服務器收集所有設備的本地梯度更新{gk(w)}，并用式(7)更新全局模型：

系統會重復上述步驟直至全局模型收斂。

1.3 空中計算模型

在每一輪次的訓練中，服務器利用空中計算技術來聚合所有的本地模型參數，其接收信號可以表示為：

其中，hk∈CN為設備k與邊緣服務器之間的上行信道向量，N為邊緣服務器中的天線數量，bk∈C為傳輸標量，s k∈C為設備k的傳輸標志，是整合過后且歸一化的本地梯度更新，假設它服從標準復正態分布，n～CN(0,σ2I)是高斯噪聲。在式(8) 中，傳輸標志sk是復數標量，因此，每個標志可以分解為本地梯度更新中的兩個元素：一個組成實部而另一個組成虛部。更多地，整個本地梯度更新向量的傳輸分布在多個資源塊中，每個資源塊中空中計算的設計是一樣的。一方面，為了加快式(7) 中的梯度聚合，需要對傳輸標志進行預處理與后處理，過程為：

（1）預處理：設備k的本地梯度更新向量中的每個元素在傳輸前被除以一個因子因此，邊緣服務器中沒有任何誤差的目標梯度聚合可以表示為：

（2）后處理：邊緣服務器收集到的聚合信號需要進行處理以得到平均梯度，表示如下：

另一方面，在邊緣服務器中，接收波束成形被用于預估式(9) 中的目標梯度聚合。加上波束成形后，邊緣服務器中的接受信號可以表示為：

其中，m∈CN為接收波束成形向量，η為歸一化因子。從而，式(7) 中的平均梯度聚合可以表示為：

2 空中計算與聯邦學習

2.1 空中計算問題構建

本文的目標為目標梯度聚合與預估梯度聚合的最小化均方誤差（MSE,Mean-Square Error），即式(9) 中的g與式(11)中的g～，可以表示為：

其中，函數MSE(x,y) 表示x與y的均方誤差。此外，每個設備有一個傳輸能量約束：

其中，P0為每個設備的最大傳輸能量。

最終問題可以概括為：

2.2 問題簡化

為了解決非凸問題P1，需要進行兩步簡化。首先，傳輸標量{bk} 可以被接收波束成形向量m表示。然后，問題可以被進一步簡化成一個二次約束二次規劃問題。第一步，根據參考文獻[7,13]，對于任意一個給定的m，可以得到最小化均方誤差的最優的傳輸標量，表示如下：

然后，考慮傳輸能量約束，可以進一步得到：

將式(17) 中的η與式(16) 中的bk帶入式(13) 中的最小均方誤差函數，可以得到：

其中，σ2是噪聲的能量。由于σ、P0都是常數，P1 中的最小均方誤差與下式是等價的：

第二步，根據參考文獻[14]，式(19)中的問題可以進一步寫成：

問題P2 是一個帶著一個非凸約束的二次約束二次規劃問題，這個問題是一個NP-hard 的問題。為了解決這個問題，文獻[13]的作者們使用了一個基于SDR 與SCA的算法。但是，文獻[13]中的算法只能得到一個次優解。本文中，參考文獻[12] 采用分支定界法（BnB,Branch and Bound），它能保證得到問題P2 的一個全局最優解。

2.3 使用分支定界算法解最小均方誤差

首先引入一個輔助變量x=[x1,x2,…,xK]T∈CK。使用該輔助變量可以將問題P2 轉化成如下的等價的形式：

（1）第一步：集合X被分割成幾個子集，每個子集對應著一個與P3 一樣形式的子問題。

（2）第二步：解出每個子問題，以求出原問題的解的上下界。

（3）第三步：如果原問題的上下界收斂到同一個值，則原問題的最優解已找到。否則，在每個子問題上重復步驟1 與步驟2。重復以上步驟直至收斂。上述三個步驟的具體過程為：

其中，為集合S的凸包。用問題P3.1i的解作為問題P3i的一個下界，因為取問題集合中 最小的下界作為當前問題P3 的下界Lt。為了得到上界Ut，可以選擇任意一個在問題P3i的可行域內的解當作一個上界。在這里，放縮問題P3.1i的最優解來生成一個可行的問題P3i的解的上界，如下所示。

把他們加入問題集合{P3i}，并且移除問題P3it。當然，也需要求得這兩個問題的解的下界。在每個迭代輪次，得到新的問題集合{P3i}，并且得到該問題集合的上界Ut與下界Lt。停止迭代直到并選擇與Lt對應的m作為問題P3 的解，否則繼續迭代。

3 實驗結果

3.1 實驗參數設定

本文使用的神經網絡模型為18 層殘差結構的脈沖神經網絡（即將標準ResNet-18 網絡的全部激活函數替換為SNN 神經元）。每個訓練輪次中，設定總的時刻數為10，并使用泊松編碼器對CIFAR10 數據集進行脈沖信號化。為了簡化問題，所有設備使用相同的學習率0.01。使用batch-SGD 訓練神經網絡，batch-size 為128。每個局部模型將會在邊緣設備本地迭代5 輪。

與此同時，通信模型使用三維模型設定。邊緣服務器被設定為位于坐標(0,0,20)，而其他設備均勻分布在以(120,20,0)為中心、半徑20 m 的圓形區域內。距離相關的大規模衰弱被設定為其中T0是參考距離為d0=1 m 時的路徑損失，d表示傳輸設備與接收器之間的距離，α是路徑損失指數。除此之外，小規模衰落被設定為Rician 因子為β的Rician 衰落。平均500 個信道模擬的結果得到最終的信道向量。設定α=3，T0=-30 dB，β=3，

3.2 分支定界算法的有效性

圖2 展示了分支定界算法的收斂性。可見，分支定界算法在300 輪迭代內達到收斂。

圖2 分支定界算法的收斂性

之后，將分支定界算法與廣泛使用的SDR 算法進行比較，如圖3 所示。使用的指標為空中計算噪音范數（NOAN,Norm of AirComp Noise），定義如下：

其中，g為目標函數，為目標函數的估計值，該估計值使用如上文所述方法生成信道向量hk，并使用分支定界算法解出波束成形向量m。如圖3 所示，分支界定算法的NOAN 一直比SDR 算法的要更小，證明分支定界算法的效果要比SDR 更好。

圖3 分支定界算法與SDR的空中計算噪聲范數比較

3.3 空中計算輔助脈沖神經網絡的聯邦學習框架的性能評估

空中計算輔助基于脈沖神經網絡的聯邦學習框架的性能如圖4 所示。

圖4 空中計算輔助基于脈沖神經網絡的聯邦學習框架性能評估

將本文使用的分支界定算法與標準模式和SDR 算法進行比較。其中，標準模式指直接使用目標函數值作為空中計算的結果，而SDR 算法指使用SDR 算法解問題P2 所得到的空中計算結果。由于分支定界算法得出的解是全局最優解，所以當σ2很小的時候訓練曲線更加接近標準模式。然而，隨著σ2增加，收斂速度變得很慢，并且最終準確率也比標準模式要低。另一方面，與只能求得局部最優解的SDR 對比，本文的算法在收斂速度與準確率上明顯表現更好。

4 結束語

本文中，研究了基于脈沖神經網絡的無線空中聯邦學習框架，并提出了相應的解決算法。由于脈沖神經網絡相較于傳統神經網絡計算效率更高且更加節省能量，所以該框架比基于傳統神經網絡的聯邦學習框架更加具有實踐意義。同時，使用的空中計算技術也很大程度上降低了通信消耗，彌補了聯邦學習在通信傳輸資源消耗方面的缺陷。大量仿真結果顯示，所提出的全局最優的分支定界算法各方面表現都比局部最優的SDR 算法要好，并且基于分支定界算法的收斂速度與準確率都優于基于SDR 的算法。在未來的工作中，可以考慮比較不同網絡模型（如CNN、SNN）在不同應用場景下的性能下界，從而根據場景選擇合適的模型。總而言之，SNN 的引入為6G 聯邦學習帶來了更多的可能性，在未來也會有越來越多的工作對基于SNN 模型的聯邦學習應用進行研究。