原標題:超越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假設特征服從高斯分布,但這項工作仍然為知識蒸餾領域帶來了重要的貢獻。
聯系作者
文章來源:新智元
作者微信:
作者簡介:智能+中國主平臺,致力于推動中國從互聯網+邁向智能+新紀元。重點關注人工智能、機器人等前沿領域發展,關注人機融合、人工智能和機器人對人類社會與文明進化的影響,領航中國新智能時代。