0. 背景

題目： FastBERT: a Self-distilling BERT with Adaptive Inference Time
機構(gòu)：北大、騰訊、北師大
作者：Weijie Liu, Peng Zhou, Zhe Zhao, Zhiruo Wang, Haotang Deng, Qi Ju
論文地址：https://arxiv.org/abs/2004.02178
收錄會議：ACL2020
代碼：https://github.com/autoliuweijie/FastBERT

摘要

預(yù)訓(xùn)練模型如 BERT 所取得的成果舉世矚目，但是如何在確保模型性能的前提下提高模型效率，是近年來的研究熱點。為此本文提出 FastBERT。FastBERT 是一種速度可調(diào)具有自適應(yīng)推理的模型。推理速度可根據(jù)不同需求靈活調(diào)整，避免樣本冗余計算。此外，該模型在微調(diào)方面采用了獨特的自精餾機制，進一步提高了計算效率，使性能損失最小。FastBERT 在 12 個中英文數(shù)據(jù)集中取得了很好的結(jié)果。如果給定不同的加速閾值以進行速度性能折衷，那么 FastBERT 能夠比 BERT 加速 1-12 倍。

介紹

預(yù)訓(xùn)練模型的成功眾所周知，從最初的 BERT，不斷升級進化，做大做強，包括 GPT2、XLNet、RoBERTa、T5、UER 等等。然而，在許多實際的場景中，這些預(yù)訓(xùn)練模型通常需要較多的計算資源。針對這一問題，常見的解決方案包括有

(1)量化，把 FP32 改成 FP16 甚至 INT8。
(2)剪枝，剪掉多余的連接、多余的注意力頭甚至是去掉部分的網(wǎng)絡(luò)層。
(3)蒸餾，用一個較小的 Student 模型來學習較大 Teacher 模型的知識。

通過查驗許多 NLP 數(shù)據(jù)集，發(fā)現(xiàn)這些樣本有不同的難度級別。較大的模型可能會對簡單的樣本輸入過多計算，而較小的模型難以應(yīng)付復(fù)雜的樣本。最近的研究表明預(yù)訓(xùn)練模型中確實存在冗余，因此設(shè)計一個適合所有樣本難度的單一化模型是有用的，該模型可以滿足不同復(fù)雜度的樣本，并在精度損失最小的情況下獲得計算效率。

基于這種目的，本文提出 FastBERT：一個具有樣本自適應(yīng)機制的預(yù)訓(xùn)練模型。

FastBERT 可以動態(tài)調(diào)整執(zhí)行層的數(shù)量，以減少計算量。換句話說，就是自適應(yīng)調(diào)整每個樣本的計算量，難度低的樣本通過一兩層就可以預(yù)測出來，難度較高的樣本則需要繼續(xù)走到更深的層級甚至是整個全程。此外，該模型也有一個獨特的自蒸餾處理機制，以盡量少改變模型結(jié)構(gòu)，在單一框架內(nèi)實現(xiàn)更快且準確的結(jié)果。與笨重的 BERT 相比，F(xiàn)astBERT 加速了 2 到 11 倍，且結(jié)果精度也極具競爭優(yōu)勢(精度只是略有損失而已)。

小結(jié)下：
(1)本文改進了 BERT，提出了一種實用的速度可調(diào)模型，即 FastBERT。
(2)首次將樣本自適應(yīng)機制和自蒸餾機制相結(jié)合，并提高了模型的推理速度。同時在 12 個 NLP 數(shù)據(jù)集上驗證了其有效性。

模型

模型架構(gòu)

FastBERT 有 2 部分組成：主干和分支。主干由 12 層的 Transformer encoder 和一個額外的 teacher-classifier 組成；而分支是附加到每個 Transformer 輸出的 student-classifiers。具體模型結(jié)構(gòu)如 Figure 2 所示：

主干由三部分組成：embedding 層、encoder(由 Transformer 堆疊起來的 blocks 組成)和 teacher classifier。embedding 層和 encoder 沒什么特別的，與 BERT 相同，是 word、position 和 segment embedding 的和。跟在 encoder 最后的是 teacher-classifier，它提取域內(nèi)特征用于下游推理。

模型訓(xùn)練

FastBERT 分別訓(xùn)練主干和分支，在訓(xùn)練另一個模塊時，一個模塊中的參數(shù)是固定住的。模型訓(xùn)練包括三個步驟:主干預(yù)訓(xùn)練、主干微調(diào)和 student classifier 的自蒸餾。

