基于類間排名相關性的解耦知識蒸餾

陳 穎,朱子奇,徐仕成,李 敏

(武漢科技大學 計算機科學與技術學院,湖北 武漢 430065)

0 引 言

卷積神經網絡(CNN)的出現革新了深度學習圖像識別領域。因為追求深層次規模架構網絡優秀的性能表現和小型網絡在生活中的適用性,知識蒸餾(KD)隨之誕生并且被廣泛應用于圖像分類[1]、目標檢測[2]和語義分割[3]等任務。其核心思想是將一個繁瑣的已優化模型(教師網絡)產生輸出概率,用于訓練另一個簡化模型(學生網絡)[4]。知識蒸餾不僅在培養學生網絡方面取得了不錯的效果,而且在通過自蒸餾改進教師網絡方面也非常有效[5-6]。

解耦知識蒸餾(Decoupled Knowledge Distillation,DKD)將教師網絡和學生網絡的logits輸出解耦成兩部分,一部分是目標類知識蒸餾(Target Class Knowledge Distillation,TCKD),另一部分是非目標類知識蒸餾(Non-Target Class Knowledge Distillation,NCKD),兩者單獨進行蒸餾[7]。通過對TCKD和NCKD的分析驗證與實驗,證明了NCKD的重要性。由于傳統知識蒸餾糅雜TCKD與NCKD一起蒸餾訓練,使用學生網絡的預測匹配教師網絡的預測,這將抑制NCKD發揮作用。即傳統知識蒸餾的高耦合性導致logit蒸餾的圖像分類效果并不理想。因此,可以通過解耦知識蒸餾解決傳統知識蒸餾的高耦合性問題。

知識蒸餾和解耦知識蒸餾都遵循精確的恢復方式,即學生損失應嚴格接近教師損失指標,教師損失應相應地改進學生網絡以達到和教師網絡近似一樣的結果[8]。然而,由于近似差距,這一假設并不總是成立的,這將給學生網絡的優化帶來偏差,導致次優結果或逆向結果。換句話說,模型不是追求評估度量的精確恢復,而是體現度量的目的,即蒸餾模型的性能。如果一個模型的損失比另一個模型小,那么它的度量應該更好。然而,當前教師網絡和學生網通常都存在代理損耗與評價指標之間的關系較為薄弱的問題。理想情況下,代理損失應保持評價指標與模型的強相關性。相比起知識的個體,知識的構成可以更好地被知識表現的關系呈現出來。而結構關系中包含更多的知識,因此保留教師和學生類間的預測關系對于優化蒸餾表現是有效的[9]。

在這些方法中,知識蒸餾以及解耦知識蒸餾被證明是一種通過模仿深層次結構網絡行為來提高小型網絡性能的有效方法。它通過在每個輸入實例的教師和學生的輸出上添加一個基于概率的點對點式精確匹配的強一致性約束,來鼓勵學生去模仿教師的行為。然而,在這樣強的約束條件下,直接優化學生損失很難得到準確的預測。該文并不只關注點對點的恢復,這忽略了類之間相對值的排名關系,而是通過顯式地學習類間排名的相關性來緩解“緊張”的約束,并獲得顯著的性能和效率的提高。

針對上述問題,通過類的排名相關性作為蒸餾的師生模型中代理損失與評價指標之間的關系,該文提出了一種排名相關一致性解耦知識蒸餾的方法來直接放大排名相關性。具體來說,該方法直接利用簡單的皮爾遜相關系數[10]。與精確恢復概率相比,基于相關性的優化更易于知識的遷移學習,并且損失函數可以與目標指標保持良好的關系一致性,與原始的解耦知識蒸餾損失相比,獲得了更好的相關性。主要貢獻如下:

(1)提出了一種新的排名相關一致性解耦知識蒸餾方法,它關注了不同類之間的相關性,解決了輸出的強一致性約束帶來的分類準確率下降問題。這是首次將類間排名相關性引入解耦知識蒸餾的方法;

(2)提出了一種能夠直接衡量類間排名相關性的新損失函數;

(3)通過充分的消融實驗和對比實驗表明,該方法在不同的任務如CIFAR-100和ImageNet-1K中都能有效提高蒸餾性能,取得良好的圖像分類精度。

1 相關工作

