自 Hinton 等人的開創(chuàng)性工作以來，基于 Kullback-Leibler 散度（KL-Div）的知識蒸餾一直占主導(dǎo)地位。 然而，KL-Div 僅比較教師和學(xué)生在相應(yīng)類別上的概率，缺乏跨類別比較的機制，應(yīng)用于中間層蒸餾時存在問題，其無法處理不重疊的分布且無法感知底層流形的幾何結(jié)構(gòu)。

為了解決這些問題，大連理工大學(xué)的研究人員提出了一種基于 Wasserstein 距離（WD）的知識蒸餾方法。所提出方法在圖像分類和目標檢測任務(wù)上均取得了當前最好的性能，論文已被 NeurIPS 2024 接受為 Poster。

論文標題： Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation

論文地址：

https://arxiv.org/abs/2412.08139

項目地址：

https://peihuali.org/WKD/

代碼地址：

https://github.com/JiamingLv/WKD

背景與動機介紹

知識蒸餾（KD）旨在將具有大容量的高性能教師模型中的知識遷移到輕量級的學(xué)生模型中。近年來，知識蒸餾在深度學(xué)習(xí)中受到了越來越多的關(guān)注，并取得了顯著進展，在視覺識別、目標檢測等多個領(lǐng)域得到了廣泛應(yīng)用。 在其開創(chuàng)性工作中，Hinton 等人引入了 Kullback-Leibler 散度（KL-Div）用于知識蒸餾，約束學(xué)生模型的類別概率預(yù)測與教師模型相似。 從那時起，KL-Div 在 Logit 蒸餾中占據(jù)主導(dǎo)地位，并且其變體方法 DKD、NKD 等也取得了令人矚目的性能。此外，這些 Logit 蒸餾方法還可以與將知識從中間層傳遞的許多先進方法相互補充。 盡管 KL-Div 取得了巨大的成功，但它存在的兩個缺點阻礙了教師模型知識的遷移。 首先，KL-Div 僅比較教師和學(xué)生在相應(yīng)類別上的概率，缺乏執(zhí)行跨類別比較的機制。 然而，現(xiàn)實世界中的類別呈現(xiàn)不同程度的視覺相似性，例如，哺乳動物物種如狗和狼彼此間的相似度較高，而與汽車和自行車等人工制品則有很大的視覺差異，如圖1所示。

不幸的是，由于 KL-Div 是類別對類別的比較，KD 和其變體方法無法顯式地利用這種豐富的跨類別知識。

▲圖1. 左圖使用t-SNE展示了100個類別的嵌入分布。可以看出，這些類別在特征空間中表現(xiàn)出豐富的相互關(guān)系 (IR)。然而，右圖中的KL散度無法顯式地利用這些相互關(guān)系。

其次，KL-Div 在用于從中間層特征進行知識蒸餾時存在局限性。圖像的深度特征通常是高維的且空間尺寸較小，因此其在特征空間中非常稀疏，不僅使得 KL-Div 在處理深度神經(jīng)網(wǎng)絡(luò)特征的分布時存在困難。 KL-Div 無法處理不重疊的離散分布，并且由于其不是一個度量，在處理連續(xù)分布時能力有限，無法感知底層流形的幾何結(jié)構(gòu)。

▲圖2. 基于Wasserstein距離（WD）的知識蒸餾方法的總覽圖

為了解決這些問題，研究人員提出了一種基于Wasserstein距離的知識蒸餾方法，稱為WKD，同時適用于Logit蒸餾（WKD-L）和Feature蒸餾（WKD-F），如圖2所示。 在WKD-L中，通過離散WD最小化教師和學(xué)生之間預(yù)測概率的差異，從而進行知識轉(zhuǎn)移。 通過這種方式，執(zhí)行跨類別的比較，能夠有效地利用類別間的相互關(guān)系（IRs），與KL-Div中的類別間比較形成鮮明對比。 對于WKD-F，研究人員利用WD從中間層特征中蒸餾知識，選擇參數(shù)化方法來建模特征的分布，并讓學(xué)生直接匹配教師的特征分布。 具體來說，利用一種最廣泛使用的連續(xù)分布（高斯分布），該分布在給定特征的1階和2階矩的情況下具有最大熵。 論文的主要貢獻可以總結(jié)如下：

