知識蒸餾的核心思想是通過遷移知識,從而通過訓練好的大模型得到更加適合推理的小模型。本文作者介紹了知識蒸餾的全過程,以及引用Hinton等人的實驗結果解釋說明,并提出了知識蒸餾的一些延伸工作方向。
0. 寫在前面
有人說過:“神經網絡用剩的logits不要扔,沾上雞蛋液,裹上面包糠...” 這兩天對知識蒸餾(Knowledge Distillation)萌生了一點興趣,正好寫一篇文章分享一下。這篇文章姑且算是一篇小科普。
1. 從模型壓縮開始
各種模型算法,最終目的都是要為某個應用服務。在買賣中,我們需要控制收入和支出。類似地,在工業應用中,除了要求模型要有好的預測(收入)以外,往往還希望它的「支出」要足夠小。具體來說,我們一般希望部署到應用中的模型使用較少的計算資源(存儲空間、計算單元等),產生較低的時延。
在深度學習的背景下,為了達到更好的預測,常常會有兩種方案:1. 使用過參數化的深度神經網絡,這類網絡學習能力非常強,因此往往加上一定的正則化策略(如dropout);2. 集成模型(ensemble),將許多弱的模型集成起來,往往可以實現較好的預測。這兩種方案無疑都有較大的「支出」,需要的計算量和計算資源很大,對部署非常不利。這也就是模型壓縮的動機:我們希望有一個規模較小的模型,能達到和大模型一樣或相當的結果。當然,從頭訓練一個小模型,從經驗上看是很難達到上述效果的,也許我們能先訓練一個大而強的模型,然后將其包含的知識轉移給小的模型呢?如何做到呢?
* 下文統一將要訓練的小模型稱為新模型,將以及訓練的大模型稱為原模型。
Rich Caruana等人在[1]中指出,可以讓新模型近似(approximate)原模型(模型即函數)。注意到,在機器學習中,我們常常假定輸入到輸出有一個潛在的函數關系,這個函數是未知的:從頭學習一個新模型就是從有限的數據中近似一個未知的函數。如果讓新模型近似原模型,因為原模型的函數是已知的,我們可以使用很多非訓練集內的偽數據來訓練新模型,這顯然要更可行。
這樣,原來我們需要讓新模型的softmax分布與真實標簽匹配,現在只需要讓新模型與原模型在給定輸入下的softmax分布匹配了。直觀來看,后者比前者具有這樣一個優勢:經過訓練后的原模型,其softmax分布包含有一定的知識——真實標簽只能告訴我們,某個圖像樣本是一輛寶馬,不是一輛垃圾車,也不是一顆蘿卜;而經過訓練的softmax可能會告訴我們,它最可能是一輛寶馬,不大可能是一輛垃圾車,但絕不可能是一顆蘿卜[2]。
2. 為什么叫「蒸餾」?
接續前面的討論,我們的目標是讓新模型與原模型的softmax輸出的分布充分接近。直接這樣做是有問題的:在一般的softmax函數中,自然指數 先拉大logits之間的差距,然后作歸一化,最終得到的分布是一個arg max的近似 (參考我之前的文章:淺談Softmax函數),其輸出是一個接近one-hot的向量,其中一個值很大,其他的都很小。這種情況下,前面說到的「可能是垃圾車,但絕不是蘿卜」這種知識的體現是非常有限的。相較類似one-hot這樣的硬性輸出,我們更希望輸出更「軟」一些。
一種方法是直接比較logits來避免這個問題。具體地,對于每一條數據,記原模型產生的某個logits是 ,新模型產生的logits是 ,我們需要最小化
文獻[2]提出了更通用的一種做法。考慮一個廣義的softmax函數
其中 是溫度,這是從統計力學中的玻爾茲曼分布中借用的概念。容易證明,當溫度 趨向于0時,softmax輸出將收斂為一個one-hot向量(證明可以參考我之前的文章:淺談Softmax函數,將 替換為 即可);溫度 趨向于無窮時,softmax的輸出則更「軟」。因此,在訓練新模型的時候,可以使用較高的 使得softmax產生的分布足夠軟,這時讓新模型的softmax輸出近似原模型;在訓練結束以后再使用正常的溫度 來預測。具體地,在訓練時我們需要最小化兩個分布的交叉熵(Cross-entropy),記新模型利用公式 產生的分布是 ,原模型產生的分布是 ,則我們需要最小化
在化學中,蒸餾是一個有效的分離沸點不同的組分的方法,大致步驟是先升溫使低沸點的組分汽化,然后降溫冷凝,達到分離出目標物質的目的。在前面提到的這個過程中,我們先讓溫度升高,然后在測試階段恢復「低溫」,從而將原模型中的知識提取出來,因此將其稱為是蒸餾,實在是妙。
當然,如果轉移時使用的是有標簽的數據,那么也可以將標簽與新模型softmax分布的交叉熵加入到損失函數中去。這里需要將式 乘上一個 ,這是為了讓損失函數的兩項的梯度大致在一個數量級上(參考公式 ),實驗表明這將大大改善新模型的表現(考慮到加入了更多的監督信號)。
3. 與直接優化logits差異相比
由公式 ,對于交叉熵損失來說,其對于新模型的某個logit 的梯度是
由于 與 是等價無窮小(時**)**,易知,當 充分大時,有
假設所有logits對每個樣本都是零均值化的,即 ,則有
所以,如果:1. 非常大,2. logits對所有樣本都是零均值化的,則知識蒸餾和最小化logits的平方差(公式 )是等價的(因為梯度大致是同一個形式)。實驗表明,溫度 不能取太大,而應該使用某個適中的值,這表明忽略極負的logits對新模型的表現很有幫助(較低的溫度產生的分布比較「硬」,傾向于忽略logits中極小的負值)。
4. 實驗與結論
Hinton等人做了三組實驗,其中兩組都驗證了知識蒸餾方法的有效性。在MNIST數據集上的實驗表明,即便有部分類別的樣本缺失,新模型也可以表現得很不錯,只需要修改相應的偏置項,就可以與原模型表現相當。在語音任務的實驗也表明,蒸餾得到的模型比從頭訓練的模型捕捉了更多數據集中的有效信息,表現僅比集成模型低了0.3個百分點。總體來說知識蒸餾是一個簡單而有效的模型壓縮/訓練方法。這大體上是因為原模型的softmax提供了比one-hot標簽更多的監督信號[3]。
知識蒸餾在后續也有很多延伸工作。在NLP方面比較有名的有Yoon Kim等人的Sequence-Level Knowledge Distillation 等。總的來說,對一些比較臃腫、不便部署的模型,可以將其「知識」轉移到小的模型上。比如,在機器翻譯中,一般的模型需要有較大的容量(capacity)才可能獲得較好的結果;現在非常流行的BERT及其變種,規模都非常大;更不用提,一些情形下我們需要將這些本身就很大的深度模型集成為一個ensemble,這時候,可以用知識蒸餾壓縮出一個較小的、「便宜」的模型。
文章地址:https://arxiv.org/abs/1606.07947
另外,在多任務的情境下,使用一般的策略訓練一個多任務模型,可能達不到比單任務更好的效果,文獻[3]探索了使用知識蒸餾,利用單任務的模型來指導訓練多任務模型的方法,很值得參考。
補充
鑒于評論區有知友對公式 有疑問,簡單補充一下這里梯度的推導(其實就是交叉熵損失對softmax輸入的梯度,LOL)。
* 這部分有一點繁瑣,能接受公式 的讀者可以跳過。
由鏈式法則,有
注意到 是原模型產生的softmax輸出,與 無關。
后一項 比較容易得到,因為 ,所以
則 是一個 維向量
前一項 是一個 的方陣,分類討論可以得到。參考公式 ,記 ,由除法的求導法則,輸出元素 對輸入 的偏導是
注意上面右側加方框部分,可以進一步展開
這樣,代入公式 ,并且將括號展開,可以得到
左側方框內偏導可以分類討論得到
帶入式 ,得到
所以 形式如下
代入式 ,可得
所以有公式 , 。
參考
[1] Caruana et al., Model Compression, 2006
[2] Hinton et al., Distilling the Knowledge in a Neural Network, 2015
[3] Kevin Clark et al., BAM! Born-Again Multi-Task Networks for Natural Language Understanding
責任編輯:xj
原文標題:知識蒸餾是什么?一份入門隨筆
文章出處:【微信公眾號:深度學習自然語言處理】歡迎添加關注!文章轉載請注明出處。
-
神經網絡
+關注
關注
42文章
4771瀏覽量
100744 -
模型
+關注
關注
1文章
3238瀏覽量
48824 -
深度學習
+關注
關注
73文章
5503瀏覽量
121137
原文標題:知識蒸餾是什么?一份入門隨筆
文章出處:【微信號:zenRRan,微信公眾號:深度學習自然語言處理】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論