導讀

在蒸餾之後對學生模型再進行微調，進一步提升能力。

隨著NLP模型的大小增加到數千億個引數，建立這些模型的更緊湊表示的重要性也隨之增加。知識蒸餾成功地實現了這一點，在一個例子中，教師模型的效能的96%保留在了一個小7倍的模型中。然而，在設計教師模型時，知識的提煉仍然被認為是事後考慮的事情，這可能會降低效率，把潛在的效能改進留給學生。

此外，在最初的提煉後對小型學生模型進行微調，而不降低他們的表現是困難的，這要求我們對教師模型進行預訓練和微調，讓他們完成我們希望學生能夠完成的任務。因此，與只訓練教師模型相比，透過知識蒸餾訓練學生模型將需要更多的訓練，這在推理的時候限制了學生模型的優點。

如果在設計和訓練教師模型的過程中，將知識的提煉放在首要位置和中心位置，會有什麼可能呢？我們能否設計併成功地訓練一個“假定”要被蒸餾的模型，蒸餾後的版本能否成功地在任何下游任務上進行微調？這是我們在本文中總結的MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices中提出的一些問題。

MobileBERT 結構

Transformer結構視覺化 (a) BERT， (b) MobileBERT 教師，(c) MobileBERT 學生。用“Linear”標記的綠色梯形稱為bottlenecks。

線性層

知識的蒸餾要求我們比較老師和學生的表示，以便將它們之間的差異最小化。當兩個矩陣或向量維數相同時，這是很直接的。因此，MobileBERT在transformer塊中引入了一個 bottleneck 層。這讓學生和老師的輸入在大小上是相等的，而他們的內部表示可以不同。這些bottleneck在上圖中用“Linear”標記為綠色梯形。在本例中，共享維度是512，而教師和學生的內部表示大小分別是1024和128。這使得我們可以使用BERT-large(340M引數)等效模型來訓練一個25M引數的學生。

此外，由於兩個模型的每個transformer塊的輸入和輸出尺寸是相同的，因此可以透過簡單的複製將嵌入引數和分類器引數從教師傳遞給學生！

多頭注意力

細心的讀者會注意到，多頭注意塊(MHA)的輸入不是先前線性投影的輸出。相反，使用初始輸入。這篇論文的設計選擇沒有任何動機，只能讓我們去推測。我相信其原因是它所允許的自由度增加了。基本上，我們將迫使模型處理資訊的方式分離為兩個單獨的流，一個流入MHA塊，另一個作為跳躍連線。(使用線性投影的輸出並不會因為初始的線性變換而改變MHA塊的行為，這也是很容易說服自己的。)

堆疊FFN

為了在這個小的學生模型中實現足夠大的容量，作者引入了他們所謂的 stacked FFN，如圖中學生模型概述中的虛線框所示。Stacked FFN只是簡單的將Feed Forward + Add & Norm blocks塊重複了4次，選擇這一方式來得到MHA和FFN block之間的良好的引數比例。本工作中的消融研究表明，當該比值在0.4-0.6範圍內時，效能最佳。

操作最佳化

由於其目標之一是在資源有限的裝置上實現快速推理，因此作者確定了他們的架構可以進一步改進的兩個方面。

把smooth GeLU的啟用函式更換為ReLU
將normalization操作轉換為element-wise的線性變換

建議知識蒸餾目標

為了實現教師和學生之間的知識轉移，作者在模型的三個階段進行了知識蒸餾：

特徵圖遷移 —— 允許學生模仿老師在每個transformer層的輸出。在上面的架構圖中，它表示為模型輸出之間的虛線箭頭。
注意力圖遷移 —— 這讓老師在不同層次上關注學生，這也是我們希望學生學習的另一個重要屬性。這是透過最小化每一層和頭部的注意力分佈(KL散度)之間的差異而實現的。
預訓練蒸餾 —— 也可以在預訓練中使用蒸餾，透過組合Masked語言建模和下一個句子預測任務的線性組合。

有了這些目標，我們就有了不止一種方法來進行知識的提煉。作者提出了三種備選方案：

輔助知識遷移。分層的知識遷移目標與主要目標(Masked語言建模和下一句預測)一起最小化。這可以被認為是最簡單的方法。
聯合知識遷移。不要試圖一次完成所有的目標，可以將知識提煉和預訓練分為兩個階段。首先對所有分層知識蒸餾損失進行訓練直到收斂，然後根據預訓練的目標進行進一步訓練。
進一步的知識轉移。兩步法還可以更進一步。如果所有層同時進行訓練，早期層沒有很好的最小化的錯誤將會傳播並影響以後層的訓練。因此，最好是一次訓練一層，同時凍結或降低前一層的學習速度。

知識遷移技術。(a) 輔助知識遷移，(b) 聯合知識遷移， (c) 進一步知識遷移

實驗結果

作者在三種配置中評估了他們提出的MobileBERT：有25M引數的主模型(MobileBERT)，沒有操作最佳化的同一個模型(MobileBERT w/o OPT)，以及只有15M引數的模型(MobileBERT-tiny)。這些模型都與基線演算法進行了比較，如ELMo，GPT和BERT-base以及相關蒸餾工作：BERT-PKD, DistilBERT。

研究發現，透過漸進式知識轉移，訓練這些不同的MobileBERT是最有效的，其效果始終顯著優於其他兩個。

在GLUE benchmark上的結果

我們發現，MobileBERT w/o OPT比大得多的BERT-base的平均GLUE得分高出0.2，而比它小4倍。另一方面，MobileBERT僅落後於BERT-base 0.6分，但推理時間快得多 —— 在Pixel 4的手機上對128個tokens序列進行推理只需要62毫秒！然而，它的效能仍然具有競爭力，因為它的效能遠遠超過GTP和ELMo。

因此，可以肯定地說，我們可以建立一個蒸餾模型，它可以在資源有限的裝置上達到既高效能又快速

與TinyBERT相比，MobileBERT-tiny的效能略好一些。然而，當你考慮到TinyBERT是如何為GLUE任務進行微調時，這一點就更令人印象深刻了。記住，在此工作之前，不可能去微調學生模型，因為他們的容量太小了。因此，在把它的知識提煉成TinyBERT之前，必須對它的老師BERT-base進行微調！但對MobileBERT來說，情況並非如此。

MobileBERT已經在GLUE上被自己進行了微調，這證明了它可以透過建議的蒸餾過程建立一個任務不可知的模型。

總結

MobileBERT在transformer模組中引入了 bottlenecks，這使得我們可以更容易地將知識從大尺寸的教師模型中提取到小尺寸的學生模型中。這種技術允許我們減少學生的寬度，而不是深度，這是已知的，以產生一個更有能力的模型。這個模型強調了這樣一個事實，它可以建立一個學生模型，它本身可以在最初的蒸餾過程後進行微調。

此外，結果還表明，這在實際應用中也是正確的，因為在Pixel 4的手機上，MobileBERT在GLUE上的引數減少了4倍，推理速度加快了5.5倍，卻能達到BERT-base效能的99.2%。

英文原文：https://medium.com/dair-ai/mobilebert-a-task-agnostic-bert-for-resource-limited-devices-️-c4fe3ee3815d

MobileBERT：一個在資源有限裝置上使用的BERT模型

導讀