TinyBERT:模型小7倍,速度快8倍,華中科大、華為出品

機器之心報導

參與：張倩、杜偉

BERT 等大模型性能強大，但很難部署到算力、內存有限的設備中。為此，來自華中科技大學、華為諾亞方舟實驗室的研究者提出了 TinyBERT，這是一種為基於 transformer 的模型專門設計的知識蒸餾方法，模型大小還不到 BERT 的 1/7，但速度是 BERT 的 9 倍還要多，而且性能沒有出現明顯下降。目前，該論文已經提交機器學習頂會 ICLR 2020。

論文：https://arxiv.org/abs/1909.10351

在自然語言處理（NLP）領域，BERT 等預訓練語言模型極大地提升了諸多 NLP 任務的性能。但是，這類預訓練語言模型通常計算開銷大，內存佔用也大，因此很難在一些資源緊張的設備上有效執行。

為了在加快推理速度和降低模型大小的同時保持準確率，來自華中科技大學和華為諾亞方舟實驗室的研究者提出了一種新穎的 transformer 蒸餾法，這是為基於 transformer 的模型專門設計的知識蒸餾（knowledge distillation，KD）方法。通過這種新的 KD 方法，大型 teacherBERT 模型中編碼的大量知識可以很好地遷移到小型 student TinyBERT 模型中。

此外，研究者還提出了一種專門用於 TinyBERT 的兩段式學習框架，從而分別在預訓練和針對特定任務的學習階段執行 transformer 蒸餾。這一框架確保 TinyBERT 可以獲取 teacherBERT 的通用知識和針對特定任務的知識。

除了提出新的 transformer 蒸餾法之外，研究者還提出了一種專門用於 TinyBERT 的兩段式學習框架，從而分別在預訓練和針對特定任務的具體學習階段執行 transformer 蒸餾。這一框架確保 TinyBERT 可以獲取 teacherBERT 的通用和針對特定任務的知識。

實證研究結果表明，TinyBERT 是有效的，在 GLUE 基準上實現了與 BERT 相當（下降 3 個百分點）的效果，並且模型大小僅為 BERT 的 13.3%（BERT 是 TinyBERT 的 7.5 倍），推理速度是 BERT 的 9.4 倍。此外，TinyBERT 還顯著優於當前的 SOTA 基準方法（BERT-PKD），但參數僅為為後者的 28％，推理時間僅為後者的 31％左右。

Transformer 蒸餾

研究者提出的 Transformer 蒸餾是專門為 Transformer 網絡設計的知識蒸餾方法，下圖 1 為本文提出的 Transformer 蒸餾方法概覽圖：

圖 1：Transformer 蒸餾概覽圖。

在這篇論文中，student 和 teacher 網絡都是通過 Transformer 層構建的。為了表述清楚，研究者在詳解 TinyBERT 之前闡述了以下問題。

假定 student 模型有 M 個 Transformer 層，teacher 模型有 N 個 Transformer 層，從 teacher 模型中選擇 M 個 Transformer 層用於 Transformer 層蒸餾。n=g(m) 是 student 層到 teacher 層的映射函數，這意味著 student 模型的第 m 層從 teacher 模型的第 n 層開始學習信息。嵌入層蒸餾和預測層蒸餾也考慮進來，將嵌入層的指數設為 0，預測層的指數設為 M+1，並且對應的層映射分別定義為 0 = g(0) 和 N + 1 = g(M + 1)。下文實驗部分將探討不同的映射函數對性能的影響。在形式上，通過最小化以下目標函數，student 模型可以獲取 teacher 模型的知識：

其中 L_layer 是給定模型層（如 Transformer 層或嵌入層）的損失函數，λ_m 是表徵第 m 層蒸餾重要度的超參數。

研究者提出的 Transformer 層蒸餾包含基於注意力的蒸餾和基於隱狀態的蒸餾，具體可參考上圖 1（b）。基於注意力的蒸餾是為了鼓勵語言知識從 teacherBERT 遷移到 student TinyBERT 模型中。具體而言，student 網絡學習如何擬合 teacher 網絡中多頭注意力的矩陣，目標函數定義如下：

其中，h 是注意力頭數。A_i ∈ R^( l×l) 是與 teacher 或 student 的第 i 個注意力頭對應的注意力矩陣。

此外，（非歸一化）注意力矩陣 A_i 用作擬合目標，而不是其 softmax 輸出的 softmax（A_i），因為實驗表明前者的設置呈現更快的收斂速度和更佳的性能。除了基於注意力的蒸餾之外，研究者還對 Transformer 層輸出的知識進行蒸餾處理（具體可參考上圖 1（b）），目標函數定義如下：

其中 H^S∈R^l×d'和 H^T∈R^l×d 分別表示 student 和 teacher 網絡的隱狀態，由方程式 4 計算得到。標量值 d 和 d'分別表示 teacher 和 student 模型的隱狀態，並且 d'通常小於 d，以獲得更小的 student 網絡。

研究者還執行了嵌入層的蒸餾，與基於隱狀態的蒸餾類似，定義如下：

其中矩陣 E^S 和 H^T 分別表示 student 和 teacher 網絡的嵌入。在論文中，這兩種矩陣的形狀與隱狀態矩陣相同。矩陣 W_e 表示線性變化，它起到與 W_h 類似的作用。

