Bert不完全手冊1. Bert推理太慢？模型蒸餾

模型蒸餾的目標主要用於模型的線上部署，解決Bert太大，推理太慢的問題。因此用一個小模型去逼近大模型的效果，實現的方式一般是Teacher-Stuent框架，先用大模型（Teacher）去對樣本進行擬合，再用小模型（Student）去模仿Teacher。

為什么蒸餾本身會比直接用小模型去擬合樣本取得更好的效果呢？這里先拋出幾個可能的方向

• 樣本降噪：大模型本身作為一個’BottleNeck‘，把難以學到的信息和噪聲樣本進行了一定程度過濾，降低了學習難度

• 樣本增強：把大模型的預測prob作為小模型的學習目標，一定程度上其實是利用大模型對原始樣本進行了數據增強，有效幫助小模型提高泛化性

• 學習目標調整：雖然小模型能力有限不能很好擬合目標，但蒸餾中核心的對齊Loss的存在使得小模型只要去模仿大模型即可，就像臨摹之於繪畫

知識蒸餾

• paper: Distilling the Knowledge in a Neural Network
• paper: Distilling Task-Specific Knowledge from BERT into Simple Neural Networks

說到蒸餾肯定要先提下Hinton在15年的paper，它並不是最早提出蒸餾概念的，但是它總結了幾點蒸餾的方案，之后被廣泛應用，主要包括以下兩點

1. Temperature Softmax

\[q_i = \frac{exp(z_i/T)}{\sum_j exp(z_j/T)} \]

想讓大模型保留更高的泛化性和樣本整體分布信息，可以用Temperature Softmax來smooth預測分布。Temperature就是以上的T參數，T越大模型輸出的probability越平滑，置信越低。因為越復雜的模型往往越傾向於在分類問題中給出更高置信度的預測結果，以二分類為例，大模型可能會給出[0.001, 0.999]這類高置信的預測結果，而加入smooth之后會得到類似 [0.2,0.8]這類熵值更大，且更多保留了樣本分布信息的預測結果。

如何理解更加smooth的分布會保留更多的樣本信息？這個用多分類可能更加直觀，例如在bert的MLM任務中，'天氣很MASK'這時對MASK位置進行預測，高置信地預測結果可能是p(MASK=好)=0.999，而Smooth之后可能會得到P(MASK=好)=0.5，P(MASK=冷)=0.3, P(MASK=熱)=0.15，很明顯后者保留的文本信息更加豐富，而豐富的文本信息能進一步幫助student模型保有更好的泛化性。

2. Distill Loss

\[Loss = L_{DS} + L_{CE} = CrossEntropy(p_i, q_i)+ \alpha CrossEntropy(p_i, y_i) \]

蒸餾的方式時讓小模型同時擬合兩個Loss。一個是蒸餾Loss，用相同的Temperature參數擬合以上大模型輸出的soft target(\(q_i\))，其實也就是最小化student和teacher輸出分布的KL散度，另一個擬合Loss是用T=1的常規softmax擬合真實的label (\(y_i\))。二者信息互為補充，大模型的soft target本身比Hard Label更容易擬合，因為大模型作為bottleNeck會過濾部分樣本噪音，同時Temperature也提供了更加smooth以及容易擬合的分布信息。而HardLabel提供了熵值更低的真實信息，幫助蒸餾模型學到正確的class分類。

KD直接應用在Bert到LSTM等小模型的蒸餾過程中，例如Distilling Task-Specific Knowledge from BERT into Simple Neural Networks。在我開發中的文本分類項目里有相關的實現SimpleClassification/knowledge_distill

PKD Bert

Paper: Patient Knowledge Distillation for BERT Model Compression
Github： https://github.com/intersun/PKD-for-BERT-Model-Compression

從PKD Bert開始，大家開始集思廣益對Bert開展瘦身行動。PKD在以上KD的基礎上主要的創新點是在蒸餾過程中除了對輸出層進行擬合之外，還加入中間層的對齊，作者稱之為patient Distill，主要包括以下幾個要點