知識蒸餾這一技術的理論最早由Hinton[4]于2015年提出,是一種基于“教師-學生”網絡的模型訓練方式。知識蒸餾是構建一個小型的學生模型,利用性能更好的大型教師模型的監督信息來訓練這個學生模型,使得小型的學生模型可以擁有更好的性能和更高的精度[11]。與直接使用One-hot標簽的訓練不同,知識蒸餾使用的概率分布可以提供更多關于樣本相似性的信息,這些信息被稱為“暗知識”。因為知識蒸餾中從教師網絡遷移到學生網絡的知識不同,所以知識蒸餾的知識可以分為基于響應的知識、基于特征的知識和基于關系的知識。

1.1 基于響應的知識

基于響應的知識一般指模型最后輸出層的logits結果,然后學生網絡直接模仿教師網絡的最終預測[12]。當選擇用來進行表示學習的神經網絡的層次越深,中間層的監督就顯得越重要,因此基于響應的知識蒸餾逐漸淡出知識蒸餾的研究視野。近幾年新提出的解耦知識蒸餾使得基于響應的知識蒸餾潛能被發現。解耦知識蒸餾是出于傳統知識蒸餾的高耦合性考慮,把logits分為目標類和非目標類,然后分別對它們進行知識蒸餾。

1.2 基于特征的知識

基于特征的知識一般指最后一層的輸出和中間層的輸出。基于特征的知識蒸餾方法通過學生網絡學習教師網絡的中間層結構進行知識的遷移。FitNet[13]首次提出了一種提取教師的中間層作為hints層和學生中間層(guided層)進行特征輸出匹配的知識蒸餾。Zagoruyko等人[14]將原始特征圖中的“注意圖”作為知識,然后提取出來進行匹配。OFD[15]引入了一個margin ReLU函數并提出了一個新的損失函數來避免對學生有不利影響的冗余信息。CRD[16]在知識蒸餾中引入了對比性學習,借用對比目標實現知識從教師到學生的傳遞。ReviewKD[17]采用溫故知新的原理,提出了新的知識蒸餾復習機制,利用教師網絡的多層次信息指導學生網絡的一級學習。

1.3 基于關系的知識

基于關系的知識研究了不同層和數據樣本之間的關系。Yim等人[18]提出了能體現任意兩個特征圖的關系的FSP矩陣來進行知識蒸餾。Lee等人[19]通過兩個特征圖之間的相關性,利用奇異值分解進行蒸餾。Park等人[20]認為遷移樣本間的關系比起遷移單個樣本的特征會更有效,提出了一種對模型輸出的結構信息進行蒸餾的方式。Peng等人[21]注意到了實例之間的相關性,提出了相關同余的方法。

值得注意的是,Peng等人提出的實例間的相關同余和Park等人提出遷移樣本間關系的方法與文中方法有一些相似之處。但是文中方法側重于類間排名的關系,著重分析了在教師網絡和學生網絡差距過大時限制解耦知識蒸餾潛力的因素。文中提出使用類間排名相關一致性來解決問題,同時提出了新損失函數來優化模型,提升任務效果。

2 文中方法

2.1 基于類間排名相關性的解耦知識蒸餾

解耦知識蒸餾雖然證明了NCKD的重要性,但它和傳統的知識蒸餾都是采用學生網絡精確匹配教師網絡輸出概率的方法,體現不了教師網絡和學生網絡內部類間的關系結構。而擁有更深層次規模架構的網絡會取得更好的表現,因此會選取一個龐大的教師網絡,這也就導致了教師網絡和學生網絡規模相差較大。在這種差距下,使用傳統的KL散度來精確地恢復預測變得更為乏力。

該文提出了一種利用皮爾遜相關系數進行關系匹配的方法。相比于KL散度的精確式匹配,提出的更為寬容的匹配方式不再關心教師中NCKD輸出的預測分數,而是注意到教師預測非目標類的相對排名的關系。如圖1所示,解耦知識蒸餾借助學生各個類的預測分數與教師各個類的預測分數一對一精確匹配從而進行知識的遷移。越大的預測分數意味著這個類與目標類越相似,最終NCKD中各個類之間會形成一個預測分數的排名。文中方法提取NCKD中各個類排名的關系,把學生網絡中NCKD的排名關系與教師網絡中NCKD的排名關系進行匹配。相比起一對一精確匹配,基于類間排名相關性的方法只需讓學生保持和教師相似的排名關系,然后進行學生和教師的關系匹配,所以文中方法包含了更多的知識,知識間的構成也可以更好地表現出來。通過這種方式,學生不再需要精確的一對一匹配教師的輸出概率,而是更多地關注整體的關系信息。

圖1 模型架構

文中方法提倡維護教師網絡與學生網絡之間的類間關系,以類間排名相關性作為知識進行傳遞,實現從優秀的教師中更好地進行知識蒸餾。

