結構剪枝：要個4層的BERT有多難？

  • 2019 年 11 月 26 日
  • 筆記

想看transformer剪枝文章的起因是美團搜索團隊的一篇文章[1]，裡面提到他們將BERT裁剪為四層，在query分類的業務場景下效果不降反升，簡直香。但是文章里也沒有透露具體的做法，我身邊也沒認識到美團搜索團隊的同學，於是決定自己動手調研一波。

首先，剪枝的主要原因是提升推理速度，另外是希望去掉在個別case上過擬合的權重，提升模型泛化能力。

我認為剪枝的研究點主要在以下兩方面：

1. 剪掉什麼？

這裡可以分為weight pruning(WP)和structured pruning(SP)。WP指將矩陣中的某個權重連接置為0，處理完後通過稀疏矩陣進行運算，或者將某一行或一列置為0（神經元），仍用矩陣進行運算。但WP的缺點是對矩陣進行改變，影響運算效率，所以不進行更深入的調研了，有興趣的同學可以參考[2]中提供的論文。SP則是對一組權重進行剪枝，比如去掉Attention heads、Attention mechanism、FFN層或者整個Transformer層，如果效果不下降的話對工程十分友好，因此下文的剪枝都指結構性剪枝。

2. 什麼時候剪？

這裡可以分為預訓練、訓練/精調、測試三個階段，是應該在訓練/預訓練/精調的時候去學習如何剪，還是預訓練之後剪掉一些再精調，還是直接在預測的時候剪呢？各位可以自己先想想～

方案一：預測時剪掉Attention Heads

在一篇NIPS2019[3]中，作者等人主要研究了注意力頭對結果的影響，評估的任務主要是encoder-decoder的NMT和BERT的MNLI。在預測階段mask掉head進行研究後可以發現：

在翻譯任務中，只有8(/96)個頭會對預測結果有顯著影響，大多數頭在預測時都是多餘的。

上圖中，作者只保留了當前層中最好的一個head。可以看出BERT中大部分層在只有一個head的情況下，對結果的影響都不大。因此通過去掉head來壓縮模型是可行的，雖然head在底層是並行計算的，在速度上可能提升不大，作者也提供了測試結果：

在剪掉50%的BERT head情況下，是有一定比例的速度提升的。不過作者同時也證明了剪掉50%的head會有效果下降：

大概降了1-2個點的感覺。所以比起我們既想壓縮模型，又想提升效果的目標還有一些距離。

方案二：訓練時加入Dropout，預測時剪掉Layer

文章[4]正在投稿ICLR2020，來自Facebook AI Research，提出了類似Dropout的LayerDrop，在訓練時隨機mask掉一組權重，讓模型適應這種「缺胳膊少腿」的情況，變得更加robust，然後在預測時直接去掉部分結構進行預測。評估的任務也非常多樣，有NMT、摘要、語言模型和NLU。

上圖左邊是剪去不同結構的效果，可見只剪layer效果就很好了。右圖是剪枝的策略，效果簡單又好的是Every other(每隔一個)。因此在之後的實驗中都是剪去了every other layer，具體的方法參見下式，比如剪枝概率為0.5，那就是剪掉除2後餘數為0的層（0, 2, 4, ..., 10)。

剪枝之後的效果要好於蒸餾和同等量級的模型：

遺憾的是，好像並沒在文中看到剪枝後模型和原BERT的效果比較。不過去BERT原論文里查了下，BERT base在上述任務的表現分別是：84.4/86.7/88.4/92.7，比剪枝成6層的模型好1-2個點，但速度卻慢一半左右。

另外，只將LayerDrop應用在訓練過程中還有提升效果的奇效，而且能提升一定的訓練速度：

看到LayerDrop這篇文章時覺得已經找到了心中最佳，後續有時間會擼一下程式碼，其他看了下但沒有細讀的論文有：

  • [5] 按權重大小進行weight pruning
  • [6] 類似方案二，在SQUAD2.0上進行實驗
  • [7] 提出了一種RPP的weight pruning方法

參考：

  1. 美團BERT的探索和實踐: https://mp.weixin.qq.com/s/433HD7gTYjYKHKrMjxpqeQ
  2. BERT 瘦身之路：Distillation，Quantization，Pruning: https://zhuanlan.zhihu.com/p/86900556
  3. Are Sixteen Heads Really Better than One?: https://arxiv.org/abs/1905.10650
  4. Reducing Transformer Depth on Demand with Structured Dropout: https://arxiv.org/abs/1909.11556
  5. Compressing BERT: Studying the Effects of Weight Pruning on Transfer Learning: https://openreview.net/forum?id=SJlPOCEKvH
  6. Pruning a BERT-based Question Answering Model: https://arxiv.org/abs/1910.06360
  7. Reweighted Proximal Pruning for Large-Scale Language Representation: https://arxiv.org/abs/1909.12486

Tags:

VirMach 便宜 VPS