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

近日，微軟亞洲研究院對 NNI 進行了更新。在最新的版本中，NNI 集成了大量前沿的剪枝算法，如 TaylorFO Weight、Movement 等?；诂F(xiàn)有的經(jīng)典預(yù)訓(xùn)練模型，研究員們通過大量實驗，發(fā)現(xiàn)了既能降低模型參數(shù)量和計算量，又能保持模型較高精度的剪枝步驟與算法組合，獲得超越 SOTA 的模型剪枝效果。

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

剪枝流程

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

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

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

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

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

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

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

1. 準備數(shù)據(jù)/模型等

在正式構(gòu)建剪枝過程之前，用戶需要加載預(yù)訓(xùn)練模型，對數(shù)據(jù)預(yù)處理并創(chuàng)建相應(yīng)的 dataloader，同時設(shè)計相應(yīng)的訓(xùn)練/評估函數(shù)，以用于后期對模型的訓(xùn)練和評估。其流程如圖2所示，共包含5步：

圖2：數(shù)據(jù)/模型準備過程的流程示意圖

具體來說，首先需要從 Transformers 庫中加載預(yù)訓(xùn)練模型，然后對數(shù)據(jù) GLUE-MNLI 進行處理，并得到相應(yīng)的 dataloader。隨后，針對模型和數(shù)據(jù)集 GLUE-MNLI，構(gòu)建相應(yīng)的訓(xùn)練/評估函數(shù)。最后將模型在 GLUE-MNLI 數(shù)據(jù)集上進行微調(diào)。

完成以上步驟就相當于完成了數(shù)據(jù)/模型等的準備工作，可以得到預(yù)訓(xùn)練模型在 MNLI 數(shù)據(jù)集上微調(diào)后的模型?？紤]到 Transformer 系列預(yù)訓(xùn)練模型的模型參數(shù)中的大頭為嵌入層，且編碼層/解碼層中包含了多頭自注意力機制和前饋神經(jīng)網(wǎng)絡(luò)。因此，在之后的步驟中需要分別對多頭自注意力機制、嵌入層和前饋神經(jīng)網(wǎng)絡(luò)剪枝，并引入動態(tài)蒸餾機制對剪枝后的模型再訓(xùn)練。

2. 多頭自注意力機制的剪枝和基于動態(tài)蒸餾機制的模型再訓(xùn)練

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

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

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

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

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

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

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

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

3. 嵌入層和前饋神經(jīng)網(wǎng)絡(luò)的剪枝，以及基于動態(tài)蒸餾機制的模型再訓(xùn)練

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

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

實驗結(jié)果

通過調(diào)整 regular_scale 參數(shù)的值和前饋神經(jīng)網(wǎng)絡(luò)的剪枝次數(shù)，研究員們得到了具有不同稀疏度和性能的模型。該過程使用了1張 A100 進行實驗，并設(shè)置 batch_size 為32。

圖5：實驗結(jié)果

從上圖實驗結(jié)果可以看出：

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

2. 隨著前饋神經(jīng)網(wǎng)絡(luò)剪枝次數(shù)的增加，模型總的稀疏度有所增加，同時模型的性能有所下降，且隨著模型總稀疏度的增加，模型的性能下降程度逐漸增大。

3. 對嵌入層剪枝3次，能夠?qū)⒛Ｐ偷木S度從768減小至561，在一定程度上提升了模型總的稀疏度。

實驗結(jié)果與平臺對比

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

（a）

(b)

(c)

圖6：NNI 在經(jīng)典預(yù)訓(xùn)練模型下的剪枝性能示意圖

三個平臺（Paper）的詳細比較結(jié)果，如表1所示?？梢钥闯?，NNI 的 Compression 模塊不僅具有完整的教程實例，同時還提供了 SpeedUp 模塊，能夠?qū)崿F(xiàn)真正意義上的減少模型參數(shù)量，而非將需要修剪的參數(shù)置為0。

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

表1：各平臺（Paper）功能對比總結(jié)

展望未來

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