基于特征重建的知識蒸餾方法

郭俊倫，彭書華，李俊杰

（北京信息科技大學自動化學院，北京 100192）

0 引言

隨著深度神經(jīng)網(wǎng)絡(luò)技術(shù)的不斷發(fā)展，深度神經(jīng)網(wǎng)絡(luò)的表現(xiàn)越來越優(yōu)異并且優(yōu)異性是傳統(tǒng)算法遠遠所不能及的。這些都是因為神經(jīng)網(wǎng)絡(luò)的深度不斷增加，從而使得結(jié)構(gòu)越來越復(fù)雜，對更高層更抽象的特征提取能力也變的越來越強。例如：寬度相同的8 層和16 層的神經(jīng)網(wǎng)絡(luò)模型用于同樣的圖片識別任務(wù)中，16 層神經(jīng)網(wǎng)絡(luò)的性能要遠高于8 層神經(jīng)網(wǎng)絡(luò)。但是隨著網(wǎng)絡(luò)的加深，所需存儲的網(wǎng)絡(luò)節(jié)點以及所需要的浮點運算次數(shù)也會成倍的增加。而這些深度神經(jīng)網(wǎng)絡(luò)有這百萬級別的網(wǎng)絡(luò)節(jié)點以及需要上億次的浮點運算，這樣會消耗很大的儲存空間并且需要強大的計算能力。但是，GPU 的出現(xiàn)使得模型訓練周期大大的縮短，計算效率大大提升。

而在生活生產(chǎn)中往往使用的使基于CPU 的嵌入式設(shè)備，CPU 設(shè)備內(nèi)存小，計算能力弱的特點，使得在嵌入式設(shè)備中應(yīng)用神經(jīng)網(wǎng)絡(luò)技術(shù)成為一大難點，這嚴重限制了神經(jīng)網(wǎng)絡(luò)技術(shù)在資源有限的嵌入式設(shè)備中應(yīng)用?，F(xiàn)在已經(jīng)有將深度神經(jīng)網(wǎng)絡(luò)壓縮技術(shù)應(yīng)用在嵌入式設(shè)備[1]中。并且在許多領(lǐng)域取得了一定的研究成果，例如無人駕駛技術(shù)[2]、醫(yī)學圖像處理[3]、人臉識別[4]等。

目前主流的網(wǎng)絡(luò)壓縮與加速方法[5]主要有以下五種：①輕量化模型的設(shè)計；②參數(shù)剪枝；③參數(shù)量化；④低秩分解；⑤知識蒸餾。

其中知識蒸餾這一方法相比于其他壓縮方法有這優(yōu)異的壓縮效果，它可以在準確率損失較小的前提下將大模型壓縮為原模型的一半以上。但是到底以何種方式“蒸餾”依然是個很大的問題。2014 年，Hinton 等人[6]第一次提出了知識蒸餾的概念，就是將復(fù)雜模型的知識遷移到另一個輕量級模型上，使輕量模型的結(jié)果近似逼近復(fù)雜模型。2015 年，Romero 等人[7]提出Fitnets 模型，選取學生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)的中間層，利用教師網(wǎng)絡(luò)中間層信息訓練學生模型前半部分參數(shù)，再利用真實標簽訓練學生模型整體。2017 年，Zagoruyko 等人[8]將注意力機制[9]引入模型壓縮中，通過選取教師網(wǎng)絡(luò)與學生網(wǎng)絡(luò)的中間卷積層生成的注意力圖作為“知識”傳遞到學生網(wǎng)絡(luò)的中，指導小模型進行訓練。并且將注意力機制與傳統(tǒng)知識蒸餾方法結(jié)合后學生網(wǎng)絡(luò)的精度會有所提高。2018 年，Zhou 等人[10]提出了一種新的知識蒸餾方式，就是學生網(wǎng)絡(luò)與教師網(wǎng)絡(luò)共用基層，這種方法不僅學生網(wǎng)絡(luò)精度有了提升，而且教師網(wǎng)絡(luò)的精度也有了一定程度的提升，這種方法的缺點就是共用基層網(wǎng)絡(luò)限制了小網(wǎng)絡(luò)的靈活性。2019 年，Mirzadeh 等人[11]提出了教師助理傳遞模式，這種研究方法發(fā)現(xiàn)大模型與小模型進行傳遞時如果加一個規(guī)模適中的中間模型，小模型會取得更好的效果。而這個中間模型就被稱為教師助理。