1. 對哪些中間層進行對齊： 作者嘗試了PKD-SKIP和PKD-LAST兩種策略，對\(Bert_{12}\)來說，前者就是對[2,4,6,8,10]層進行對齊，后者就是對[7,8,9,10,11]層進行對齊。這背后主要是考慮Bert的信息分布，究竟是相對均勻的分布在各個層，還是頂層信息包含了底層信息。從效果上看是PKD- SKIP效果更好，之后的Distill 和TinyBert也都多少借鑒了這里的方案

2. 對中間層的那些信息進行對齊：這里作者只選擇了對齊CLS token

3. 用什么方式進行對齊：作者用了正則化之后的CLS token Embedding之間的歐式距離來進行對齊

以上的擬合方式作者稱為PT Loss，蒸餾的過程是在KD的基礎上加入了PT Loss，如下

\[L_{PKD} = (1-\alpha)L_{CE} + \alpha L_{DS} + \beta L_{PT} \]

PKD的局限性，是只針對下游任務進行蒸餾，也就是teacher模型是fine-tune Bert。並且對初始化student的方式也相對簡單，直接使用了\(Bert_3\) \(Bert_6\)的預訓練模型進行初始化，這里的初始化方式和以上對齊使用的PKD-Skip/LAST策略存在一定的不一致性。

Distill Bert

paper: DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
github：https://huggingface.co/docs/transformers/model_doc/distilbert

從DistillBert開始，蒸餾被提前到了預訓練階段。Distill Bert Base縮減了一半的層數，對比Bert Base，實現用40%更少的參數，在部分任務達到Bert97%左右的效果，比Bert預測快60%。

預訓練

Distill Bert的模型結構也是6層的Transformer，在訓練目標上Distill Bert使用了3個訓練目標的線性組合。包括

1. \(L_{mlm}\): Mask LM Loss

student模型用和Bert相同的MLM方式直接去學習預訓練任務。在后續的對比中，是否加入MLM對效果影響最小。這也印證了最初對蒸餾為何有效的討論，蒸餾弱化了student去直接學習任務的重要性，而是更多通過模仿teacher模型去學習有效信息

2. \({L_{ce}}\): 輸出層對齊

這里和KD相同，都是用了temperature softmax，除初始化之外，\(L_{ce}\)對模型效果的影響最大，所以模擬teacher輸出才是distill的制勝秘籍

3. \(L_{cos}\): 隱藏層距離

Distill並沒有非常強調要對內部參數進行對齊，在paper中也只是簡單提及加入隱藏層之間的cosine Loss會對效果有提升

DistillBert的初始化策略借鑒了PKD-Skip，用\(Bert_{12}\)每兩層取一層參數來初始化Distill，對比PKD直接用\(Bert_{6}\)的參數來初始化，這種初始化策略一致性更好。之后的預訓練過程可以說是對已有參數的微調，因此初始化對整體效果的影響程度甚至超過以上3個loss function。這里其實已經能發現student和teacher能否在內部結構上保持一致對最終效果的影響很大。以下是分別去除初始化以及三種loss function對DistillBert效果影響的評估結果

下游遷移

在預訓練任務上訓練好的student模型，可以選擇直接微調或者遷移到下游任務。不過如果再進行一次蒸餾，也就是在下游任務進行過微調的大模型上繼續進行蒸餾，會得到更好的效果，在SQuad數據集上可以再有1～2個點的提升。因為小模型的擬合能力有限，所以感覺這里下游任務的復雜程度越高，進行二次蒸餾帶來的效果提升會越明顯。

Tiny Bert

• paper: TinyBERT: Distilling BERT for Natural Language Understanding
• github：https://github.com/huawei-noah/Pretrained-Language-Model/tree/master/TinyBERT

Tiny Bert在Distill的基礎上完善了預訓練和微調過程中蒸餾方式。作者分別給出4層和6層的兩種大小，對比相同層數的DistillBert，TinyBert表現顯著更好，並且進一步縮減了2/3的參數，推理速度再提升3倍

預訓練

和Distill不同的是，Tiny並沒有直接使用Teacher模型的參數來初始化student模型，這個差異讓Tiny在參數壓縮上有更大的空間。Distill只能壓縮層數，而不能壓縮Embedding或者Hidden Size，因為這部分是直接從Teacher初始化來的，而Tiny可以更自由的選擇隱藏層的大小。

