NLP中的預訓練語言模型（四）—— 小型化bert（DistillBert, ALBERT, TINYBERT）

　　bert之類的預訓練模型在NLP各項任務上取得的效果是顯著的，可是由於bert的模型參數多，推斷速度慢等緣由，致使bert在工業界上的應用很難普及，針對預訓練模型作模型壓縮是促進其在工業界應用的關鍵，今天介紹三篇小型化bert模型——DistillBert, ALBERT, TINYBERT。

一，DistillBert

　　論文：DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter 

　　GitHub：暫無

　　DistillBert是在bert的基礎上用知識蒸餾技術訓練出來的小型化bert。總體上來講這篇論文仍是很是簡單的，只是引入了知識蒸餾技術來訓練一個小的bert。具體作法以下：

　　1）給定原始的bert-base做爲teacher網絡。

　　2）在bert-base的基礎上將網絡層數減半（也就是從原來的12層減小到6層）。

　　3）利用teacher的軟標籤和teacher的隱層參數來訓練student網絡。

　　訓練時的損失函數定義爲三種損失函數的線性和，三種損失函數分別爲：

　　1）$L_{ce}$。這是teacher網絡softmax層輸出的機率分佈和student網絡softmax層輸出的機率分佈的交叉熵（注：MLM任務的輸出）。

　　2）$L_{mlm}$。這是student網絡softmax層輸出的機率分佈和真實的one-hot標籤的交叉熵

　　3）$L_{cos}$。這是student網絡隱層輸出和teacher網絡隱層輸出的餘弦類似度值，在上面咱們說student的網絡層數只有6層，teacher網絡的層數有12層，所以我的認爲這裏在計算該損失的時候是用student的第1層對應teacher的第2層，student的第2層對應teacher的第4層，以此類推。

　　做者對student的初始化也作了些工做，做者用teacher的參數來初始化student的網絡參數，作法和上面相似，用teacher的第2層初始化student的第1層，teacher的第4層初始化student的第2層。

　　做者也解釋了爲何減少網絡的層數，而不減少隱層大小，做者認爲在現代線性代數框架中，在張量計算中，下降最後一維（也就是隱層大小）的維度對計算效率提高不大，反卻是減少層數，也提高計算效率。

　　另外做者在這裏移除了句子向量和pooler層，在這裏也沒有看到NSP任務的損失函數，所以我的認爲做者也去除了NSP任務（主要是不少人證實該任務並無什麼效果）。

　　總體上來講雖然方法簡單，可是效果仍是很不錯的，模型大小減少了40%（66M），推斷速度提高了60%，但性能只下降了約3%。

 

二，ALBERT 

　　論文：ALBERT: A LITE BERT FOR SELF-SUPERVISEDLEARNING OF LANGUAGE REPRESENTATIONS 

　　GitHub：https://github.com/brightmart/albert_zh

　　ALBERT主要是從模型架構上作的改變，能極大的減少模型大小，可是沒有提高推斷速度。ALBERT主要作了三點改變：Factorized embedding parameterization ，Cross-layer parameter sharing ，Inter-sentence coherence loss 。咱們先來介紹這三個改變：

　　Factorized embedding parameterization

　　在bert中採用的是embedding層的嵌入詞向量大小E等於隱層大小H，可是做者認爲embedding層只是作了詞嵌入，在這一層詞與詞之間是相互獨立的，詞嵌入後獲得的向量包含的信息也僅僅只有當前詞的信息，所以向量長度不須要那麼大（咱們經常使用的word2vec向量長度通常不超過300），可是隱層由於會和其餘詞計算self-attention，所以隱層詞對應的向量是包含了上下文信息的，此時含有的信息是很是豐富的，用小的向量容易丟失信息，須要將隱層大小設大一點。針對這樣的分析，做者認爲像bert中embedding層和隱層大小設置爲相等是不合理的（bert-base中，E=H=768）。並且在embedding層，有一個大的詞嵌入矩陣V x E。在這裏V爲vocab size（一般比較大，bert中就是20000多），所以當E很大時，這裏的參數就很是多，又基於上面的分析，E能夠不用設這麼大，所以做者在這裏作了一個矩陣分解，將矩陣V x H(E)分解爲兩個小的矩陣V x E，E x H，E << H。在這裏再也不將E=H，而是將E設置爲一個遠小於H的值，而後再通過一個矩陣E x H將詞向量維度映射到H。

　　Cross-layer parameter sharing

　　再Transformer結構中，咱們也能夠選擇共享一些參數，不如共享self-attention中的參數，或者是共享FFN層的參數，在這裏做者共享了encoder中全部的參數（包括self attention和FFN），結合上面的矩陣分解，模型參數獲得了極大的減少，具體結果以下圖：

　　　　

　　另外這種層之間的參數共享機制，也讓各層的輸入輸出的L2距離和餘弦距離變化的更加平滑，以下圖所示：

　　　　

　　Inter-sentence coherence loss

