【NLP】模型壓縮與蒸餾！BERT的忒修斯船
作者?|?許明???
整理 | NewBeeNLP公眾號(hào)
如果忒修斯的船上的木頭被逐漸替換，直到所有的木頭都不是原來(lái)的木頭，那這艘船還是原來(lái)的那艘船嗎？
-普魯塔克
最近遇到一個(gè)需要對(duì)算法加速的場(chǎng)景，了解到了一個(gè)比較簡(jiǎn)潔實(shí)用的方法：Bert-of-theseus[1], 了解了原理后參考代碼實(shí)驗(yàn)后，驗(yàn)證了其有效性，總結(jié)一下。
模型壓縮
模型在設(shè)計(jì)之初都是過(guò)參數(shù)化的，這是因?yàn)槟Ｐ偷膮?shù)量與復(fù)雜度代表著模型的容量與學(xué)習(xí)能力，但當(dāng)我們實(shí)際使用時(shí)，我們需要更好的部署他（低資源），更快的響應(yīng)（快速推理），常常需要進(jìn)行模型壓縮。
模型壓縮就是簡(jiǎn)化大的模型，得到推理快資源占用低的小模型，而想"即要馬而跑又不用吃草"通常是很難的，所以壓縮后的模型常常也會(huì)有不同程度的犧牲，如模型性能下降。此外，模型壓縮是作用在推理階段，帶來(lái)的常常是訓(xùn)練時(shí)間的增加。
模型壓縮又分為幾種方式：一種是剪枝(Pruning)與量化(Quantization),一種是知識(shí)蒸餾(Knowledge Distillation),
還有一種是權(quán)重共享（Sharing）與因數(shù)分解（Factorization）。該部分內(nèi)容推薦一篇博客：All The Ways You Can Compress BERT[2]
剪枝
剪枝技術(shù)是通過(guò)將大模型中一些"不重要"的連接剪斷，得到一個(gè)"稀疏"結(jié)構(gòu)的模型。剪枝又分為"結(jié)構(gòu)性剪枝"與"非結(jié)構(gòu)性剪枝".剪枝可以作用在權(quán)重粒度，
也可以作用在attention heads / layer粒度上。不過(guò)剪枝技術(shù)感覺(jué)會(huì)逐步被
量化
量化不改變模型的網(wǎng)絡(luò)結(jié)構(gòu)，而是改變模型的參數(shù)的數(shù)據(jù)格式，通常模型在建立與訓(xùn)練時(shí)使用的是 float32 格式的，量化就是將格式轉(zhuǎn)換為 low-bit, 如 float16 甚至二值化，如此即提速又省顯存。
知識(shí)蒸餾
知識(shí)蒸餾是訓(xùn)練一個(gè)小模型(student)來(lái)學(xué)習(xí)大模型(teacher)，由于大模型是之前已經(jīng)fine-tuning的，所以此時(shí)學(xué)習(xí)的目標(biāo)已經(jīng)轉(zhuǎn)換為對(duì)應(yīng)的logit而不再是one-hot編碼了，所以student有可能比teacher的性能更好。這樣即小又準(zhǔn)的模型實(shí)在太好了。
不過(guò)為了達(dá)到這樣的效果，通常設(shè)計(jì)小模型時(shí)不光要學(xué)習(xí)大模型的輸出，還要學(xué)習(xí)各個(gè)中間層結(jié)果，相關(guān)矩陣等，這就需要仔細(xì)設(shè)計(jì)模型的結(jié)構(gòu)與loss及l(fā)oss融合方案了。一種簡(jiǎn)單的方法是只學(xué)習(xí)大模型的logit，這與對(duì)label做embedding有點(diǎn)類(lèi)似，不過(guò)我沒(méi)做過(guò)實(shí)驗(yàn)還。
權(quán)重共享
將部分權(quán)重在多個(gè)層中共享以達(dá)到壓縮模型的效果，如ALBERT中共享self-attention中的參數(shù)
權(quán)重分解
將權(quán)重矩陣進(jìn)行因數(shù)分解，形成兩個(gè)低秩的矩陣相乘的形式，從而降低計(jì)算量
模型壓縮的必要性
看了上面模型壓縮的方法，每一個(gè)都有種"脫褲子放屁"的感覺(jué)，與其訓(xùn)練一個(gè)大模型，再費(fèi)力把它變小，為何不直接開(kāi)始就弄個(gè)小的呢？
首先，模型在設(shè)計(jì)之初是都是會(huì)或多或少的過(guò)參數(shù)化，因?yàn)槟Ｐ偷膮?shù)量與復(fù)雜度代表著模型的容量與學(xué)習(xí)能力； 其次，開(kāi)始就用一個(gè)小模型，那這個(gè)小模型也是需要設(shè)計(jì)的，不能隨便拿來(lái)一個(gè)，而設(shè)計(jì)一個(gè)性能高參數(shù)規(guī)模小的小模型難度是非常大的，往往是模型小了性能也低了； 第三點(diǎn)，大模型壓縮后與小模型雖然參數(shù)規(guī)模相當(dāng)，但是對(duì)應(yīng)的模型空間并不相同 此外，為了更好的部署，如手機(jī)或FPGA等，得到精度更高模型更小(distillation)或者利用硬件加速(low-bit)，模型壓縮都是值得試一試的手段。
更詳細(xì)的討論，可以參考為什么要壓縮模型，而不直接訓(xùn)練一個(gè)小的CNN[3]
Bert of theseus
論文：BERT-of-Theseus: Compressing BERT by Progressive Module Replacing 地址：https://arxiv.org/abs/2002.02925 arxiv訪(fǎng)問(wèn)不方便的同學(xué)公眾號(hào)后臺(tái)回復(fù)『0024』直接獲取
Bert of theseus 方法屬于上面提到的知識(shí)蒸餾，知識(shí)蒸餾中我們提到，在蒸餾時(shí)，我們不光要學(xué)習(xí)teacher的輸出，對(duì)中間層我們也希望他們直接盡量相似。
那想象一個(gè)這種狀態(tài)對(duì)應(yīng)的理想情況：中間層的結(jié)果一致，最終的結(jié)果一致,既然我們的期望中間結(jié)果一致，那也就意味著兩者可以互相替換。正如開(kāi)頭提到的忒修斯之船一樣。所以核心思想是：「與其設(shè)計(jì)復(fù)雜的loss來(lái)讓中間層結(jié)果相似不如直接用小模型替換大模型來(lái)訓(xùn)練」
通過(guò)復(fù)雜loss來(lái)達(dá)到與中間層結(jié)果相似可以看作是一種整體漸進(jìn)式的逼近，讓小模型一點(diǎn)點(diǎn)去學(xué)習(xí)，而直接替換可以看作是一種簡(jiǎn)單粗暴的方式，
但是他不需要設(shè)計(jì)各種loss，優(yōu)化目標(biāo)也是同一個(gè)，就只有一個(gè)下游任務(wù)相關(guān)的loss，突出一個(gè)簡(jiǎn)潔。
這就好比高中上學(xué)一樣，即使花高價(jià)也要讓孩子去一所好高中，因?yàn)閷W(xué)校的"氛圍"能讓孩子的學(xué)習(xí)成績(jī)進(jìn)步，其實(shí)是因?yàn)橹車(chē)暮⒆訋е黄饘W(xué)，弱雞也能學(xué)的比平時(shí)更多一點(diǎn)。bert-of-theseus也是類(lèi)似的道理，跟著大佬（teacher）總比單獨(dú)fine-tuning效果好。
具體流程
如果直接將小模型替換大模型，那其實(shí)是在對(duì)小模型進(jìn)行微調(diào)，與大模型就脫離了，也達(dá)不到對(duì)應(yīng)的效果，所以作者采用了一種概率替換的方式。
首先呢，想象我們現(xiàn)在已經(jīng)訓(xùn)練好了一個(gè)6層的BERT，我們成為Predecessor（前輩）, 而我們需要訓(xùn)練一個(gè)三層的bert，
他的結(jié)果近似12層BERT的效果，我們成為Successor(傳承者),那 bert-of-theseus的模型結(jié)構(gòu)如下圖[4]所示：
在bert-of-theseus中，首先固定predecessor的權(quán)重，然后將6層的Bert分為3個(gè)block，每個(gè)block與successor的一層對(duì)應(yīng)，訓(xùn)練過(guò)程分為兩個(gè)stage：
首先用successor中的層概率替換predecessor中對(duì)應(yīng)的block，在下游任務(wù)中直接fine-tuning（只訓(xùn)練successor）， 然后將successor從bert-of-theseus中分離出來(lái)，單獨(dú)在下游任務(wù)中進(jìn)行fine-tuning，直到指標(biāo)不再上升。
所謂替換，就是輸出的替換，在進(jìn)入下一層前在predecessor和successor的輸出中二選一。替換概率作者也給出了兩種方式，一種是固定 0.5,一種是線(xiàn)性從0-1,如下圖所示：
實(shí)驗(yàn)效果
實(shí)驗(yàn)代碼主要參考bert-of-theseus[5], 實(shí)驗(yàn)主要做了三組，一組文本分類(lèi)兩組ner-crf，結(jié)果如下：
文本分類(lèi)：CLUE的iflytek數(shù)據(jù)集
ner-crf: 公司數(shù)據(jù)
可以看到，相比直接那前幾層微調(diào)，bert-of-theseus的效果確實(shí)更好，此外，我還嘗試了線(xiàn)性策略的替換概率，效果上差別不大。
實(shí)驗(yàn)代碼：
classification_ifytek_bert_of_theseus[6] sequence_labeling_ner_bert_of_theseus[7]
本文參考資料
Bert-of-theseus: https://arxiv.org/abs/2002.02925
[2]All The Ways You Can Compress BERT: http://mitchgordon.me/machine/learning/2019/11/18/all-the-ways-to-compress-BERT.html
[3]為什么要壓縮模型，而不直接訓(xùn)練一個(gè)小的CNN: https://www.zhihu.com/question/303922732
[4]下圖: https://spaces.ac.cn/archives/7575
[5]bert-of-theseus: https://github.com/bojone/bert-of-theseus
[6]classification_ifytek_bert_of_theseus: https://github.com/xv44586/toolkit4nlp/blob/master/examples/classification_ifytek_bert_of_theseus.py
[7]sequence_labeling_ner_bert_of_theseus: https://github.com/xv44586/toolkit4nlp/blob/master/examples/sequence_labeling_ner_bert_of_theseus.py
-?END?-
往期精彩回顧
本站qq群704220115，加入微信群請(qǐng)掃碼：