2.2 NCKD的類間排名相關性匹配

知識蒸餾的目標是在訓練有素的大型教師網絡的指導下培養更便于部署的小型學生網絡。而在解耦知識蒸餾中,通過最小化教師網絡與學生網絡的TCKD和NCKD的預測分數之間的差異實現知識從教師網絡到學生網絡的轉移。在蒸餾中,訓練有素的教師網絡提供了超出通常數據監督的額外語義知識,而挑戰在于如何對教師的知識進行編碼和轉換,使學生的表現可以達到最大化。

對于來自第t類的訓練樣本,學生網絡和教師網絡的logits輸出分別為Z(T)∈R1×C和Z(S)∈R1×C,其中[Z(T),Z(S)]∈Z。分類概率可以表示為P=[p1,p2,…,pT,…,pC]∈R1×C,其中pi表示第i個類的概率,C表示類的個數。分類概率p中的每個元素都可以通過softmax()函數和溫度因子T進行評估:

(1)

非目標類∧P是眾多類的集合(除去目標類),即∧P=[∧p1,…,∧pT-1,∧pT+1,…,∧pC]∈R1×(C-1),其中每個類的概率可表示為:

(2)

對于解耦知識蒸餾,需要解耦目標類和非目標類。因此定義b=[pT,pN]∈R1×2,b代表目標類和所有非目標類的二值概率,pT代表目標類的概率,pN代表非目標類的概率:

(3)

(4)

根據式1～式4可以推斷出pi=∧pi×pN,所以KD的loss可表示為:

(5)

LKD=αTCKD+βNCKD

(6)

對于預測向量∧p(T)和∧p(S),若x=y,則x和y之間可以得到一個精確的匹配距離為0,即d(x,y)=0。該文使用一種較為“輕松”的匹配方式,所以引入關系映射f()和g():

d(x,y)=d(f(x),g(y))

(7)

文中方法注重x,y之間的內部關系,所以并不要求x和y完全一樣。為了不影響NCKD向量所包含的信息,映射f()和g()應該是等值的,因此選擇恒等變換:

d(x,y)=d(a1x+b1,a2y+b2)

(8)

其中,a1,a2,b1,b2都是常數。

該文使用皮爾遜相關系數。皮爾遜相關系數廣泛用于衡量兩個變量之間的相關程度,反映了兩個變量之間的線性關系和相關性的方向。最重要的一點是皮爾遜相關系數有一個很重要的數學特性:當變量x和y的位置發生變化時,皮爾遜相關系數不會發生變化。也就是說把變量x和y移動為ax+b和cy+d(其中a,b,c和d皆為常數)并不會引起x和y相關系數的變化,而這一點數學特性恰好完美地契合了文中方法所追求的特性d(x,y)=d(a1x+b1,a2y+b2)。皮爾遜相關系數ρ∧p(x,y)可表示為:

(9)

2.3 損失函數

皮爾遜相關系數的變化范圍為[-1,1],當相關系數值為0時代表兩個變量之間沒有任何線性關系,故皮爾遜距離應當為1。所以,皮爾遜距離d(x,y)和皮爾遜相關系數ρ∧p(x,y)的關系可以表示為:

d∧p(x,y)=1-ρ∧p(x,y)

(10)

所以,該文的損失函數是:

LKD=αTCKD+βd∧p(x,y)

(11)

文中方法通過最大線性相關來保留教師網絡和學生網絡在每個實例的概率分布上的關系,所以可以更好地泛化到輸出的整個分布,是一種更為穩健有效的方法。

3 實驗與分析

3.1 數據集介紹及評價標準

CIFAR-100[22]是知識蒸餾分類領域應用最為廣泛的一個數據集,也是文中實驗所采用的一個數據集。CIFAR-100數據集包含50 000張訓練圖像和100個類別。每個類別有600張大小為32×32的彩色圖像,其中500張作為訓練集,100張作為測試集。對于每一張圖像,它有fine_labels和coarse_labels兩個標簽,分別代表圖像的細粒度和粗粒度標簽。

ImageNet-1K[23]是ImageNet的子集,來自于斯坦福大學的課程項目,是一個極具挑戰性的數據庫,也是圖像分類任務中使用次數較多的大型數據庫。ImageNet-1K共有1 000個類別,訓練集有1 281 167張圖片且附帶標簽,驗證集有50 000張圖片并且附帶標簽,測試集有100 000張圖片。

