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

郭俊倫，彭書華，李俊杰

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

0 引言

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

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

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

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

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

1 理論方法

1.1 殘差網絡

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

圖1 殘差網絡的殘差塊結構

殘差網絡在普通卷積網絡中引入了殘差塊，由公式（1）可見引入殘差塊后會對原先網絡的訓練結果進行微調，從而抑制原網絡在訓練中產生的錯誤，修正深度神經網絡的訓練結果。

其中，Xl、Xl+1為殘差網絡第l 層和第l+1 層的卷積層運算結果；W 為殘差塊結構的網絡參數；F（·）為殘差塊中的卷積運算函數。

因為網絡壓縮實驗中需要使用深度神經網絡，并且為了方便與其他方法的實驗結果進行對比，選用了寬度殘差網絡結構（Wide Residual Net，WRN）即以模型的寬度作為基準的殘差網絡，進行相關實驗。

1.2 知識蒸餾

用于分類任務的深度神經網絡往往使用Softmax層來計算出類概率。模型的訓練往往采用的是人工標注的標簽，又叫硬標簽。知識蒸餾則是采用教師網絡Softmax 層的輸出作為小模型訓練的標簽，又叫軟標簽，并且在Softmax 中引入溫度參數T，此時基于知識蒸餾的Softmax 公式變為：

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

圖2 知識蒸餾架構

由圖2 可知，引入參數T 后的學生模型損失函數變為：

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

1.3 特征重建算法

Johnson 等人[13]在圖片風格遷移實驗中發現，兩張圖片輸入同一個網絡時，圖片的內容越相似則網絡的深度特征越相似。于是在圖片風格遷移實驗中提出了特征重建損失函數。此損失函數是由已經訓練好的VGGNet[14]計算得到的，將VGGNet 作為損失函數產生網絡，用來訓練圖片風格遷移網絡。風格遷移實驗內容遷移（即特征重方法）的訓練架構如圖3 所示。

圖3 圖片風格遷移之內容遷移訓練架構

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

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

2 使用特征重建算法設計知識蒸餾新架構

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

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

圖4 基于學生網絡特征重建的知識蒸餾架構設計

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

那么其中的利用特征重建生成損失函數的網絡框架的結構如圖5 所示。

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

使用此損失函數生成網絡生成損失函數對學生網絡進行訓練。此時結合傳統知識蒸餾訓練學生網絡的損失函數為：

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

3 實驗結果及分析

3.1 在Cifar10上的實驗結果及分析

（1）利用知識蒸餾新架構進行知識蒸餾的可行性分析

在使用數據集Cifar10 的實驗中，分別選用學生網絡為 16 層（WRN-16-1 和 WRN-16-2），教師網絡為40 層（WRN-40-1 和 WRN-40-2）的寬度殘差網絡結構，其中1 和2 表示網絡寬度為基準寬度的倍數。

第一組實驗學生網絡選用WRN-16-1，教師網絡選用WRN-40-1，特征重建算法（Feature Reconstruction，FR）公式（5）中選用β為 0.01；若結合傳統知識蒸餾（Knowledge Distillation，KD）算法（即 FR+KD 算法），則選用β為 0.1，公式（3）中選用參數α為 0.9，T 為 4。

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

表1 為對比實驗的實驗結果。（其中base 和teacher 分別問為不使用任何蒸餾算法的淺層神經網絡和深層神經網絡的錯誤率）。

表1 知識蒸餾新架構的錯誤率（單位：%）

由表1 可見，在Cifar10 數據集上使用基于特征重建算法的新蒸餾框架將學生網絡的準確率提升了將近0.5%，而結合傳統知識蒸餾算法后，新架構的準確率更是提高到1%。

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

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

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

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

3.2 在Cifar100上的實驗結果及分析

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

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

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

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

4 結語

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

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

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