高精度壓縮Transformer，NNI剪枝一站式指南

無論在學術界還是產業界，今年人工智能大模型都是爆款話題。但面對這些動不動就數十億級別參數的模型，使用傳統方法微調，宛如水中撈月、海底撈針。作為微軟亞洲研究院為科研人員和算法工程師量身定制的一站式 AutoML（自動機器學習）工具， NNI（Neural Network Intelligence）在過去的三年間不斷迭代更新，加強了對各種分布式訓練環境的支持，成為了最熱門的 AutoML 開源項目之一。

近日，微軟亞洲研究院對 NNI 進行了更新。在最新的版本中，NNI 集成了大量前沿的剪枝算法，如 TaylorFO Weight、Movement 等。基于現有的經典預訓練模型，研究員們通過大量實驗，發現了既能降低模型參數量和計算量，又能保持模型較高精度的剪枝步驟與算法組合，獲得超越 SOTA 的模型剪枝效果。

今天我們就以 Transformer 系列的預訓練模型和數據集 GLUE-MNLI 為例，為大家介紹一下 NNI 的 pruner 剪枝流程和使用的剪枝算法組合。

剪枝流程

在正式介紹剪枝流程前，我們需要先了解什么是 pruner，mask 和 SpeedUp。

• pruner：使用具體的剪枝算法實例化的剪枝器。

• mask：在剪枝過程中，pruner 會生成一個和目標子模塊大小相同的 mask（全1）矩陣，并在 mask 矩陣中將目標子模塊中需要剪掉的部分的對應位置置為0。最后通過將目標子模塊和對應的 mask 矩陣相乘，即可得到模擬剪枝后的模型效果。

• SpeedUp：從上述描述可以看出，在剪枝過程中，實際上只是將需要剪枝的部分用0進行了替換，因此使用 SpeedUp 模塊是修剪上述目標子模塊中需要剪掉的參數，而不是用0替代，從而實現真正意義上的減少參數量。

在使用 NNI Compression 模塊中的 pruner 進行剪枝操作時，用戶只需完成數據/模型等的準備、pruner 的構建，以及模型剪枝和再訓練，即可為模型構建一個剪枝的 pipeline。

以 Transformer 系列的預訓練模型為例，其剪枝流程共包含4步：首先準備數據/模型等，接著針對多頭自注意力機制（Multi-head Attention）、嵌入層（embedding）和前饋神經網絡（FFN）分別剪枝和再訓練模型。

圖1：Transformer 系列模型的剪枝流程示意圖

1. 準備數據/模型等

在正式構建剪枝過程之前，用戶需要加載預訓練模型，對數據預處理并創建相應的 dataloader，同時設計相應的訓練/評估函數，以用于后期對模型的訓練和評估。其流程如圖2所示，共包含5步：

圖2：數據/模型準備過程的流程示意圖

具體來說，首先需要從 Transformers 庫中加載預訓練模型，然后對數據 GLUE-MNLI 進行處理，并得到相應的 dataloader。隨后，針對模型和數據集 GLUE-MNLI，構建相應的訓練/評估函數。最后將模型在 GLUE-MNLI 數據集上進行微調。

完成以上步驟就相當于完成了數據/模型等的準備工作，可以得到預訓練模型在 MNLI 數據集上微調后的模型。考慮到 Transformer 系列預訓練模型的模型參數中的大頭為嵌入層，且編碼層/解碼層中包含了多頭自注意力機制和前饋神經網絡。因此，在之后的步驟中需要分別對多頭自注意力機制、嵌入層和前饋神經網絡剪枝，并引入動態蒸餾機制對剪枝后的模型再訓練。

2. 多頭自注意力機制的剪枝和基于動態蒸餾機制的模型再訓練

多頭自注意力模塊的剪枝和模型再訓練分為3步，如圖3所示：首先要構建 pruner，接著對多頭自注意力模塊進行剪枝，最后使用動態蒸餾機制再訓練模型。

圖3：多頭自注意力機制的剪枝和再訓練流程示意圖

在進行剪枝前，用戶需要選定一個剪枝算法并實例化相應的 pruner。所有的剪枝算法均需向模型中傳入 config_list 參數，因為其定義了需要剪枝的運算名、運算類別及稀疏度等。具體到 Movement 剪枝算法，還需要設置其他的一些參數，如：evaluator 參數，用于訓練感知的模型壓縮過程；movement_mode 參數，共有“soft“和”hard“兩種模式，若為”soft”，則難以精確地控制模型剪枝后的稀疏度，但是可以得到性能更好的模型。參數 regular_scale 用于控制剪枝的稀疏度，regular_scale 越大，模型剪枝后的稀疏度越高。更多其他參數可參閱

https://nni.readthedocs.io/zh/stable/reference/compression/pruner.html#movement-pruner

接下來，要使用構造的剪枝算法實例 pruner 對多頭自注意力模塊進行剪枝。用戶只需調用 pruner.compress() 即可執行對模型的剪枝過程，并得到剪枝后的模型和 attention_mask。其中 attention_mask 給出了需要剪枝的子模塊的參數剪枝范圍，0代表該位置被剪掉，1代表該位置被保留。