綜上，傳統(tǒng)的知識蒸餾忽略了教師模型的互相交互，而一些新模型Zhou[10]、Mirzadeh[11]等人提出的模型雖然有了交互過程，但是網(wǎng)絡(luò)交互方式大大限制了學生網(wǎng)絡(luò)的靈活性。由此本文結(jié)合圖片風格遷移實驗中的特征重建方法提出了一種新的知識蒸餾方法。

1 理論方法

1.1 殘差網(wǎng)絡(luò)

研究發(fā)現(xiàn)神經(jīng)網(wǎng)絡(luò)的層次越深，網(wǎng)絡(luò)就會有越強的特征提取能力，從而就能有更好的使用效果。但是人們發(fā)現(xiàn)隨著網(wǎng)絡(luò)的不斷加深會出現(xiàn)梯度彌散、梯度消失甚至神經(jīng)網(wǎng)絡(luò)退化等問題，這些問題使得深度神經(jīng)網(wǎng)絡(luò)的發(fā)展變得困難，在2015 年ImageNet 比賽中，He 等人[12]提出殘差網(wǎng)絡(luò)（Residual Network），網(wǎng)絡(luò)的深度達到了152 層，為加深神經(jīng)網(wǎng)絡(luò)提供了新的思路。殘差網(wǎng)絡(luò)解決了層數(shù)加深后出現(xiàn)的一系列問題。殘差網(wǎng)絡(luò)的殘差塊結(jié)構(gòu)如圖1 所示。

圖1 殘差網(wǎng)絡(luò)的殘差塊結(jié)構(gòu)

殘差網(wǎng)絡(luò)在普通卷積網(wǎng)絡(luò)中引入了殘差塊，由公式（1）可見引入殘差塊后會對原先網(wǎng)絡(luò)的訓練結(jié)果進行微調(diào)，從而抑制原網(wǎng)絡(luò)在訓練中產(chǎn)生的錯誤，修正深度神經(jīng)網(wǎng)絡(luò)的訓練結(jié)果。

其中，Xl、Xl+1為殘差網(wǎng)絡(luò)第l 層和第l+1 層的卷積層運算結(jié)果；W 為殘差塊結(jié)構(gòu)的網(wǎng)絡(luò)參數(shù)；F（·）為殘差塊中的卷積運算函數(shù)。

因為網(wǎng)絡(luò)壓縮實驗中需要使用深度神經(jīng)網(wǎng)絡(luò)，并且為了方便與其他方法的實驗結(jié)果進行對比，選用了寬度殘差網(wǎng)絡(luò)結(jié)構(gòu)（Wide Residual Net，WRN）即以模型的寬度作為基準的殘差網(wǎng)絡(luò)，進行相關(guān)實驗。

1.2 知識蒸餾

用于分類任務(wù)的深度神經(jīng)網(wǎng)絡(luò)往往使用Softmax層來計算出類概率。模型的訓練往往采用的是人工標注的標簽，又叫硬標簽。知識蒸餾則是采用教師網(wǎng)絡(luò)Softmax 層的輸出作為小模型訓練的標簽，又叫軟標簽，并且在Softmax 中引入溫度參數(shù)T，此時基于知識蒸餾的Softmax 公式變?yōu)椋?/p>

