1. 程式人生 > 其它 >模型壓縮 -- 知識蒸餾

模型壓縮 -- 知識蒸餾

輕量化網路已經是一個熱點,主要的技術路線如下:

  • 1、壓縮已經訓練好的大模型:知識蒸餾、剪枝(權重剪枝、通道剪枝)、權值量化、注意力遷移
  • 2、重新設計輕量化模型:SqueezeNet、MobileNet系列、ShuffleNet系列、EfficientNet、EfficientDet等
  • 3、加速卷積運算:im2col+GEMM、winograd、低秩分解
  • 4、硬體部署:TensorRT、OpenVIno、JetSon、Tensorflow-Lite、FPGA、積體電路、嵌入式開發、無人駕駛

追求模型在引數量、計算量、記憶體訪問量等方面耗時、能耗盡可能的減小,以便在邊緣裝置上進行部署


本節來了解下蒸餾的開山之作:知識蒸餾是由Hinton老爺子在2015年首次在《Distilling the Knowledge in a Neural Network》論文中提出的,論文地址:

https://arxiv.org/pdf/1503.0253

知識蒸餾就是把一個大模型或者多個模型ensemble學到的知識遷移到另一個輕量級的單模型上,最主要的目的是為了方便線上部署。知識蒸餾主要有兩個方面:第一個是將大而深的模型遷移到一個輕量級的小模型上。這就像我們線上把大而深的BERT模型學到的知識遷移到輕量級的TextCNN小模型上;另一個就是將多個模型ensemble學到的知識遷移到單個輕量級的模型。

問題:

  1. 知識蒸餾要解決什麼問題
  2. 什麼是知識蒸餾?它具有什麼優勢?和其他網路的區別?
  3. 具體的創新點在哪裡
  4. 效果如何

現在來回到上面這幾個問題:

1、知識蒸餾要解決什麼問題?或者為什麼要提出知識蒸餾?

簡單來說,它就是為了實現更加輕量化的網路而提出的。

在訓練過程中,我們需要使用複雜的模型,大量的計算資源,以便從非常大、高度冗餘的資料集中提取出資訊。在實驗中,效果最好的模型往往規模很大,甚至由多個模型整合得到。而大模型不方便部署到服務中去,常見的瓶頸如下:
• 推斷速度慢
• 對部署資源要求高(記憶體,視訊記憶體等) • 在部署時,我們對延遲以及計算資源都有著嚴格的限制。

因此,模型壓縮(在保證效能的前提下減少模型的引數量)成為了一個重要的問題。而”模型蒸餾“屬於模型壓縮的一種方法

一般地,大模型往往都是單個複雜網路或者是若干網路的集合,具有良好的效能和泛化能力;而小模型因為網路規模較小,表達能力有限。因此,可以利用大模型學習到的知識去指導小模型訓練,使得小模型具有與大模型相當的效能,但是引數數量大幅降低,從而實現模型壓縮與加速,這就是知識蒸餾與遷移學習在模型優化中的應用。

舉個例子:
BERT這一類模型優點在於效果好,但是如果用於線上推理就比較麻煩了,因為基礎版本的BERT模型接近330M包含一億的引數,你想讓一個一億引數的模型完成線上10ms內的線上推理基本有點不現實。而傳統的文字分類演算法比如TextCNN可以輕鬆滿足線上推理的需求,但是效果相比BERT還是有點不如人意。知識蒸餾通俗的理解就是BERT當老師,TextCNN當學生,讓BERT這個老師把學到的知識傳授給TextCNN這個學生,這樣就能讓TextCNN達到和BERT媲美的效果,最後我們線上去部署TextCNN,就能做到模型效果和線上推理速度兼得。這就是知識蒸餾的作用。

2、什麼是知識蒸餾?它具有什麼優勢?它為什麼有效?

知識蒸餾本質上是一種對映關係,將老師學到的東西傳遞給學生網路