NNI 的 SpeedUp 模塊可以將被 mask 住的參數和計算從模型中刪除，具體的刪除邏輯如圖4所示，以 Query Linear 層的 weight（記作Q）為例，其維度為[768,768]，那么 Q 的 weight 的 mask 矩陣維度也為[768, 768]，將其記作 mask。首先將該 mask 矩陣的維度進行變換，第一維是多頭數目8，其余的則是第二維，將變換后的 mask 矩陣記作 reshaped mask 矩陣。接著，對 reshaped mask 矩陣在第二維度上求和，并判斷求和后的值是否為0，此時的 mask 矩陣維度變為[8]，每個位置對應著一個多頭。對于變換后的 mask 矩陣，若位置 i 的值為0，則代表在 Q 中的第 i 個多頭需要被剪掉。在圖中，位置0、3、7的值均為0，因此，在Q中的第0、3、7個多頭需要被剪掉。最后，將[0,3,7]作為參數傳入 prune_heads 函數中，對 Q 進行修剪。修剪后，Q 的維度為[576,768]。對 SpeedUp 更加全面的介紹可以參考發表于 OSDI 2022 的論文 SparTA。在即將發布的 NNI 3.0 中 SpeedUp 會對更多模型提供更加完善的支持。

圖4：利用 prune_heads 函數修剪自注意力模塊的過程示意圖

在對多頭自注意力模塊剪枝后，以微調后的模型作為教師模型，以剪枝后的模型作為學生模型，然后借鑒 CoFi 中的動態蒸餾機制 [1] 對模型進行再訓練，就可以得到新的模型。這里的動態蒸餾機制，是指教師模型的層和學生模型的層之間不是一個靜態對應關系，每次蒸餾教師都可以選擇從自身的高層動態蒸餾信息到學生模型低層中的一層里。

3. 嵌入層和前饋神經網絡的剪枝，以及基于動態蒸餾機制的模型再訓練

嵌入層和前饋神經網絡的剪枝過程與多頭自注意力模塊的剪枝過程類似。此處使用 Taylor 剪枝算法 （https://nni.readthedocs.io/zh/stable/reference/compression/pruner.html#taylor-fo-weight-pruner ） 對嵌入層和前饋神經網絡進行剪枝。同樣地，研究員們定義了 config_list、evaluator 參數及 taylor_pruner_steps 參數。由于嵌入層的維度與后續模型中的維度具有相關性。因此，基于上述參數，在嵌入層的剪枝過程中研究員們將剪枝模式 mode 設置為了“dependency-aware”模式，并傳入模型的輸入 dummy_input，以幫助 pruner 捕捉和嵌入層維度具有依賴關系的子模型。

接下來，使用分別構造的 pruner 對前饋神經網絡和嵌入層進行剪枝。和多頭自注意力模塊的剪枝不同的是，此處使用了迭代式剪枝法，即在模型基于動態蒸餾的再訓練過程中，每2000步分別使用 pruner 對前饋神經網絡和嵌入層剪枝一次，其中，前饋神經網絡共剪枝19/24次，嵌入層共剪枝3次。每次剪枝后，使用 ModelSpeedUp 對前饋神經網絡層進行剪枝，以實現真正意義上的修剪參數，而不是將需要修剪的參數用0替換。

實驗結果

通過調整 regular_scale 參數的值和前饋神經網絡的剪枝次數，研究員們得到了具有不同稀疏度和性能的模型。該過程使用了1張 A100 進行實驗，并設置 batch_size 為32。

圖5：實驗結果

從上圖實驗結果可以看出：

1. 隨著 regular_scale 的增加，模型總的稀疏度有所增加。當 regular_scale 大于等于10時，模型總的稀疏度超過了69%，性能損失超過1%。

2. 隨著前饋神經網絡剪枝次數的增加，模型總的稀疏度有所增加，同時模型的性能有所下降，且隨著模型總稀疏度的增加，模型的性能下降程度逐漸增大。

3. 對嵌入層剪枝3次，能夠將模型的維度從768減小至561，在一定程度上提升了模型總的稀疏度。

實驗結果與平臺對比

進一步分析實驗結果可以發現，使用 NNI 對 BERT 在 MNLI 數據集上剪枝后的性能好于 nn pruning 框架（圖6(a)），且當模型總的稀疏度低于65%時，NNI 和 CoFi 對 BERT 在 MNLI 數據集上剪枝的性能差距較小，當模型總的稀疏度大于65%時，使用 NNI 對 BERT 在 MNLI 數據集上剪枝后的性能好于 CoFi。圖6(b)和圖6(c)分別展示了 NNI 在 T5 和 ViT 模型上的剪枝性能。從圖中可以看出，當模型相應部分的稀疏度超過了75%后，模型性能下降約為3%，當模型相應部分的稀疏度低于50%時，模型性能下降較少。

（a）

(b)

(c)

圖6：NNI 在經典預訓練模型下的剪枝性能示意圖

三個平臺（Paper）的詳細比較結果，如表1所示。可以看出，NNI 的 Compression 模塊不僅具有完整的教程實例，同時還提供了 SpeedUp 模塊，能夠實現真正意義上的減少模型參數量，而非將需要修剪的參數置為0。

同時，NNI 支持 BERT、RoBerta、GPT、BART、T5、ViT 等主流模型，并提供了 Taylor、Movement、ADMM、Slim、AGP、Activation APoZ、Activation Mean 等16種前沿剪枝算法，能夠更好地滿足用戶的需求，具有較強的通用性。

表1：各平臺（Paper）功能對比總結

展望未來

在 NNI 3.0 版本中，微軟亞洲研究院的研究員們還將引入蒸餾模塊，更好地為用戶提供集剪枝、蒸餾為一體的壓縮工具，同時 SpeedUp 模塊也將更全面地支持對 Transformer 的修剪。敬請期待！
關于最新版 NNI 的完整代碼和 tutorial，請參見：https://nni.readthedocs.io/zh/stable/tutorials/pruning_bert_glue.html
如有任何問題，歡迎在評論區提問、交流！