zi為進入Softmax 輸出層前的標簽；qi為知識蒸餾后計算出的每一類的概率；T 為知識蒸餾溫度參數(shù)，當T 為1 時，公式便是神經(jīng)網(wǎng)絡(luò)分類任務(wù)中最常用的Softmax 公式，由公式（2）可見溫度T 越高神經(jīng)網(wǎng)絡(luò)計算出的軟標簽分布越均勻，這也是“蒸餾”名字的由來。典型的知識蒸餾架構(gòu)如圖2 所示。

圖2 知識蒸餾架構(gòu)

由圖2 可知，引入?yún)?shù)T 后的學生模型損失函數(shù)變?yōu)椋?/p>

其中，α為軟標簽在學生模型訓練中所占的比例；T為溫度參數(shù)，與公式（2）中T 相同；、分別為學生網(wǎng)絡(luò)和教師網(wǎng)絡(luò)含參數(shù)T 的Softmax 公式；qs為學生網(wǎng)絡(luò)中的 Softmax 公式；Cr（·）為交叉熵函數(shù)；y 為模型訓練的真實標簽。

1.3 特征重建算法

Johnson 等人[13]在圖片風格遷移實驗中發(fā)現(xiàn)，兩張圖片輸入同一個網(wǎng)絡(luò)時，圖片的內(nèi)容越相似則網(wǎng)絡(luò)的深度特征越相似。于是在圖片風格遷移實驗中提出了特征重建損失函數(shù)。此損失函數(shù)是由已經(jīng)訓練好的VGGNet[14]計算得到的，將VGGNet 作為損失函數(shù)產(chǎn)生網(wǎng)絡(luò)，用來訓練圖片風格遷移網(wǎng)絡(luò)。風格遷移實驗內(nèi)容遷移（即特征重方法）的訓練架構(gòu)如圖3 所示。

圖3 圖片風格遷移之內(nèi)容遷移訓練架構(gòu)

其中y 為圖片生成網(wǎng)絡(luò)生成的新圖片，yc為圖片生成網(wǎng)絡(luò)生成新圖片的目標圖片。實驗中VGGNet 的深度特征歐氏距離越小，則圖片生成網(wǎng)絡(luò)生成的圖片內(nèi)容越接近目標圖片。由此也可以得出特征重建損失函數(shù)。此損失函數(shù)如下：

Cj×Hj×Wj為學生網(wǎng)絡(luò)第 j 層特征圖的尺寸；、為圖片y 和yc分別輸入VGGNet 后第j 層的特征向量。

2 使用特征重建算法設(shè)計知識蒸餾新架構(gòu)

傳統(tǒng)知識蒸餾架構(gòu)都是利用教師網(wǎng)絡(luò)與學生網(wǎng)絡(luò)中間層特征向量作為知識傳遞介質(zhì)進行知識蒸餾。這種做法的缺點就是教師網(wǎng)絡(luò)和學生網(wǎng)絡(luò)相對獨立，沒有充分發(fā)掘出模型間的內(nèi)在聯(lián)系。目前，已經(jīng)存在多種知識蒸餾框架，這些框架大都還是保持教師網(wǎng)絡(luò)與學生網(wǎng)絡(luò)相對獨立的結(jié)構(gòu)，只是尋求不同的知識傳播形式。所以到底以何種架構(gòu)的知識傳播形式更加高效需要進行不斷的實驗。

在圖片風格遷移實驗中，是使用已經(jīng)訓練好的網(wǎng)絡(luò)作為損失函數(shù)生成網(wǎng)絡(luò)，去訓練圖片風格遷移網(wǎng)絡(luò)。而在知識蒸餾中，也是已知一個訓練好的教師網(wǎng)絡(luò)，去訓練一個規(guī)模更小的學生網(wǎng)絡(luò)，由此，參照風格遷移實驗，使用教師網(wǎng)絡(luò)作為損失函數(shù)生成網(wǎng)絡(luò)，并且使用風格遷移實驗中的特征重建損失函數(shù)對學生網(wǎng)絡(luò)進行訓練。蒸餾框架總設(shè)計如圖4 所示。

