BERT 模型蒸餾 TinyBERT

BERT 的效果好，但是模型太大且速度慢，因此需要有一些模型壓縮的方法。TinyBERT 是一種對 BERT 壓縮後的模型，由華中科技和華為的研究人員提出。TinyBERT 主要用來模型蒸餾的方法進行壓縮，在 GLUE 實驗中可以保留 BERT-base 96% 的性能，但體積比 BERT 小了 7 倍，速度快了 9 倍。

1.前言

模型蒸餾 Distillation 是一種常用的模型壓縮方法，首先訓練一個大的 teacher 模型，然後使用 teacher 模型輸出的預測值訓練小的 student 模型。student 模型學習 teacher 模型的預測結果 (概率值) 從而學習到 teacher 模型的泛化能力。

有不少利用模型蒸餾壓縮 BERT 的研究，例如將 BERT 蒸餾到 BiLSTM，還有 huggingface 的 DistilBERT 等，在之前的文章BERT 模型蒸餾 Distillation BERT有比較詳細的介紹，對模型蒸餾不熟悉的童鞋可以參考一下。

這裡主要介紹另一種 BERT 蒸餾的模型 TinyBERT，之前蒸餾模型的損失函數主要是針對 teacher 模型輸出的預測概率值，而 TinyBERT 的損失函數包括四個部分：Embedding 層的損失，Transformer 層 attention 的損失，Transformer 層 hidden state 的損失和最後預測層的損失。即 student 模型不僅僅學習 teacher 模型的預測概率，也學習其 Embedding 層和 Transformer 層的特性。

上面的圖片展示了 TinyBERT (studet) 和 BERT (teacher) 的結構，可以看到 TinyBERT 減少了 BERT 的層數，並且減小了 BERT 隱藏層的維度。

2.TinyBERT

TinyBERT 蒸餾過程中的損失函數主要包含以下四個：

Embedding 層損失函數Transformer 層 attention 損失函數Transformer 層 hidden state 損失函數預測層損失函數 我們先看一下 TinyBERT 蒸餾時候每一層的映射方法。

2.1 TinyBERT 蒸餾的映射方法

假設 TinyBERT 有 M 個 Transformer 層，而 BERT 有 N 個 Transformer 層。TinyBERT 蒸餾主要涉及的層有 embedding 層 (編號為0)、Transformer 層 (編號為1到M) 和輸出層 (編號 M+1)。

我們需要將 TinyBERT 每一層和 BERT 中要學習的層對應起來，然後再蒸餾。對應的函數為 g(m) = n，m 是 TinyBERT 層的編號，n 是 BERT 層的編號。

對於 embedding 層，TinyBERT 蒸餾的時候 embedding 層 (0) 對應了 BERT 的 embedding 層 (0)，即 g(0) = 0。

對於輸出層，TinyBERT 的輸出層 (M+1) 對應了 BERT 的輸出層 (N+1)，即 g(M+1) = N+1。

對於中間的 Transformer 層，TinyBERT 採用 k 層蒸餾的方法，即 g(m) = m × N / M。例如 TinyBERT 有 4 層 Transformer，BERT 有 12 層 Transformer，則 TinyBERT 第 1 層 Transformer 學習的是 BERT 的第 3 層；而TinyBERT 第 2 層學習 BERT 的第 6 層。

2.2 Embedding 層損失函數

Embedding

ES是 TinyBERT 的 embedding，ET 是 BERT 的 embedding，l 是句子序列的長度，而 d『 是 TinyBERT embedding 維度，d 是 BERT embedding 維度。因為是要壓縮 BERT 模型，所以 d' < d，TinyBERT 希望模型學到的 embedding 與 BERT 原來的 embedding 具有相似的語義，因此採用了上面的損失函數，減少兩者 embedding 的差異。

embedding 維度不同，不能直接計算 loss，因此 TInyBERT 增加了一個映射矩陣 We(d'×d) 的矩陣，ES 乘以映射矩陣後維度與 ET 一樣。embedding loss 就是二者的均方誤差 MSE。

2.3 Transformer 層 attention 損失函數

TinyBERT 在 Transformer 層損失函數有兩個，第一個是 attention loss，如下圖所示。

Transformer 層的 attention loss

attention loss 主要是希望 TinyBERT Multi-Head Attention 部分輸出的 attention score 矩陣 能夠接近 BERT 的 attention score 矩陣。因為有研究發現 BERT 學習到的 attention score 矩陣能夠包含語義知識，例如語法和相互關係等，具體可參考論文《What Does BERT Look At? An Analysis of BERT’s Attention》。TinyBERT 通過下面的損失函數學習 BERT attention 的功能，h 表示 Multi-Head Attention 中 head 的個數。

Attention loss

2.4 Transformer 層 hidden state 損失函數

TinyBERT 在 Transformer 層的第二個損失函數是 hidden loss，如下圖所示。

Transformer 層的 hidden loss

hidden state loss 和 embedding loss 類似，計算公式如下，也需要經過一個映射矩陣。

Hidden loss

2.5 預測層損失函數

預測層的損失函數採用了交叉熵，計算公式如下，其中 t 是模型蒸餾的 temperature value，zT是 BERT 的預測概率，而 zS 是 TinyBERT 的預測概率。

Prediction loss

3.TinyBERT 兩階段訓練方法

TinyBERT 兩階段訓練

BERT 有兩個訓練階段，第一個訓練階段是訓練一個預訓練模型 (預訓練 BERT)，第二個訓練階段是針對具體下遊任務微調 (微調後的 BERT)。TinyBERT 的蒸餾過程也分為兩個階段。

第一個階段，TinyBERT 在大規模的 general 數據集上，利用預訓練 BERT 蒸餾出一個 General TinyBERT。第二個階段，TinyBERT 採用數據增強，利用微調後的 BERT 訓練一個 task specific 模型。4.參考文獻

TINYBERT: DISTILLING BERT FOR NATURAL LANGUAGE UNDERSTANDING