最簡單的模型輕量化方法：20行程式碼為BERT剪枝

• 2019 年 11 月 22 日
• 筆記

| 導語 BERT模型在多種下游任務表現優異，但龐大的模型結果也帶來了訓練及推理速度過慢的問題，難以滿足對實時響應速度要求高的場景，模型輕量化就顯得非常重要。因此，筆者對BERT系列模型進行剪枝，並部署到實際項目中，在滿足準確率的前提下提高推理速度。

一. 模型輕量化

    模型輕量化是業界一直在探索的一個課題，尤其是當你使用了BERT系列的預訓練語言模型，inference速度始終是個繞不開的問題，而且訓練平台可能還會對訓練機器、速度有限制，訓練時長也是一個難題。

   目前業界上主要的輕量化方法如下：

• 蒸餾：將大模型蒸餾至小模型，思路是先訓練好一個大模型，輸入原始數據得到logits作為小模型的soft label，而原始數據的標籤則為hard label，使用soft label和hard label訓練小模型，旨在將大模型的能力教給小模型。
• 剪枝：不改變模型結構，減小模型的維度，以減小模型量級。
• 量化：將高精度的浮點數轉化為低精度的浮點數，例如4-bit、8-bit等。
• OP重建：合併底層操作，加速矩陣運算。
• 低秩分解：將原始的權重張量分解為多個張量，並對分解張量進行優化。

    我們團隊對這些輕量化方法都進行了嘗試，簡單總結如下：

• 蒸餾：可以很好地將大模型的能力教給小模型，將12層BERT蒸餾至2層BERT，可以達到非常接近的效果。但這種方法需要先訓練出一個大模型。
• 剪枝：速度有非常顯著的提升，結合蒸餾，可以達到很好的效果；即使不結合蒸餾，也能達到不錯的效果。
• 量化：主要用於模型壓縮，可以將大文件壓縮成小文件存儲，方便部署於移動端，但是在速度上無明顯提升。
• OP重建：有明顯加速功能，但是操作較為複雜，需要修改底層C++程式碼。
• 低秩分解：基於PCA演算法，有一倍多的加速作用，但是效果也下降了許多。

    在這些方法中，剪枝顯得非常簡單又高效，如果你想快速得對BERT模型進行輕量化，不僅inference快，還希望訓練快，模型文件小，效果基本維持，那麼剪枝將是一個非常好的選擇，本文將介紹如何為BERT系列模型剪枝，並附上程式碼，教你十分鐘剪枝。

二. BERT剪枝

    本節先重溫BERT[1]及其變體AL-BERT[2]的模型結構，分析在哪裡地方參數量大，再介紹如何為這類結構進行剪枝。

1. BERT模型主要組件

• Input Embedding：詞嵌入，包含token、segment、position三種嵌入方式；
• Multi-Head Attention：多頭注意力機制，共12頭；
• Feed Forward：全連接層，對注意力的輸出向量做進一步映射；
• Output pooler：對hidden向量進行平均/或取cls，得到輸出向量，用於下游任務。

按照默認的維度配置，得到的模型參數大小如下（此處僅展示一層）：

    可以看到BERT模型的參數維度都比較大，都是768起步，而在每一層的結構中，全連接層的3072維，是造成該層參數爆炸的主要原因。單層的參數量已經比普通模型大了許多，當該層參數量再乘以12，殺傷指數更是暴增。

    海量的參數加上海量的無監督訓練數據，BERT模型取得奇效，但我們在訓練我們的下游任務時，是否真的需要這麼大的模型呢？

    可以看到，AL-BERT對Embedding參數進行了因式分解，分解成了2個小矩陣，先將Embedding矩陣投射到一個更小的矩陣E，再投影到隱藏空間H中，減少了參數量（註：同時AL-BERT進行了跨層參數共享，所以保存的參數量少，得到的模型文件非常小），大大加快了模型的訓練速度，但遺憾的是AL-BERT並沒有提高inference速度。

2. 剪枝方法

  基於以上分析，針對BERT系列模型的結構，可採取的剪枝方法如下：

