PaperWeekly - ECCV 2022 | FedX: 在無監督聯邦學習中進行知識蒸餾－鑽石舞台

©PaperWeekly 原創 · 作者 | 張成蹊

單位 | Freewheel機器學習工程師

研究方向 | 自然語言處理

論文的立意是在聯邦學習過程中，結合對比學習，用局部與全局的兩種蒸餾方式來無監督地學習樣本的向量表示。FedX 不僅是一種創新的無監督學習算法，更是一種可以熱插拔用於傳統對比學習，使之效果提升的無監督損失函數。

關鍵詞：聯邦學習（Federated Learning），對比學習（Contrastive Learning），知識蒸餾（Knowledge Distillation），無監督學習（Unsupervised Learning）。

論文標題：

FedX: Unsupervised Federated Learning with Cross Knowledge Distillation

論文鏈接：

https://arxiv.org/abs/2207.09158

代碼鏈接：

https://github.com/Sungwon-Han/FEDX

相關背景

聯邦學習假設數據分散在不同的客戶端中，且不能與中央服務器共享；截止目前，已經有很多有監督聯邦學習用在下游任務的相關研究，如目標檢測、圖像分割、人員重識別等，但對於無監督聯邦學習的需求也在逐步增加，以應對一些本地客戶端缺失數據的場景。其中比較具有挑戰性的研究點為：

數據在不同客戶端是離散且異構的，不滿足數據獨立同分布的條件，所以聯邦學習算法必須考慮到獨立客戶端與全局數據分布不同時的有偏情況；

在數據有偏到一定程度，客戶端可能只擁有較少的數據，此時需要引入無監督/半監督的聯邦學習方法，使得每一個客戶端都從全局模型中獲益，達到較好利用本地數據的目的。

聯邦學習的整體流程如下圖所示：

其中：

Local Update：每個client各自學習自己的樣本表示，獲得各自獨立的本地參數模型；

Local model update：將本地參數模型上傳到中心服務器；

Global aggregation：中心服務器對所有 client 上傳的模型進行聚合，最簡單的方式是（假設本地模型都是同構的）根據 client 各自擁有樣本數量的占比進行加權平均；

Global model download：每個 client 將 global model 各自下載到本地。

論文主要聚焦討論的是步驟 1 中的獨立學習部分，也是數據偏差最大的部分。

總的來說，論文具有以下幾方面的貢獻：

提出了無監督的聯邦學習算法 FedX，使用了局部與全局兩個階段的蒸餾來學習不同客戶端的數據表示；兩階段的知識蒸餾能夠發現有意義的本地數據向量表示，同時通過全局的知識蒸餾消除不同客戶端之間的偏差；

FedX 是一個可以被用於現有聯邦學習擴展的算法，提升現有聯邦學習算法的性能（1.58-5.52pp），同時對不同客戶端的數據嚴格保密，不（像部分現有的無監督聯邦學習算法一樣）直接在不同客戶端之間進行數據的共享；

FedX 已開源。

局部與全局蒸餾

以 SimCLR 為例，傳統對比學習的損失函數如下所示：

其中，與分別表示原批次的訓練樣本，和基於原樣本進行了數據增強後的訓練樣本，即從一個真實樣本中採用窗口看到的某個切片圖。表示某種對兩個向量的相似性度量方法（比如餘弦相似度，值得注意的是，作者在論文中也提到了餘弦相似度，不過在實際的代碼實現里用的是兩個正則化（torch.F.normalize）後的向量內積），類似知識蒸餾中用於控制程度的超參（temperature）。

可以容易地看出，對比學習損失函數的目標是：對於樣本，已知模型最後一個隱層對於它的向量表示，及其增強後的樣本，希望兩者的向量更加接近，其相似度更高；而對於兩個批次中其他個負樣本，希望其向量相似度越低。

2.1 局部蒸餾

作為一個獨立的本地客戶端，其學習的目標是基於自己擁有的本地數據，獲得數據有意義的向量表示。

作者認為，局部模型可以通過提取不同樣本間的結構化知識，實現局部知識傳遞的目標。在整個過程中，從自身學習到的知識逐步地由不同的局部樣本傳遞給自身。在本地模型更新時，除了每個樣本和自己的增強後樣本進行對比學習，兩者還可以從彼此的分布中進一步學習，謂之為關係損失（relational loss）。下式是關係損失的形式化定義：

該公式表示了正樣本的隱層向量對於一個隨機採樣的批次中任意一個樣本的隱層向量的概率分布。與之對應，正樣本在對比學習中的增強樣本隱層向量與後者也會有一個對應的向量分布：

這兩個分布代表了一組相似的正樣本關於其它負樣本的相對分布情況。容易發現，如果模型真的能辨認出增強前後的樣本歸屬於同一個樣本（或者標籤），那麼，除了對比學習本身設計的學習目標 —— 這兩個樣本間的相似度高於同批次中其他負樣本的相似度——以外，這兩個樣本對於某個隨機的負樣本的相似度分布也應該相似。