圖4 基于學生網(wǎng)絡(luò)特征重建的知識蒸餾架構(gòu)設(shè)計

由圖4 可見，該架構(gòu)主要分成兩部分，一部分為利用特征重建算法的損失函數(shù)生成網(wǎng)絡(luò)；另一部分是傳統(tǒng)的知識蒸餾算法。損失函數(shù)生成網(wǎng)絡(luò)是將學生網(wǎng)絡(luò)的中間層特征分組嵌入到教師網(wǎng)絡(luò)中，使用教師網(wǎng)絡(luò)的深度特征對學生網(wǎng)絡(luò)插入教師網(wǎng)絡(luò)之后生成的新特征進行特征重建。

那么其中的利用特征重建生成損失函數(shù)的網(wǎng)絡(luò)框架的結(jié)構(gòu)如圖5 所示。

圖5 利用特征重建生成損失函數(shù)的框架

使用此損失函數(shù)生成網(wǎng)絡(luò)生成損失函數(shù)對學生網(wǎng)絡(luò)進行訓練。此時結(jié)合傳統(tǒng)知識蒸餾訓練學生網(wǎng)絡(luò)的損失函數(shù)為：

其中j 為重建特征層的分組數(shù)；β為特征重建損失函數(shù)組在訓練中所占的比重；Lkd和lfeat分別為上文所提到的傳統(tǒng)知識蒸餾損失函數(shù)以及特征重建損失函數(shù)。

3 實驗結(jié)果及分析

3.1 在Cifar10上的實驗結(jié)果及分析

（1）利用知識蒸餾新架構(gòu)進行知識蒸餾的可行性分析

在使用數(shù)據(jù)集Cifar10 的實驗中，分別選用學生網(wǎng)絡(luò)為 16 層（WRN-16-1 和 WRN-16-2），教師網(wǎng)絡(luò)為40 層（WRN-40-1 和 WRN-40-2）的寬度殘差網(wǎng)絡(luò)結(jié)構(gòu)，其中1 和2 表示網(wǎng)絡(luò)寬度為基準寬度的倍數(shù)。

第一組實驗學生網(wǎng)絡(luò)選用WRN-16-1，教師網(wǎng)絡(luò)選用WRN-40-1，特征重建算法（Feature Reconstruction，F(xiàn)R）公式（5）中選用β為 0.01；若結(jié)合傳統(tǒng)知識蒸餾（Knowledge Distillation，KD）算法（即 FR+KD 算法），則選用β為 0.1，公式（3）中選用參數(shù)α為 0.9，T 為 4。

第二組實驗學生網(wǎng)絡(luò)選用WRN-16-2，教師網(wǎng)絡(luò)選用WRN-40-2，特征重建算法中選用β為0.01；結(jié)合傳統(tǒng)知識蒸餾算法后選用β為0.2，其中公式（3）中仍然選用參數(shù)α為 0.9，T 為 4。

表1 為對比實驗的實驗結(jié)果。（其中base 和teacher 分別問為不使用任何蒸餾算法的淺層神經(jīng)網(wǎng)絡(luò)和深層神經(jīng)網(wǎng)絡(luò)的錯誤率）。

表1 知識蒸餾新架構(gòu)的錯誤率（單位：%）

由表1 可見，在Cifar10 數(shù)據(jù)集上使用基于特征重建算法的新蒸餾框架將學生網(wǎng)絡(luò)的準確率提升了將近0.5%，而結(jié)合傳統(tǒng)知識蒸餾算法后，新架構(gòu)的準確率更是提高到1%。

而其中WRN-16-1 的參數(shù)量為0.2M，WRN-40-1的參數(shù)量為 0.6M；WRN-16-2 的參數(shù)量為 0.7M，WRN-40-2 的參數(shù)量為2.2M。由此可見蒸餾后在保證準確率的前提下，學生網(wǎng)絡(luò)的參數(shù)量降低到了教師網(wǎng)絡(luò)參數(shù)量的1/3，甚至更低。并且網(wǎng)絡(luò)的深度還降低到了原教師網(wǎng)絡(luò)的一半多。

