引言

上一節介紹了$\text{DistilBERT}$模型，本節將繼續介紹優化性更強的知識蒸餾$\text{BERT}$模型——$\text{TinyBERT}$模型。

回顧：DistilBERT模型

$\text{DistilBERT}$模型是一種基于$\text{BERT-base}$的知識蒸餾版本，其模型結構表示如下。單從模型結構的角度觀察，學生模型神經元的維度沒有發生變化$(768)$，僅是$\text{Encoder}$層數減少為$\text{BERT-base}$的一半；并且各層的初始化繼承了一部分$\text{BERT-base}$：從教師模型的$\text{Encoder}$層中每兩層選擇一層作為學生模型$\text{Encoder}$層的初始化。

$\text{DistilBERT}$訓練出的學生模型($\text{param:66M}$)依然可以達到$\text{BERT-base}$模型幾乎$97$%的準確度。能夠達到這個效果離不開$\text{DistilBERT}$的三個核心策略：

• 掩碼語言模型策略$(\text{Masked?Language?Model})$：根據$\text{RoBERTa}$中的描述，摒棄掉下句預測$(\text{Next?Sentence?Prediction,NSP})$策略，并使用動態掩碼替代靜態掩碼作為$\text{BERT}$模型的訓練策略；
• 蒸餾策略$(\text{Distillation?loss})$：通過使用$\text{Softmax}$溫度函數將教師模型$\text{BERT-base}$與學生模型$\text{DistilBERT}$輸出層的解空間盡可能地相似：
  其中$N$表示教師模型和學生模型的輸出層維度，在$\text{DistilBERT}$模型中，兩者的維度相同，均為$768$。
  $$\begin{array}{rl}& \mathcal{T}(x)=(t_1,t_2,?,t_N)\\ & \mathcal{S}(x)=(s_1,s_2,?,s_N)\\ & {\mathcal{L}}_{cross}=?\sum _{i=1}^N{t}_i?\log (s_i)\end{array}$$
• 余弦嵌入策略$(\text{Cosine?Embedding?loss})$：通過計算輸出層分布向量之間夾角的余弦值$\cos [\mathcal{T}(x),\mathcal{S}(x)]$，當該值為$1$時，對應的${\mathcal{L}}_{cosine}$達到最小。此時兩向量的方向為同一方向，教師和學生模型輸出的解空間已被對齊：
  $${\mathcal{L}}_{cosine}=1?\cos [\mathcal{T}(x),\mathcal{S}(x)]$$

總結：
在$\text{ALBERT}$模型中介紹過，雖然$\text{ALBERT}$也是$\text{BERT}$的簡化版本，但它們的解空間并不相同；

與此相反，$\text{DistilBERT}$中除了繼承了$\text{BERT}$中的掩碼語言模型策略外，剩余的兩條策略均是圍繞牢牢綁定教師模型與學生模型的解空間而設計。

重新觀察學生模型$\text{DistilBERT}$，它能夠達到如此精煉的模型結構$(\text{param:110M?->?66M})$，但依然保持極高的準確性，沒有出現欠擬合的情況。這至少意味著：$\text{DistilBERT}$模型中的神經元被利用得更加充分。

在這種情況下，是否可以百尺竿頭更進一步呢$?\text{?TinyBERT}$模型給了我們一個更精進的答案。

TinyBERT模型結構

相比于$\text{DistilBERT}$模型中$\text{Encoder}$層數減半的嚴肅操作，$\text{TinyBERT}$模型可以自定義學生模型的層數。并且還可以設置隱藏層單元中神經元的維度，從而使模型更加精簡。那么它是如何實現在如此精簡的模型結構下，不僅沒有欠擬合，而且還能保持優秀的訓練結果呢$?$ 自然是依靠更加嚴苛的策略作為約束。

$\text{TinyBERT}$模型的教師-學生模型結構表示如下：

其中索引$0$表示嵌入層，$1$表示第一個$\text{Encoder}$，以此類推。最后$\text{N+1,M+1}$分別表示教師、學生模型的預測層。

該蒸餾結構與$\text{DistilBERT}$之間沒有太大區別，只不過沒有$\text{DistilBERT}$中的初始化操作。教師與學生模型中各層的遷移過程可以表示為如下式子：
$$n=\mathcal{G}(m)$$
其表達的含義是：將教師模型中的第$n$層遷移到學生模型的第$m$層。例如：

• $0=\mathcal{G}(0)$表示將教師模型的嵌入層知識遷移到學生模型的嵌入層；
• $N+1=\mathcal{G}(M+1)$表示將教師模型的預測層知識遷移到學生模型的預測層；
• $n=\mathcal{G}(m)$表示將教師模型的第$n$個$\text{Encoder}$層知識遷移到學生模型的第$m$個$\text{Encoder}$層。

