【NLP】模型壓縮與蒸餾！BERT的忒修斯船

作者?|?許明???

整理 | NewBeeNLP公眾號

如果忒修斯的船上的木頭被逐漸替換，直到所有的木頭都不是原來的木頭，那這艘船還是原來的那艘船嗎？

-普魯塔克

最近遇到一個需要對算法加速的場景，了解到了一個比較簡潔實用的方法：Bert-of-theseus^[1], 了解了原理后參考代碼實驗后，驗證了其有效性，總結(jié)一下。

模型壓縮

模型在設(shè)計之初都是過參數(shù)化的，這是因為模型的參數(shù)量與復(fù)雜度代表著模型的容量與學(xué)習(xí)能力，但當(dāng)我們實際使用時，我們需要更好的部署他（低資源），更快的響應(yīng)（快速推理），常常需要進行模型壓縮。

模型壓縮就是簡化大的模型，得到推理快資源占用低的小模型，而想"即要馬而跑又不用吃草"通常是很難的，所以壓縮后的模型常常也會有不同程度的犧牲，如模型性能下降。此外，模型壓縮是作用在推理階段，帶來的常常是訓(xùn)練時間的增加。

模型壓縮又分為幾種方式：一種是剪枝(Pruning)與量化(Quantization),一種是知識蒸餾(Knowledge Distillation), 還有一種是權(quán)重共享（Sharing）與因數(shù)分解（Factorization）。該部分內(nèi)容推薦一篇博客：All The Ways You Can Compress BERT^[2]

剪枝

剪枝技術(shù)是通過將大模型中一些"不重要"的連接剪斷，得到一個"稀疏"結(jié)構(gòu)的模型。剪枝又分為"結(jié)構(gòu)性剪枝"與"非結(jié)構(gòu)性剪枝".剪枝可以作用在權(quán)重粒度， 也可以作用在attention heads / layer粒度上。不過剪枝技術(shù)感覺會逐步被NAS（Neural Architecture Search）取。

量化

量化不改變模型的網(wǎng)絡(luò)結(jié)構(gòu)，而是改變模型的參數(shù)的數(shù)據(jù)格式，通常模型在建立與訓(xùn)練時使用的是 float32 格式的，量化就是將格式轉(zhuǎn)換為 low-bit, 如 float16 甚至二值化，如此即提速又省顯存。

知識蒸餾

知識蒸餾是訓(xùn)練一個小模型(student)來學(xué)習(xí)大模型(teacher)，由于大模型是之前已經(jīng)fine-tuning的，所以此時學(xué)習(xí)的目標已經(jīng)轉(zhuǎn)換為對應(yīng)的logit而不再是one-hot編碼了，所以student有可能比teacher的性能更好。這樣即小又準的模型實在太好了。

不過為了達到這樣的效果，通常設(shè)計小模型時不光要學(xué)習(xí)大模型的輸出，還要學(xué)習(xí)各個中間層結(jié)果，相關(guān)矩陣等，這就需要仔細設(shè)計模型的結(jié)構(gòu)與loss及l(fā)oss融合方案了。一種簡單的方法是只學(xué)習(xí)大模型的logit，這與對label做embedding有點類似，不過我沒做過實驗還。

權(quán)重共享

將部分權(quán)重在多個層中共享以達到壓縮模型的效果，如ALBERT中共享self-attention中的參數(shù)

權(quán)重分解

將權(quán)重矩陣進行因數(shù)分解，形成兩個低秩的矩陣相乘的形式，從而降低計算量

模型壓縮的必要性

看了上面模型壓縮的方法，每一個都有種"脫褲子放屁"的感覺，與其訓(xùn)練一個大模型，再費力把它變小，為何不直接開始就弄個小的呢？

