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

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

論文標(biāo)題： Wasserstein Distance Rivals Kullback-Leibler Divergence for Knowledge Distillation

論文地址：

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é)生模型中。近年來，知識(shí)蒸餾在深度學(xué)習(xí)中受到了越來越多的關(guān)注，并取得了顯著進(jìn)展，在視覺識(shí)別、目標(biāo)檢測(cè)等多個(gè)領(lǐng)域得到了廣泛應(yīng)用。 在其開創(chuàng)性工作中，Hinton 等人引入了 Kullback-Leibler 散度（KL-Div）用于知識(shí)蒸餾，約束學(xué)生模型的類別概率預(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)類別上的概率，缺乏執(zhí)行跨類別比較的機(jī)制。 然而，現(xiàn)實(shí)世界中的類別呈現(xiàn)不同程度的視覺相似性，例如，哺乳動(dòng)物物種如狗和狼彼此間的相似度較高，而與汽車和自行車等人工制品則有很大的視覺差異，如圖1所示。

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)分析和比較 研究人員在ImageNet和CIFAR-100上評(píng)估了WKD在圖像分類任務(wù)中的表現(xiàn)，還評(píng)估了WKD在自蒸餾（Self-KD）中的有效，并且將WKD擴(kuò)展到目標(biāo)檢測(cè)，并在MS-COCO上進(jìn)行了實(shí)驗(yàn)。 在ImageNet上的圖像分類 研究人員在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也超過當(dāng)前的最佳方法ReviewKD；最后，WKD-L和WKD-F的結(jié)合進(jìn)一步提升了性能，超越了強(qiáng)有力的競(jìng)爭方法。

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

在CIFAR-100上的圖像分類 研究人員在教師模型為CNN、學(xué)生為Transformer或反之的設(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的圖像分類結(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è)頭中的分類分支進(jìn)行Logit蒸餾。對(duì)于WKD-F，直接從輸入到分類分支的特征中進(jìn)行知識(shí)遷移，即從RoIAlign層輸出的特征來計(jì)算高斯分布。 對(duì)于Logit蒸餾，WKD-L顯著優(yōu)于經(jīng)典的KD，并略微優(yōu)于DKD。對(duì)于特征蒸餾，WKD-F在兩個(gè)設(shè)置中均顯著超過之前的最佳特征蒸餾方法ReviewKD。最后，通過結(jié)合WKD-L和WKD-F，表現(xiàn)超過了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ì)算開銷更高，而WKD-F假設(shè)特征服從高斯分布。