該文引入了圖像分類任務中的準確率(Accuracy)作為實驗的性能評價指標,該基本指標代表在所有樣本中預測正確的概率,是圖像分類任務中最直觀的性能評價指標。Top-1 Accuracy代表排名第一的類別與實際結果相符合的準確率。

(12)

其中,TP代表被判定為正樣本,實際上也是正樣本的樣本數;TN代表被判定為負樣本,實際上也是負樣本的樣本數;FP代表被判定為正樣本,但實際上是負樣本的樣本數;FN代表被判定為負樣本,但實際上是正樣本的樣本數。

3.2 實驗參數

實驗皆在Linux系統上進行并基于Pytorch框架,使用了五張NVIDIA GeForce RTX 2080 Ti GPU。實驗中使用ResNet[24],VGG[25],ShuffleNet[26-27],MobileNet[28]和Wide ResNet[29]網絡。對于CIFAR100數據集,實驗的Batchsize大小為128,不同的學生網絡設置了不同的學習率初始值,如表1所示。實驗所有模型進行240個Epoch的訓練,在150個Epoch之后,每30個Epoch的學習率衰減0.1。對于ImageNet-1K數據集,實驗使用標準的訓練過程,將Batchsize大小設置為256,所有模型進行100個Epoch的訓練,實驗將學習率初始化為0.1,然后每30個Epoch衰減學習率。文中方法的總損失函數中的超參數α和β分別設置為1.0和8.0。

表1 在CIFAR-100數據集上不同網絡設置的學習率

3.3 消融實驗

為了研究損失函數對于蒸餾模型性能的影響,使用ResNet32×4作為教師網絡,ShuffleNetV2作為學生網絡,對三種方法分別訓練240次,在CIFAR-100數據集上實驗得到的精度結果如表2所示。傳統的知識蒸餾圖像分類達到了74.07%的準確率。接下來選取基于解耦的知識蒸餾作為基準網絡,解耦目標類和非目標類進行蒸餾,圖像分類達到了76.45%的準確率,解決了因為傳統知識蒸餾的高耦合性帶來的精度降低問題。最后是該文提出的基于非目標類類之間預測排序關系的知識蒸餾,改善了強一致性約束后,模型達到了77.38%的準確率,相比基準網絡提升了0.93百分點,解決了解耦知識蒸餾強一致性約束對精度帶來負面影響的問題,證明了該方法的有效性。消融實驗結果表明,該方法可以更好地體現類之間的相關性并泛化到輸出的整個分布,降低了教師和學生因模型架構懸殊和強一致性約束所帶來的影響,具有更優秀的分類準確率。

表2 不同方法在CIFAR-100數據集上的實驗結果

同時,文中方法和基準方法的loss變化對比如圖2所示,文中方法的loss不僅波動更為平緩,而且比基準方法的loss更小一些(差值在0.01～0.05之間波動)。

圖2 loss變化對比

3.4 對比實驗

為了驗證文中方法的先進性,在圖像分類任務上將文中方法與其它先進的方法進行比較。文中使用的基準方法是解耦知識蒸餾,加入基于類間排名相關性模塊后,在數據集CIFAR-100和ImageNet-1K提升了0.2～1百分點。實驗結果表明,與基準方法和傳統知識蒸餾相比,文中方法在教師學生網絡組中一致取得了改善。基于CIFAR-100數據集,文中方法在同系列教師學生網絡組中獲得了0.2～0.8百分點的提升,在不同系列教師學生網絡組中獲得了0.2～1百分點的提升,這進一步證明了文中方法的先進性。

表3 同系列教師學生網絡實驗結果 %

續表3

表4 不同系列教師學生網絡實驗結果 %

由于CIFAR-100的圖像數量較少,因此在ImageNet-1K上進行了實驗來驗證文中方法的可擴展性。實驗了從ResNet34到ResNet18的蒸餾設置,文中方法再次優于所有其它方法,如表5所示。

表5 ImageNet-1K數據集的對比實驗結果

4 結束語

在解耦知識蒸餾中,針對架構規模懸殊的教師網絡和學生網絡使用點對點式強一致性約束精確匹配預測概率的問題,提出了一種相關一致性解耦知識蒸餾的方法,即類間排名相關性。該方法針對NCKD采用了類間排名相關性進行一致匹配來緩解logit蒸餾強一致性約束。通過在同系列教師網絡和學生網絡以及不同系列教師網絡和學生網絡上的訓練測試,該模型的圖像分類準確率顯著提高,有效提高了解耦知識蒸餾的圖像分類能力。