華為劉群團隊構造兩階段知識蒸餾模型TinyBERT，模型壓縮7.5倍，推理時間快9.4倍

2019 年 11 月 1 日
筆記

作者 | 劉旺旺

編輯 | 唐里

來源 | ICLR 2020 在審

鏈接 | https://arxiv.org/pdf/1909.10351.pdf

代碼 | 暫未公布

‍‍‍‍動機預訓練的語言模型+微調的模式提升了很多nlp任務的baseline，但是預訓練的語言模型太過龐大，參數都是百萬級別，因此很難應用到實際的應用場景中，尤其是一些資源有限的設備上。

解決方法

文獻(https://arxiv.org/pdf/1901.10430.pdf)等證明了預訓練語言模型的參數是冗餘的。因此論文提出了一種基於transformer結構的知識蒸餾方法，用於bert的壓縮和加速。最終將模型大小減小為原來的1/7.5, inference的時間減少為原來的1/9.4，並且可以達到與原有bert模型相當的效果。

預備知識

transformer layer 包含兩個部分：Multi-Head Attention (MHA) 和 Position-wise Feed-Forward Network (FFN)，bert主要的結構也是多層transformer layer的堆疊。

MHA：

FFN：

MHA層和FFN層之間，以及FFN層之後還包含一個殘差連接和norm層，詳細可見文獻(https://arxiv.org/pdf/1706.03762.pdf)

Knowledge Distillation(KD)

知識蒸餾是一種模型壓縮常見方法，指的是在teacher-student框架中，將複雜、學習能力強的網絡(teacher)學到的特徵表示「知識」蒸餾出來，傳遞給參數量小、學習能力弱的網絡(student)。teacher網絡中學習到的特徵表示可作為監督信息，訓練student網絡以模仿teacher網絡的行為。整個知識蒸餾過程的誤差函數為：

‍‍‍‍‍‍‍‍‍‍其中，x是網絡輸入，L(⋅)是衡量在輸入x下，teacher網絡和student的網絡的差異性，

是訓練集。

和

分別表示teacher和student網絡的行為函數，可以通俗理解成網絡中特徵表示。從公式可看出，知識蒸餾的過程關鍵在於，如何定義網絡的差異性loss，以及如何選取網絡的行為函數。

‍‍

基於transformer的知識蒸餾模型壓縮

先前也有一些工作使用知識蒸餾的方法來做bert模型的壓縮：

通過表格可以看出，論文tinybert的創新點在於學習了teacher Bert中更多的層數的特徵表示，蒸餾的特徵表示包括：

詞向量層的輸出，
Transformer layer的輸出以及注意力矩陣
預測層輸出(僅在微調階段使用)

上圖描述了bert知識蒸餾的過程，左邊的圖整體概括了知識蒸餾的過程：左邊是Teacher BERT，右邊是Student TinyBERT，論文的目的是將Teacher BERT學習到的知識遷移到TinyBERT中；右邊的圖描述了知識遷移的細節，在訓練過程中選用Teacher BERT中每一層transformer layer的attention矩陣和輸出作為監督信息。

假設 Teacher BERT 有M層，TinyBERT有N層(N<M), n = g(m) 是TinyBERT到Teacher BERT 映射函數，表示TinyBERT的第m層是從Teacher BERT中的第n層學習知識的。Teacher BERT 詞向量層和預測輸出層的特徵表示也被用於知識蒸餾。TinyBERT詞向量層和預測輸出層是從Teacher BERT相應的層學習知識的。將Teacher BERT 詞向量層看作第0層，預測輸出層看作M+1層，映射函數可表示為 0 = g(0), n = g(m), N+1=g(M + 1)

將Teacher BERT學習到的特徵表示作為TinyBERT的監督信息，從而訓練TinyBERT，訓練的loss可表示為如下公式：

‍‍

‍‍‍‍其中

為當前層衡量Teacher BERT絡和TinyBERT的差異性的loss函數,

是超參數，代表當前層的重要程度。

對應不同的層，論文採用了不同的loss函數‍‍‍‍

‍‍Transformer層：

attention矩陣

論文第一次將attention矩陣作為知識蒸餾過程中teacher網絡的監督信息。因為文獻（https://arxiv.org/pdf/1906.04341.pdf）證明了attention舉證中包含了語法和共指信息，通過將attention矩陣作為teacher網絡中的監督信息，可以將這些信息遷移到student網絡。採用的是均方誤差函數‍‍‍‍‍‍‍‍‍‍，h是atteniton的頭數，每層共‍‍有h個注‍‍意力矩‍‍陣‍‍A‍‍。

是teacher Bert或者tinyBert中的注意力矩陣，l是輸入文本的長度。

‍‍

Transformer層輸出

同樣使用均方誤差函數，使‍‍用Wh進行維度‍‍轉換

詞向量層：

同樣使用均方誤差函數，使‍‍用W‍‍e進行維‍‍度‍‍轉換

預測輸出層：

‍

‍‍僅在微調階段的知識蒸餾過程中使用‍‍‍‍‍‍，

和

分別表示teacher Bert和tinyBert網絡在具體下游任務中預測層輸出。t是知識蒸餾中的溫度參數，實驗中被設置為1.

‍

整體可表示為：

Pre-training和Fine-tuning兩個階段知識蒸餾

知識蒸餾的過程也被分為兩個階段， General distillation (Pre-training階段)，使用大規模無監督的數據，幫助student網絡TinyBERT學習到尚未微調的teacher BERT中的知識，有利於提高TinyBERT的泛化能力。此時，由於詞向量維度的減小，隱層神經元的減少，以及網絡層數的減少，tinybert的表現遠不如teacher bert。（注意：在Pre-training階段，蒸餾的特徵表示未使用預測層輸出）

task-specific distillation(Fine-tuning階段)，使用具體任務的數據，幫助TinyBERT學習到更多任務相關的具體知識。值得注意的是，在Fine-tuning階段，論文採用了數據增強的策略，從後面的實驗中可以看出，數據增強起到了很重要的作用。數據擴充的過程如下：對於特定任務的數據中每一條文本，首先使用bert自帶的方式進行bpe分詞，bpe分詞之後是完整單詞（single-piece word），用[MASK]符號代替，然後使用bert進行預測並選擇其對應的候選詞N個；如果bpe分詞之後不是完整單詞，則使用Glove詞向量以及餘弦相似度來選擇對應的N個候選詞，最後以‍‍‍‍概率

選擇是否替換這個單詞，從而產生更多的文本數據。

‍‍‍‍ 實驗結果

實驗設置

TinyBERT(參數個數14.5M):‍‍‍‍‍‍‍‍‍‍

‍‍‍‍‍‍‍‍‍‍‍‍網絡層數‍‍M=4, 隱層維度

,FFN層維度

,h=12

Techer BERT(參數個數109M):‍‍‍‍‍‍‍‍‍‍‍‍

‍‍

網‍‍‍‍‍‍‍‍‍‍絡‍‍‍‍‍‍層數‍‍‍‍N=12, 隱層維度‍‍d=768‍‍‍‍‍‍‍‍, FFN層維度

,h=12

‍‍

‍‍‍‍‍‍‍‍‍‍映射函數‍‍‍‍‍‍‍‍：g(m)=3×m‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

‍‍

每層的重要性參數λ設‍‍置‍‍‍‍‍‍為1‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

‍‍

實驗結果‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍‍

在GLUE數據集上的實驗結果如下表所示：

‍‍

由table2和table3可得：

由於參數的減小，Bert small 與 Bert base相比，效果相差很大
TinyBERT 效果在所有任務上都超越了 Bert small，最多平均提升了6.3%
TinyBERT 的效果也優於先前的基於知識蒸餾的模型壓縮方法BERT-PKD 和 DistillBERT
TinyBERT模型大小減小為原來的7.5分之一，Inference的時間減少為原來的9.4分之一
TinyBERT模型大小比Distilled BiLSTM大。但是，Inference 速度要比Distilled BiLSTM快，並且在所展示的任務上效果都超過了Distilled BiLSTM
對於CoLA這個數據集，所有壓縮模型的效果都不如Bert base，但是相比於其他壓縮模型，TinyBERT有最好的效果。CoLA是判斷一句話是否語法正確的數據集，需要更多語言學知識，需要更大的網絡去擬合，隨着TinyBERT參數增大，也能提高TinyBERT在該數據集上的效果，Table 4 也證明了這個結論。

通過嘗試更深更寬的網絡實驗結果如Table 4 所以，可以得到如下結論：

論文提出的知識蒸餾的方法，適用於任何大小的TinyBERT
對於Cola數據集，加深加寬網絡可以能帶來效果的大幅提升 (f49.7 -> 54.0),僅僅加寬或加深帶來的提升不大
4層的TinyBERT比6層的其他壓縮的模型效果要好，這也證明了論文提出模型的有效性

模型簡化測試

分析兩個階段的知識蒸餾 TD (Task-specific Distillation)和GD (General Distillation)，以及數據擴充DA (Data Augmentation) 對TinyBERT整體效果的作用

總體可看出，去掉TD和DA對整體結果影響較大，去掉GD對整體的結果作用較小
去掉GD對CoLA的作用大於MNLI和MRPC(CoLA在沒有GD的情況下降了9%)，CoLA是判斷一句話是否語法正確的數據集，需要更多語言學知識，而GD的過程正是捕獲這種知識的手段。

分析知識蒸餾過程中，選取的不同的特徵表示對整體結果的作用

沒有Transformer層對模型的影響最大，Transformer層是整個模型的主要構成部分
Transformer層中attention矩陣相比隱層輸出的作用要大
整體來說，Transformer層，embeding層，預測輸出層，對於提高模型的整體效果都是有效的。

分析知識蒸餾的過程中，使用不同的映射函數, 對整體結果的作用

整體而言Uniform-strategy結果最好
Top-strategy和Bottom-strategy相比，在不同的任務上，效果優劣不一。

總結

論文提出一個基於知識蒸餾進行模型的壓縮的方法，分別對bert的預訓練階段以及微調階段進行知識蒸餾。創新部分在於，使用更多的teacher bert的學習到的知識(不同層的特徵表示)作為student網絡的監督信息。實驗結果表明，與先前的知識蒸餾用於模型壓縮的方法的相比，論文提出的方法更有效，提升效果的同時更加快了inference速度。實驗翔實，具有參考意義。

但是從table5可以看出數據增強的對模型最終的效果影響很大。在CoLA數據集上，tinyBert不用數據增強（No DA）的結果低於table2中DistilBERT的結果；在MRPC數據集上，則結果相當。table2 中用於對比的方法並未使用該數據增強的方法，若能增加一組實驗會更有說服力。

延伸閱讀：

《加速 BERT 模型有多少種方法？從架構優化、模型壓縮到模型蒸餾最新進展詳解！》

‍‍‍‍‍‍蒸餾

另一個有趣的模型壓縮方法是蒸餾，這是一種將大型「teacher」網絡的知識轉移到較小的「student」網絡的技術，訓練學生網絡來模仿教師網絡的行為。

Rich Caruana 及其合作者率先採用了這種策略。在他們先驅性的論文中，他們提供了令人信服的證明：大型集成模型所獲得的知識可以轉移到單個小型的模型中。

Geoffrey Hinton 等人在他們的「Distilling the Knowledge in a Neural Network」論文中證明了這種技術可以應用於神經網絡模型。

DistilBERT

從 Hinton 開始，蒸餾的方法逐漸被應用到了不同的神經網絡中，當然你可能也聽過 HuggingFace 提出的 DistilBERT，這是一種對 BERT 的蒸餾。這項工作出自論文「DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter」，目前已經提交到 NeurIPS 2019。

DistilBERT 是一種較小的語言模型，受 BERT 的監督而訓練。在該模型中，作者刪除了令牌類型嵌入和合併器（用於下一個句子分類任務），並保持體系架構其餘部分不變，同時將層數減少了兩倍。

您可以在 HuggingFace（以前叫做 pytorch-transformers 和 pytorch-pretrained-bert）的 translators python 軟件包的幫助下，使用現成的 DistilBERT。該軟件包的 2.0.0 版本支持 TensorFlow 2.0 / PyTorch 互操作。

DistilBERT 的作者還使用了最新 RoBERTa 論文中的一些訓練技巧，這些技巧表明 BERT 的訓練方式對其最終性能至關重要。

DistilBERT 與 BERT 相比具有出人意料的結果：作者保留了 95％以上的性能，但參數卻減少了 40％。

比較 GLUE 基準測試的開發集

在推斷時間方面，DistilBERT 比 BERT 快了 60％以上，比 ELMo + BiLSTM 快 120％。

推理速度

TinyBERT

就在幾天前，出現了一種新的 BERT 蒸餾方法，來自華為諾亞方舟實驗室的劉群團隊提出了 TinyBERT。

為了構建一個具有競爭力的 TinyBERT，作者首先提出了一種新的 Transformer 蒸餾方法，來蒸餾嵌入 BERT 的知識。

具體來說就是，他們設計了幾個損失函數來適合 BERT 層的不同表示形式：

1、嵌入層的輸出；

2、Transformer 層派生的隱藏狀態和注意力矩陣;

3、預測層輸出的 logits 輸出。

論文中基於注意力的擬合則得益於最近的發現，即 BERT 學習的注意力權重可以捕獲大量的語言知識，這意味着語言知識可以很好地從教師 BERT 轉移到學生 TinyBERT。而在 BERT 的現有知識蒸餾的方法（如 Distilled BiLSTM_SOFT，BERT-PKD 和 DistilBERT）中卻忽略了這一點。

在這項工作中，作者還提出了一種新穎的兩階段學習框架，包括通用蒸餾和特定任務蒸餾。在通用蒸餾階段，未經微調的原始 BERT 充當教師模型，學生 TinyBERT 通過在通用領域對大型語料庫執行通常的 Transformer 蒸餾來學習模仿教師的行為。他們獲得了可以對各種下游任務進行微調的通用 TinyBERT。在特定任務蒸餾階段，他們將數據進行擴充，來提供更多與任務相關的材料供教師-學生學習，然後在增強的數據上重新執行 Tranformer 蒸餾。

這個兩階段方法對於提升 TinyBERT 的性能和泛化能力是必不可少的。

TinyBERY 在實驗上取得了非常的成績，相對於 GLUE 數據集的 BERT-base，其性能並沒有下降多少，而推理參數小了 7.5 倍，推理時間快了 9.4 倍。

TinyBERT 與其他基準的比較

我們期待他們能夠將這種方法應用到 BERT-large 和 XLNet 等大型模型中，同樣也期待他們開放源碼。

其他蒸餾方法

除了 DistilBERT 和 TinyBERT 外，還有其他一些為大家所熟知的蒸餾方法。

（2019/03）「Distilling Task-Specific Knowledge from BERT into Simple Neural Networks」

這篇論文將 BERT 蒸餾到單層 BiLSTM 中，取得了與 ELMo 可比的結果，同時使用的參數減少了大約 100 倍，推理時間減少了 15 倍。

BiLSTM_SOF 是 TBiLSTM 的蒸餾，後者是在 soft logit 目標上訓練出來的。

（2019/08）「Patient Knowledge Distillation for BERT Model Compression」

這篇論文提出了一種耐心知識蒸餾的方法，這是首次嘗試使用教師的隱藏狀態，而不僅僅是最後一層的輸出。他們的學生模型從教師模型的多個中間層「耐心」地學習來獲得更多知識。在他們的耐心蒸餾知識框架中，只訓練學生模仿中間層的 [CLS] 令牌的表示形式。代碼已公開。

（2019/09）「Extreme Language Model Compression with Optimal Subwords and Shared Projections」

這是最近提交到 ICLR 2020 的一篇論文，這篇論文專註於一種用於訓練詞彙量顯著較小、嵌入和隱藏狀態維度較低的學生模型的知識蒸餾技術。作者採用了雙重訓練機制，可以同時訓練教師和學生模型，從而能夠獲得針對學生詞彙的最佳詞嵌入。該方法能夠將 BERT-base 模型壓縮 60 倍以上，而下游任務指標只有很小的下降，從而使得語言模型佔用的空間只有不到 7MB。

TinyBERT 的結果似乎更好，但一個 7MB 的類 BERT 模型簡直爽的不能再爽！