Hinton等人在《Distilling the Knowledge in a Neural Network》一文中提出的知識蒸餾概念,其核心思想是先訓練一個複雜網路模型,然後使用這個複雜網路的輸出和資料的真實標籤去訓練一個更小的網路,因此知識蒸餾框架通常包含了一個複雜模型(被稱為Teacher模型)和一個小模型(被稱為Student模型)。

2.1 它具有什麼優勢?

一個好的模型最重要的是通過訓練資料獲得一定的泛化能力,不僅僅是擬合訓練資料,最重要的是在新資料集上能有一定的泛化識別能力。而知識蒸餾的目的是讓學生去學習老師的這種泛化能力,所以從理論上來說學生比老師單純的去擬合訓練資料能獲得更多的知識。

(1)提升模型精度

如果對目前的網路模型A的精度不是很滿意,那麼可以先訓練一個更高精度的teacher模型B(通常引數量更多,時延更大),然後用這個訓練好的teacher模型B對student模型A進行知識蒸餾,得到一個更高精度的A模型。

(2)降低模型時延,壓縮網路引數

如果對目前的網路模型A的時延不滿意,可以先找到一個時延更低,引數量更小的模型B,通常來講,這種模型精度也會比較低,然後通過訓練一個更高精度的teacher模型C來對這個引數量小的模型B進行知識蒸餾,使得該模型B的精度接近最原始的模型A,從而達到降低時延的目的。

(3)標籤之間的域遷移

假如使用狗和貓的資料集訓練了一個teacher模型A,使用香蕉和蘋果訓練了一個teacher模型B,那麼就可以用這兩個模型同時蒸餾出一個可以識別狗、貓、香蕉以及蘋果的模型,將兩個不同域的資料集進行整合和遷移。

2.2 為什麼有效


1) 對於複雜的模型,理論搜尋空間要大於較小網路的搜尋空間,但如果使用較小的網路可以實現相同(甚至相似)的收斂,則教師網路的收斂空間應與學生網路的解空間重疊,如上圖。

但僅此一項並不能保證學生網路在同一位置收斂,學生網路的收斂可能與教師網路的收斂大不相同。但如果指導學生網路複製教師網路的行為(教師網路已經在更大的解空間中進行搜尋了),則可以預期其收斂空間與原始教師網路收斂空間重疊。

2) softmax層的輸出,除了正例之外,負標籤也帶有大量的資訊,比如某些負標籤對應的概率遠遠大於其他負標籤(寶馬,兔子和垃圾車)。而在傳統的訓練過程(hard target)中,所有負標籤都被統一對待。也就是說,KD的訓練方式使得每個樣本給Net-S帶來的資訊量大於傳統的訓練方式。

3、創新體的創新點在哪裡

3.0 Hard-target和Soft-target

在Hinton這篇論文中,Hinton將問題限定在分類問題,分類問題的共同點是模型最後會有一個softmax層,其輸出值對應了相應類別的概率值。在知識蒸餾時,由於我們已經有了一個泛化能力較強的Teacher模型,我們在利用Teacher模型來蒸餾訓練Student模型時,可以直接讓Student模型去學習Teacher模型的泛化能力。一個很直白且高效的遷移泛化能力的方法就是:使用softmax層輸出的類別的概率來作為“Soft-target” 。

傳統的神經網路訓練方法是定義一個損失函式,目標是使預測值儘可能接近於真實值(Hard- target),損失函式就是使神經網路的損失值和儘可能小。這種訓練過程是對ground truth求極大似然。在知識蒸餾中,是使用大模型的類別概率作為Soft-target的訓練過程。

  • Hard-target:原始資料集標註的 one-shot 標籤,除了正標籤為 1,其他負標籤都是 0。
  • Soft-target:Teacher模型softmax層輸出的類別概率,每個類別都分配了概率,正標籤的概率最高。

知識蒸餾用Teacher模型預測的 Soft-target 來輔助 Hard-target 訓練 Student模型的方式為什麼有效呢?