(1)主干預(yù)訓(xùn)練：
與 BERT 的各種版本模型一樣，比如 BERT-WWM、RoBERTa 等等模型的預(yù)訓(xùn)練方法和結(jié)果都可以直接拿來用。所以 FastBERT 可以不需要自己進行預(yù)訓(xùn)練，直接加載各種已經(jīng)公布的高質(zhì)量預(yù)訓(xùn)練模型。需要注意的是，由于 teacher-classifier 僅用于推理，因此預(yù)訓(xùn)練過程不受影響。
(2)主干微調(diào)：
對于每個下游任務(wù)，用對應(yīng)的任務(wù)數(shù)據(jù)輸入到模型，對主干進行和 teacher-classifier 微調(diào)。注意，此時尚未啟用 student classifier。
(3)分支自蒸餾：

經(jīng)過上述預(yù)訓(xùn)練和微調(diào)的主干對知識具有良好的表征和抽取，其輸出是高質(zhì)量的 soft-label，既包含原始 embedding，又包含通用知識。在此基礎(chǔ)上蒸餾出 student classifiers。由于 student classifier 之間是相互獨立的，他們的預(yù)測結(jié)果分別與 teacher-classifier 的 soft-label 進行比較，其差異用 KL 散度來衡量。FastBERT 中一共有 L-1 個 student classifier，所以自蒸餾的總損失即為它們 KL 散度的總和。

預(yù)訓(xùn)練和微調(diào)階段都只更新主干參數(shù)，蒸餾階段固定住主干參數(shù)，用 student classifier 蒸餾主干中 teacher classifier 的概率分布。由于在這個階段僅僅需要 teacher classifier 的輸出，那么可以自由地使用數(shù)量無限的未標記數(shù)據(jù)，而無需局限于有標記的數(shù)據(jù)。這意味著只要 teacher classifier 允許，總是可以提升 student classifier，以盡可能的擬合 teacher。此外，以前的蒸餾方法是用兩個模型去做，一個模型學習另一個模型的知識，而 FastBERT 的自蒸餾與之不同，因為 teacher 和 student 在同一個模型。此所謂“自蒸餾”便是這個意思，自己(分支)蒸餾自己(主干)的知識。

模型自適應(yīng)推理

在推理階段，F(xiàn)astBERT 可以根據(jù)樣本的復(fù)雜性調(diào)整模型中執(zhí)行的 encoding 層的數(shù)量。在每一 Transformer 層，用歸一化的熵衡量每個樣本的當前推理是否可靠到足以終止。歸一化熵定義如下：

實驗結(jié)果

在 12 個 NLP 數(shù)據(jù)集(6 個英文和 6 個中文)上對 FastBERT 進行評測。

分析 FLOPs： 所謂的 FLOPs(floating point operations)是指每秒所執(zhí)行的浮點運算次數(shù)，可以用來衡量模型計算復(fù)雜度。FLOPs 與模型的運行環(huán)境(CPU、GPU 或 TPU)無關(guān)，反映模型計算的復(fù)雜性。一般來說，模型的 FLOPs 越大，推理時間就會越長。在同樣的精度下，更低 FLOPs 的模型效率更高，更適合生產(chǎn)環(huán)境。Table 1 中列出了 Transformer 和分類器這兩種結(jié)構(gòu)的 FLOPs，從中可以看出，分類器的計算負荷(FLOPs)要比 Transformer 輕得多。這是 FastBERT 加速的基礎(chǔ)，雖然它增加了額外的分類器，但減少 Transformer 的計算量來實現(xiàn)加速。

性能對比：BERT、DistilBERT 和 FastBERT 在 12 個數(shù)據(jù)集上對比結(jié)果如 Table 2 所示。

從 Table 2 可以看出，Speed=0.1 時，F(xiàn)astBERT 在多數(shù)數(shù)據(jù)集上比 BERT 加速 2-5 倍且無損精度。如果允許精確度稍有損失，F(xiàn)astBERT 比 BERT 快 7 到 11 倍。FastBERT甚至在部分任務(wù)上(如 Weibo)還有細微提升。相比之下 DistillBERT 的性價比就不是這么高了，損失的性能換得的速度提升不如 FastBERT 多。

總結(jié)

本文提出一個快速版 BERT，即 FastBERT。具體來說，F(xiàn)astBERT 在訓(xùn)練階段采用一種自蒸餾機制，在推理階段采用一種自適應(yīng)機制，提高了推理效率且準確度損失極少。此外，F(xiàn)astBERT 在工業(yè)場景中有一個非常實用的特性，即推理速度可調(diào)。實驗表明，在不降低性能的情況下，F(xiàn)astBERT 可以比 BERT 快 2 到 3 倍。若允許精度略有損失，模型推理加速在 1 到 12 倍。此外，F(xiàn)astBERT 仍然與其他類 BERT 模型的參數(shù)設(shè)置兼容，即可以自由加載其他如 BERT-WWM、ERNIE 和 RoBERTa 等預(yù)訓(xùn)練的模型。