自Hinton等人的開(kāi)創(chuàng)性工作以來(lái)，基于Kullback-Leibler散度（KL-Div）的知識(shí)蒸餾一直占主導(dǎo)地位。

然而，KL-Div僅比較教師和學(xué)生在相應(yīng)類(lèi)別上的概率，缺乏跨類(lèi)別比較的機(jī)制，應(yīng)用于中間層蒸餾時(shí)存在問(wèn)題，其無(wú)法處理不重疊的分布且無(wú)法感知底層流形的幾何結(jié)構(gòu)。

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

論文地址：https://arxiv.org/abs/2412.08139

項(xiàng)目地址：https://peihuali.org/WKD/

代碼地址：https://github.com/JiamingLv/WKD

背景與動(dòng)機(jī)介紹

知識(shí)蒸餾（KD）旨在將具有大容量的高性能教師模型中的知識(shí)遷移到輕量級(jí)的學(xué)生模型中。近年來(lái)，知識(shí)蒸餾在深度學(xué)習(xí)中受到了越來(lái)越多的關(guān)注，并取得了顯著進(jìn)展，在視覺(jué)識(shí)別、目標(biāo)檢測(cè)等多個(gè)領(lǐng)域得到了廣泛應(yīng)用。

在其開(kāi)創(chuàng)性工作中，Hinton等人引入了Kullback-Leibler散度（KL-Div）用于知識(shí)蒸餾，約束學(xué)生模型的類(lèi)別概率預(yù)測(cè)與教師模型相似。

從那時(shí)起，KL-Div在Logit蒸餾中占據(jù)主導(dǎo)地位，并且其變體方法DKD、NKD等也取得了令人矚目的性能。此外，這些Logit蒸餾方法還可以與將知識(shí)從中間層傳遞的許多先進(jìn)方法相互補(bǔ)充。

盡管KL-Div取得了巨大的成功，但它存在的兩個(gè)缺點(diǎn)阻礙了教師模型知識(shí)的遷移。

首先，KL-Div僅比較教師和學(xué)生在相應(yīng)類(lèi)別上的概率，缺乏執(zhí)行跨類(lèi)別比較的機(jī)制。

然而，現(xiàn)實(shí)世界中的類(lèi)別呈現(xiàn)不同程度的視覺(jué)相似性，例如，哺乳動(dòng)物物種如狗和狼彼此間的相似度較高，而與汽車(chē)和自行車(chē)等人工制品則有很大的視覺(jué)差異，如圖1所示。

不幸的是，由于KL-Div是類(lèi)別對(duì)類(lèi)別的比較，KD和其變體方法無(wú)法顯式地利用這種豐富的跨類(lèi)別知識(shí)。

圖1 左圖使用t-SNE展示了100個(gè)類(lèi)別的嵌入分布?？梢钥闯觯@些類(lèi)別在特征空間中表現(xiàn)出豐富的相互關(guān)系 (IR)。然而，右圖中的KL散度無(wú)法顯式地利用這些相互關(guān)系

其次，KL-Div在用于從中間層特征進(jìn)行知識(shí)蒸餾時(shí)存在局限性。圖像的深度特征通常是高維的且空間尺寸較小，因此其在特征空間中非常稀疏，不僅使得KL-Div在處理深度神經(jīng)網(wǎng)絡(luò)特征的分布時(shí)存在困難。

KL-Div無(wú)法處理不重疊的離散分布，并且由于其不是一個(gè)度量，在處理連續(xù)分布時(shí)能力有限，無(wú)法感知底層流形的幾何結(jié)構(gòu)。

圖2 基于Wasserstein距離（WD）的知識(shí)蒸餾方法的總覽圖

為了解決這些問(wèn)題，研究人員提出了一種基于Wasserstein距離的知識(shí)蒸餾方法，稱(chēng)為WKD，同時(shí)適用于Logit蒸餾（WKD-L）和Feature蒸餾（WKD-F），如圖2所示。

在WKD-L中，通過(guò)離散WD最小化教師和學(xué)生之間預(yù)測(cè)概率的差異，從而進(jìn)行知識(shí)轉(zhuǎn)移。

通過(guò)這種方式，執(zhí)行跨類(lèi)別的比較，能夠有效地利用類(lèi)別間的相互關(guān)系（IRs），與KL-Div中的類(lèi)別間比較形成鮮明對(duì)比。

對(duì)于WKD-F，研究人員利用WD從中間層特征中蒸餾知識(shí)，選擇參數(shù)化方法來(lái)建模特征的分布，并讓學(xué)生直接匹配教師的特征分布。