（2）不同知識蒸餾算法間的對比實驗

為了驗證上述知識蒸餾架構(gòu)的有效性性，針對基于注意力轉(zhuǎn)移機制[8]（Attention Transfer，AT）的知識蒸餾方法以及傳統(tǒng)知識蒸餾算法做對比實驗。實驗結(jié)果如表2 所示。在實驗中使用了控制變量法，保證了不同算法間使用相同的神經(jīng)網(wǎng)絡(luò)以及相同的訓練方法（隨機梯度下降法），等等。

表2 各種知識蒸餾算法的錯誤率（單位：%）

3.2 在Cifar100上的實驗結(jié)果及分析

Cifar100 相比于Cifar10 就是同樣的數(shù)據(jù)集細分為了100 類，增加了分類任務(wù)的難度。由此，可以實驗出在其它的復(fù)雜分類任務(wù)中此框架是否還可適用。在Cifar100 上同樣選用學生網(wǎng)絡(luò)為 16 層（WRN-16-1），教師網(wǎng)絡(luò)為40 層（WRN-40-1）的寬度殘差網(wǎng)絡(luò)結(jié)構(gòu)。這兩個網(wǎng)絡(luò)的深度以及參數(shù)量都是與上述在Cifar10 中使用的網(wǎng)絡(luò)是相同的。

其中，特征重建算法中選用β為0.01；結(jié)合傳統(tǒng)知識蒸餾算法后，則選用為0.1，公式（3）中選用參數(shù)α為0.9，T 為4。其中公式（3）中參數(shù)的選擇仍參照經(jīng)典知識蒸餾方法中的實驗結(jié)果。表3 為各種算法的實驗結(jié)果。

表3 在Cifar100 中各種算法的錯誤率（單位：%）

由表3 可得，在復(fù)雜分類任務(wù)中，此框架的表現(xiàn)效果優(yōu)異，不僅將學生網(wǎng)絡(luò)的準確率提高到了10%以上，而且基于此框架訓練出的學生網(wǎng)絡(luò)的準確率要高于目前較常用的傳統(tǒng)知識蒸餾框架。相比與未壓縮的原網(wǎng)絡(luò)而言，模型大小僅為原模型的1/3，而準確率僅僅損失了3%。相比與未使用蒸餾模型訓練出的小模型準確率有了極大的提升。

4 結(jié)語

本文所提出的新知識蒸餾框架更進一步發(fā)揮了教師網(wǎng)絡(luò)強大的特征提取能力，對學生網(wǎng)絡(luò)的特征進行重建，從而使學生網(wǎng)絡(luò)的準確率方面超越了大部分傳統(tǒng)知識蒸餾架構(gòu)。該構(gòu)架在學生網(wǎng)絡(luò)的準確率方面，與原教師網(wǎng)絡(luò)的準確率相比損失了1%-3%；在壓縮效率方面，學生網(wǎng)絡(luò)的參數(shù)量以及網(wǎng)絡(luò)深度都降低為了原教師網(wǎng)絡(luò)的一半以上。由此，新的蒸餾框架在準確率以及壓縮效率上都有了一定的提高。

雖然新框架取得了一定的實驗成果，但是由于學生網(wǎng)絡(luò)的特征是通過人工分組嵌入教師網(wǎng)絡(luò)中進行特征重建，這在一定程度上限制了學生網(wǎng)絡(luò)的結(jié)構(gòu)，學生網(wǎng)絡(luò)只能采用與教師網(wǎng)絡(luò)寬度相同的網(wǎng)絡(luò)結(jié)構(gòu)，這也是此架構(gòu)存在的缺點。

盡管新架構(gòu)存在一定的缺點，但是本實驗為知識蒸餾的研究提供了一種新思路，并且往后將對此架構(gòu)進行進一步的完善。