softmax層的輸出,除了正例之外,負標籤也帶有Teacher模型歸納推理的大量資訊,比如某些負標籤對應的概率遠遠大於其他負標籤,則代表 Teacher模型在推理時認為該樣本與該負標籤有一定的相似性。而在傳統的訓練過程(Hard-target)中,所有負標籤都被統一對待。也就是說,知識蒸餾的訓練方式使得每個樣本給Student模型帶來的資訊量大於傳統的訓練方式。

如在MNIST資料集中做手寫體數字識別任務,假設某個輸入的“2”更加形似"3",softmax的輸出值中"3"對應的概率會比其他負標籤類別高;而另一個"2"更加形似"7",則這個樣本分配給"7"對應的概率會比其他負標籤類別高。這兩個"2"對應的Hard-target的值是相同的,但是它們的Soft-target卻是不同的,由此我們可見Soft-target蘊含著比Hard-target更多的資訊。

在使用 Soft-target 訓練時,Student模型可以很快學習到 Teacher模型的推理過程;而傳統的 Hard-target 的訓練方式,所有的負標籤都會被平等對待。因此,Soft-target 給 Student模型帶來的資訊量要大於 Hard-target,並且Soft-target分佈的熵相對高時,其Soft-target蘊含的知識就更豐富。同時,使用 Soft-target 訓練時,梯度的方差會更小,訓練時可以使用更大的學習率,所需要的樣本也更少。這也解釋了為什麼通過蒸餾的方法訓練出的Student模型相比使用完全相同的模型結構和訓練資料只使用Hard-target的訓練方法得到的模型,擁有更好的泛化能力。

3.1 帶“溫度”的softmax

將要訓練的小模型稱為新模型,將以及訓練的大模型稱為原模型。

我們的目標是讓新模型與原模型的softmax輸出的分佈充分接近。直接這樣做是有問題的:在一般的softmax函式中,自然指數先拉大logits之間的差距,然後作歸一化,最終得到的分佈是一個arg max的近似 ,其輸出是一個接近one-hot的向量,其中一個值很大,其他的都很小。這種情況下,前面說到的「可能是垃圾車,但絕不是蘿蔔」這種知識的體現是非常有限的。相較類似one-hot這樣的硬性輸出,我們更希望輸出更「軟」一些。一種方法是直接比較logits來避免這個問題。具體地,對於每一條資料,記原模型產生的某個logits是,新模型產生的logits是 ,我們需要最小化

該《Distilling the Knowledge in a Neural Network》論文中提出了帶有“溫度”的softmax,作為一種廣義是softmax函式。
原始的softmax函式:

但要是直接使用softmax層的輸出值作為soft target, 這又會帶來一個問題: 當softmax輸出的概率分佈熵相對較小時,負標籤的值都很接近0,對損失函式的貢獻非常小,小到可以忽略不計。因此"溫度"這個變數就派上了用場

其中 T 是溫度,這是從統計力學中的玻爾茲曼分佈中借用的概念。容易證明,當溫度 T 趨向於0時,softmax輸出將收斂為一個one-hot向量;溫度 T 趨向於無窮時,softmax的輸出則更「軟」。因此,在訓練新模型的時候,可以使用較高的 T 使得softmax產生的分佈足夠軟,這時讓新模型(同樣溫度下)的softmax輸出近似原模型;在訓練結束以後再使用正常的溫度 T = 1 來預測。具體地,在訓練時我們需要最小化兩個分佈的交叉熵(Cross-entropy),記新模型利用公式 (2) 產生的分佈是 q ,原模型產生的分佈是 p ,則我們需要最小化

在化學中,蒸餾是一個有效的分離沸點不同的組分的方法,大致步驟是先升溫使低沸點的組分汽化,然後降溫冷凝,達到分離出目標物質的目的。在前面提到的這個過程中,我們先讓溫度 T 升高,然後在測試階段恢復「低溫」,從而將原模型中的知識提取出來,因此將其稱為是蒸餾,實在是妙。