1）層數剪枝

   在BERT模型的應用中，我們一般取第12層的hidden向量用於下游任務。而低層向量基本上包含了基礎資訊，我們可以取低層的輸出向量接到任務層，進行微調。

（跟許老闆討論過一個論文，BERT的低層向量可以學習到一些基礎的詞法資訊，高層向量可以學到更多跟任務相關的特徵，暫時找不到這篇論文了，找到會補上）

2）維度剪枝

    接下來對每一層的維度進行剪枝，ok，全連接層的3072維，在一堆768中成功引起了我們的注意：

    intermediate層的參數量 =（768+1）*3072 *2 = 4724736

    假設我們剪到768維，全連接層的參數量可以減少75%，假如剪到384維，全連接的參數量可以減少87.5%！

3）Attention剪枝

    在12頭注意力中，每頭維度是64，最終疊加註意力向量共768維。

    相關研究[3]表明：

• 在inference階段，大部分head在被單獨去掉的時候，效果不會損失太多；
• 將某一層的head只保留1個，其餘的head去掉，對效果基本不會有什麼影響。

    因此，我們可以嘗試只保留1-2層模型，裁剪ffn維度，減少head個數，在裁剪大量參數的同時維持精度不會下降太多。

三. 工程實現

首先我們看下市面上有沒有啥方便的工具可以剪枝：

• Tensorflow Pruning API：tensorflow官方剪枝工具，該工具基於Keras，如果要用在Tensorflow的模型中，需要將Tensorflow模型轉化為Keras模型，諸多不便。
• Pocketflow Pruning API：騰訊開源的模型壓縮框架，基於tensorflow，為卷積層提供通道剪枝，無法用於BERT結構。
• PaddlePaddle Pruning API：基於百度自家研發的深度學習框架。

這些工具都不適合使用，那就讓我們自己來動手剪枝吧：

• 簡單方法：直接改配置文件的參數設置，不載入Googlepretrain好的語言模型，使用自己的數據重新pretrain語言模型，再載入該模型進行task-specific fine-tune；
• 進階方法：在fine-tune的時候，首先隨機初始化參數，假設從原始的m維裁剪到了n維，那麼取預訓練BERT模型相應的前n維賦值給剪枝後的參數。
• 終極方法：在pretrain階段，取通用BERT模型前n維參數進行賦值再train一遍；在fine-tune階段，就可以直接載入train好的模型進行微調。

下面進入了超級簡單的程式碼環節！關鍵程式碼僅20行！

1）首先，將Googlepretrain的模型參數預存好，保存到一個json文件中：

2）參數賦值，在model_fn_builder函數中，載入預存的參數進行剪枝賦值：

是的！剪枝就是如此簡單！從前筆者為了多方面做對比實驗（例如，第一層剪到768維，第2層剪到384維），強行修改了BERT的模型程式碼，傳入一個字典進行剪枝，遷移到另一個BERT變體模型就不太方便。

      最後附上部分實驗結果（時間可能會有所波動）：

• 不要懷疑，為什麼BERT效果這麼差，因為這份結果是拿口語化badcase測試的，與訓練集相符合的驗證集可以到達99%的準確率~
• AL-BERT訓練速度起飛，在同等訓練數據、模型層數、維度基本等同的前提下，1層AL-BERT 1.5小時即可收斂，而1層BERT模型需要4個小時！在本次場景下，BERT模型收斂得比較慢，這一戰，AL-BERT勝！
• 取前n維向量的剪枝方法是否過於粗暴？是有點，我們也簡單嘗試過，對權重根據絕對值進行排序裁剪，但結果相差不大。或許可以繼續優化~

小結：對BERT系列模型來說，剪枝是一個非常不錯的輕量化方法，很多下游任務可以不需要這麼龐大的模型，也能達到很好的效果。

References

• Devlin J , Chang M W , Lee K , et al. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding[J]. 2018.
• Lan Z , Chen M , Goodman S , et al. ALBERT: A Lite BERT for Self-supervised Learning of Language Representations[J]. 2019.
• Michel P, Levy O, Neubig G. Are Sixteen Heads Really Better than One?[J]. arXiv preprint arXiv:1905.10650, 2019.