緩解模型訓(xùn)練成本過高的問題

我們都知道，為了讓以深度神經(jīng)網(wǎng)絡(luò)為基礎(chǔ)的模型更快地訓(xùn)練，人們提出了單機(jī)多卡、多機(jī)多卡等分布式訓(xùn)練的方式，那么，在模型預(yù)測推理階段，有什么方法可以加速推理呢？遺憾的是，并行/分布式的加速方法并不適用于模型推理階段。

但這并不意味著沒有方法可以加速模型的推理。既然多核的方式不行，可以考慮在單核上面的加速，比如減小模型大?。茨Ｐ蛪嚎s）。

模型壓縮可以分為模型剪枝（pruning）和模型蒸餾（distillation）。由于模型中的參數(shù)對模型推理的貢獻(xiàn)天生就是不平等的，我們可以利用剪枝將貢獻(xiàn)度不高的模型參數(shù)剪去，從而減小模型大小，但剪枝帶來的加速比并不高，最多只有2~3倍的速度提升；蒸餾的方法可以帶來較大的加速比，推理精度也不會有太大的損失，但通常情況下，蒸餾會用到大量無標(biāo)簽的數(shù)據(jù)預(yù)訓(xùn)練學(xué)生模型（student），然后用任務(wù)相關(guān)的帶標(biāo)簽數(shù)據(jù)進(jìn)一步微調(diào)或蒸餾學(xué)生模型，但是，預(yù)訓(xùn)練階段又可能會花費大量的時間。

那么有沒有什么方法既能獲得較大的加速比和較低的精度損失，又可以緩解模型訓(xùn)練成本過高的問題呢？

最近丹琦女神組提出了解決這一問題的方法，讓我們一起來看看吧。

論文標(biāo)題：

Structured Pruning Learns Compact and Accurate Models

論文鏈接：

https://arxiv.org/pdf/2204.00408.pdf

github地址：

https://github.com/princeton-nlp/CoFiPruning

背景介紹

在介紹這篇論文提出的方法之前，還需要簡單說明下作者研究的背景。

作者提出的模型壓縮方法針對的是原模型（教師模型）為transformer 架構(gòu)的模型壓縮。眾所周知， transformer 由多個塊組成，每個塊由一個多頭自注意力（multi-head self-attention，MHA）和兩個前饋神經(jīng)網(wǎng)絡(luò)（FFN）組成。其中， MHAs 和 FFNs 的參數(shù)量比為 1:2。在 GPU 上，兩者的推理時間基本相同，而在 CPU 上， FFNs 則會耗費更多的推理時間。

正如前面所述，模型壓縮可以歸納為兩種方法：知識蒸餾和剪枝。知識蒸餾在通常情況下，需要預(yù)先定義一個結(jié)構(gòu)固定的學(xué)生網(wǎng)絡(luò)（當(dāng)然，也有一些嘗試動態(tài)學(xué)生網(wǎng)絡(luò)的研究），通過用大量無標(biāo)簽數(shù)據(jù)預(yù)訓(xùn)練學(xué)生網(wǎng)絡(luò)的方式進(jìn)行模型參數(shù)初始化，然后用任務(wù)特定的帶標(biāo)簽數(shù)據(jù)微調(diào)學(xué)生模型，當(dāng)然，學(xué)生模型的初始化方式可以有很多種，例如用教師模型的某些層初始化學(xué)生模型等，但基本的解決思想是一致的。

剪枝則指的是從原模型中去除冗余的參數(shù)，按照剪枝的粒度可以分為以下幾種：

層剪枝（layer pruning）：從模型中刪去整個 block 塊（包括 MHA 和 FFN ），一些研究表明，去除 50% 的層并不會有太多的精度下降，而且還可以獲得 2X 的加速比；

頭剪枝（head pruning）：通過 mask 矩陣只保留一部分 head，但研究表明，這種做法并不能帶來較大的加速比，當(dāng)僅保留一個 head 的時候，加速比為 1.4X；

前饋神經(jīng)網(wǎng)絡(luò)剪枝（FFN pruning）：去除整個 FFN 層或者去除 FFN 層的某些維度；

更加細(xì)粒度的塊和非結(jié)構(gòu)化的剪枝：去除 MHA 和 FFN 中更小的塊或者去除某些參數(shù)權(quán)重，目前這種做法很難優(yōu)化模型，也難以獲得推理加速

除此之外，還可以將剪枝和蒸餾融合起來，但目前，該方法具體實現(xiàn)尚不清晰。

結(jié)構(gòu)化剪枝方法 CoFi

為獲得較大的加速比和較低的精度損失，以及緩解模型訓(xùn)練成本過高的問題，作者提出了結(jié)構(gòu)化剪枝方法CoFi（Coarse- andFine-grained Pruning），方法由兩個部分組成：粗粒度和細(xì)粒度的剪枝以及從原模型（未剪枝）到剪枝模型的逐層蒸餾