當然,如果轉移時使用的是有標籤的資料,那麼也可以將標籤與新模型softmax分佈的交叉熵加入到損失函式中去。這裡需要將式 (3) 乘上一個\(T^2\),這是為了讓損失函式的兩項的梯度大致在一個數量級上(參考公式 (9) ),實驗表明這將大大改善新模型的表現(考慮到加入了更多的監督訊號)。

與直接優化logits差異相比
由公式(2)(3),對於交叉熵損失來說,其對於新模型的某個logit 的梯度是 \(z_i\)的梯度是

由於\(e^2 - 2\)\(x\)等價無窮小(\(x -> 0\) 時),易知,當 \(T\) 充分大時,有

假設所有logits對每個樣本都是零均值化的,即\(\sum_{j}z_j=\sum_j v_j = 0\),則有

PS:
溫度:
原來的softmax函式是T = 1的特例。 T越高,softmax的output probability distribution越趨於平滑,其分佈的熵越大,負標籤攜帶的資訊會被相對地放大,模型訓練將更加關注負標籤

3.2 通用的知識蒸餾方法

知識蒸餾訓練的具體方法如下圖所示,主要包括以下幾個步驟:

  • 第一步是訓練Net-T;
  • 第二步是在高溫T下,蒸餾Net-T的知識到Net-S
  • 訓練Net-T的過程很簡單,下面詳細講講第二步:高溫蒸餾的過程。高溫蒸餾過程的目標函式由distill loss(對應soft target)和student loss(對應hard target)加權得到。示意圖如

知識蒸餾示意圖(來自:https://nervanasystems.github.io/distiller/knowledge_distillation.html)

Net-T和Net-S同時輸入transfer set(這裡可以這裡可以直接複用訓練Net-T用到的training set),用Net-T產生的softmax distribution(with high temperature)來作為soft target,Net-S在相同溫度T條件下的softmax輸出和soft target的cross entropy就是Loss函式的一部分$$L_{soft}$$

Net-S在T=1的條件下的softmax輸出和ground truth的cross entropy就是Loss函式的第二部分$$L_{hard}$$

第二部分Loss必要性其實很好理解:Net-T也有一定的錯誤率,使用round truth可以有效降低錯誤被傳播給Net-S的可能。打個比方,老師雖然學識遠遠超過學生,但是他仍然有出錯的可能,而這時候如果學生在老師的教授之外,可以同時參考到標準答案,

【注意】
在Net-S訓練完畢後,做inference時其Softmax的溫度T要恢復到1

最後,\(\alpha\)\(\beta\) 是關於\({L_{soft}}\)\({L_{hard}}\)的權重,實驗發現,當\({L_{hard}}\) 權重較小時,能產生最好的效果,這是一個經驗性的結論。直接給出結論:由於\({L_{soft}}\) 貢獻的梯度大約為\({L_{hard}}\)\(\frac{1}{T^2}\),因此在同時使用Soft-target和Hard-target的時候,需要在\({L_{soft}}\)的權重上乘以\({T^2}\)的係數,這樣才能保證Soft-target和Hard-target貢獻的梯度量基本一致。

4、效果如何

Hinton等人做了三組實驗,其中兩組都驗證了知識蒸餾方法的有效性。在MNIST資料集上的實驗表明,即便有部分類別的樣本缺失,新模型也可以表現得很不錯,只需要修改相應的偏置項,就可以與原模型表現相當。在語音任務的實驗也表明,蒸餾得到的模型比從頭訓練的模型捕捉了更多資料集中的有效資訊,表現僅比整合模型低了0.3個百分點。總體來說知識蒸餾是一個簡單而有效的模型壓縮/訓練方法。這大體上是因為原模型的softmax提供了比one-hot標籤更多的監督訊號

參考:
https://www.bilibili.com/video/BV1gS4y1k7vj/?spm_id_from=333.788
https://zhuanlan.zhihu.com/p/90049906