除了模擬中間層的行為之外，研究者還利用知識蒸餾來擬合 teacher 模型的預測結果。具體而言，他們對 student 網絡 logits 和 teacher 網絡 logits 之間的 soft 交叉熵損失進行懲罰：

其中 z^S 和 z^T 分別表示 student 和 teacher 模型預測的 logits 向量，log_softmax() 表示 log 似然，t 表示溫度值。實驗表明，t=1 時運行良好。

通過以上幾個蒸餾目標函數（即方程式 7、8、9 和 10），可以整合 teacher 和 student 網絡之間對應層的蒸餾損失：

在實驗中，研究者首先執行的是中間層蒸餾（M ≥ m ≥ 0），其次是預測層蒸餾（m = M + 1）。

兩段式學習框架

BERT 的應用通常包含兩個學習階段：預訓練和微調。BERT 在預訓練階段學到的大量知識非常重要，並且遷移的時候也應該包含在內。因此，研究者提出了一個兩段式學習框架，包含通用蒸餾和特定於任務的蒸餾，如下圖 2 所示：

圖 2：TinyBERT 的兩段式學習圖示。

通用蒸餾可以幫助 student TinyBERT 學習到 teacher BERT 中嵌入的豐富知識，對於提升 TinyBERT 的泛化能力至關重要。特定於任務的蒸餾賦予 student 模型特定於任務的知識。這種兩段式蒸餾可以縮小 teacher 和 student 模型之間的差距。

通用蒸餾

在通用蒸餾中，研究者使用原始 BERT 作為 teacher 模型，而且不對其進行微調，利用大規模文本語料庫作為學習數據。通過在通用領域文本上執行 Transformer 蒸餾，他們獲取了一個通用 TinyBERT，可以針對下遊任務進行微調。然而，由於隱藏/嵌入層大小及層數顯著降低，通用 TinyBERT 的表現不如 BERT。

針對特定任務的蒸餾

研究者提出通過針對特定任務的蒸餾來獲得有競爭力的微調 TinyBERT 模型。而在蒸餾過程中，他們在針對特定任務的增強數據集上（如圖 2 所示）重新執行了提出的 Transformer 蒸餾。具體而言，微調的 BERT 用作 teacher 模型，並提出以數據增強方法來擴展針對特定任務的訓練集。

此外，上述兩個學習階段是相輔相成的：通用蒸餾為針對特定任務的蒸餾提供良好的初始化，而針對特定任務的蒸餾通過專注於學習針對特定任務的知識來進一步提升 TinyBERT 的效果。

實驗效果

為了驗證 TinyBERT 的效果，研究者在多個任務上將其與其他模型進行了比較。

在 GLUE 上的實驗效果

研究者在 GLUE 基準上評估了 TinyBERT 的性能，結果如下表 2 所示。模型大小和推理時間的效率見下表 3。

表 2：在 GLUE 基準上的評估結果。

表 3：基線模型和 TinyBERT 的模型大小和推理時間。層數量不包含嵌入和預測層。

實驗結果表明：1）TinyBERT 在所有 GlUE 任務中的表現都優於 BERTSMALL，平均性能提升了 6.3%，表明本文提出的 KD 學習框架可以有效地提升小模型在下遊任務中的性能；2）TinyBERT 顯著超越了 KD SOTA 基線（即 BERT-PKD 和 DistillBERT），比 BERT-PKD 高出 3.9%（見圖 2），但參數只有基線模型的 28%，推理時間只有基線模型的 31% 左右（見圖 3）；3）與 teacher BERTBASE 相比，TinyBERT 的大小僅為前者的 13.3%，但速度卻是前者的 9.4 倍，而且性能損失不大。4）TinyBERT 與 Distilled BiLSTM_SOFT 模型效率相當，在 BiLSTM 基線公開的所有任務中均顯示出了明顯的性能提升。5）對於具有挑戰性的 CoLA 數據集，所有的蒸餾小模型與 teacher 模型的性能差距都比較大。TinyBERT 與基線模型相比實現了顯著的性能提升，如果利用更深、更寬的模型來捕獲更複雜的語言信息，它的性能還能進一步提升。

模型大小的影響

為了測試模型大小對性能的影響，研究者在幾個典型的 GLUE 任務中測試了不同大小 TinyBERT 模型的性能。結果如下表 4 所示：

表 4：提升寬度、深度之後的 TinyBERT 變體與基線的性能比較結果。

不同學習步驟的影響

本文提出的兩段式 TinyBERT 學習框架包含三個關鍵步驟：TD（特定於任務的蒸餾）、GD（通用蒸餾）和 DA（數據蒸餾）。每個學習步驟的影響如下表 5 所示：

表 5：不同學習步驟的控制變量研究。

不同蒸餾目標的影響

研究者還探索了不同目標對 TinyBERT 學習的影響，結果如下表 6 所示：

表 6：不同蒸餾目標對 TinyBERT 學習的影響。

映射函數對模型性能的影響

研究者探究了不同映射函數 n = g(m) 對於 TinyBERT 學習的影響，比較結果見下表 7：

表 7：不同映射策略的對比結果。

本文為機器之心報導，轉載請聯繫本公眾號獲得授權。

------------------------------------------------