大連理工提出基于Wasserstein距離(WD)的知識蒸餾方法

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

為了解決這些問題，大連理工大學的研究人員提出了一種基于 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）旨在將具有大容量的高性能教師模型中的知識遷移到輕量級的學生模型中。近年來，知識蒸餾在深度學習中受到了越來越多的關(guān)注，并取得了顯著進展，在視覺識別、目標檢測等多個領(lǐng)域得到了廣泛應(yīng)用。 在其開創(chuàng)性工作中，Hinton 等人引入了 Kullback-Leibler 散度（KL-Div）用于知識蒸餾，約束學生模型的類別概率預測與教師模型相似。 從那時起，KL-Div 在 Logit 蒸餾中占據(jù)主導地位，并且其變體方法 DKD、NKD 等也取得了令人矚目的性能。此外，這些 Logit 蒸餾方法還可以與將知識從中間層傳遞的許多先進方法相互補充。 盡管 KL-Div 取得了巨大的成功，但它存在的兩個缺點阻礙了教師模型知識的遷移。 首先，KL-Div 僅比較教師和學生在相應(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最小化教師和學生之間預測概率的差異，從而進行知識轉(zhuǎn)移。 通過這種方式，執(zhí)行跨類別的比較，能夠有效地利用類別間的相互關(guān)系（IRs），與KL-Div中的類別間比較形成鮮明對比。 對于WKD-F，研究人員利用WD從中間層特征中蒸餾知識，選擇參數(shù)化方法來建模特征的分布，并讓學生直接匹配教師的特征分布。 具體來說，利用一種最廣泛使用的連續(xù)分布（高斯分布），該分布在給定特征的1階和2階矩的情況下具有最大熵。 論文的主要貢獻可以總結(jié)如下：

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

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

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

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

。類似地，學生的分布記為。 兩者之間的連續(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)，其中教師和學生網(wǎng)絡(luò)分別為ResNet34和ResNet18；設(shè)置（b）涉及異質(zhì)架構(gòu)，在該設(shè)置中，教師網(wǎng)絡(luò)為ResNet50，學生網(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、學生為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ù)主導地位，目前尚不清楚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è)特征服從高斯分布。