具體來(lái)說(shuō)，利用一種最廣泛使用的連續(xù)分布（高斯分布），該分布在給定特征的1階和2階矩的情況下具有最大熵。

論文的主要貢獻(xiàn)可以總結(jié)如下：

1. 提出了一種基于離散WD的Logit蒸餾方法（WKD-L），可以通過(guò)教師和學(xué)生預(yù)測(cè)概率之間的跨類(lèi)別比較，利用類(lèi)別間豐富的相互關(guān)系，克服KL-Div無(wú)法進(jìn)行類(lèi)別間比較的缺點(diǎn)。
2. 將連續(xù)WD引入中間層進(jìn)行Feature蒸餾（WKD-F），可以有效地利用高斯分布的Riemann空間幾何結(jié)構(gòu)，優(yōu)于無(wú)法感知幾何結(jié)構(gòu)的KL-Div。
3. 在圖像分類(lèi)和目標(biāo)檢測(cè)任務(wù)中，WKD-L優(yōu)于非常強(qiáng)的基于KL-Div的Logit蒸餾方法，而WKD-F在特征蒸餾中優(yōu)于KL-Div的對(duì)比方法和最先進(jìn)的方法。WKD-L和WKD-F的結(jié)合可以進(jìn)一步提高性能。

用于知識(shí)遷移的WD距離

用于Logit蒸餾的離散WD距離

類(lèi)別之間的相互關(guān)系（IRs）

如圖1所示，現(xiàn)實(shí)世界中的類(lèi)別在特征空間中表現(xiàn)出復(fù)雜的拓?fù)潢P(guān)系。相同類(lèi)別的特征會(huì)聚集并形成一個(gè)分布，而相鄰類(lèi)別的特征有重疊且不能完全分離。

因此，研究人員提出基于CKA量化類(lèi)別間的相互關(guān)系（IRs），CKA是一種歸一化的Hilbert-Schmidt獨(dú)立性準(zhǔn)則（HSIC），通過(guò)將兩個(gè)特征集映射到再生核希爾伯特空間（RKHS）來(lái)建模統(tǒng)計(jì)關(guān)系。

首先將每個(gè)類(lèi)別中所有訓(xùn)練樣本的特征構(gòu)成一個(gè)特征矩陣，之后通過(guò)計(jì)算任意兩個(gè)類(lèi)別特征矩陣之間的CKA得到類(lèi)間相互關(guān)系（IR）。計(jì)算IR的成本可以忽略，因?yàn)樵谟?xùn)練前僅需計(jì)算一次。

由于教師模型通常包含更豐富的知識(shí)，因此使用教師模型來(lái)計(jì)算類(lèi)別間的相互關(guān)系。

損失函數(shù)

用和分別表示教師模型和學(xué)生模型的預(yù)測(cè)類(lèi)別概率，其通過(guò)softmax函數(shù)和溫度對(duì)Logit計(jì)算得到。將離散的WD表示為一種熵正則化的線(xiàn)性規(guī)劃：

其中和分別表示每單位質(zhì)量的運(yùn)輸成本和在將概率質(zhì)量從移動(dòng)到

時(shí)的運(yùn)輸量；是正則化參數(shù)。

定義運(yùn)輸成本與相似度度量成負(fù)相關(guān)。

因此，WKD-L的損失函數(shù)可以定義為：

用于Feature蒸餾的連續(xù)WD距離

特征分布建模

將模型某個(gè)中間層輸出的特征圖重塑為一個(gè)矩陣，其中第i列表示一個(gè)空間特征。

之后，估計(jì)這些特征的一階矩和二階矩，并將二者作為高斯分布的參數(shù)來(lái)建模輸入圖像特征的分布。

損失函數(shù)

設(shè)教師的特征分布為高斯分布。類(lèi)似地，學(xué)生的分布記為。

兩者之間的連續(xù)Wasserstein距離（WD）定義為：

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

實(shí)驗(yàn)分析和比較

研究人員在ImageNet和CIFAR-100上評(píng)估了WKD在圖像分類(lèi)任務(wù)中的表現(xiàn)，還評(píng)估了WKD在自蒸餾（Self-KD）中的有效，并且將WKD擴(kuò)展到目標(biāo)檢測(cè)，并在MS-COCO上進(jìn)行了實(shí)驗(yàn)。

在ImageNet上的圖像分類(lèi)