提出了一種基于離散WD的Logit蒸餾方法（WKD-L），可以通過教師和學(xué)生預(yù)測概率之間的跨類別比較，利用類別間豐富的相互關(guān)系，克服KL-Div無法進行類別間比較的缺點。

將連續(xù)WD引入中間層進行Feature蒸餾（WKD-F），可以有效地利用高斯分布的Riemann空間幾何結(jié)構(gòu)，優(yōu)于無法感知幾何結(jié)構(gòu)的KL-Div。

在圖像分類和目標檢測任務(wù)中，WKD-L優(yōu)于非常強的基于KL-Div的Logit蒸餾方法，而WKD-F在特征蒸餾中優(yōu)于KL-Div的對比方法和最先進的方法。WKD-L和WKD-F的結(jié)合可以進一步提高性能。

用于知識遷移的WD距離

用于Logit蒸餾的離散WD距離 類別之間的相互關(guān)系（IRs） 如圖1所示，現(xiàn)實世界中的類別在特征空間中表現(xiàn)出復(fù)雜的拓撲關(guān)系。相同類別的特征會聚集并形成一個分布，而相鄰類別的特征有重疊且不能完全分離。 因此，研究人員提出基于CKA量化類別間的相互關(guān)系（IRs），CKA是一種歸一化的Hilbert-Schmidt獨立性準則（HSIC），通過將兩個特征集映射到再生核希爾伯特空間（RKHS）來建模統(tǒng)計關(guān)系。 首先將每個類別中所有訓(xùn)練樣本的特征構(gòu)成一個特征矩陣，之后通過計算任意兩個類別特征矩陣之間的CKA得到類間相互關(guān)系（IR）。計算IR的成本可以忽略，因為在訓(xùn)練前僅需計算一次。 由于教師模型通常包含更豐富的知識，因此使用教師模型來計算類別間的相互關(guān)系。 損失函數(shù) 用和分別表示教師模型和學(xué)生模型的預(yù)測類別概率，其通過softmax函數(shù)和溫度對Logit計算得到。將離散的WD表示為一種熵正則化的線性規(guī)劃：

其中和分別表示每單位質(zhì)量的運輸成本和在將概率質(zhì)量從移動到 時的運輸量；是正則化參數(shù)。 定義運輸成本與相似度度量成負相關(guān)。 因此，WKD-L的損失函數(shù)可以定義為：

用于Feature蒸餾的連續(xù)WD距離 特征分布建模 將模型某個中間層輸出的特征圖重塑為一個矩陣，其中第i列表示一個空間特征。 之后，估計這些特征的一階矩和二階矩

，并將二者作為高斯分布的參數(shù)來建模輸入圖像特征的分布。 損失函數(shù) 設(shè)教師的特征分布為高斯分布

。類似地，學(xué)生的分布記為。 兩者之間的連續(xù)Wasserstein距離（WD）定義為：

其中，和是高斯變量，q表示聯(lián)合分布。最小化上式可以得到閉集形式的WD距離。此外，為了平衡均值和協(xié)方差的作用，引入了一個均值-協(xié)方差比率γ，最后損失定義為：

實驗分析和比較 研究人員在ImageNet和CIFAR-100上評估了WKD在圖像分類任務(wù)中的表現(xiàn)，還評估了WKD在自蒸餾（Self-KD）中的有效，并且將WKD擴展到目標檢測，并在MS-COCO上進行了實驗。 在ImageNet上的圖像分類 研究人員在ImageNet的在兩種設(shè)置下與現(xiàn)有工作進行了比較。設(shè)置（a）涉及同質(zhì)架構(gòu)，其中教師和學(xué)生網(wǎng)絡(luò)分別為ResNet34和ResNet18；設(shè)置（b）涉及異質(zhì)架構(gòu)，在該設(shè)置中，教師網(wǎng)絡(luò)為ResNet50，學(xué)生網(wǎng)絡(luò)為MobileNetV1。 對于Logit蒸餾，WKD-L在兩種設(shè)置下均優(yōu)于經(jīng)典的KD及其所有變體。對于特征蒸餾，WKD-F也超過當前的最佳方法ReviewKD；最后，WKD-L和WKD-F的結(jié)合進一步提升了性能，超越了強有力的競爭方法。

