TINYBERT: DISTILLING BERT FOR NATURAL LANGUAGE UNDERSTANDING

TINYBERT: DISTILLING BERT FOR NATURAL LANGUAGE UNDERSTANDING

來源：ICLR 2020 在審 鏈接：https://arxiv.org/pdf/1909.10351.pdf 代碼：暫未公佈

動機

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

模型	參數
ELMo	90M
BERT	110M
BERT-large	340M
GPT	110M
GPT-2	1.5B

解決方法

文獻 等證明了預訓練語言模型的參數是冗餘的。因此論文提出了一種基於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層，詳細可見文獻

Knowledge Distillation(KD)

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

其中，$x$是網絡輸入，$L(\cdot)$是衡量在輸入$x$下，teacher網絡和student的網絡的差異性，$\mathcal{X}$ 是訓練集。 $f^{T}$和 $f^{S}$分別表示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可表示爲如下公式：

其中$\mathcal{L}_{\text {layer }}$爲當前層衡量Teacher BERT絡和TinyBERT的差異性的loss函數, $\lambda_{m}$是超參數，代表當前層的重要程度。

對應不同的層，論文采用了不同的loss函數
Transformer層：

attention矩陣

論文第一次將attention矩陣作爲知識蒸餾過程中teacher網絡的監督信息。因爲 文獻 證明了attention舉證中包含了語法和共指信息，通過將attention矩陣作爲teacher網絡中的監督信息，可以將這些信息遷移到student網絡。採用的是均方誤差函數，$h$是atteniton的頭數，每層共$h$個注意力矩陣$A$。$\boldsymbol{A}_{i}\in \mathbb{R}^{l \times l}$是teacher Bert或者tinyBert中的注意力矩陣，$l$是輸入文本的長度。

Transformer層輸出

同樣使用均方誤差函數，使用$\mathbf{W}_{h}$進行維度轉換

詞向量層：

同樣使用均方誤差函數，使用$\mathbf{W}_{e}$進行維度轉換

預測輸出層：

僅在微調階段的知識蒸餾過程中使用，$z^{T}$和$z^{S}$分別表示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個候選詞，最後以概率$p_{t}$選擇是否替換這個單詞，從而產生更多的文本數據。

實驗結果

實驗設置

TinyBERT(參數個數14.5M):

網絡層數$M=4$, 隱層維度$d_{i}^{\prime}=1200$, FFN層維度$d_{i}^{\prime}=1200$, h=12

Techer BERT(參數個數109M):

網絡層數$N=12$, 隱層維度$d=768$, FFN層維度$d_{i}=3072$,$h=12$

映射函數：$g(m)=3 \times m$

每層的重要性參數$\lambda$設置爲1

實驗結果

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

由table2和table3可得：

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

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

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

模型簡化測試

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

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

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

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

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

Uniform-strategy: g(m) = m + N - M; 0 < m <= M
Top-strategy: g(m) = m; 0 < m <= M
Bottom-strategy: g(m) = 3 * m; 0 < m <= M

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

總結

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

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