研究人員在ImageNet的在兩種設(shè)置下與現(xiàn)有工作進(jì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。

對(duì)于Logit蒸餾，WKD-L在兩種設(shè)置下均優(yōu)于經(jīng)典的KD及其所有變體。對(duì)于特征蒸餾，WKD-F也超過(guò)當(dāng)前的最佳方法ReviewKD；最后，WKD-L和WKD-F的結(jié)合進(jìn)一步提升了性能，超越了強(qiáng)有力的競(jìng)爭(zhēng)方法。

表1 在ImageNet上的圖像分類(lèi)結(jié)果

在CIFAR-100上的圖像分類(lèi)

研究人員在教師模型為CNN、學(xué)生為T(mén)ransformer或反之的設(shè)置下評(píng)估了WKD方法，使用的CNN模型包括ResNet（RN）、MobileNetV2（MNV2）和ConvNeXt；Transformer模型包括ViT、DeiT和Swin Transformer。

對(duì)于Logit蒸餾，WKD-L在從Transformer到CNN遷移知識(shí)或反之的設(shè)置下始終優(yōu)于最新的OFA方法。對(duì)于特征蒸餾，WKD-F在所有實(shí)驗(yàn)設(shè)置中排名第一；

研究人員認(rèn)為，對(duì)于跨CNN和Transformer的知識(shí)轉(zhuǎn)移，考慮到兩者特征差異較大，WKD-F比像FitNet和CRD這樣直接對(duì)原始特征進(jìn)行對(duì)齊的方法更為合適。

表2  CIFAR-100上跨CNN和Transformer的圖像分類(lèi)結(jié)果（Top-1準(zhǔn)確率）

在ImageNet上的自蒸餾

研究人員在Born-Again Network（BAN）框架中將WKD方法用于自蒸餾任務(wù)（Self-KD）。

使用ResNet18在ImageNet上進(jìn)行實(shí)驗(yàn)，結(jié)果如表3所示，WKD-L取得了最佳結(jié)果，比BAN的Top-1準(zhǔn)確率高出約0.9%，比第二高的USKD方法高出0.6%。這一比較表明，WKD方法可以很好地推廣到自蒸餾任務(wù)中。

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

在MS-COCO上的目標(biāo)檢測(cè)

研究人員將WKD擴(kuò)展到Faster-RCNN框架下的目標(biāo)檢測(cè)中。對(duì)于WKD-L，使用檢測(cè)頭中的分類(lèi)分支進(jìn)行Logit蒸餾。對(duì)于WKD-F，直接從輸入到分類(lèi)分支的特征中進(jìn)行知識(shí)遷移，即從RoIAlign層輸出的特征來(lái)計(jì)算高斯分布。

對(duì)于Logit蒸餾，WKD-L顯著優(yōu)于經(jīng)典的KD，并略微優(yōu)于DKD。對(duì)于特征蒸餾，WKD-F在兩個(gè)設(shè)置中均顯著超過(guò)之前的最佳特征蒸餾方法ReviewKD。最后，通過(guò)結(jié)合WKD-L和WKD-F，表現(xiàn)超過(guò)了DKD+ReviewKD。當(dāng)使用額外的邊框回歸進(jìn)行知識(shí)遷移時(shí)，WKD-L+WKD-F進(jìn)一步提高并超越了之前的最先進(jìn)方法FCFD。

表4 在MS-COCO上的目標(biāo)檢測(cè)結(jié)果。?：使用了額外的邊框回歸進(jìn)行知識(shí)遷移

結(jié)論

Wasserstein距離（WD）在生成模型等多個(gè)領(lǐng)域中已展現(xiàn)出相較于KL散度（KL-Div）的顯著優(yōu)勢(shì)。

然而，在知識(shí)蒸餾領(lǐng)域，KL散度仍然占據(jù)主導(dǎo)地位，目前尚不清楚Wasserstein距離能否實(shí)現(xiàn)更優(yōu)的表現(xiàn)。

研究人員認(rèn)為，早期基于Wasserstein距離的知識(shí)蒸餾研究未能充分發(fā)揮該度量的潛力。

因此，文中提出了一種基于Wasserstein距離的全新知識(shí)蒸餾方法，能夠從Logit和Feature兩個(gè)方面進(jìn)行知識(shí)遷移。

大量的實(shí)驗(yàn)表明，離散形式的Wasserstein距離在Logit蒸餾中是當(dāng)前主流KL散度的極具潛力的替代方案，而連續(xù)形式的Wasserstein距離在中間層特征遷移中也取得了令人信服的性能表現(xiàn)。

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