▲表1. 在ImageNet上的圖像分類結(jié)果

在CIFAR-100上的圖像分類 研究人員在教師模型為CNN、學(xué)生為Transformer或反之的設(shè)置下評估了WKD方法，使用的CNN模型包括ResNet（RN）、MobileNetV2（MNV2）和ConvNeXt；Transformer模型包括ViT、DeiT和Swin Transformer。 對于Logit蒸餾，WKD-L在從Transformer到CNN遷移知識或反之的設(shè)置下始終優(yōu)于最新的OFA方法。對于特征蒸餾，WKD-F在所有實驗設(shè)置中排名第一； 研究人員認為，對于跨CNN和Transformer的知識轉(zhuǎn)移，考慮到兩者特征差異較大，WKD-F比像FitNet和CRD這樣直接對原始特征進行對齊的方法更為合適。

▲表2. CIFAR-100上跨CNN和Transformer的圖像分類結(jié)果（Top-1準確率） 在ImageNet上的自蒸餾 研究人員在Born-Again Network（BAN）框架中將WKD方法用于自蒸餾任務(wù)（Self-KD）。 使用ResNet18在ImageNet上進行實驗，結(jié)果如表3所示，WKD-L取得了最佳結(jié)果，比BAN的Top-1準確率高出約0.9%，比第二高的USKD方法高出0.6%。這一比較表明，WKD方法可以很好地推廣到自蒸餾任務(wù)中。

▲表3. 在ImageNet上自蒸餾（Self-KD）的結(jié)果

在MS-COCO上的目標檢測 研究人員將WKD擴展到Faster-RCNN框架下的目標檢測中。對于WKD-L，使用檢測頭中的分類分支進行Logit蒸餾。對于WKD-F，直接從輸入到分類分支的特征中進行知識遷移，即從RoIAlign層輸出的特征來計算高斯分布。 對于Logit蒸餾，WKD-L顯著優(yōu)于經(jīng)典的KD，并略微優(yōu)于DKD。對于特征蒸餾，WKD-F在兩個設(shè)置中均顯著超過之前的最佳特征蒸餾方法ReviewKD。最后，通過結(jié)合WKD-L和WKD-F，表現(xiàn)超過了DKD+ReviewKD。當使用額外的邊框回歸進行知識遷移時，WKD-L+WKD-F進一步提高并超越了之前的最先進方法FCFD。

▲表4. 在MS-COCO上的目標檢測結(jié)果。?：使用了額外的邊框回歸進行知識遷移

總結(jié)

Wasserstein距離（WD）在生成模型等多個領(lǐng)域中已展現(xiàn)出相較于KL散度（KL-Div）的顯著優(yōu)勢。 然而，在知識蒸餾領(lǐng)域，KL散度仍然占據(jù)主導(dǎo)地位，目前尚不清楚Wasserstein距離能否實現(xiàn)更優(yōu)的表現(xiàn)。 研究人員認為，早期基于Wasserstein距離的知識蒸餾研究未能充分發(fā)揮該度量的潛力。 因此，文中提出了一種基于Wasserstein距離的全新知識蒸餾方法，能夠從Logit和Feature兩個方面進行知識遷移。 大量的實驗表明，離散形式的Wasserstein距離在Logit蒸餾中是當前主流KL散度的極具潛力的替代方案，而連續(xù)形式的Wasserstein距離在中間層特征遷移中也取得了令人信服的性能表現(xiàn)。

盡管如此，該方法仍存在一定局限性：WKD-L相比基于KL散度的Logit蒸餾方法計算開銷更高，而WKD-F假設(shè)特征服從高斯分布。