ACL2020 | FastBERT：放飛BERT的推理速度
本文介紹的是 ACL 2020 論文《FastBERT: a Self-distilling BERT with Adaptive Inference Time》，論文作者來自北大、騰訊以及北師大。
項目地址：//github.com/autoliuweijie/FastBERT
自從BERT問世以來，大多數NLP任務的效果都有了一次質的飛躍。BERT Large在GLUE test上甚至提升了7個點之多。但BERT同時也開啟了模型的「做大做深」之路，普通玩家根本訓不起，高端玩家雖然訓得起但也不一定用得起。
所以BERT之後的發展也比較清晰，一部分壕大佬們繼續搞預訓練提升效果，當你對BERT Large望而卻步的時候，又出了GPT2，又雙出了威震天Megatron-LM，又雙叒出了T5，又雙叒叕出了DeepSpeed。。。每次都是照著一個數量級去加，剩下的人只能默默觀望，翻翻《顯示記憶體不夠，如何訓練大型神經網路？》看哪個trick可以用上。
另一部分大佬著力於給BERT瘦身提升速度。比如剪枝，剪掉多餘的連接、多餘的注意力頭、甚至LayerDrop[1]直接砍掉一半Transformer層；再比如量化，把FP32改成FP16或者INT8；還有蒸餾，用一個學生模型來學習大模型的知識，不僅要學logits，還要學attention score。。。
然而，大部分減肥方法都會帶來精度的下降。剪枝會直接降低模型的擬合能力，量化雖然有提升但也有瓶頸，蒸餾的不確定性最大，很難預知你的BERT教出來怎樣的學生。
昨天刷到了一篇讓我眼前一亮的文章《FastBERT: a Self-distilling BERT with Adaptive Inference Time》[2]，是北大+騰訊+北師大的ACL2020。作者提出了一種新的inference速度提升方式，相比單純的student蒸餾有更高的確定性，且可以自行權衡效果與速度，簡單實用。
FastBERT的創新點很容易理解，就是在每層Transformer後都去預測樣本標籤，如果某樣本預測結果的置信度很高，就不用繼續計算了。論文把這個邏輯稱為樣本自適應機制（Sample-wise adaptive mechanism），就是自適應調整每個樣本的計算量，容易的樣本通過一兩層就可以預測出來，較難的樣本則需要走完全程。
那麼問題來了，用什麼去預測中間層的結果呢？作者的解決方案是給每層後面接一個分類器，畢竟分類器比Transformer需要的成本小多了：
註：FLOPs (floating point operations)是Tensorflow中提供的浮點數計算量統計
於是模型的整體結構就呼之欲出了：
作者將原BERT模型稱為主幹（Backbone），每個分類器稱為分支（Branch）。
要注意的是，這裡的分支Classifier都是最後一層的分類器蒸餾來的，作者將這稱為自蒸餾（Self-distillation）。就是在預訓練和精調階段都只更新主幹參數，精調完後freeze主幹參數，用分支分類器（圖中的student）蒸餾主幹分類器（圖中的teacher）的概率分布。
之所以叫自蒸餾，是因為之前的蒸餾都是用兩個模型去做，一個模型學習另一個模型的知識，而FastBERT是自己（分支）蒸餾自己（主幹）的知識。值得注意的是，蒸餾時需要freeze主幹部分，保證pretrain和finetune階段學習的知識不被影響，僅用brach 來儘可能的擬合teacher的分布。
那為什麼不直接用標註數據訓分支分類器呢？因為直接訓效果不好唄（攤手～下面是作者在消融實驗給出的結果：
可以看到，非蒸餾的結果沒有蒸餾要好。個人認為是合理的，因為這兩種方式在精調階段的目標不一樣。非自蒸餾是在精調階段訓練所有分類器，目標函數有所改變，迫使前幾層編碼器抽取更多的任務feature。但BERT強大的能力與網路深度的相關性很大，所以過早地判斷不一定準確，致使效果下降。
同時，使用自蒸餾還有一點重要的好處，就是不再依賴於標註數據。蒸餾的效果可以通過源源不斷的無標籤數據來提升。
了解模型結構之後，訓練與推理也就很自然了。只比普通的BERT模型多了自蒸餾這個步驟：
-
Pre-training：同BERT系模型是一樣的，網上那麼多開源的模型也可以隨意拿來～
-
Fine-tuning for Backbone：主幹精調，也就是給BERT加上分類器，用任務數據訓練，這裡也用不到分支分類器，可以盡情地優化
-
Self-distillation for branch：分支自蒸餾，用無標籤任務數據就可以，將主幹分類器預測的概率分布蒸餾給分支分類器。這裡使用KL散度衡量分布距離，loss是所有分支分類器與主幹分類器的KL散度之和
-
Adaptive inference：自適應推理，及根據分支分類器的結果對樣本進行層層過濾，簡單的直接給結果，困難的繼續預測。這裡作者定義了新的不確定性指標，用預測結果的熵來衡量，熵越大則不確定性越大：
對於每層分類結果，作者用「Speed」代表不確定性的閾值，和推理速度是正比關係。因為閾值越小 => 不確定性越小 => 過濾的樣本越少 => 推理速度越慢。
模型最終在12個數據集（6個中文的和6個英文的）上的表現還是很好的：
可以看到，在Speed=0.2時速度可以提升1-10倍，且精度下降全部在0.11個點之內，甚至部分任務上還有細微提升。相比之下HuggingFace的DistillBERT的波動就比較劇烈了，6層模型速度只提升2倍，但精度下降最高會達到7個點。
FastBERT是一個在工程上十分實用的模型，通過提前輸出簡單樣本的預測結果，減少模型的計算負擔，從而提高推理速度。雖然每層都多了一個分類器，但分類器的計算量也比Transformer小了兩個數量級，對速度影響較小。後續的分支自蒸餾也設計的比較巧妙，可以利用無監督數據不斷提升分支分類器的效果。
另外，Sample-wise adaptive mechanism和Self-distillation這兩個idea也是在本文首次提出，相信會起到拋玉引玉的作用，衍生更多此類工作。論文本身也還有一些想像空間，比如分別優化每層分類器，因為在主幹被freeze的情況下各個分支是獨立的；或者自蒸餾unfreeze主幹，再加上數據自蒸餾一把，說不定還會有性能提升。
值得一提的是，本文的一作劉偉傑（北大）正是K-BERT[3]的作者，那也是一篇我很喜歡的文章，把知識融入到BERT的方式比較優雅，真心期待作者有更多的idea～
最後再回來誇一下，FastBERT著實很實用，而且完全不會影響到手頭調好的BERT，只需要蒸餾幾個淺層分類器，再把判斷機制加上就可以了。而且比起不太穩定的蒸餾來說放在線上也更有底，穩穩的幸福。
[1] Reducing Transformer Depth on Demand with Structured Dropout: //arxiv.org/abs/1909.11556
[2] FastBERT: a Self-distilling BERT with Adaptive Inference Time: //arxiv.org/abs/2004.02178
[3] K-BERT: Enabling Language Representation with Knowledge Graph: //arxiv.org/abs/1909.07606