但是初始化的缺失，導致Tiny模仿Bert的難度直線上升。因此Tiny不只對隱藏層輸出和輸出層進行對齊，同時還加入了對Embedding和Attention矩陣的對齊。

具體的Loss Function包括以下四個部分

1. \(L_{attn}\): attention矩陣對齊

有研究表明attention矩陣包含了豐富的語義語法信息，所以作者加入了未正則化的attention矩陣的對齊（對比softmax歸一化后收斂更快），分別計算每個head，attention矩陣(seq_len * seq_len)之間的歐式距離

\[L_{attn} = \frac{1}{h}\sum_{i=1}^h MSE(A_i^S,A_i^T) \]

2. \(L_{hidn}\)：隱藏層距離

這里和Distill的差異在於，Tiny允許隱藏層大小的差異，所以加入了一個線性映射\(W_h\)來對齊student和teacher隱藏層的維度。

\[L_{hidn} = MSE(H^SW_h, H^T) \]

3. \(L_{emb}\): 輸入層對齊

Tiny還加入了對輸入向量的對齊，因為允許輸入層的維度差異，所以也加入了對應的線性映射\(W_e\)

\[L_{emb} = MSE(E^SW_e, E^T) \]

4. \(L_{pred}\):輸出層對齊

這里和KD相同，都是用了temperature softmax，對齊student和teacher的輸出層概率。

在pretrain階段，作者評估\(L_{pred}\) 完全沒有收益，因此預訓練階段只用了對齊內部參數的3個Loss，並且只train了3個epochs。個人感覺pretrain階段只加入對齊的loss，效果類似於對大模型每個Block的hidden state進行了PCA降維，從相似的Attention信息里，只保留更重要的hidden信息輸出到下一個Block。pretrain的過程只是在對齊必要的attention信息后，學習最佳的降維矩陣。所以總感覺這里好像還有進一步優化的空間，因為PCA只對hidden和emb做了，有沒有可能對Attention也做個壓縮？？？

在下游遷移的蒸餾中使用了以上4個loss的線性組合，在消融實驗中，對模型效果的重要性影響是Atten>pred>hidden>Emb，其中Atten+hidden的效果是1+1>>2的，可見每個transformer block內部進行整體對齊是很重要的。

下游遷移

在遷移到下游任務時，Tiny選擇先對中間部分（不包括\(L_{pred}\)）蒸餾10～50個epochs。因為相同的輸出概率底層transformer參數分布不一定一致，但是底層transformer一致輸出概率是一定一致的，所以是希望優先對齊transformer內部的參數分布。

除了蒸餾之外，TinyBert還加入數據增強。作者用預訓練的Bert，以及Glove詞向量來進行同義詞樣本增強，有P的概率，詞會被GLove中的Top K同義詞替換，或者會被MASK並用Bert預測得到Top K Token替換。這里作者用了P=0.4， K=15，每個原始樣本生成最多20個增強樣本。在日常使用中也會發現數據增強在越復雜的模型上收益越小，但是用復雜模型進行增強往往會在小模型上有更大的收益，可以更有效的提高小模型的泛化性。

在消融實驗中，對遷移任務的影響是遷移蒸餾>數據增強>>預訓練蒸餾，其中預訓練蒸餾影響非常小，不過這里的評估數據相對有限，感覺不太能直接得出預訓練蒸餾沒啥用這種結論。感覺預訓練的蒸餾在更復雜，樣本多樣性更高的任務上效果應該會更顯著。畢竟對於簡單任務本身Finetune對底層layer的影響就很有限，對finetune還是pretrain模型蒸餾不會有太大差異。對於Bert Finetune到底對參數有啥影響可以看下這篇博客 博觀約取系列 ～ 探測Bert Finetune對向量空間的影響]

BERT手冊相關論文和博客詳見BertManual

---

Reference

1. https://mp.weixin.qq.com/s/tKfHq49heakvjM0EVQPgHw
2. https://medium.com/huggingface/distilbert-8cf3380435b5
3. https://www.jiqizhixin.com/articles/2019-12-05-3
4. https://huggingface.co/docs/transformers/model_doc/distilbert