粗粒度和細(xì)粒度的剪枝

在頭剪枝（head pruning）中，經(jīng)常通過由 組成的 mask 矩陣來保留一部分 head ，但是，當(dāng) mask 矩陣全為 0，即去除掉所有的 head 時，會使模型優(yōu)化變得困難。為此，作者為每一層的 MHA 和 FFN 引入兩個掩碼變量 和 ，多頭自注意力和前饋神經(jīng)網(wǎng)絡(luò)可以表示為：

其中， 是輸入向量， 是 head 數(shù)量， ， ， ， 分別是 query、key、value 和輸出的權(quán)重矩陣， 是 attention 函數(shù)， 和 分別是 FFN 的兩個權(quán)重， 和 分別是掩碼矩陣變量。

作者用 和 控制每一層的 MHA 和 FFN 是否剪枝以及 和 控制每一層的 MHA 和 FFN 中的哪些 head 和哪些維度需要剪枝。

除此之外，作者還對 和 的輸出在維度上作剪枝操作。具體做法是將掩碼變量 應(yīng)用到模型中所有的權(quán)重矩陣，掩碼跨層共享的原因是作者考慮到模型中的殘差使得隱向量中的每個維度都可以連接到下一層相應(yīng)的維度。

此外，作者定義了預(yù)期稀疏度：

其中， 是整個模型大小， 是 block 層數(shù)， 是隱藏層維度， 是多頭自注意力的每個 head 的向量維度， 是前饋網(wǎng)絡(luò)的維度，一般情況下，。

模型訓(xùn)練階段，所有的掩碼元素的值處于之間，推理階段，會將低于閾值的掩碼變量映射為 0，得到最終的剪枝模型，其中，閾值由每個權(quán)重矩陣的預(yù)期稀疏度確定。

從原模型到剪枝模型的蒸餾

考慮到將剪枝和蒸餾融合可以提高性能，作者提出了用于剪枝的逐層蒸餾方法。與一般的蒸餾做法不同，作者沒有預(yù)先定義從教師網(wǎng)絡(luò)到學(xué)生網(wǎng)絡(luò)的固定的層映射，而是動態(tài)地搜索兩者之間的層映射。具體來說，假設(shè) 表示準(zhǔn)備將知識蒸餾到學(xué)生網(wǎng)絡(luò)的教師網(wǎng)絡(luò)的層的集合， 是層映射函數(shù)，表示從教師網(wǎng)絡(luò)的第 層映射到的學(xué)生網(wǎng)絡(luò)的層，那么，隱藏層的蒸餾 可以定義為：

其中， 是線性變換矩陣， 和 分別是第 層學(xué)生網(wǎng)絡(luò)和第 層教師網(wǎng)絡(luò)的隱藏層表示。 定義如下：

其中，計算兩個層集合之間的距離（MSE）的操作是可以并行執(zhí)行的。通過上面的層映射函數(shù)，教師網(wǎng)絡(luò)和學(xué)生網(wǎng)絡(luò)之間的層映射總是按照最有利于剪枝的方向進(jìn)行。

最后，作者將逐層蒸餾和來自預(yù)測層（模型輸出）的蒸餾結(jié)合起來，得到最終的loss：

其中， 是超參， 和 分別是學(xué)生模型和教師模型的輸出概率分布。

實驗

數(shù)據(jù)集

作者使用的是 GLUE 數(shù)據(jù)集和 SQuAD v1.1 數(shù)據(jù)集，其中， GLUE 數(shù)據(jù)集包括 SST2、MNLI、QQP、QNLI、MRPC、CoLA、STS-B 和 RTE 八個數(shù)據(jù)集。

在四個相對較大的 GLUE 數(shù)據(jù)集（包括 MNLI、QNLI、SST-2 和 QQP）以及 SQuAD 數(shù)據(jù)集上，作者訓(xùn)練了 20 個 epoch，并對最終的學(xué)生網(wǎng)絡(luò)微調(diào)了額外 20 個 epoch。在前 20 個 epoch 中，作者使用蒸餾的目標(biāo)函數(shù)對模型進(jìn)行微調(diào) 1 個 epoch，然后在 2 個 epoch 之內(nèi)使模型達(dá)到期望的目標(biāo)稀疏度。而對于四個較小的 GLUE 數(shù)據(jù)集，作者訓(xùn)練了 100 個 epoch 并微調(diào) 20 個 epoch。作者用蒸餾目標(biāo)函數(shù)微調(diào)模型 4 個epoch，并在接下來的 20 個 epoch 內(nèi)將模型剪枝到期望的目標(biāo)稀疏度。在達(dá)到目標(biāo)稀疏度后，還會在剩余的訓(xùn)練 epoch 階段繼續(xù)剪枝以搜索性能更好的網(wǎng)絡(luò)結(jié)構(gòu)。

