DynaBERT: Dynamic BERT with Adaptive Width and Depth
論文中作者提出了新的訓練算法，同時對不同尺寸的子網絡進行訓練，通過該方法訓練后可以在推理階段直接對模型裁剪。依靠新的訓練算法，本文在效果上超越了眾多壓縮模型，比如DistillBERT、TinyBERT以及LayerDrop后的模型。
論文對於BERT的壓縮流程是這樣的：
- 訓練時，對寬度和深度進行裁剪，訓練不同的子網絡
- 推理時，根據速度需要直接裁剪，用裁剪后的子網絡進行預測
整體的訓練分為兩個階段，先進行寬度自適應訓練，再進行寬度+深度自適應訓練。
寬度自適應 Adaptive Width
寬度自適應的訓練流程是：
- 得到適合裁剪的teacher模型，並用它初始化student模型
- 裁剪得到不同尺寸的子網絡作為student模型，對teacher進行蒸餾
在MHA中，我們認為不同的head抽取到了不同的特征，因此每個head的作用和權重肯定也是不同的，intermediate中的神經元連接也是。如果直接按照粗暴裁剪的話，大概率會丟失重要的信息，因此作者想到了一種方法，對head和神經元進行排序，每次剪枝掉不重要的部分，並稱這種方法為Netword Rewiring。重要程度的計算核心思想是計算去掉head之前和之后的loss變化，變化越大則越重要。[Analyzing multi-head self-attention]
利用Rewiring機制，便可以對注意力頭和神經元進行排序，得到第一步的teacher模型，如圖：
要注意的是，雖然隨着參數更新，注意力頭和神經元的權重會變化，但teacher模型只初始化一次（在后文有驗證增加頻率並沒帶來太大提升）。之后，每個batch會訓練mw=[1.0, 0.75, 0.5, 0.25]四種student模型，如圖：
蒸餾的最終loss來源於三方面：logits、embedding和每層的hidden state。
(soft cross-entropy loss)
（computed over all L Transformer layers）
最終的loss公式：
論文中取(lambda1, lambda2) = (1, 0.1)
emb和hidden的loss權重是因為他們有相同的維度和相似的規模。
深度自適應 Adaptive Depth
訓好了width-adaptive的模型之后，就可以訓自適應深度的了。淺層BERT模型的優化其實比較成熟了，主要的技巧就是蒸餾。作者直接使用訓好的DynaBERTw作為teacher，蒸餾裁剪深度后的小版本BERT。
對於深度，系數md = [1.0, 0,75, 0,5]，設層的深度為[1,12]，作者根據mod(d+1, 1/md) ≡ 0去掉深度為d的層。之所以取d+1是因為研究表明最后一層比較重要[Minilm]。
例如，對於BERT，包含12個transformer層1,2,...,12，當md=0.75時，被drop的層是深度為3，7，11的。（因為用了d+1，所以第12層沒有被drop）
最終被保留的層的loss：
整體loss：
最后，為了避免災難性遺忘，作者繼續對寬度進行剪枝訓練，第二階段的訓練方式如圖：
根據訓練時寬度和深度的裁剪系數，作者最終可得到12個大小不同的BERT模型，在GLUE上的效果如下：
可以看到，剪枝的BERT效果並沒有太多下降，並且在9個任務中都超越了BERT-base。同時，這種靈活的訓練方式也給BERT本身的效果帶來了提升，在與BERT和RoBERTa的對比中都更勝一籌：
另外，作者還和DistillBERT、TinyBERT、LayerDrop進行了實驗對比，DynaBERT均獲得了更好的效果。
在消融實驗中，作者發現在加了rewiring機制后准確率平均提升了2個點之多：
結論
本篇論文的創新點主要在於Adaptive width的訓練方式，考慮到后續的裁剪，作者對head和neuron進行了排序，並利用蒸餾讓子網絡學習大網絡的知識。
總體來說還是有些點可以挖的，比如作者為什么選擇先對寬度進行自適應，再寬度+深度自適應？這樣的好處可能是在第二階段的蒸餾中學習到寬度自適應過的子網絡知識。但直接進行同時訓練不可以嗎？還是希望作者再驗證一下不同順序的差距。
為了簡化，作者在寬度上所做的壓縮比較簡單，之后可以繼續嘗試壓縮hidden dim。另外，ALBERT相比原始BERT其實更適合淺層Transformer，也可以作為之后的嘗試方向。
參考：
別再蒸餾3層BERT了！變矮又能變瘦的DynaBERT了解一下