但是這個相似度要怎麼度量呢？作者提出將這兩個概率向量相加作為兩者要學習的目標，即：

其中指的是 KL 散度，也即相對熵（其實就是真實分布不是 one-hot 前提下的交叉熵）。

正如上面所述，這個損失函數最小化的目標是兩個概率分布對於其分布均值之間的距離。舉個具體的例子，比如正樣本及其增強分別代表下圖中的同一隻狗的兩種不同的視角：

那麼，對於另一個隨機採樣的樣本（比如貓的某個視角），正樣本的向量和隨機採樣樣本的向量的相似度，預期上應該是低於正樣本的向量和另一張狗的隨機採樣樣本的相似度的；而作為正樣本的增強樣本，它對於不同的隨機樣本，也應該具有這樣相似的特性。

最後，作者將對比學習的損失函數與關係損失函數加和，得到了聯邦學習的步驟 1 中的最終損失函數：

2.2 全局損失

除了局部的損失函數外，作者也引入了全局的損失函數。其出發點為：局部模型不但要從自己獨立的樣本中來學習分布，自然也要從全局模型中學習信息，所以在全局上提出了與局部損失類似的全局損失函數。唯一的差別在於，全局損失函數的關係損失函數中，隨機採樣的負樣本的表示向量從全局模型中取得（從作者的代碼中可以發現，此時全局模型的參數是凍結的，只用於提供特徵向量）。

全局模型主要目的是解決局部數據和全局數據的分布不一致所造成的的偏差，從而避免局部模型從上述過程中往局部最優收斂。

在聯邦學習的過程中，步驟 4 將全局聚合模型下載的不同客戶端本地，通過該全局模型來試圖對自身擁有的局部模型糾偏。值得注意的是，全局模型的結構可能與局部模型不一致，為了使得知識蒸餾能夠有效進行，在全局蒸餾過程中，作者在局部模型中添加了一個額外的預測層（在代碼中體現為一個兩層的全連接，維度均為）。

基於此，全局的對比學習損失表示為：

其中上標表示該樣本的特徵表示經過了局部模型的額外預測層，而上標表示該樣本經過了全局模型的最終預測層（softmax 的前一層輸出）。在全局的對比學習損失中，旨在通過全局模型對局部樣本進行糾偏。

與此對應，樣本對於其它負樣本的相對分布也分別有局部模型和全局模型產生：

最終的 FedX 損失函數由局部損失和全局損失加和得到，即：

FedX效果評估

作者在 CIFAR-10，SVHN，F-MNIST 三個數據集上進行了實驗，使用 Dirichlet分布模擬聯邦學習中各個 client 分得的樣本分布，採用了 SimCLR，MoCo，BYOL，ProtoCL，FedU 一共 5 個方法作為基線模型，將各自加入關係損失後的 + FedX 方法作為提出的方法，在三個數據集上分別取得了 4.29pp，5.52pp 和1.58pp 的增長。

作者也進行了一系列的有效性實驗，以證明 FedX 在不同場景，不同設定下的有效性。

作者首先通過了不同模塊的消融實驗來驗證模塊有效性（本地關係損失函數+全局關係損失函數+一個額外添加的兩層全連接層變換）。下圖的結果表明關係損失函數、全局蒸餾，以及為了匹配全局蒸餾而增加的額外全連接層都對模型有着積極的影響。而去除全局知識蒸餾過程會使得整個聯邦學習最終準確率有着非常大的下降，從而從側面證明了在局部樣本有偏的情況下，不通過全局模型進行糾偏，會大概率使得局部模型陷入局部最優，從而進一步對全局模型產生比較大的負面影響。

在參數敏感性實驗中，作者控制了不同的數據量，不同的 client 數量與不同的本地-全局模型交互次數，以探究 FedX 在不同的實驗設定下是否會出現不穩定的情況 —— 實驗結果表明其在不同的數據量、不同的訓練輪次，以及不同的客戶端數量下，都能給現有的聯邦學習算法帶來比較顯著的提升。

除此以外，作者還進行了一個比較新穎的 Embed 空間實驗：通過引入反餘弦函數來量化全局模型與局部模型向量間的分布情況，預期上類內更加相似，類間更加不同。FedX 的表現如預期一樣，使得類內的角度偏向更大的夾角。

值得注意的是，在論文的最後，作者也提出了將 FedX 擴展到半監督學習的場景中的設定，在這個場景下，首先用未標註數據對模型進行參數的初始化，然後再在模型上層添加一個全連接層，使用標註數據再額外訓練一定的輪次做微調。發現 FedX 在半監督的算法場景下也能帶來一定的提升，進一步證明了 FedX 不僅能夠用於無監督場景下的聯邦學習，也能用於少數據場景下的半監督學習。

更多閱讀