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

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

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

剪枝流程

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

實(shí)驗(yàn)結(jié)果

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

圖5：實(shí)驗(yàn)結(jié)果

從上圖實(shí)驗(yàn)結(jié)果可以看出：

1. 隨著 regular_scale 的增加，模型總的稀疏度有所增加。當(dāng) regular_scale 大于等于10時(shí)，模型總的稀疏度超過(guò)了69%，性能損失超過(guò)1%。

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

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

實(shí)驗(yàn)結(jié)果與平臺(tái)對(duì)比

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

（a）

(b)

(c)

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

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

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

表1：各平臺(tái)（Paper）功能對(duì)比總結(jié)

展望未來(lái)

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