作者在每個數(shù)據(jù)集上依次訓(xùn)練并微調(diào)不同目標(biāo)稀疏度的模型。此外，作者在實驗中發(fā)現(xiàn)，訓(xùn)練結(jié)束后的微調(diào)可以有效地保持模型不會有太大的精度損失。

模型對比

作者將 、、、、 作為 baseline，對比了 CoFi 和 baseline 在不同的加速比和模型大小的情況下準(zhǔn)確率（acc）和 f1 的變化

可以看到，在相同的加速比或相同的模型大小下，作者提出的方法 CoFi 都可以獲得更高的準(zhǔn)確率（acc）或 f1

此外，作者還展示了 CoFi 和 的對比結(jié)果。

其中， 使用大量無標(biāo)簽數(shù)據(jù)通過預(yù)訓(xùn)練的方式初始化學(xué)生網(wǎng)絡(luò)，和大規(guī)模預(yù)訓(xùn)練語言模型一樣，這種從大規(guī)模語料庫中獲取通用知識的方式對學(xué)生模型性能有著十分重要的作用，但預(yù)訓(xùn)練會花費太多的時間，如下圖所示：

從表 2 可以看到，CoFi 可以獲得和 基本一致的加速比和略高的模型預(yù)測準(zhǔn)確度，與此同時，CoFi 模型的訓(xùn)練時間較 得到了大幅的縮減，也證明了以蒸餾目標(biāo)函數(shù)訓(xùn)練的剪枝方法在模型壓縮方面是經(jīng)濟(jì)且高效的。

除此之外，作者還使用相同的任務(wù)特定的數(shù)據(jù)分別為 和 CoFi 做數(shù)據(jù)增強(qiáng)，并對比了數(shù)據(jù)增強(qiáng)后的模型性能。

可以看到，數(shù)據(jù)增強(qiáng)后，CoFi 也基本有著高于 的模型性能。

▲sota剪枝和蒸餾方法的性能對比

消融實驗

作者做了一系列的消融實驗，證明了 CoFi 各個模塊的有效性。實驗結(jié)果如下所示：

▲CoFi消融實驗-剪枝單元

▲CoFi消融實驗-蒸餾目標(biāo)函數(shù)

剪枝后的模型的結(jié)構(gòu)

作者研究了經(jīng)過 CoFi 剪枝后得到的模型的結(jié)構(gòu)，作者分別在五個數(shù)據(jù)集上在不同的目標(biāo)稀疏度下訓(xùn)練微調(diào)得到剪枝后的模型，并對模型的 FFN 層的平均中間維度以及 MHA 層的 head 平均剩余數(shù)量做了統(tǒng)計，結(jié)果如下：

可以看到：

幾乎在所有的稀疏度下，前饋神經(jīng)網(wǎng)絡(luò)層都存在明顯的剪枝，這表明 FFN 層比 MHA 層存在更多的冗余信息；

CoFi 傾向于更多地剪枝上層網(wǎng)絡(luò)結(jié)構(gòu)

作者還詳細(xì)地展示了不同數(shù)據(jù)集下經(jīng) CoFi 剪枝后的模型的具體結(jié)構(gòu)，下圖展示的是每個模型每層 MHA 和 FFN 模塊的保留情況。盡管模型大小基本一致，但不同的數(shù)據(jù)集卻訓(xùn)練出結(jié)構(gòu)存在明顯差異的模型，從側(cè)面也證明了不同的數(shù)據(jù)集存在不同的最優(yōu)學(xué)生模型。

▲經(jīng)CoFi剪枝后的模型不同層子模塊保留情況

總結(jié)

作者提出的結(jié)構(gòu)化剪枝方法 CoFi 在幾乎沒有太多精度損失的情況下，達(dá)到了 10 倍以上的加速比，同時，和常規(guī)的蒸餾做法相比，避免了因使用大量無標(biāo)簽數(shù)據(jù)預(yù)訓(xùn)練模型而帶來的訓(xùn)練成本過高的問題，按作者的話來說，該方法可以是蒸餾的一個有效替代品。

當(dāng)然，作者也指出，盡管 CoFi 可以應(yīng)用到任務(wù)無關(guān)的模型的剪枝中，比如大規(guī)模預(yù)訓(xùn)練語言模型，但由于上游剪枝方案設(shè)計的復(fù)雜性，作者還是將 CoFi 的應(yīng)用場景限制到任務(wù)相關(guān)的模型壓縮中。試想下，如果可以將任務(wù)無關(guān)的模型壓縮到可以部署到移動設(shè)備或者可穿戴智能設(shè)備上，那么，世界又會是一番怎樣的景象呢？不知道這是不是丹琦女神給我們新開的一個“坑”呢？

審核編輯 ：李倩