　　上面兩種改變主要是在減少模型的參數，而這裏的改變主要是提升模型在下游任務上的性能，做者認爲在bert中的NSP任務其實是有問題的，當初設計這個任務時，是但願能學到句子一致性的，但事實上因爲當初在構造句子對負例時，是從不一樣的document中選擇的句子，所以網絡並不須要學到句子的一致性，而只要學到句子的主題（由於不一樣的document，主題極可能不同），就能夠判斷兩個句子是不是負例，所以做者認爲NSP任務只學到了句子的主題，而主題分類是一個淺層語義的NLP任務。

　　所以在這裏做者提出了一種新的任務SOP（句子順序預測），在這裏構造正例和bert中同樣，而負例就是將正例中兩個句子的順序顛倒，這裏應爲負例中的句子都來自於同一文檔，因此經過句子的主題是沒法區分正例和負例的，須要理解句子的深層語義才能區分，做者也給出了實驗證實：

　　　　

 　　如上面表中所示，學習SOP（86.5）任務，也能在NSP（78.9）任務上取得不錯的效果，可是學習NSP（90.5）任務，是沒法在SOP（52.0）任務上取得效果的。

　　以上就爲這篇論文改變的三個點，所以矩陣分解和參數共享，模型參數減少了不少，可是由於層數沒有變，在推斷時的計算量並無降低，所以推斷速度沒啥提高，可是訓練時的速度有提高。引入SOP任務也能提高模型的效果。此外，ALBERT中的mask方式，採用的是n-gram mask，其實和SpanBert中的span mask是同樣的，只是span中將mask的最大長度控制在10，而這裏將最大長度控制在3。另外優化算法採用了LAMB優化器。

　　最後上一個ALBERT的模型效果：

　　　　

　　　　

 　　這裏的ALBERT是用的ALBERT-xxlarge，其參數以下：

　　　　

 　　增大了隱層大小到4096，但實際上在這裏做者證實了並非全部的模型增大隱層大小，模型的效果都會提高，好比做者在bert-large上證實了提高隱層大小到2048時，模型降低不少：

　　　　

 

三，TINYBERT

　　論文：TINYBERT: DISTILLING BERT FOR NATURAL LANGUAGE UNDERSTANDING

　　GitHub：暫無

　　TINYBERT也是採用了知識蒸餾的方法來壓縮模型的，只是在設計上叫distillBert作了更多的工做，做者提出了兩個點：針對Transformer結構的知識蒸餾和針對pre-training和fine-tuning兩階段的知識蒸餾。

　　做者在這裏構造了四類損失函數來對模型中各層的參數進行約束來訓練模型，具體模型結構以下：

　　　　

 　　做者構造了四類損失，分別針對embedding layer，attention 權重矩陣，隱層輸出，predict layer。能夠將這個統一到一個損失函數中：

　　　　

 　　上面式子中$\lambda_m$表示每一層對應的係數，$S_m$表示studnet網絡的第m層，$T_{g(m)}$表示teacher網絡的第n層，其中$n = g(m)$。而且有$g(0) = 0$，$g(M+1) = N+1$，0表示embedding layer，M+1和N+1表示perdict layer。

　　針對上面四層具體的損失函數表達式以下：

　　attention 權重矩陣

　　　　

 　　h爲multi attention中頭數

　　隱層輸出

　　　　

 　　由於student網絡的隱層大小一般會設置的比teacher的小，所以爲了在計算時維度一致，這裏用一個矩陣$W_h$將student的隱層向量線性映射到和teacher一樣的空間下。

　　embedding layer

　　　　

 　　$W_s$同理上。

　　以上三種損失函數都採用了MSE，主要是爲了將模型的各項參數對齊。

　　predict layer

　　　　

 　　predict layer也就是softmax層，在這裏的損失函數是交叉熵，t是溫度參數，在這裏設置爲1。

　　以上四種損失函數是做者針對transformer提出的知識蒸餾方法。除此以外做者認爲除了對pre-training蒸餾以外，在fine-tuning時也利用teacher的知識來訓練模型能夠取得在下游任務更好的效果。所以做者提出了兩階段知識蒸餾，以下圖所示：

　　　　

 

 　　本質上就是在pre-training蒸餾一個general TinyBERT，而後再在general TinyBERT的基礎上利用task-bert上再蒸餾出fine-tuned TinyBERT。

　　做者給出了TinyBERT的效果：

　　　　

 　　另外做者也給出了四種損失對最終結果的貢獻：

　　　　

 　　還有就是關於$n = g(m)$這個式子中$g(m)$怎麼選擇，假設student的層數爲4層，這裏的$n = g(m) = 3m$，做者將這種稱爲Uniform-strategy。另外做者還和其餘的$g(m)$作了對比：

　　　　

 　　Top-strategy指用teacher最後4層，Bottom-strategy指用前面4層，其實這裏的映射函數，我感受可能還有更優的方案，例如取平均，或者用attention來作，可能效果會更好。