TinyBERT模型策略

那么$\text{TinyBERT}$是如何制定策略的呢$?$ 主要圍繞三個部分制定策略：

• $\text{Transformer}$層$(\text{Encoder})$
• 嵌入層$(\text{Embedding?Layer})$
• 預測層$(\text{Predict?Layer})$

Transformer層蒸餾

在$\text{Transformer}$層也就是編碼器層，需要使用多頭注意力機制計算注意力矩陣，再使用$\text{FeedForward?Network}$進行一個前饋計算，并將最終計算得到的隱藏狀態特征作為該編碼器的輸出。在$\text{TinyBERT}$中除了將教師模型中$\text{Encoder}$內的注意力矩陣遷移到學生模型相應的$\text{Encoder}$中，也同時將相應的隱藏狀態特征遷移到學生模型中。因而$\text{Transformer}$層蒸餾包括兩次知識蒸餾：

• 基于注意力的蒸餾
  通過最小化對應學生$\text{Encoder}$和教師$\text{Encoder}$內注意力矩陣的均方誤差來訓練對應學生$\text{Encoder}$層：
  • 其中$h$表示注意力機制頭的數量;${\mathcal{A}}_i^{\mathcal{S}}$表示學生$\text{Encoder}$內第$i$個頭的注意力矩陣;${\mathcal{A}}_i^{\mathcal{T}}$表示教師$\text{Encoder}$內第$i$個頭的注意力矩陣;$\text{MSE}$表示均方誤差操作。
  • 個人疑問：當學生模型隱藏層維度變化的時候${\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}}$是一樣大的嗎？但書中并沒有解釋。
    $${\mathcal{L}}_{attn}=\frac{1}{h}\sum _{i=1}^h\text{MSE}({\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}})$$
    需要注意的是，這里的注意力矩陣${\mathcal{A}}_i^{\mathcal{S}},{\mathcal{A}}_i^{\mathcal{T}}$使用的是執行$\text{Layer?Norm}$映射前的矩陣，這樣做的目的是保證信息的完整性，并且更快地收斂。
• 基于隱藏狀態的蒸餾
  隱藏狀態是當前$\text{Encoder}$的輸出，我們同樣需要將教師$\text{Encoder}$的隱藏層知識遷移到學生$\text{Encoder}$的隱藏層狀態中：
  其中${\mathcal{H}}_{\mathcal{S}}$表示學生$\text{Encoder}$內的隱藏層狀態;${\mathcal{H}}_{\mathcal{T}}$表示教師$\text{Encoder}$內的隱藏層狀態。同樣使用均方誤差使${\mathcal{H}}_{\mathcal{S}}$向${\mathcal{H}}_{\mathcal{T}}$方向擬合。
  $${\mathcal{L}}_{hidn}=\text{MSE}({\mathcal{H}}_{\mathcal{S}},{\mathcal{H}}_{\mathcal{T}})$$
  但需要注意的是：當學生$\text{Encoder}$隱藏層維度發生變化時，${\mathcal{H}}_{\mathcal{S}}$和${\mathcal{H}}_{\mathcal{T}}$兩者之間的維度之間存在差異，因而需要訓練一個新的權重矩陣${\mathcal{W}}_h$使兩者處于同一級別的維度空間：
  相當于作用在損失函數上的權重矩陣，反向傳播過程中同樣存在梯度更新。
  $${\mathcal{L}}_{hidn}=\text{MSE}({\mathcal{H}}_{\mathcal{S}}{\mathcal{W}}_h,{\mathcal{H}}_{\mathcal{T}})$$

嵌入層蒸餾

關于嵌入層的蒸餾與隱藏狀態的蒸餾相似，當學生模型設置的隱藏層維度與教師模型維度不同時，兩者對應的$\text{Embedding}$也不同。同樣在損失函數中添加一個新的權重參數${\mathcal{W}}_{\mathcal{E}}$，使兩個$\text{Embedding}$處于同一級別的維度空間：
其中${\mathcal{E}}_{\mathcal{S}}$表示學生模型的$\text{Embedding}$矩陣;${\mathcal{E}}_{\mathcal{T}}$表示教師模型的$\text{Embedding}$矩陣。$\text{MSE}$表示均方誤差。

$${\mathcal{L}}_{emb}=\text{MSE}({\mathcal{E}}_{\mathcal{S}}{\mathcal{W}}_{\mathcal{E}},{\mathcal{E}}_{\mathcal{T}})$$

