超越KL！大連理工發布Wasserstein距離知識蒸餾新方法｜NeurIPS 2024

原標題：超越KL！大連理工發布Wasserstein距離知識蒸餾新方法｜NeurIPS 2024
文章來源：新智元
內容字數：7110字

基于Wasserstein距離的知識蒸餾方法：超越KL散度的知識遷移

本文總結了大連理工大學研究人員發表在NeurIPS 2024上的論文，該論文提出了一種基于Wasserstein距離 (WD) 的知識蒸餾 (KD) 方法，在圖像分類和目標檢測任務上取得了顯著成果，并超越了傳統的基于Kullback-Leibler散度 (KL-Div) 的方法。

1. 傳統KL-Div方法的局限性

傳統的基于KL-Div的知識蒸餾方法，盡管在Logit蒸餾中取得了成功，但存在兩個主要缺陷：首先，KL-Div只比較教師和學生模型在相同類別上的概率，忽略了類別間的相互關系；其次，在Feature蒸餾中，KL-Div難以處理高維稀疏的深度特征，無法有效處理不重疊的分布，也無法捕捉底層流形的幾何結構。

2. 基于Wasserstein距離的知識蒸餾方法 (WKD)

為了克服KL-Div的局限性，研究人員提出了WKD方法，該方法包含兩個部分：WKD-L (Logit蒸餾) 和 WKD-F (Feature蒸餾)。

2.1 WKD-L: 基于離散WD的Logit蒸餾

WKD-L利用離散WD最小化教師和學生模型預測概率的差異。通過這種方式，WKD-L能夠進行跨類別比較，有效利用類別間的相互關系 (IRs)，這與KL-Div的類別內比較形成對比。研究人員使用CKA來量化類別間的相互關系，并將其融入損失函數中。

2.2 WKD-F: 基于連續WD的Feature蒸餾

WKD-F利用連續WD進行中間層特征的知識蒸餾。它通過高斯分布建模特征的分布，并最小化教師和學生模型特征分布之間的WD距離。這種方法能夠有效利用高斯分布的Riemann空間幾何結構，克服了KL-Div無法感知幾何結構的缺點。

3. 實驗結果與分析

研究人員在ImageNet、CIFAR-100和MS-COCO數據集上進行了大量的實驗，結果表明：

在ImageNet圖像分類任務中，WKD-L優于各種基于KL-Div的Logit蒸餾方法；WKD-F優于現有最先進的Feature蒸餾方法；WKD-L和WKD-F結合后，性能進一步提升。
在CIFAR-100圖像分類任務中，WKD在跨CNN和Transformer的知識遷移中表現出色，尤其是在Feature蒸餾方面。
在ImageNet自蒸餾任務中，WKD-L取得了最佳結果。
在MS-COCO目標檢測任務中，WKD-L和WKD-F均顯著優于現有方法，結合后性能進一步提升，甚至超越了最先進的方法。

4. 結論

這項研究證明了Wasserstein距離在知識蒸餾領域的有效性。WKD方法在Logit和Feature蒸餾方面均取得了優異的性能，為知識蒸餾研究提供了新的方向。雖然WKD-L的計算成本略高于基于KL-Div的方法，且WKD-F假設特征服從高斯分布，但這項工作仍然為知識蒸餾領域帶來了重要的貢獻。