• 首先，模型在設(shè)計之初是都是會或多或少的過參數(shù)化，因為模型的參數(shù)量與復(fù)雜度代表著模型的容量與學(xué)習(xí)能力；
• 其次，開始就用一個小模型，那這個小模型也是需要設(shè)計的，不能隨便拿來一個，而設(shè)計一個性能高參數(shù)規(guī)模小的小模型難度是非常大的，往往是模型小了性能也低了；
• 第三點，大模型壓縮后與小模型雖然參數(shù)規(guī)模相當(dāng)，但是對應(yīng)的模型空間并不相同
• 此外，為了更好的部署，如手機或FPGA等，得到精度更高模型更小(distillation)或者利用硬件加速(low-bit)，模型壓縮都是值得試一試的手段。

更詳細的討論，可以參考為什么要壓縮模型，而不直接訓(xùn)練一個小的CNN^[3]

Bert of theseus

• 論文：BERT-of-Theseus: Compressing BERT by Progressive Module Replacing
• 地址：https://arxiv.org/abs/2002.02925
• arxiv訪問不方便的同學(xué)公眾號后臺回復(fù)『0024』直接獲取
  

Bert of theseus 方法屬于上面提到的知識蒸餾，知識蒸餾中我們提到，在蒸餾時，我們不光要學(xué)習(xí)teacher的輸出，對中間層我們也希望他們直接盡量相似。

那想象一個這種狀態(tài)對應(yīng)的理想情況：中間層的結(jié)果一致，最終的結(jié)果一致,既然我們的期望中間結(jié)果一致，那也就意味著兩者可以互相替換。正如開頭提到的忒修斯之船一樣。所以核心思想是：「與其設(shè)計復(fù)雜的loss來讓中間層結(jié)果相似不如直接用小模型替換大模型來訓(xùn)練」

通過復(fù)雜loss來達到與中間層結(jié)果相似可以看作是一種整體漸進式的逼近，讓小模型一點點去學(xué)習(xí)，而直接替換可以看作是一種簡單粗暴的方式， 但是他不需要設(shè)計各種loss，優(yōu)化目標也是同一個，就只有一個下游任務(wù)相關(guān)的loss，突出一個簡潔。

這就好比高中上學(xué)一樣，即使花高價也要讓孩子去一所好高中，因為學(xué)校的"氛圍"能讓孩子的學(xué)習(xí)成績進步，其實是因為周圍的孩子帶著一起學(xué)，弱雞也能學(xué)的比平時更多一點。bert-of-theseus也是類似的道理，跟著大佬（teacher）總比單獨fine-tuning效果好。

具體流程

如果直接將小模型替換大模型，那其實是在對小模型進行微調(diào)，與大模型就脫離了，也達不到對應(yīng)的效果，所以作者采用了一種概率替換的方式。

首先呢，想象我們現(xiàn)在已經(jīng)訓(xùn)練好了一個6層的BERT，我們成為Predecessor（前輩）, 而我們需要訓(xùn)練一個三層的bert， 他的結(jié)果近似12層BERT的效果，我們成為Successor(傳承者),那 bert-of-theseus的模型結(jié)構(gòu)如下圖^[4]所示：

在bert-of-theseus中，首先固定predecessor的權(quán)重，然后將6層的Bert分為3個block，每個block與successor的一層對應(yīng)，訓(xùn)練過程分為兩個stage：

• 首先用successor中的層概率替換predecessor中對應(yīng)的block，在下游任務(wù)中直接fine-tuning（只訓(xùn)練successor），
• 然后將successor從bert-of-theseus中分離出來，單獨在下游任務(wù)中進行fine-tuning，直到指標不再上升。

所謂替換，就是輸出的替換，在進入下一層前在predecessor和successor的輸出中二選一。替換概率作者也給出了兩種方式，一種是固定 0.5,一種是線性從0-1,如下圖所示：

實驗效果

實驗代碼主要參考bert-of-theseus^[5], 實驗主要做了三組，一組文本分類兩組ner-crf，結(jié)果如下：

文本分類：CLUE的iflytek數(shù)據(jù)集

ner-crf: 公司數(shù)據(jù)

可以看到，相比直接那前幾層微調(diào)，bert-of-theseus的效果確實更好，此外，我還嘗試了線性策略的替換概率，效果上差別不大。

實驗代碼：

• classification_ifytek_bert_of_theseus^[6]
• sequence_labeling_ner_bert_of_theseus^[7]

本文參考資料