預測層蒸餾

在預測層蒸餾部分，遷移的是輸出層的知識信息。這里和$\text{DistilBERT}$模型關于預測層的損失類似。對于教師模型的輸出${\mathcal{Z}}^{\mathcal{T}}$和學生模型的輸出${\mathcal{Z}}^{\mathcal{S}}$：

• 使用$\text{Softmax}$溫度函數分別獲取對應的軟目標${\mathcal{P}}^{\mathcal{T}}$和軟預測${\mathcal{P}}^{\mathcal{S}}$結果：
  同理,${\mathcal{Z}}^{\mathcal{T}}$對應的軟目標結果${\mathcal{P}}^{\mathcal{T}}$不再贅述。
  $$?\begin{array}{l}\begin{array}{r}{\mathcal{P}}_i^{\mathcal{S}}=\frac{\exp ({\mathcal{Z}}_i^{\mathcal{S}}/\mathcal{T})}{{\sum }_j\exp ({\mathcal{Z}}_j^{\mathcal{S}})/\mathcal{T}}\end{array}\\ \\ {\mathcal{P}}^{\mathcal{S}}=({\mathcal{P}}_1^{\mathcal{S}},{\mathcal{P}}_2^{\mathcal{S}},?,{\mathcal{P}}_N^{\mathcal{S}})\end{array}$$
• 再使用交叉熵損失函數對${\mathcal{P}}^{\mathcal{S}}$與${\mathcal{P}}^{\mathcal{T}}$進行描述：
  $${\mathcal{L}}_{pred}=?{\mathcal{P}}^{\mathcal{T}}?\log \left({\mathcal{P}}^{\mathcal{S}}\right)$$

最終，$\text{TinyBERT}$包含所有層的損失函數表示如下：
這里$[{\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)}]$表示學生模型的第$m$層與教師模型第$\mathcal{G}(m)$之間的遷移關系。
$$\mathcal{L}\left[{\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)}\right]=?\begin{array}{l}{\mathcal{L}}_{emb}({\mathcal{S}}_0,{\mathcal{T}}_0)m=0\\ {\mathcal{L}}_{hidn}({\mathcal{S}}_m,{\mathcal{T}}_{\mathcal{G}(m)})M\ge m>0\\ {\mathcal{L}}_{pred}({\mathcal{S}}_{M+1},{\mathcal{T}}_{N+1})m=M+1\end{array}$$
可以看出：

• $\text{TinyBERT}$損失函數數量是不確定的。它取決于設計學生模型$(\text{TinyBERT})$的層的數量；
• 相比于$\text{DistilBERT}$，$\text{TinyBERT}$需要為削減隱藏層狀態維度和層數付出相應的代價——設計的策略需要與教師模型關系更加緊密，并精確到注意力矩陣和隱藏層狀態，從而得到一個與教師模型關聯更加緊密的、學生模型的解空間。

TinyBERT模型的訓練

在文章中作者描述的訓練流程表示如下：

在$\text{TinyBERT}$模型中，使用兩個階段進行訓練：

• 通用蒸餾：在該階段，使用$\text{BERT-base}$預訓練模型作為教師，并使用$\text{BERT-base}$的訓練集對學生模型$(\text{TinyBERT})$進行蒸餾。并將該模型稱作通用$\text{TinyBERT}$模型。
• 特定任務蒸餾：在微調階段，將基于一項具體任務對通用$\text{TinyBERT}$模型進行微調$(\text{fine-tuning})$。具體微調過程方式為：
  • 使用預訓練$\text{BERT-base}$模型針對具體任務進行微調，并將這個微調后的$\text{BERT-base}$模型作為教師；
  • 將上述經過通用蒸餾得到的通用$\text{TinyBERT}$模型作為學生，經過蒸餾，得到的$\text{TinyBERT}$模型稱作微調的$\text{TinyBERT}$模型。

效果展示

論文中關于$\text{TinyBERT}$對于各下游任務中，與各模型比較結果如下：

其中，${\text{DistilBERT}}_4$表示學生模型包含$4$層$\text{Encoder}$，其他同理。可以發現：

• 相比于${\text{DistilBERT}}_4$，${\text{TinyBERT}}_4$使用不到其$30$%，但準確率卻遠高于$\text{DistilBERT}$模型；
• ${\text{TinyBERT}}_6$參數數量是$\text{BERT-base}$的$60$%左右，但其準確性基本與$\text{BERT-base}$持平。

$\text{Reference}$：
論文鏈接
《BERT基礎教程——Transformer大模型實戰》